moved backbone models

app.py

@@ -20,8 +20,8 @@ import time
 import threading
 import os
 
-from backbone import extract_features, get_model
-from backbone import MODEL_DICT, LAYER_DICT, RES_DICT
+from ncut_pytorch.backbone import extract_features, load_model
+from ncut_pytorch.backbone import MODEL_DICT, LAYER_DICT, RES_DICT
 from ncut_pytorch import NCUT, eigenvector_to_rgb
 
 DATASET_TUPS = [
@@ -218,7 +218,7 @@ def ncut_run(
         
     start = time.time()
     features = extract_features(
-        images, model, model_name=model_name, node_type=node_type, layer=layer-1, batch_size=BATCH_SIZE
+        images, model, 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"
@@ -407,7 +407,7 @@ def run_fn(
     images = [transform_image(image, resolution=resolution) for image in images]
     images = torch.stack(images)
         
-    model = get_model(model_name)
+    model = load_model(model_name)
     
     kwargs = {
         "model_name": model_name,
@@ -585,8 +585,8 @@ def make_output_images_section():
 
 def make_parameters_section():
     gr.Markdown("### Parameters <a style='color: #0044CC;' href='https://ncut-pytorch.readthedocs.io/en/latest/how_to_get_better_segmentation/' target='_blank'>Help</a>")
-    from backbone import get_all_model_names
-    model_names = get_all_model_names()
+    from backbone import get_demo_model_names
+    model_names = get_demo_model_names()
     model_dropdown = gr.Dropdown(model_names, label="Backbone", value="DiNO(dino_vitb8)", elem_id="model_name")
     layer_slider = gr.Slider(1, 12, step=1, label="Backbone: Layer index", value=10, 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?")

app_text.py

@@ -22,8 +22,8 @@ import numpy as np
 
 from ncut_pytorch import NCUT, eigenvector_to_rgb
 
-from backbone_text import MODEL_DICT as TEXT_MODEL_DICT
-from backbone_text import LAYER_DICT as TEXT_LAYER_DICT
+from ncut_pytorch.backbone_text import MODEL_DICT as TEXT_MODEL_DICT
+from ncut_pytorch.backbone_text import LAYER_DICT as TEXT_LAYER_DICT
 
 def compute_ncut(
     features,
@@ -41,8 +41,6 @@ def compute_ncut(
     metric="cosine",
 ):        
     logging_str = ""
-    print("running ncut")
-    print(features.shape)
     num_nodes = np.prod(features.shape[:-1])
     if num_nodes / 2 < num_eig:
         # raise gr.Error("Number of eigenvectors should be less than half the number of nodes.")
@@ -197,8 +195,10 @@ def ncut_run(
     )
     logging_str += _logging_str
     
+    start = time.time() 
     title = f"{model_name}, Layer {layer}, {node_type}"
     fig = make_plot(token_texts, rgb, title=title)
+    logging_str += f"Plotting time: {time.time() - start:.2f}s\n"
     return fig, logging_str
 
 def _ncut_run(*args, **kwargs):
@@ -302,3 +302,5 @@ if __name__ == "__main__":
     with gr.Blocks() as demo:
         make_demo()
     demo.launch(share=True)
+
+# %%

backbone.py

@@ -1,881 +0,0 @@
-# Author: Huzheng Yang
-# %%
-from typing import Optional, Tuple
-from einops import rearrange
-import requests
-import torch
-import torch.nn.functional as F
-import timm
-from torch import nn
-import numpy as np
-import os
-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 = '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)
-
-
-class MobileSAM(nn.Module):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-        from mobile_sam import sam_model_registry
-
-        url = "https://raw.githubusercontent.com/ChaoningZhang/MobileSAM/master/weights/mobile_sam.pt"
-        model_type = "vit_t"
-        sam_checkpoint = "mobile_sam.pt"
-        if not os.path.exists(sam_checkpoint):
-            import requests
-
-            r = requests.get(url)
-            with open(sam_checkpoint, "wb") as f:
-                f.write(r.content)
-
-        mobile_sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
-
-        def new_forward_fn(self, x):
-            shortcut = x
-
-            x = self.conv1(x)
-            x = self.act1(x)
-
-            x = self.conv2(x)
-            x = self.act2(x)
-
-            self.attn_output = rearrange(x.clone(), "b c h w -> b h w c")
-
-            x = self.conv3(x)
-
-            self.mlp_output = rearrange(x.clone(), "b c h w -> b h w c")
-
-            x = self.drop_path(x)
-
-            x += shortcut
-            x = self.act3(x)
-
-            self.block_output = rearrange(x.clone(), "b c h w -> b h w c")
-
-            return x
-
-        setattr(
-            mobile_sam.image_encoder.layers[0].blocks[0].__class__,
-            "forward",
-            new_forward_fn,
-        )
-
-        def new_forward_fn2(self, x):
-            H, W = self.input_resolution
-            B, L, C = x.shape
-            assert L == H * W, "input feature has wrong size"
-            res_x = x
-            if H == self.window_size and W == self.window_size:
-                x = self.attn(x)
-            else:
-                x = x.view(B, H, W, C)
-                pad_b = (self.window_size - H % self.window_size) % self.window_size
-                pad_r = (self.window_size - W % self.window_size) % self.window_size
-                padding = pad_b > 0 or pad_r > 0
-
-                if padding:
-                    x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b))
-
-                pH, pW = H + pad_b, W + pad_r
-                nH = pH // self.window_size
-                nW = pW // self.window_size
-                # window partition
-                x = (
-                    x.view(B, nH, self.window_size, nW, self.window_size, C)
-                    .transpose(2, 3)
-                    .reshape(B * nH * nW, self.window_size * self.window_size, C)
-                )
-                x = self.attn(x)
-                # window reverse
-                x = (
-                    x.view(B, nH, nW, self.window_size, self.window_size, C)
-                    .transpose(2, 3)
-                    .reshape(B, pH, pW, C)
-                )
-
-                if padding:
-                    x = x[:, :H, :W].contiguous()
-
-                x = x.view(B, L, C)
-
-            hw = np.sqrt(x.shape[1]).astype(int)
-            self.attn_output = rearrange(x.clone(), "b (h w) c -> b h w c", h=hw)
-
-            x = res_x + self.drop_path(x)
-
-            x = x.transpose(1, 2).reshape(B, C, H, W)
-            x = self.local_conv(x)
-            x = x.view(B, C, L).transpose(1, 2)
-
-            mlp_output = self.mlp(x)
-            self.mlp_output = rearrange(
-                mlp_output.clone(), "b (h w) c -> b h w c", h=hw
-            )
-
-            x = x + self.drop_path(mlp_output)
-            self.block_output = rearrange(x.clone(), "b (h w) c -> b h w c", h=hw)
-            return x
-
-        setattr(
-            mobile_sam.image_encoder.layers[1].blocks[0].__class__,
-            "forward",
-            new_forward_fn2,
-        )
-
-        mobile_sam.eval()
-        self.image_encoder = mobile_sam.image_encoder
-
-    @torch.no_grad()
-    def forward(self, x):
-        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_layer in range(len(self.image_encoder.layers)):
-            for i_block in range(len(self.image_encoder.layers[i_layer].blocks)):
-                blk = self.image_encoder.layers[i_layer].blocks[i_block]
-                attn_outputs.append(blk.attn_output)
-                mlp_outputs.append(blk.mlp_output)
-                block_outputs.append(blk.block_output)
-        return {
-            'attn': attn_outputs,
-            'mlp': mlp_outputs,
-            'block': block_outputs
-        }
-
-MODEL_DICT["MobileSAM"] = partial(MobileSAM)
-LAYER_DICT["MobileSAM"] = 12
-RES_DICT["MobileSAM"] = (1024, 1024)
-
-
-class DiNOv2(torch.nn.Module):
-    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):
-                return self.ls1(self.attn(self.norm1(x)))
-
-            def ffn_residual_func(x):
-                return self.ls2(self.mlp(self.norm2(x)))
-
-            attn_output = attn_residual_func(x)
-
-            hw = np.sqrt(attn_output.shape[1] - num_reg).astype(int)
-            self.attn_output = rearrange(
-                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()[:, 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()[:, num_reg:], "b (h w) c -> b h w c", h=hw
-            )
-            return x
-
-        setattr(self.dinov2.blocks[0].__class__, "forward", new_block_forward)
-
-    @torch.no_grad()
-    def forward(self, x):
-
-        out = self.dinov2(x)
-
-        attn_outputs, mlp_outputs, block_outputs = [], [], []
-        for i, blk in enumerate(self.dinov2.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["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)
-
-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]
-    patch_embeddings = embeddings[1:]  # [14*14, 768]
-    hw = np.sqrt(patch_embeddings.shape[0]).astype(int)
-    patch_embeddings = rearrange(patch_embeddings, "(h w) c -> c h w", h=hw)
-    patch_embeddings = F.interpolate(patch_embeddings.unsqueeze(0), size=(h, w), mode="nearest").squeeze(0)
-    patch_embeddings = rearrange(patch_embeddings, "c h w -> (h w) c")
-    embeddings = torch.cat([cls_embeddings.unsqueeze(0), patch_embeddings], dim=0)
-    return embeddings
-
-# class CLIP(torch.nn.Module):
-#     def __init__(self, ver="openai/clip-vit-base-patch16"):
-#         super().__init__()
-
-#         from transformers import CLIPProcessor, CLIPModel
-
-#         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, 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)
-        
-        def new_forward(
-                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(model.visual.transformer.resblocks[0].__class__, "forward", new_forward)
-        
-        self.model = model
-        self.model.eval()
-    
-    def forward(self, x):
-        out = self.model(x)
-        attn_outputs, mlp_outputs, block_outputs = [], [], []
-        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_base_patch14_448.mim_in22k_ft_in1k',
-            pretrained=True,
-            num_classes=0,  # remove classifier nn.Linear
-        )
-        model = model.eval()
-
-        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["CLIP(eva02_base_patch14_448/mim_in22k_ft_in1k)"] = partial(EVA02)
-LAYER_DICT["CLIP(eva02_base_patch14_448/mim_in22k_ft_in1k)"] = 12
-RES_DICT["CLIP(eva02_base_patch14_448/mim_in22k_ft_in1k)"] = (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
-
-        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):
-    def __init__(self, **kwargs):
-        super(MAE, self).__init__(**kwargs)
-
-        sd = torch.hub.load_state_dict_from_url(
-            "https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_base.pth"
-        )
-
-        checkpoint_model = sd["model"]
-        state_dict = self.state_dict()
-        for k in ["head.weight", "head.bias"]:
-            if (
-                k in checkpoint_model
-                and checkpoint_model[k].shape != state_dict[k].shape
-            ):
-                print(f"Removing key {k} from pretrained checkpoint")
-                del checkpoint_model[k]
-
-        # load pre-trained model
-        msg = self.load_state_dict(checkpoint_model, strict=False)
-        print(msg)
-        
-        # resample the patch embeddings to 56x56, take 896x896 input
-        pos_embed = self.pos_embed[0]
-        pos_embed = resample_position_embeddings(pos_embed, 42, 42)
-        self.pos_embed = nn.Parameter(pos_embed.unsqueeze(0))
-        self.img_size = (672, 672)
-        self.patch_embed.img_size = (672, 672)
-
-        self.requires_grad_(False)
-        self.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(self.blocks[0].__class__, "forward", forward)
-        
-    def forward(self, x):
-        out = super().forward(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.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': 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)
-
-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}")
-        try:
-            model = MODEL_DICT[model_name]()
-        except Exception as e:
-            print(f"Error downloading {model_name}: {e}")
-            continue
-        
-def get_all_model_names():
-    return list(MODEL_DICT.keys())
-        
-def get_model(model_name):
-    return MODEL_DICT[model_name]()
-
-@torch.no_grad()
-def extract_features(images, model, model_name, node_type, layer, batch_size=8):
-    use_cuda = torch.cuda.is_available()
-    
-    if use_cuda:
-        model = model.cuda()
-
-    chunked_idxs = torch.split(torch.arange(images.shape[0]), batch_size)
-    
-    outputs = []
-    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]
-        if out is None:
-            raise ValueError(f"Node type {node_type} not found in model {model_name}")
-        out = out[layer]
-        # normalize
-        out = F.normalize(out, dim=-1)
-        outputs.append(out.cpu().float())
-    outputs = torch.cat(outputs, dim=0)
-
-    return outputs
-
-    
-if __name__ == '__main__':
-    inp = torch.rand(1, 3, 1024, 1024)
-    model = MAE()
-    out = model(inp)
-    print(out[0][0].shape, out[0][1].shape, out[0][2].shape)

backbone_text.py

@@ -1,239 +0,0 @@
-# %%
-# 
-from typing import List, Union
-import torch
-import os
-from torch import nn
-from typing import Optional, Tuple
-
-from functools import partial
-
-MODEL_DICT = {}
-LAYER_DICT = {}
-
-class Llama(nn.Module):
-    def __init__(self, model_id="meta-llama/Meta-Llama-3.1-8B"):
-        super().__init__()
-        
-        import transformers
-        
-        access_token = os.getenv("HF_ACCESS_TOKEN")
-        if access_token is None:
-            raise ValueError("HF_ACCESS_TOKEN environment variable must be set")
-        
-        pipeline = transformers.pipeline(
-            "text-generation",
-            model=model_id,
-            model_kwargs={"torch_dtype": torch.bfloat16},
-            token=access_token,
-            device='cpu',
-        )
-        
-        tokenizer = pipeline.tokenizer
-        model = pipeline.model
-        
-        def new_forward(
-            self,
-            hidden_states: torch.Tensor,
-            attention_mask: Optional[torch.Tensor] = None,
-            position_ids: Optional[torch.LongTensor] = None,
-            past_key_value = None,
-            output_attentions: Optional[bool] = False,
-            use_cache: Optional[bool] = False,
-            cache_position: Optional[torch.LongTensor] = None,
-            position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.45
-            **kwargs,
-        ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
-            residual = hidden_states
-
-            hidden_states = self.input_layernorm(hidden_states)
-
-            # Self Attention
-            hidden_states, self_attn_weights, present_key_value = self.self_attn(
-                hidden_states=hidden_states,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_value=past_key_value,
-                output_attentions=output_attentions,
-                use_cache=use_cache,
-                cache_position=cache_position,
-                position_embeddings=position_embeddings,
-                **kwargs,
-            )
-            
-            self.attn_output = hidden_states.clone()
-            
-            hidden_states = residual + hidden_states
-
-            # Fully Connected
-            residual = hidden_states
-            hidden_states = self.post_attention_layernorm(hidden_states)
-            hidden_states = self.mlp(hidden_states)
-            
-            self.mlp_output = hidden_states.clone()
-            
-            hidden_states = residual + hidden_states
-            
-            self.block_output = hidden_states.clone()
-
-            outputs = (hidden_states,)
-
-            if output_attentions:
-                outputs += (self_attn_weights,)
-
-            if use_cache:
-                outputs += (present_key_value,)
-
-            return outputs
-        
-        # for layer in model.model.layers:
-            # setattr(layer.__class__, "forward", new_forward)
-            # setattr(layer.__class__, "__call__", new_forward)
-        setattr(model.model.layers[0].__class__, "forward", new_forward)
-        setattr(model.model.layers[0].__class__, "__call__", new_forward)
-        
-        self.model = model
-        self.tokenizer = tokenizer
-    
-    @torch.no_grad()
-    def forward(self, text: str):
-        encoded_input = self.tokenizer(text, return_tensors='pt')
-        device = next(self.model.parameters()).device
-        encoded_input = {k: v.to(device) for k, v in encoded_input.items()}
-        output = self.model(**encoded_input, output_hidden_states=True)
-    
-        attn_outputs, mlp_outputs, block_outputs = [], [], []
-        for i, blk in enumerate(self.model.model.layers):
-            attn_outputs.append(blk.attn_output)
-            mlp_outputs.append(blk.mlp_output)
-            block_outputs.append(blk.block_output)
-        
-        token_ids = encoded_input['input_ids']
-        token_texts = [self.tokenizer.decode([token_id]) for token_id in token_ids[0]]
-        
-        return {"attn": attn_outputs, "mlp": mlp_outputs, "block": block_outputs, "token_texts": token_texts}
-
-MODEL_DICT["meta-llama/Meta-Llama-3.1-8B"] = partial(Llama, model_id="meta-llama/Meta-Llama-3.1-8B")
-LAYER_DICT["meta-llama/Meta-Llama-3.1-8B"] = 32
-MODEL_DICT["meta-llama/Meta-Llama-3-8B"] = partial(Llama, model_id="meta-llama/Meta-Llama-3-8B")
-LAYER_DICT["meta-llama/Meta-Llama-3-8B"] = 32
-
-class GPT2(nn.Module):
-    def __init__(self):
-        super().__init__()
-        from transformers import GPT2Tokenizer, GPT2Model
-        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
-        model = GPT2Model.from_pretrained('gpt2')
-        
-        def new_forward(
-            self,
-            hidden_states: Optional[Tuple[torch.FloatTensor]],
-            layer_past: Optional[Tuple[torch.Tensor]] = None,
-            attention_mask: Optional[torch.FloatTensor] = None,
-            head_mask: Optional[torch.FloatTensor] = None,
-            encoder_hidden_states: Optional[torch.Tensor] = None,
-            encoder_attention_mask: Optional[torch.FloatTensor] = None,
-            use_cache: Optional[bool] = False,
-            output_attentions: Optional[bool] = False,
-        ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
-            residual = hidden_states
-            hidden_states = self.ln_1(hidden_states)
-            attn_outputs = self.attn(
-                hidden_states,
-                layer_past=layer_past,
-                attention_mask=attention_mask,
-                head_mask=head_mask,
-                use_cache=use_cache,
-                output_attentions=output_attentions,
-            )
-            attn_output = attn_outputs[0]  # output_attn: a, present, (attentions)
-            outputs = attn_outputs[1:]
-            # residual connection
-            self.attn_output = attn_output.clone()
-            hidden_states = attn_output + residual
-
-            if encoder_hidden_states is not None:
-                # add one self-attention block for cross-attention
-                if not hasattr(self, "crossattention"):
-                    raise ValueError(
-                        f"If `encoder_hidden_states` are passed, {self} has to be instantiated with "
-                        "cross-attention layers by setting `config.add_cross_attention=True`"
-                    )
-                residual = hidden_states
-                hidden_states = self.ln_cross_attn(hidden_states)
-                cross_attn_outputs = self.crossattention(
-                    hidden_states,
-                    attention_mask=attention_mask,
-                    head_mask=head_mask,
-                    encoder_hidden_states=encoder_hidden_states,
-                    encoder_attention_mask=encoder_attention_mask,
-                    output_attentions=output_attentions,
-                )
-                attn_output = cross_attn_outputs[0]
-                # residual connection
-                hidden_states = residual + attn_output
-                outputs = outputs + cross_attn_outputs[2:]  # add cross attentions if we output attention weights
-
-            residual = hidden_states
-            hidden_states = self.ln_2(hidden_states)
-            feed_forward_hidden_states = self.mlp(hidden_states)
-            # residual connection
-            self.mlp_output = feed_forward_hidden_states.clone()
-            hidden_states = residual + feed_forward_hidden_states
-
-            if use_cache:
-                outputs = (hidden_states,) + outputs
-            else:
-                outputs = (hidden_states,) + outputs[1:]
-
-            self.block_output = hidden_states.clone()
-            return outputs  # hidden_states, present, (attentions, cross_attentions)
-        
-        setattr(model.h[0].__class__, "forward", new_forward)
-        
-        self.model = model
-        self.tokenizer = tokenizer
-
-    @torch.no_grad()
-    def forward(self, text: str):
-        encoded_input = self.tokenizer(text, return_tensors='pt')
-        device = next(self.model.parameters()).device
-        encoded_input = {k: v.to(device) for k, v in encoded_input.items()}
-        output = self.model(**encoded_input, output_hidden_states=True)
-    
-        attn_outputs, mlp_outputs, block_outputs = [], [], []
-        for i, blk in enumerate(self.model.h):
-            attn_outputs.append(blk.attn_output)
-            mlp_outputs.append(blk.mlp_output)
-            block_outputs.append(blk.block_output)
-            
-        token_ids = encoded_input['input_ids']
-        token_texts = [self.tokenizer.decode([token_id]) for token_id in token_ids[0]]
-        
-        return {"attn": attn_outputs, "mlp": mlp_outputs, "block": block_outputs, "token_texts": token_texts}
-    
-MODEL_DICT["gpt2"] = GPT2
-LAYER_DICT["gpt2"] = 12
-
-
-def download_all_models():
-    for model_name in MODEL_DICT:
-        print(f"Downloading {model_name}")
-        try:
-            model = MODEL_DICT[model_name]()
-        except Exception as e:
-            print(f"Error downloading {model_name}: {e}")
-            continue
-
-
-if __name__ == '__main__':
-
-    model = MODEL_DICT["meta-llama/Meta-Llama-3-8B"]()
-    # model = MODEL_DICT["gpt2"]()
-    text = """
-    1. The majestic giraffe, with its towering height and distinctive long neck, roams the savannas of Africa. These gentle giants use their elongated tongues to pluck leaves from the tallest trees, making them well-adapted to their environment. Their unique coat patterns, much like human fingerprints, are unique to each individual.
-    """
-    model = model.cuda()
-    output = model(text)
-    print(output["block"][1].shape)
-    print(output["token_texts"])