update data and model

app.py

@@ -1,26 +1,51 @@
 # Author: Huzheng Yang
 # %%
 USE_SPACES = True
+BATCH_SIZE = 4
 
-if USE_SPACES:  # huggingface ZeroGPU
+if USE_SPACES:  # huggingface ZeroGPU, dynamic GPU allocation
     try:
         import spaces
     except ImportError:
-        USE_SPACES = False  # run on standard GPU
+        USE_SPACES = False  # run on local machine
+        BATCH_SIZE = 1
 
 import os
 import gradio as gr
 
 import torch
+import torch.nn.functional as F
 from PIL import Image
 import numpy as np
 import time
 
 import gradio as gr
 
-from backbone import extract_features
+from backbone import extract_features, download_all_models, get_model
+from backbone import MODEL_DICT, LAYER_DICT, RES_DICT
 from ncut_pytorch import NCUT, eigenvector_to_rgb
 
+DATASET_TUPS = [
+    # (name, num_classes)
+    ('UCSC-VLAA/Recap-COCO-30K', None),
+    ('nateraw/pascal-voc-2012', None),
+    ('johnowhitaker/imagenette2-320', 10),
+    ('jainr3/diffusiondb-pixelart', None),
+    ('nielsr/CelebA-faces', None),
+    ('JapanDegitalMaterial/Places_in_Japan', None),
+    ('Borismile/Anime-dataset', None),
+    ('Multimodal-Fatima/CUB_train', 200),
+    ('mrm8488/ImageNet1K-val', 1000),
+]
+DATASET_NAMES = [tup[0] for tup in DATASET_TUPS]
+DATASET_CLASSES = [tup[1] for tup in DATASET_TUPS]
+
+from datasets import load_dataset
+
+def download_all_datasets():
+    for name in DATASET_NAMES:
+        print(f"Downloading {name}")
+        load_dataset(name, trust_remote_code=True)
 
 def compute_ncut(
     features,
@@ -55,6 +80,7 @@ def compute_ncut(
         knn=knn_ncut,
         sample_method=sampling_method,
         distance=metric,
+        normalize_features=False,
     ).fit_transform(features.reshape(-1, features.shape[-1]))
     # print(f"NCUT time: {time.time() - start:.2f}s")
     logging_str += f"NCUT time: {time.time() - start:.2f}s\n"
@@ -147,6 +173,7 @@ example_items = downscaled_images[:3] + downscaled_outputs[:3]
 
 
 def ncut_run(
+    model,
     images,
     model_name="SAM(sam_vit_b)",
     layer=-1,
@@ -176,15 +203,15 @@ def ncut_run(
         logging_str += f"Perplexity/n_neighbors must be less than the number of samples.\n" f"Setting Perplexity to {num_sample_tsne-1}.\n"
         perplexity = num_sample_tsne - 1
         n_neighbors = num_sample_tsne - 1
-    
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
         
     node_type = node_type.split(":")[0].strip()
-    
-    images = [image[0] for image in images]     # remove the label
-    
+        
     start = time.time()
     features = extract_features(
-        images, model_name=model_name, node_type=node_type, layer=layer
+        images, model, model_name=model_name, node_type=node_type, layer=layer-1, batch_size=BATCH_SIZE
     )
     # print(f"Feature extraction time (gpu): {time.time() - start:.2f}s")
     logging_str += f"Backbone time: {time.time() - start:.2f}s\n"
@@ -213,6 +240,8 @@ def ncut_run(
             rgb = dont_use_too_much_green(rgb)
             rgbs.append(to_pil_images(rgb))
             inp = eigvecs.reshape(*features.shape[:3], -1)
+            if recursion_metric == "cosine":
+                inp = F.normalize(inp, dim=-1)
         return rgbs[0], rgbs[1], rgbs[2], logging_str
     
     if old_school_ncut:  # individual images
@@ -273,6 +302,8 @@ def _ncut_run(*args, **kwargs):
         return ret
     except Exception as e:
         gr.Error(str(e))
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
         return [], "Error: " + str(e)
 
 if USE_SPACES:
@@ -318,6 +349,15 @@ def extract_video_frames(video_path, max_frames=100):
     # return as list of PIL images
     return [(Image.fromarray(frames[i]), "") for i in range(frames.shape[0])]
 
+def transform_image(image, resolution=(1024, 1024)):
+    image = image.convert('RGB').resize(resolution, Image.LANCZOS)
+    # Convert to torch tensor
+    image = torch.tensor(np.array(image).transpose(2, 0, 1)).float()
+    image = image / 255
+    # Normalize
+    image = (image - 0.5) / 0.5
+    return image
+
 def run_fn(
     images,
     model_name="SAM(sam_vit_b)",
@@ -341,7 +381,7 @@ def run_fn(
     recursion_l3_n_eigs=20,
     recursion_metric="euclidean",
 ):
-    # print("Running...")
+    
     if images is None:
         gr.Warning("No images selected.")
         return [], "No images selected."
@@ -353,6 +393,14 @@ def run_fn(
     
     if sampling_method == "fps":
         sampling_method = "farthest"
+        
+    # resize the images before acquiring GPU
+    resolution = RES_DICT[model_name]
+    images = [tup[0] for tup in images]
+    images = [transform_image(image, resolution=resolution) for image in images]
+    images = torch.stack(images)
+        
+    model = get_model(model_name)
     
     kwargs = {
         "model_name": model_name,
@@ -379,25 +427,25 @@ def run_fn(
     # print(kwargs)
     num_images = len(images)
     if num_images > 100:
-        return super_duper_long_run(images, **kwargs)
+        return super_duper_long_run(model, images, **kwargs)
     if recursion:
-        return longer_run(images, **kwargs)
+        return longer_run(model, images, **kwargs)
     if num_images > 50:
-        return longer_run(images, **kwargs)
+        return longer_run(model, images, **kwargs)
     if old_school_ncut:
-        return longer_run(images, **kwargs)
+        return longer_run(model, images, **kwargs)
     if num_images > 10:
-        return long_run(images, **kwargs)
+        return long_run(model, images, **kwargs)
     if embedding_method == "UMAP":
         if perplexity >= 250 or num_sample_tsne >= 500:
-            return longer_run(images, **kwargs)
-        return long_run(images, **kwargs)
+            return longer_run(model, images, **kwargs)
+        return long_run(model, images, **kwargs)
     if embedding_method == "t-SNE":
         if perplexity >= 250 or num_sample_tsne >= 500:
-            return long_run(images, **kwargs)
-        return quick_run(images, **kwargs)
+            return long_run(model, images, **kwargs)
+        return quick_run(model, images, **kwargs)
     
-    return quick_run(images, **kwargs)
+    return quick_run(model, images, **kwargs)
 
 
 
@@ -435,23 +483,42 @@ def make_example_video_section():
     return load_video_button
 
 def make_dataset_images_section(open=False):
+    
     with gr.Accordion("➡️ Click to expand: Load from dataset", open=open):
-        dataset_names = [
-            'UCSC-VLAA/Recap-COCO-30K',
-            'nateraw/pascal-voc-2012',
-            'johnowhitaker/imagenette2-320',
-            'jainr3/diffusiondb-pixelart',
-            'nielsr/CelebA-faces',
-            'JapanDegitalMaterial/Places_in_Japan',
-            'Borismile/Anime-dataset',
-        ]
-        dataset_dropdown = gr.Dropdown(dataset_names, label="Dataset name", value="UCSC-VLAA/Recap-COCO-30K", elem_id="dataset")
-        num_images_slider = gr.Slider(1, 200, step=1, label="Number of images", value=9, elem_id="num_images")
-        # random_seed_slider = gr.Number(0, label="Random seed", elem_id="random_seed")
-        random_seed_slider = gr.Slider(0, 1000, step=1, label="Random seed", value=1, elem_id="random_seed")
+        dataset_names = DATASET_NAMES
+        dataset_classes = DATASET_CLASSES
+        dataset_dropdown = gr.Dropdown(dataset_names, label="Dataset name", value="mrm8488/ImageNet1K-val", elem_id="dataset")
+        num_images_slider = gr.Number(10, label="Number of images", elem_id="num_images")
+        filter_by_class_checkbox = gr.Checkbox(label="Filter by class", value=True, elem_id="filter_by_class_checkbox")
+        filter_by_class_text = gr.Textbox(label="Class to select", value="0,33,99", elem_id="filter_by_class_text", info=f"e.g. `0,1,2`. (1000 classes)", visible=True)
+        is_random_checkbox = gr.Checkbox(label="Random shuffle", value=False, elem_id="random_seed_checkbox")
+        random_seed_slider = gr.Slider(0, 1000, step=1, label="Random seed", value=1, elem_id="random_seed", visible=False)
         load_dataset_button = gr.Button("Load Dataset", elem_id="load-dataset-button")
-    def load_dataset_images(dataset_name, num_images=10, random_seed=42):
-        from datasets import load_dataset
+    
+    def change_filter_options(dataset_name):
+        idx = dataset_names.index(dataset_name)
+        num_classes = dataset_classes[idx]
+        if num_classes is None:
+            return (gr.Checkbox(label="Filter by class", value=False, elem_id="filter_by_class_checkbox", visible=False),
+                    gr.Textbox(label="Class to select", value="0,1,2", elem_id="filter_by_class_text", info="e.g. `0,1,2`. This dataset has no class label", visible=False))
+        return (gr.Checkbox(label="Filter by class", value=True, elem_id="filter_by_class_checkbox", visible=True),
+                gr.Textbox(label="Class to select", value="0,1,2", elem_id="filter_by_class_text", info=f"e.g. `0,1,2`. ({num_classes} classes)", visible=True))
+    dataset_dropdown.change(fn=change_filter_options, inputs=dataset_dropdown, outputs=[filter_by_class_checkbox, filter_by_class_text])
+    
+    def change_filter_by_class(is_filter, dataset_name):
+        idx = dataset_names.index(dataset_name)
+        num_classes = dataset_classes[idx]
+        return gr.Textbox(label="Class to select", value="0,1,2", elem_id="filter_by_class_text", info=f"e.g. `0,1,2`. ({num_classes} classes)", visible=is_filter)
+    filter_by_class_checkbox.change(fn=change_filter_by_class, inputs=[filter_by_class_checkbox, dataset_dropdown], outputs=filter_by_class_text)
+    
+    def change_random_seed(is_random):
+        return gr.Slider(0, 1000, step=1, label="Random seed", value=1, elem_id="random_seed", visible=is_random)
+    is_random_checkbox.change(fn=change_random_seed, inputs=is_random_checkbox, outputs=random_seed_slider)
+    
+        
+    def load_dataset_images(dataset_name, num_images=10, 
+                            is_filter=True, filter_by_class_text="0,1,2", 
+                            is_random=False, seed=1):
         try:
             dataset = load_dataset(dataset_name, trust_remote_code=True)
             key = list(dataset.keys())[0]
@@ -461,11 +528,38 @@ def make_dataset_images_section(open=False):
             return None
         if num_images > len(dataset):
             num_images = len(dataset)
-        image_idx = np.random.RandomState(random_seed).choice(len(dataset), num_images, replace=False)
-        image_idx = image_idx.tolist()
+        if is_filter:
+            classes = list(map(int, filter_by_class_text.split(",")))
+            labels = np.array(dataset['label'])
+            unique_labels = np.unique(labels)
+            valid_classes = [i for i in classes if i in unique_labels]
+            if len(valid_classes) == 0:
+                gr.Error(f"Classes {classes} not found in the dataset.")
+                return None
+            # shuffle each class
+            chunk_size = num_images // len(valid_classes)
+            image_idx = []
+            for i in valid_classes:
+                idx = np.where(labels == i)[0]
+                if is_random:
+                    idx = np.random.RandomState(seed).choice(idx, chunk_size, replace=False)
+                else:
+                    idx = idx[:chunk_size]
+                image_idx.extend(idx.tolist())
+        if not is_filter:
+            if is_random:
+                image_idx = np.random.RandomState(seed).choice(len(dataset), num_images, replace=False).tolist()
+            else:
+                image_idx = list(range(num_images))
         images = [dataset[i]['image'] for i in image_idx]
         return images   
-    load_dataset_button.click(load_dataset_images, inputs=[dataset_dropdown, num_images_slider, random_seed_slider], outputs=[input_gallery])
+    
+    load_dataset_button.click(load_dataset_images, 
+                        inputs=[dataset_dropdown, num_images_slider,
+                                filter_by_class_checkbox, filter_by_class_text, 
+                                is_random_checkbox, random_seed_slider],
+                        outputs=[input_gallery])
+    
     return dataset_dropdown, num_images_slider, random_seed_slider, load_dataset_button
     
 def make_output_images_section():
@@ -475,25 +569,21 @@ def make_output_images_section():
 
 def make_parameters_section():
     gr.Markdown('### Parameters')
-    model_names = [
-        "SAM(sam_vit_b)",
-        "MobileSAM",
-        "DiNO(dinov2_vitb14_reg)",
-        "CLIP(openai/clip-vit-base-patch16)",
-        "MAE(vit_base)",
-        "SAM2(sam2_hiera_b+)",
-        "SAM2(sam2_hiera_t)",
-    ]
+    from backbone import get_all_model_names
+    model_names = get_all_model_names()
     model_dropdown = gr.Dropdown(model_names, label="Backbone", value="SAM2(sam2_hiera_t)", elem_id="model_name")
-    layer_slider = gr.Slider(0, 11, step=1, label="Backbone: Layer index", value=11, elem_id="layer")
+    layer_slider = gr.Slider(1, 12, step=1, label="Backbone: Layer index", value=12, elem_id="layer")
     node_type_dropdown = gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Backbone: Layer type", value="block: sum of residual", elem_id="node_type", info="which feature to take from each layer?")
     num_eig_slider = gr.Slider(1, 1000, step=1, label="NCUT: Number of eigenvectors", value=100, elem_id="num_eig", info='increase for more clusters')
 
     def change_layer_slider(model_name):
-        if model_name == "SAM2(sam2_hiera_b+)":
-            return gr.Slider(0, 23, step=1, label="Backbone: Layer index", value=23, elem_id="layer", visible=True)
+        layer_dict = LAYER_DICT
+        if model_name in layer_dict:
+            value = layer_dict[model_name]
+            return gr.Slider(1, value, step=1, label="Backbone: Layer index", value=value, elem_id="layer", visible=True)
         else:
-            return gr.Slider(0, 11, step=1, label="Backbone: Layer index", value=11, elem_id="layer", visible=True)
+            value = 12
+            return gr.Slider(1, value, step=1, label="Backbone: Layer index", value=value, elem_id="layer", visible=True)
     model_dropdown.change(fn=change_layer_slider, inputs=model_dropdown, outputs=layer_slider)
     
     with gr.Accordion("➡️ Click to expand: more parameters", open=False):
@@ -728,75 +818,68 @@ with gr.Blocks() as demo:
         gr.Markdown('![ncut](https://ncut-pytorch.readthedocs.io/en/latest/images/gallery/llama3/llama3_layer_31.jpg)')
     
     with gr.Tab('Compare'): 
+        def add_one_model(i_model=1):
+            with gr.Column(scale=5, min_width=200) as col:
+                gr.Markdown(f'### Output Model {i_model}')
+                output_gallery = gr.Gallery(value=[], label="NCUT Embedding", show_label=False, elem_id=f"ncut{i_model}", columns=[3], rows=[1], object_fit="contain", height="auto")
+                submit_button = gr.Button("🔴RUN", elem_id=f"submit_button{i_model}")
+                [
+                    model_dropdown, layer_slider, node_type_dropdown, num_eig_slider, 
+                    affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
+                    embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
+                    perplexity_slider, n_neighbors_slider, min_dist_slider, 
+                    sampling_method_dropdown
+                ] = make_parameters_section()
+                # logging text box
+                logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information")
+                submit_button.click(
+                    run_fn,
+                    inputs=[
+                        input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
+                        affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
+                        embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
+                        perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown
+                    ],
+                    outputs=[output_gallery, logging_text]
+                )
+                
+                return col
 
         with gr.Row():
             with gr.Column(scale=5, min_width=200):
                 input_gallery, submit_button, clear_images_button = make_input_images_section()
+                clear_images_button.click(lambda x: ([],), outputs=[input_gallery])
                 submit_button.visible = False
                 load_images_button, example_gallery, hide_button = make_example_images_section()
                 example_gallery.visible = False
                 hide_button.visible = False
                 dataset_dropdown, num_images_slider, random_seed_slider, load_dataset_button = make_dataset_images_section(open=True)
                 load_images_button.click(lambda x: default_images, outputs=input_gallery)
+
                 
-            with gr.Column(scale=5, min_width=200):
-                gr.Markdown('### Output Model1')
-                output_gallery1 = gr.Gallery(value=[], label="NCUT Embedding", show_label=False, elem_id="ncut1", columns=[3], rows=[1], object_fit="contain", height="auto")
-                submit_button1 = gr.Button("🔴RUN", elem_id="submit_button1")
-                [
-                    model_dropdown1, layer_slider1, node_type_dropdown1, num_eig_slider1,
-                    affinity_focal_gamma_slider1, num_sample_ncut_slider1, knn_ncut_slider1,
-                    embedding_method_dropdown1, num_sample_tsne_slider1, knn_tsne_slider1,
-                    perplexity_slider1, n_neighbors_slider1, min_dist_slider1,
-                    sampling_method_dropdown1
-                ] = make_parameters_section()
-                model_dropdown1.value = 'DiNO(dinov2_vitb14_reg)'
-                layer_slider1.value = 11
-                node_type_dropdown1.value = 'block: sum of residual'
-                # logging text box
-                logging_text1 = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information")
+            for i in range(1, 3):
+                add_one_model(i)
+                
+        with gr.Row():
+            for i in range(1, 4):
+                with gr.Column(scale=5, min_width=200):
+                    slot = gr.Button("Add model", elem_id=f"add_model_button{i}")
+                    col = add_one_model(i+2)
+                    col.visible = False
+                    
+                    slot.click(
+                        fn=lambda x: gr.update(visible=True),
+                        outputs=col
+                    )
+                    slot.click(
+                        fn=lambda x: gr.update(visible=False),
+                        outputs=slot
+                    )
+
 
-            with gr.Column(scale=5, min_width=200):
-                gr.Markdown('### Output Model2')
-                output_gallery2 = gr.Gallery(value=[], label="NCUT Embedding", show_label=False, elem_id="ncut2", columns=[3], rows=[1], object_fit="contain", height="auto")
-                submit_button2 = gr.Button("🔴RUN", elem_id="submit_button2")
-                [
-                    model_dropdown2, layer_slider2, node_type_dropdown2, num_eig_slider2,
-                    affinity_focal_gamma_slider2, num_sample_ncut_slider2, knn_ncut_slider2,
-                    embedding_method_dropdown2, num_sample_tsne_slider2, knn_tsne_slider2,
-                    perplexity_slider2, n_neighbors_slider2, min_dist_slider2,
-                    sampling_method_dropdown2
-                ] = make_parameters_section()
-                model_dropdown2.value = 'DiNO(dinov2_vitb14_reg)'
-                layer_slider2.value = 9
-                node_type_dropdown2.value = 'attn: attention output'
-                # logging text box
-                logging_text2 = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information")
             
-        clear_images_button.click(lambda x: ([], [], []), outputs=[input_gallery, output_gallery1, output_gallery2])
-        submit_button1.click(
-            run_fn,
-            inputs=[
-                input_gallery, model_dropdown1, layer_slider1, num_eig_slider1, node_type_dropdown1, 
-                affinity_focal_gamma_slider1, num_sample_ncut_slider1, knn_ncut_slider1, 
-                embedding_method_dropdown1, num_sample_tsne_slider1, knn_tsne_slider1, 
-                perplexity_slider1, n_neighbors_slider1, min_dist_slider1, sampling_method_dropdown1
-            ],
-            outputs=[output_gallery1, logging_text1]
-        )
-        
-        submit_button2.click(
-            run_fn,
-            inputs=[
-                input_gallery, model_dropdown2, layer_slider2, num_eig_slider2, node_type_dropdown2, 
-                affinity_focal_gamma_slider2, num_sample_ncut_slider2, knn_ncut_slider2, 
-                embedding_method_dropdown2, num_sample_tsne_slider2, knn_tsne_slider2, 
-                perplexity_slider2, n_neighbors_slider2, min_dist_slider2, sampling_method_dropdown2
-            ],
-            outputs=[output_gallery2, logging_text2]
-        )
-        
-        
+download_all_models()
+download_all_datasets()
 demo.launch(share=True)
 
 # %%

backbone.py

@@ -12,19 +12,193 @@ import time
 
 import gradio as gr
 
+from functools import partial
+
 MODEL_DICT = {}
+LAYER_DICT = {}
+RES_DICT = {}
 
+class SAM2(nn.Module):
+
+    def __init__(self, model_cfg='sam2_hiera_b+',):
+        super().__init__()
+
+        try:
+            from sam2.build_sam import build_sam2
+        except ImportError:
+            print("Please install segment_anything_2 from https://github.com/facebookresearch/segment-anything-2.git")
+            return
+        
+        config_dict = {
+            'sam2_hiera_l': ("sam2_hiera_large.pt", "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt"),
+            'sam2_hiera_b+': ("sam2_hiera_base_plus.pt", "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_base_plus.pt"),
+            'sam2_hiera_s': ("sam2_hiera_small.pt", "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt"),
+            'sam2_hiera_t': ("sam2_hiera_tiny.pt", "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt"),
+        }
+        filename, url = config_dict[model_cfg]
+        if not os.path.exists(filename):
+            print(f"Downloading {url}")
+            r = requests.get(url)
+            with open(filename, 'wb') as f:
+                f.write(r.content)
+        sam2_checkpoint = filename
+        
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        sam2_model = build_sam2(model_cfg, sam2_checkpoint, device=device)
+
+        image_encoder = sam2_model.image_encoder
+        image_encoder.eval()
+        
+        from sam2.modeling.backbones.hieradet import do_pool
+        from sam2.modeling.backbones.utils import window_partition, window_unpartition
+        def new_forward(self, x: torch.Tensor) -> torch.Tensor:
+            shortcut = x  # B, H, W, C
+            x = self.norm1(x)
+
+            # Skip connection
+            if self.dim != self.dim_out:
+                shortcut = do_pool(self.proj(x), self.pool)
+
+            # Window partition
+            window_size = self.window_size
+            if window_size > 0:
+                H, W = x.shape[1], x.shape[2]
+                x, pad_hw = window_partition(x, window_size)
+
+            # Window Attention + Q Pooling (if stage change)
+            x = self.attn(x)
+            if self.q_stride:
+                # Shapes have changed due to Q pooling
+                window_size = self.window_size // self.q_stride[0]
+                H, W = shortcut.shape[1:3]
+
+                pad_h = (window_size - H % window_size) % window_size
+                pad_w = (window_size - W % window_size) % window_size
+                pad_hw = (H + pad_h, W + pad_w)
+
+            # Reverse window partition
+            if self.window_size > 0:
+                x = window_unpartition(x, window_size, pad_hw, (H, W))
+
+            self.attn_output = x.clone()
+            
+            x = shortcut + self.drop_path(x)
+            # MLP
+            mlp_out = self.mlp(self.norm2(x))
+            self.mlp_output = mlp_out.clone()
+            x = x + self.drop_path(mlp_out)
+            self.block_output = x.clone()
+            return x        
+        
+        setattr(image_encoder.trunk.blocks[0].__class__, 'forward', new_forward)
+        
+        self.image_encoder = image_encoder
+        
+        
+        
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        output = self.image_encoder(x)
+        attn_outputs, mlp_outputs, block_outputs = [], [], []
+        for block in self.image_encoder.trunk.blocks:
+            attn_outputs.append(block.attn_output)
+            mlp_outputs.append(block.mlp_output)
+            block_outputs.append(block.block_output)
+        return {
+            'attn': attn_outputs,
+            'mlp': mlp_outputs,
+            'block': block_outputs
+        }
+
+MODEL_DICT["SAM2(sam2_hiera_t)"] = partial(SAM2, model_cfg='sam2_hiera_t')
+LAYER_DICT["SAM2(sam2_hiera_t)"] = 12
+RES_DICT["SAM2(sam2_hiera_t)"] = (1024, 1024)
+MODEL_DICT["SAM2(sam2_hiera_s)"] = partial(SAM2, model_cfg='sam2_hiera_s')
+LAYER_DICT["SAM2(sam2_hiera_s)"] = 16
+RES_DICT["SAM2(sam2_hiera_s)"] = (1024, 1024)
+MODEL_DICT["SAM2(sam2_hiera_b+)"] = partial(SAM2, model_cfg='sam2_hiera_b+')
+LAYER_DICT["SAM2(sam2_hiera_b+)"] = 24
+RES_DICT["SAM2(sam2_hiera_b+)"] = (1024, 1024)
+MODEL_DICT["SAM2(sam2_hiera_l)"] = partial(SAM2, model_cfg='sam2_hiera_l')
+LAYER_DICT["SAM2(sam2_hiera_l)"] = 48
+RES_DICT["SAM2(sam2_hiera_l)"] = (1024, 1024)
+
+
+class SAM(torch.nn.Module):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        from segment_anything import sam_model_registry, SamPredictor
+        from segment_anything.modeling.sam import Sam
+
+        checkpoint = "sam_vit_b_01ec64.pth"
+        if not os.path.exists(checkpoint):
+            checkpoint_url = (
+                "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
+            )
+            import requests
+
+            r = requests.get(checkpoint_url)
+            with open(checkpoint, "wb") as f:
+                f.write(r.content)
+
+        sam: Sam = sam_model_registry["vit_b"](checkpoint=checkpoint)
+
+        from segment_anything.modeling.image_encoder import (
+            window_partition,
+            window_unpartition,
+        )
+
+        def new_block_forward(self, x: torch.Tensor) -> torch.Tensor:
+            shortcut = x
+            x = self.norm1(x)
+            # Window partition
+            if self.window_size > 0:
+                H, W = x.shape[1], x.shape[2]
+                x, pad_hw = window_partition(x, self.window_size)
+
+            x = self.attn(x)
+            # Reverse window partition
+            if self.window_size > 0:
+                x = window_unpartition(x, self.window_size, pad_hw, (H, W))
+            self.attn_output = x.clone()
+
+            x = shortcut + x
+            mlp_outout = self.mlp(self.norm2(x))
+            self.mlp_output = mlp_outout.clone()
+            x = x + mlp_outout
+            self.block_output = x.clone()
+
+            return x
+
+        setattr(sam.image_encoder.blocks[0].__class__, "forward", new_block_forward)
+
+        self.image_encoder = sam.image_encoder
+        self.image_encoder.eval()
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        with torch.no_grad():
+            x = torch.nn.functional.interpolate(x, size=(1024, 1024), mode="bilinear")
+        out = self.image_encoder(x)
+
+        attn_outputs, mlp_outputs, block_outputs = [], [], []
+        for i, blk in enumerate(self.image_encoder.blocks):
+            attn_outputs.append(blk.attn_output)
+            mlp_outputs.append(blk.mlp_output)
+            block_outputs.append(blk.block_output)
+        attn_outputs = torch.stack(attn_outputs)
+        mlp_outputs = torch.stack(mlp_outputs)
+        block_outputs = torch.stack(block_outputs)
+        return {
+            'attn': attn_outputs,
+            'mlp': mlp_outputs,
+            'block': block_outputs
+        }
+
+MODEL_DICT["SAM(sam_vit_b)"] = partial(SAM)
+LAYER_DICT["SAM(sam_vit_b)"] = 12
+RES_DICT["SAM(sam_vit_b)"] = (1024, 1024)
 
-def transform_image(image, resolution=(1024, 1024), use_cuda=False):
-    image = image.convert('RGB').resize(resolution, Image.Resampling.NEAREST)
-    # Convert to torch tensor
-    image = torch.tensor(np.array(image).transpose(2, 0, 1)).float()
-    if use_cuda:
-        image = image.cuda()
-    image = image / 255
-    # Normalize
-    image = (image - 0.5) / 0.5
-    return image
 
 class MobileSAM(nn.Module):
     def __init__(self, **kwargs):
@@ -152,181 +326,24 @@ class MobileSAM(nn.Module):
                 attn_outputs.append(blk.attn_output)
                 mlp_outputs.append(blk.mlp_output)
                 block_outputs.append(blk.block_output)
-        return attn_outputs, mlp_outputs, block_outputs
-
-
-MODEL_DICT["MobileSAM"] = MobileSAM()
-
-
-class SAM(torch.nn.Module):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        from segment_anything import sam_model_registry, SamPredictor
-        from segment_anything.modeling.sam import Sam
-
-        checkpoint = "sam_vit_b_01ec64.pth"
-        if not os.path.exists(checkpoint):
-            checkpoint_url = (
-                "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
-            )
-            import requests
-
-            r = requests.get(checkpoint_url)
-            with open(checkpoint, "wb") as f:
-                f.write(r.content)
-
-        sam: Sam = sam_model_registry["vit_b"](checkpoint=checkpoint)
-
-        from segment_anything.modeling.image_encoder import (
-            window_partition,
-            window_unpartition,
-        )
-
-        def new_block_forward(self, x: torch.Tensor) -> torch.Tensor:
-            shortcut = x
-            x = self.norm1(x)
-            # Window partition
-            if self.window_size > 0:
-                H, W = x.shape[1], x.shape[2]
-                x, pad_hw = window_partition(x, self.window_size)
-
-            x = self.attn(x)
-            # Reverse window partition
-            if self.window_size > 0:
-                x = window_unpartition(x, self.window_size, pad_hw, (H, W))
-            self.attn_output = x.clone()
-
-            x = shortcut + x
-            mlp_outout = self.mlp(self.norm2(x))
-            self.mlp_output = mlp_outout.clone()
-            x = x + mlp_outout
-            self.block_output = x.clone()
-
-            return x
-
-        setattr(sam.image_encoder.blocks[0].__class__, "forward", new_block_forward)
-
-        self.image_encoder = sam.image_encoder
-        self.image_encoder.eval()
-
-    @torch.no_grad()
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        with torch.no_grad():
-            x = torch.nn.functional.interpolate(x, size=(1024, 1024), mode="bilinear")
-        out = self.image_encoder(x)
-
-        attn_outputs, mlp_outputs, block_outputs = [], [], []
-        for i, blk in enumerate(self.image_encoder.blocks):
-            attn_outputs.append(blk.attn_output)
-            mlp_outputs.append(blk.mlp_output)
-            block_outputs.append(blk.block_output)
-        attn_outputs = torch.stack(attn_outputs)
-        mlp_outputs = torch.stack(mlp_outputs)
-        block_outputs = torch.stack(block_outputs)
-        return attn_outputs, mlp_outputs, block_outputs
-
-
-MODEL_DICT["SAM(sam_vit_b)"] = SAM()
-
-
-class SAM2(nn.Module):
-
-    def __init__(self, model_cfg='sam2_hiera_b+',):
-        super().__init__()
-
-        try:
-            from sam2.build_sam import build_sam2
-        except ImportError:
-            print("Please install segment_anything_2 from https://github.com/facebookresearch/segment-anything-2.git")
-            return
-        
-        config_dict = {
-            'sam2_hiera_large': ("sam2_hiera_large.pt", "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt"),
-            'sam2_hiera_b+': ("sam2_hiera_base_plus.pt", "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_base_plus.pt"),
-            'sam2_hiera_s': ("sam2_hiera_small.pt", "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt"),
-            'sam2_hiera_t': ("sam2_hiera_tiny.pt", "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt"),
+        return {
+            'attn': attn_outputs,
+            'mlp': mlp_outputs,
+            'block': block_outputs
         }
-        filename, url = config_dict[model_cfg]
-        if not os.path.exists(filename):
-            print(f"Downloading {url}")
-            r = requests.get(url)
-            with open(filename, 'wb') as f:
-                f.write(r.content)
-        sam2_checkpoint = filename
-        
-        device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        sam2_model = build_sam2(model_cfg, sam2_checkpoint, device=device)
-
-        image_encoder = sam2_model.image_encoder
-        image_encoder.eval()
-        
-        from sam2.modeling.backbones.hieradet import do_pool
-        from sam2.modeling.backbones.utils import window_partition, window_unpartition
-        def new_forward(self, x: torch.Tensor) -> torch.Tensor:
-            shortcut = x  # B, H, W, C
-            x = self.norm1(x)
 
-            # Skip connection
-            if self.dim != self.dim_out:
-                shortcut = do_pool(self.proj(x), self.pool)
+MODEL_DICT["MobileSAM"] = partial(MobileSAM)
+LAYER_DICT["MobileSAM"] = 12
+RES_DICT["MobileSAM"] = (1024, 1024)
 
-            # Window partition
-            window_size = self.window_size
-            if window_size > 0:
-                H, W = x.shape[1], x.shape[2]
-                x, pad_hw = window_partition(x, window_size)
-
-            # Window Attention + Q Pooling (if stage change)
-            x = self.attn(x)
-            if self.q_stride:
-                # Shapes have changed due to Q pooling
-                window_size = self.window_size // self.q_stride[0]
-                H, W = shortcut.shape[1:3]
-
-                pad_h = (window_size - H % window_size) % window_size
-                pad_w = (window_size - W % window_size) % window_size
-                pad_hw = (H + pad_h, W + pad_w)
-
-            # Reverse window partition
-            if self.window_size > 0:
-                x = window_unpartition(x, window_size, pad_hw, (H, W))
-
-            self.attn_output = x.clone()
-            
-            x = shortcut + self.drop_path(x)
-            # MLP
-            mlp_out = self.mlp(self.norm2(x))
-            self.mlp_output = mlp_out.clone()
-            x = x + self.drop_path(mlp_out)
-            self.block_output = x.clone()
-            return x        
-        
-        setattr(image_encoder.trunk.blocks[0].__class__, 'forward', new_forward)
-        
-        self.image_encoder = image_encoder
-        
-        
-        
-    @torch.no_grad()
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        output = self.image_encoder(x)
-        attn_outputs, mlp_outputs, block_outputs = [], [], []
-        for block in self.image_encoder.trunk.blocks:
-            attn_outputs.append(block.attn_output)
-            mlp_outputs.append(block.mlp_output)
-            block_outputs.append(block.block_output)
-        return attn_outputs, mlp_outputs, block_outputs
-
-
-MODEL_DICT["SAM2(sam2_hiera_b+)"] = SAM2(model_cfg='sam2_hiera_b+')
-MODEL_DICT["SAM2(sam2_hiera_t)"] = SAM2(model_cfg='sam2_hiera_t')
 
 class DiNOv2(torch.nn.Module):
-    def __init__(self, ver="dinov2_vitb14_reg"):
+    def __init__(self, ver="dinov2_vitb14_reg", num_reg=5):
         super().__init__()
         self.dinov2 = torch.hub.load("facebookresearch/dinov2", ver)
         self.dinov2.requires_grad_(False)
         self.dinov2.eval()
+        self.num_reg = num_reg
 
         def new_block_forward(self, x: torch.Tensor) -> torch.Tensor:
             def attn_residual_func(x):
@@ -337,20 +354,20 @@ class DiNOv2(torch.nn.Module):
 
             attn_output = attn_residual_func(x)
 
-            hw = np.sqrt(attn_output.shape[1] - 5).astype(int)
+            hw = np.sqrt(attn_output.shape[1] - num_reg).astype(int)
             self.attn_output = rearrange(
-                attn_output.clone()[:, 5:], "b (h w) c -> b h w c", h=hw
+                attn_output.clone()[:, num_reg:], "b (h w) c -> b h w c", h=hw
             )
 
             x = x + attn_output
             mlp_output = ffn_residual_func(x)
             self.mlp_output = rearrange(
-                mlp_output.clone()[:, 5:], "b (h w) c -> b h w c", h=hw
+                mlp_output.clone()[:, num_reg:], "b (h w) c -> b h w c", h=hw
             )
             x = x + mlp_output
             block_output = x
             self.block_output = rearrange(
-                block_output.clone()[:, 5:], "b (h w) c -> b h w c", h=hw
+                block_output.clone()[:, num_reg:], "b (h w) c -> b h w c", h=hw
             )
             return x
 
@@ -370,10 +387,64 @@ class DiNOv2(torch.nn.Module):
         attn_outputs = torch.stack(attn_outputs)
         mlp_outputs = torch.stack(mlp_outputs)
         block_outputs = torch.stack(block_outputs)
-        return attn_outputs, mlp_outputs, block_outputs
+        return {
+            'attn': attn_outputs,
+            'mlp': mlp_outputs,
+            'block': block_outputs
+        }
 
+MODEL_DICT["DiNOv2reg(dinov2_vitb14_reg)"] = partial(DiNOv2, ver="dinov2_vitb14_reg", num_reg=5)
+LAYER_DICT["DiNOv2reg(dinov2_vitb14_reg)"] = 12
+RES_DICT["DiNOv2reg(dinov2_vitb14_reg)"] = (672, 672)
+MODEL_DICT["DiNOv2(dinov2_vitb14)"] = partial(DiNOv2, ver="dinov2_vitb14", num_reg=1)
+LAYER_DICT["DiNOv2(dinov2_vitb14)"] = 12
+RES_DICT["DiNOv2(dinov2_vitb14)"] = (672, 672)
 
-MODEL_DICT["DiNO(dinov2_vitb14_reg)"] = DiNOv2()
+class DiNO(nn.Module):
+    def __init__(self, ver="dino_vitb8"):
+        super().__init__()
+        model = torch.hub.load('facebookresearch/dino:main', ver)
+        model = model.eval()
+
+        def remove_cls_and_reshape(x):
+            x = x.clone()
+            x = x[:, 1:]
+            hw = np.sqrt(x.shape[1]).astype(int)
+            x = rearrange(x, "b (h w) c -> b h w c", h=hw)
+            return x
+
+        def new_forward(self, x, return_attention=False):
+            y, attn = self.attn(self.norm1(x))
+            self.attn_output = remove_cls_and_reshape(y.clone())
+            if return_attention:
+                return attn
+            x = x + self.drop_path(y)
+            mlp_output = self.mlp(self.norm2(x))
+            self.mlp_output = remove_cls_and_reshape(mlp_output.clone())
+            x = x + self.drop_path(mlp_output)
+            self.block_output = remove_cls_and_reshape(x.clone())
+            return x
+
+        setattr(model.blocks[0].__class__, "forward", new_forward)
+        
+        self.model = model
+        self.model.eval()
+        self.model.requires_grad_(False)
+        
+    def forward(self, x):
+        out = self.model(x)
+        attn_outputs = [block.attn_output for block in self.model.blocks]
+        mlp_outputs = [block.mlp_output for block in self.model.blocks]
+        block_outputs = [block.block_output for block in self.model.blocks]
+        return {
+            'attn': attn_outputs,
+            'mlp': mlp_outputs,
+            'block': block_outputs
+        }
+            
+MODEL_DICT["DiNO(dino_vitb8)"] = partial(DiNO)
+LAYER_DICT["DiNO(dino_vitb8)"] = 12
+RES_DICT["DiNO(dino_vitb8)"] = (448, 448)
 
 def resample_position_embeddings(embeddings, h, w):
     cls_embeddings = embeddings[0]
@@ -385,87 +456,264 @@ def resample_position_embeddings(embeddings, h, w):
     embeddings = torch.cat([cls_embeddings.unsqueeze(0), patch_embeddings], dim=0)
     return embeddings
 
-class CLIP(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
+# class CLIP(torch.nn.Module):
+#     def __init__(self, ver="openai/clip-vit-base-patch16"):
+#         super().__init__()
 
-        from transformers import CLIPProcessor, CLIPModel
+#         from transformers import CLIPProcessor, CLIPModel
 
-        model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
+#         model = CLIPModel.from_pretrained(ver)
                 
-        # resample the patch embeddings to 56x56, take 896x896 input
-        embeddings = model.vision_model.embeddings.position_embedding.weight
-        embeddings = resample_position_embeddings(embeddings, 56, 56)
-        model.vision_model.embeddings.position_embedding.weight = nn.Parameter(embeddings)
-        model.vision_model.embeddings.position_ids = torch.arange(0, 1+56*56)
+#         # resample the patch embeddings to 56x56, take 896x896 input
+#         embeddings = model.vision_model.embeddings.position_embedding.weight
+#         embeddings = resample_position_embeddings(embeddings, 42, 42)
+#         model.vision_model.embeddings.position_embedding.weight = nn.Parameter(embeddings)
+#         model.vision_model.embeddings.position_ids = torch.arange(0, 1+56*56)
+        
+#         # processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")
+#         self.model = model.eval()
+
+#         def new_forward(
+#             self,
+#             hidden_states: torch.Tensor,
+#             attention_mask: torch.Tensor,
+#             causal_attention_mask: torch.Tensor,
+#             output_attentions: Optional[bool] = False,
+#         ) -> Tuple[torch.FloatTensor]:
+
+#             residual = hidden_states
+
+#             hidden_states = self.layer_norm1(hidden_states)
+#             hidden_states, attn_weights = self.self_attn(
+#                 hidden_states=hidden_states,
+#                 attention_mask=attention_mask,
+#                 causal_attention_mask=causal_attention_mask,
+#                 output_attentions=output_attentions,
+#             )
+#             hw = np.sqrt(hidden_states.shape[1] - 1).astype(int)
+#             self.attn_output = rearrange(
+#                 hidden_states.clone()[:, 1:], "b (h w) c -> b h w c", h=hw
+#             )
+#             hidden_states = residual + hidden_states
+
+#             residual = hidden_states
+#             hidden_states = self.layer_norm2(hidden_states)
+#             hidden_states = self.mlp(hidden_states)
+#             self.mlp_output = rearrange(
+#                 hidden_states.clone()[:, 1:], "b (h w) c -> b h w c", h=hw
+#             )
+
+#             hidden_states = residual + hidden_states
+
+#             outputs = (hidden_states,)
+
+#             if output_attentions:
+#                 outputs += (attn_weights,)
+
+#             self.block_output = rearrange(
+#                 hidden_states.clone()[:, 1:], "b (h w) c -> b h w c", h=hw
+#             )
+#             return outputs
+
+#         setattr(
+#             self.model.vision_model.encoder.layers[0].__class__, "forward", new_forward
+#         )
+
+#     @torch.no_grad()
+#     def forward(self, x):
+
+#         out = self.model.vision_model(x)
+
+#         attn_outputs, mlp_outputs, block_outputs = [], [], []
+#         for i, blk in enumerate(self.model.vision_model.encoder.layers):
+#             attn_outputs.append(blk.attn_output)
+#             mlp_outputs.append(blk.mlp_output)
+#             block_outputs.append(blk.block_output)
+
+#         attn_outputs = torch.stack(attn_outputs)
+#         mlp_outputs = torch.stack(mlp_outputs)
+#         block_outputs = torch.stack(block_outputs)
+#         return attn_outputs, mlp_outputs, block_outputs
+
+
+# MODEL_DICT["CLIP(openai/clip-vit-base-patch16)"] = partial(CLIP, ver="openai/clip-vit-base-patch16")
+# LAYER_DICT["CLIP(openai/clip-vit-base-patch16)"] = 12
+# RES_DICT["CLIP(openai/clip-vit-base-patch16)"] = (896, 896)
+
+
+class OpenCLIPViT(nn.Module):
+    def __init__(self, version='ViT-B-16', pretrained='laion2b_s34b_b88k'):
+        super().__init__()
+        try:
+            import open_clip
+        except ImportError:
+            print("Please install open_clip to use this class.")
+            return
+        
+        model, _, _ = open_clip.create_model_and_transforms(version, pretrained=pretrained)
+        
+        positional_embedding = resample_position_embeddings(model.visual.positional_embedding, 42, 42)
+        model.visual.positional_embedding = nn.Parameter(positional_embedding)
         
-        # processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")
-        self.model = model.eval()
-
         def new_forward(
-            self,
-            hidden_states: torch.Tensor,
-            attention_mask: torch.Tensor,
-            causal_attention_mask: torch.Tensor,
-            output_attentions: Optional[bool] = False,
-        ) -> Tuple[torch.FloatTensor]:
-
-            residual = hidden_states
-
-            hidden_states = self.layer_norm1(hidden_states)
-            hidden_states, attn_weights = self.self_attn(
-                hidden_states=hidden_states,
-                attention_mask=attention_mask,
-                causal_attention_mask=causal_attention_mask,
-                output_attentions=output_attentions,
-            )
-            hw = np.sqrt(hidden_states.shape[1] - 1).astype(int)
-            self.attn_output = rearrange(
-                hidden_states.clone()[:, 1:], "b (h w) c -> b h w c", h=hw
-            )
-            hidden_states = residual + hidden_states
-
-            residual = hidden_states
-            hidden_states = self.layer_norm2(hidden_states)
-            hidden_states = self.mlp(hidden_states)
-            self.mlp_output = rearrange(
-                hidden_states.clone()[:, 1:], "b (h w) c -> b h w c", h=hw
-            )
-
-            hidden_states = residual + hidden_states
-
-            outputs = (hidden_states,)
-
-            if output_attentions:
-                outputs += (attn_weights,)
-
-            self.block_output = rearrange(
-                hidden_states.clone()[:, 1:], "b (h w) c -> b h w c", h=hw
-            )
-            return outputs
+                self,
+                q_x: torch.Tensor,
+                k_x: Optional[torch.Tensor] = None,
+                v_x: Optional[torch.Tensor] = None,
+                attn_mask: Optional[torch.Tensor] = None,
+        ):
+            def remove_cls_and_reshape(x):
+                x = x.clone()
+                x = x[1:]
+                hw = np.sqrt(x.shape[0]).astype(int)
+                x = rearrange(x, "(h w) b c -> b h w c", h=hw)
+                return x
+            
+            k_x = self.ln_1_kv(k_x) if hasattr(self, "ln_1_kv") and k_x is not None else None
+            v_x = self.ln_1_kv(v_x) if hasattr(self, "ln_1_kv") and v_x is not None else None
+
+            attn_output = self.attention(q_x=self.ln_1(q_x), k_x=k_x, v_x=v_x, attn_mask=attn_mask)
+            self.attn_output = remove_cls_and_reshape(attn_output.clone())
+            x = q_x + self.ls_1(attn_output)
+            mlp_output = self.mlp(self.ln_2(x))
+            self.mlp_output = remove_cls_and_reshape(mlp_output.clone())
+            x = x + self.ls_2(mlp_output)    
+            self.block_output = remove_cls_and_reshape(x.clone())
+            return x
 
-        setattr(
-            self.model.vision_model.encoder.layers[0].__class__, "forward", new_forward
-        )
 
-    @torch.no_grad()
+        setattr(model.visual.transformer.resblocks[0].__class__, "forward", new_forward)
+        
+        self.model = model
+        self.model.eval()
+    
     def forward(self, x):
-
-        out = self.model.vision_model(x)
-
+        out = self.model(x)
         attn_outputs, mlp_outputs, block_outputs = [], [], []
-        for i, blk in enumerate(self.model.vision_model.encoder.layers):
-            attn_outputs.append(blk.attn_output)
-            mlp_outputs.append(blk.mlp_output)
-            block_outputs.append(blk.block_output)
+        for block in self.model.visual.transformer.resblocks:
+            attn_outputs.append(block.attn_output)
+            mlp_outputs.append(block.mlp_output)
+            block_outputs.append(block.block_output)
+        return {
+            'attn': attn_outputs,
+            'mlp': mlp_outputs,
+            'block': block_outputs
+        }
+        
+MODEL_DICT["CLIP(ViT-B-16/openai)"] = partial(OpenCLIPViT, version='ViT-B-16', pretrained='openai')
+LAYER_DICT["CLIP(ViT-B-16/openai)"] = 12
+RES_DICT["CLIP(ViT-B-16/openai)"] = (672, 672)
+MODEL_DICT["CLIP(ViT-B-16/laion2b_s34b_b88k)"] = partial(OpenCLIPViT, version='ViT-B-16', pretrained='laion2b_s34b_b88k')
+LAYER_DICT["CLIP(ViT-B-16/laion2b_s34b_b88k)"] = 12
+RES_DICT["CLIP(ViT-B-16/laion2b_s34b_b88k)"] = (672, 672)
+
+class EVA02(nn.Module):
+    
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        
+        model = timm.create_model(
+            'eva02_large_patch14_448.mim_m38m_ft_in22k_in1k',
+            pretrained=True,
+            num_classes=0,  # remove classifier nn.Linear
+        )
+        model = model.eval()
 
-        attn_outputs = torch.stack(attn_outputs)
-        mlp_outputs = torch.stack(mlp_outputs)
-        block_outputs = torch.stack(block_outputs)
-        return attn_outputs, mlp_outputs, block_outputs
+        def new_forward(self, x, rope: Optional[torch.Tensor] = None, attn_mask: Optional[torch.Tensor] = None):
+        
+            def remove_cls_and_reshape(x):
+                x = x.clone()
+                x = x[:, 1:]
+                hw = np.sqrt(x.shape[1]).astype(int)
+                x = rearrange(x, "b (h w) c -> b h w c", h=hw)
+                return x
+            
+            if self.gamma_1 is None:
+                attn_output = self.attn(self.norm1(x), rope=rope, attn_mask=attn_mask)
+                self.attn_output = remove_cls_and_reshape(attn_output.clone())
+                x = x + self.drop_path1(attn_output)
+                mlp_output = self.mlp(self.norm2(x))
+                self.mlp_output = remove_cls_and_reshape(mlp_output.clone())
+                x = x + self.drop_path2(mlp_output)
+            else:
+                attn_output = self.attn(self.norm1(x), rope=rope, attn_mask=attn_mask)
+                self.attn_output = remove_cls_and_reshape(attn_output.clone())
+                x = x + self.drop_path1(self.gamma_1 * attn_output)
+                mlp_output = self.mlp(self.norm2(x))
+                self.mlp_output = remove_cls_and_reshape(mlp_output.clone())
+                x = x + self.drop_path2(self.gamma_2 * mlp_output)
+            self.block_output = remove_cls_and_reshape(x.clone())
+            return x
+
+        setattr(model.blocks[0].__class__, "forward", new_forward)
+        
+        self.model = model
+        
+    def forward(self, x):
+        out = self.model(x)
+        attn_outputs = [block.attn_output for block in self.model.blocks]
+        mlp_outputs = [block.mlp_output for block in self.model.blocks]
+        block_outputs = [block.block_output for block in self.model.blocks]
+        return {
+            'attn': attn_outputs,
+            'mlp': mlp_outputs,
+            'block': block_outputs
+        }
+        
+MODEL_DICT["EVA-CLIP(eva02_large_patch14_448)"] = partial(EVA02)
+LAYER_DICT["EVA-CLIP(eva02_large_patch14_448)"] = 24
+RES_DICT["EVA-CLIP(eva02_large_patch14_448)"] = (448, 448)
 
+class CLIPConvnext(nn.Module):
+    def __init__(self):
+        super().__init__()
+        try:
+            import open_clip
+        except ImportError:
+            print("Please install open_clip to use this class.")
+            return
+    
+        model, _, _ = open_clip.create_model_and_transforms('convnext_base_w_320', pretrained='laion_aesthetic_s13b_b82k')
+        
+        def new_forward(self, x):
+            shortcut = x
+            x = self.conv_dw(x)
+            if self.use_conv_mlp:
+                x = self.norm(x)
+                x = self.mlp(x)
+            else:
+                x = x.permute(0, 2, 3, 1)
+                x = self.norm(x)
+                x = self.mlp(x)
+                x = x.permute(0, 3, 1, 2)
+            if self.gamma is not None:
+                x = x.mul(self.gamma.reshape(1, -1, 1, 1))
+
+            x = self.drop_path(x) + self.shortcut(shortcut)
+            self.block_output = rearrange(x.clone(), "b c h w -> b h w c")
+            return x
 
-MODEL_DICT["CLIP(openai/clip-vit-base-patch16)"] = CLIP()
+        setattr(model.visual.trunk.stages[0].blocks[0].__class__, "forward", new_forward)
+        
+        self.model = model
+        self.model.eval()
+    
+    def forward(self, x):
+        out = self.model(x)
+        block_outputs = []
+        for stage in self.model.visual.trunk.stages:
+            for block in stage.blocks:
+                block_outputs.append(block.block_output)
+        return {
+            'attn': None,
+            'mlp': None,
+            'block': block_outputs
+        }
+        
+    
+MODEL_DICT["CLIP(convnext_base_w_320/laion_aesthetic_s13b_b82k)"] = partial(CLIPConvnext)
+LAYER_DICT["CLIP(convnext_base_w_320/laion_aesthetic_s13b_b82k)"] = 36
+RES_DICT["CLIP(convnext_base_w_320/laion_aesthetic_s13b_b82k)"] = (960, 960)
 
 
 class MAE(timm.models.vision_transformer.VisionTransformer):
@@ -492,10 +740,10 @@ class MAE(timm.models.vision_transformer.VisionTransformer):
         
         # resample the patch embeddings to 56x56, take 896x896 input
         pos_embed = self.pos_embed[0]
-        pos_embed = resample_position_embeddings(pos_embed, 56, 56)
+        pos_embed = resample_position_embeddings(pos_embed, 42, 42)
         self.pos_embed = nn.Parameter(pos_embed.unsqueeze(0))
-        self.img_size = (896, 896)
-        self.patch_embed.img_size = (896, 896)
+        self.img_size = (672, 672)
+        self.patch_embed.img_size = (672, 672)
 
         self.requires_grad_(False)
         self.eval()
@@ -519,45 +767,105 @@ class MAE(timm.models.vision_transformer.VisionTransformer):
             x = rearrange(x, "b (h w) c -> b h w c", h=hw)
             return x
         
-        attn_nodes = [remove_cls_and_reshape(block.saved_attn_node) for block in self.blocks]
-        mlp_nodes = [remove_cls_and_reshape(block.saved_mlp_node) for block in self.blocks]
+        attn_outputs = [remove_cls_and_reshape(block.saved_attn_node) for block in self.blocks]
+        mlp_outputs = [remove_cls_and_reshape(block.saved_mlp_node) for block in self.blocks]
         block_outputs = [remove_cls_and_reshape(block.saved_block_output) for block in self.blocks]
-        return attn_nodes, mlp_nodes, block_outputs
+        return {
+            'attn': attn_outputs,
+            'mlp': mlp_outputs,
+            'block': block_outputs
+        }
+        
 
+MODEL_DICT["MAE(vit_base)"] = partial(MAE)
+LAYER_DICT["MAE(vit_base)"] = 12
+RES_DICT["MAE(vit_base)"] = (672, 672)
 
-MODEL_DICT["MAE(vit_base)"] = MAE()
+class ImageNet(nn.Module):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
 
+        model = timm.create_model(
+            'vit_base_patch16_224.augreg2_in21k_ft_in1k',
+            pretrained=True,
+            num_classes=0,  # remove classifier nn.Linear
+        )
+        
+        # resample the patch embeddings to 56x56, take 896x896 input
+        pos_embed = model.pos_embed[0]
+        pos_embed = resample_position_embeddings(pos_embed, 42, 42)
+        model.pos_embed = nn.Parameter(pos_embed.unsqueeze(0))
+        model.img_size = (672, 672)
+        model.patch_embed.img_size = (672, 672)
+
+        model.requires_grad_(False)
+        model.eval()
+        
+        def forward(self, x):
+            self.saved_attn_node = self.ls1(self.attn(self.norm1(x)))
+            x = x + self.saved_attn_node.clone()
+            self.saved_mlp_node = self.ls2(self.mlp(self.norm2(x)))
+            x = x + self.saved_mlp_node.clone()
+            self.saved_block_output = x.clone()
+            return x
+        
+        setattr(model.blocks[0].__class__, "forward", forward)
+        
+        self.model = model
+        
+    def forward(self, x):
+        out = self.model(x)
+        def remove_cls_and_reshape(x):
+            x = x.clone()
+            x = x[:, 1:]
+            hw = np.sqrt(x.shape[1]).astype(int)
+            x = rearrange(x, "b (h w) c -> b h w c", h=hw)
+            return x
+        
+        attn_outputs = [remove_cls_and_reshape(block.saved_attn_node) for block in self.model.blocks]
+        mlp_outputs = [remove_cls_and_reshape(block.saved_mlp_node) for block in self.model.blocks]
+        block_outputs = [remove_cls_and_reshape(block.saved_block_output) for block in self.model.blocks]
+        return {
+            'attn': attn_outputs,
+            'mlp': mlp_outputs,
+            'block': block_outputs
+        }
+            
+MODEL_DICT["ImageNet(vit_base)"] = partial(ImageNet)
+LAYER_DICT["ImageNet(vit_base)"] = 12
+RES_DICT["ImageNet(vit_base)"] = (672, 672)
+
+def download_all_models():
+    for model_name in MODEL_DICT:
+        print(f"Downloading {model_name}")
+        model = MODEL_DICT[model_name]()
+        
+def get_all_model_names():
+    return list(MODEL_DICT.keys())
+        
+def get_model(model_name):
+    return MODEL_DICT[model_name]()
 
-def extract_features(images, model_name, node_type, layer):
+@torch.no_grad()
+def extract_features(images, model, model_name, node_type, layer, batch_size=8):
     use_cuda = torch.cuda.is_available()
     
-    resolution = (1024, 1024)
-    resolution_dict = {
-        "DiNO(dinov2_vitb14_reg)": (896, 896),
-        'CLIP(openai/clip-vit-base-patch16)': (896, 896),
-        'MAE(vit_base)': (896, 896),
-    }
-    if model_name in resolution_dict:
-        resolution = resolution_dict[model_name]
-
-    model = MODEL_DICT[model_name]
-
     if use_cuda:
         model = model.cuda()
 
+    chunked_idxs = torch.split(torch.arange(images.shape[0]), batch_size)
+    
     outputs = []
-    for i in range(len(images)):
-        image = transform_image(images[i], resolution=resolution, use_cuda=use_cuda)
-        inp = image.unsqueeze(0)
-        attn_output, mlp_output, block_output = model(inp)
-        out_dict = {
-            "attn": attn_output,
-            "mlp": mlp_output,
-            "block": block_output,
-        }
-        out = out_dict[node_type]
+    for idxs in chunked_idxs:
+        inp = images[idxs]
+        if use_cuda:
+            inp = inp.cuda()
+        out = model(inp)  # {'attn': [B, H, W, C], 'mlp': [B, H, W, C], 'block': [B, H, W, C]}
+        out = out[node_type]
         out = out[layer]
-        outputs.append(out)
+        # normalize
+        out = F.normalize(out, dim=-1)
+        outputs.append(out.cpu().float())
     outputs = torch.cat(outputs, dim=0)
 
     return outputs

requirements.txt

@@ -11,4 +11,5 @@ pillow==9.4.0
 SAM-2 @ git+https://github.com/huzeyann/segment-anything-2.git
 segment-anything @ git+https://github.com/facebookresearch/segment-anything.git@6fdee8f
 mobile-sam @ git+https://github.com/ChaoningZhang/MobileSAM.git@c12dd83
-timm
+timm
+open-clip-torch==2.20.0