initial commit

app.py

@@ -0,0 +1,42 @@
+import gradio as gr
+import PIL
+import numpy as np
+from models.maskclip import MaskClip
+from models.dino import DINO
+import torchvision.transforms as T
+import torch.nn.functional as F
+from lposs import lposs, lposs_plus
+import torch
+
+device = "cpu"
+if torch.cuda.is_available():
+    print("Using GPU")
+    device = "cuda"
+# elif torch.backends.mps.is_available():
+#     device = "mps"
+
+print(f"Using device: {device}")
+
+maskclip = MaskClip().to(device)
+dino = DINO().to(device)
+to_torch_tensor = T.Compose([T.Resize(size=448, max_size=2048), T.ToTensor()])
+
+def segment_image(img: PIL.Image.Image, classnames: str, use_lposs_plus: bool | None) -> tuple[np.ndarray | PIL.Image.Image | str, list[tuple[np.ndarray | tuple[int, int, int, int], str]]]:
+    img_tensor = to_torch_tensor(PIL.Image.fromarray(img)).unsqueeze(0).to(device)
+    classnames = [c.strip() for c in classnames.split(",")]
+    num_classes = len(classnames)
+    
+    preds = lposs(maskclip, dino, img_tensor, classnames)
+    if use_lposs_plus:
+        preds = lposs_plus(img_tensor, preds)
+    preds = F.interpolate(preds, size=img.shape[:-1], mode="bilinear", align_corners=False)
+    preds = F.softmax(preds * 100, dim=1).cpu().numpy()
+    return (img, [(preds[0, i, :, :], classnames[i]) for i in range(num_classes)])
+
+demo = gr.Interface(
+    fn=segment_image,
+    inputs=["image", "text", "checkbox"],
+    outputs=["annotatedimage"],
+)
+
+demo.launch()

lposs.py

@@ -0,0 +1,293 @@
+from itertools import product
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+try:
+    import cupy as cp
+    from cupyx.scipy.sparse import csr_matrix as cp_csr_matrix, eye as cp_eye, diags as cp_diags
+    from cupyx.scipy.sparse import linalg as cp_s_linalg
+except ImportError:
+    print("Cupy not installed")
+import numpy as np
+from scipy.sparse import csr_matrix, eye, diags
+from scipy.sparse import linalg as s_linalg
+from kornia.color import rgb_to_lab
+
+
+def make_input_divisible(x: torch.Tensor, patch_size=16) -> torch.Tensor:
+    """Pad some pixels to make the input size divisible by the patch size."""
+    B, _, H_0, W_0 = x.shape
+    pad_w = (patch_size - W_0 % patch_size) % patch_size
+    pad_h = (patch_size - H_0 % patch_size) % patch_size
+
+    x = nn.functional.pad(x, (0, pad_w, 0, pad_h), value=0)
+
+    return x
+
+
+def reshape_windows(x):
+    height_width = [(y.shape[0], y.shape[1]) for y in x]
+    dim = x[0].shape[-1]
+    x = [torch.reshape(y, (-1, dim)) for y in x]
+    
+    return torch.cat(x, dim=0), height_width
+
+
+def normalize_connection_graph_cupy(G):
+    W = cp_csr_matrix(G)
+    W = W - cp_diags(W.diagonal(), 0)
+    S = W.sum(axis=1)
+    # breakpoint()
+    S[S == 0] = 1
+    D = cp.array(1.0 / cp.sqrt(S))
+    D[cp.isnan(D)] = 0
+    D[cp.isinf(D)] = 0
+    D_mh = cp_diags(D.reshape(-1), 0)
+    Wn = D_mh * W * D_mh
+    return Wn
+
+
+def normalize_connection_graph(G):
+    W = csr_matrix(G)
+    W = W - diags(W.diagonal(), 0)
+    S = W.sum(axis=1)
+    S[S == 0] = 1
+    D = np.array(1.0 / np.sqrt(S))
+    D[np.isnan(D)] = 0
+    D[np.isinf(D)] = 0
+    D_mh = diags(D.reshape(-1), 0)
+    Wn = D_mh * W * D_mh
+    return Wn
+
+
+def cp_dfs_search(L, Y, tol=1e-6, maxiter=10):
+    out = cp_s_linalg.cg(L, Y, tol=tol, maxiter=maxiter)[0]
+
+    return out
+
+
+def dfs_search(L, Y, tol=1e-6, maxiter=10):
+    out = s_linalg.cg(L, Y, rtol=tol, maxiter=maxiter)[0]
+
+    return out
+
+
+def perform_lp(L, preds):
+    if torch.cuda.is_available():
+        lp_preds = cp.zeros(preds.shape)
+        preds = cp.asarray(preds)
+        for cls_idx, y_cls in enumerate(preds.T):
+            Y = y_cls
+            lp_preds[:, cls_idx] = cp_dfs_search(L, Y)
+        lp_preds = torch.as_tensor(lp_preds, device="cuda")
+    else:
+        lp_preds = np.zeros(preds.shape)
+        for cls_idx, y_cls in enumerate(preds.T):
+            Y = y_cls
+            lp_preds[:, cls_idx] = dfs_search(L, Y)
+        lp_preds = torch.as_tensor(lp_preds, device="cpu")
+
+    return lp_preds
+
+
+def get_lposs_laplacian(feats, locations, height_width, sigma=0.0, pix_dist_pow=2, k=100, gamma=1.0, alpha=0.95, patch_size=16):
+    idx_window = torch.cat([window * torch.ones((h*w, ), device=feats.device, dtype=torch.int64) for window, (h, w) in enumerate(height_width)])
+    idx_h = torch.cat([torch.arange(h).view(-1,1).repeat(1, w).flatten() for h, w in height_width]).to(feats.device)
+    idx_w = torch.cat([torch.arange(w).view(1,-1).repeat(h, 1).flatten() for h, w in height_width]).to(feats.device)
+    loc_h = locations[idx_window, 0] + (patch_size // 2) + idx_h * patch_size
+    loc_w = locations[idx_window, 2] + (patch_size // 2) + idx_w * patch_size
+    locs = torch.stack((loc_h, loc_w), 1)
+    locs = torch.unsqueeze(locs, 0)
+    dist = torch.cdist(locs, locs, p=2)
+    dist = dist[0, ...]
+    dist = dist ** pix_dist_pow
+    geometry_affinity = torch.exp(-sigma * dist)
+
+    N = feats.shape[0]
+    
+    affinity = feats @ feats.T
+    sims, ks = torch.topk(affinity, k=k, dim=1)
+
+    sims[sims < 0] = 0
+    sims = sims ** gamma
+    geometry_affinity = geometry_affinity.gather(1, ks).flatten()
+    sims = sims.flatten()
+    sims = sims * geometry_affinity
+    ks = ks.flatten()
+    rows = torch.arange(N).repeat_interleave(k)
+    
+    if torch.cuda.is_available():
+        W = cp_csr_matrix(
+            (cp.asarray(sims), (cp.asarray(rows), cp.asarray(ks))),
+            shape=(N, N),
+        )
+        W = W + W.T
+        Wn = normalize_connection_graph_cupy(W)
+        L = cp_eye(Wn.shape[0]) - alpha * Wn
+    else:
+        W = csr_matrix(
+            (sims.cpu().numpy(), (rows.cpu().numpy(), ks.cpu().numpy())),
+            shape=(N, N),
+        )
+        W = W + W.T
+        Wn = normalize_connection_graph(W)
+        L = eye(Wn.shape[0]) - alpha * Wn
+
+    return L
+
+
+def lposs(clip, dino, img, classnames, window_size=(224,224), window_stride=(112, 112), sigma=0.01, pix_dist_pow=1, lp_k_image=400, lp_gamma=3.0, lp_alpha=0.95):
+    h_stride, w_stride = window_stride
+    h_crop, w_crop = window_size
+    batch_size, _, h_img, w_img = img.size()
+    h_grids = max(h_img - h_crop + h_stride - 1, 0) // h_stride + 1
+    w_grids = max(w_img - w_crop + w_stride - 1, 0) // w_stride + 1
+
+    clf = clip.get_classifier(classnames)
+        
+    locations = img.new_zeros((h_grids*w_grids, 4))
+    dino_feats = []
+    clip_feats = []
+    for h_idx in range(h_grids):
+        for w_idx in range(w_grids):
+            y1 = h_idx * h_stride
+            x1 = w_idx * w_stride
+            y2 = min(y1 + h_crop, h_img)
+            x2 = min(x1 + w_crop, w_img)
+            y1 = max(y2 - h_crop, 0)
+            x1 = max(x2 - w_crop, 0)
+            crop_img = img[:, :, y1:y2, x1:x2]
+
+            img_dino_feats, (h_dino, w_dino) = dino(make_input_divisible(crop_img, dino.patch_size)) # (1, 768, N)
+            img_dino_feats = img_dino_feats.reshape((batch_size, -1, h_dino, w_dino)).permute(0, 2, 3, 1) # (1, h_dino, w_dino, 768)
+            img_clip_feats = clip(make_input_divisible(crop_img, clip.patch_size)) # (1, 512, h, w)
+
+            if img_clip_feats.shape[1] != img_dino_feats.shape[1] or img_clip_feats.shape[2] != img_dino_feats.shape[2]:
+                img_clip_feats = F.interpolate(img_clip_feats, size=(img_dino_feats.shape[1], img_dino_feats.shape[2]), mode='bilinear', align_corners=False)
+                
+            img_clip_feats = img_clip_feats.permute(0, 2, 3, 1) # (1, h, w, 512)
+
+            dino_feats.append(img_dino_feats[0, ...])
+            clip_feats.append(img_clip_feats[0, ...])
+            locations[h_idx*w_grids + w_idx, 0] = y1
+            locations[h_idx*w_grids + w_idx, 1] = y2
+            locations[h_idx*w_grids + w_idx, 2] = x1
+            locations[h_idx*w_grids + w_idx, 3] = x2
+
+    num_classes = clf.shape[0]
+
+    patch_size = dino.patch_size
+        
+    dino_feats, height_width = reshape_windows(dino_feats)
+    clip_feats, _ = reshape_windows(clip_feats)
+    dino_feats = F.normalize(dino_feats, p=2, dim=-1)
+    clip_feats = F.normalize(clip_feats, p=2, dim=-1)
+
+    L = get_lposs_laplacian(dino_feats, locations, height_width, sigma=sigma, pix_dist_pow=pix_dist_pow, k=lp_k_image, gamma=lp_gamma, alpha=lp_alpha, patch_size=patch_size)
+    clip_preds = clip_feats @ clf.T
+    
+    lp_preds = perform_lp(L, clip_preds)
+    
+    preds = img.new_zeros((batch_size, num_classes, h_img, w_img))
+    count_mat = img.new_zeros((batch_size, 1, h_img, w_img))
+    idx_window = torch.cat([window * torch.ones((h*w, ), device=dino_feats.device, dtype=torch.int64) for window, (h, w) in enumerate(height_width)])
+    for h_idx in range(h_grids):
+        for w_idx in range(w_grids):
+            y1 = h_idx * h_stride
+            x1 = w_idx * w_stride
+            y2 = min(y1 + h_crop, h_img)
+            x2 = min(x1 + w_crop, w_img)
+            y1 = max(y2 - h_crop, 0)
+            x1 = max(x2 - w_crop, 0)
+            win_id = h_idx*w_grids + w_idx
+            crop_seg_logit = lp_preds[torch.where(idx_window == win_id)[0], :]
+            crop_seg_logit = torch.reshape(crop_seg_logit, height_width[win_id]+(num_classes, ))
+            crop_seg_logit = torch.unsqueeze(crop_seg_logit, 0)
+            crop_seg_logit = torch.permute(crop_seg_logit, (0, 3, 1, 2))
+            crop_seg_logit = F.interpolate(
+                input=crop_seg_logit,
+                size=(y2-y1, x2-x1),
+                mode='bilinear',
+                align_corners=False
+            )
+            assert crop_seg_logit.shape[2] == (y2 - y1) and crop_seg_logit.shape[3] == (x2 - x1)
+            preds += F.pad(crop_seg_logit,
+                            (int(x1), int(preds.shape[3] - x2), int(y1),
+                            int(preds.shape[2] - y2)))
+
+            count_mat[:, :, y1:y2, x1:x2] += 1
+    assert (count_mat == 0).sum() == 0
+    preds = preds / count_mat
+
+    return preds
+
+
+def get_pixel_connections(img, neigh=1):
+    img = img[0, ...]
+    img_lab = rgb_to_lab(img)
+    img_lab = img_lab.permute((1, 2, 0))
+    img_lab /= torch.tensor([100, 128, 128], device=img.device) # project Lab values to 0-1 range
+    img_h, img_w, _ = img_lab.shape
+    img_lab = img_lab.reshape((img_h*img_w, -1))
+
+    idx = torch.arange(img_h * img_w).to(img.device)
+    loc_h = idx // img_w
+    loc_w = idx % img_w
+    locs = torch.stack((loc_h, loc_w), 1)
+    
+    rows, cols = [], []
+
+    for mov in product(range(-neigh, neigh+1), range(-neigh, neigh+1)):
+        if mov[0] == 0 and mov[1] == 0:
+            continue
+        new_locs = locs + torch.tensor(mov).to(img.device)
+        mask = torch.logical_and(torch.logical_and(torch.logical_and(new_locs[:, 0] >= 0, new_locs[:, 1] >= 0), new_locs[:, 0] < img_h), new_locs[:, 1] < img_w)
+        rows.append(torch.where(mask)[0])
+        col = new_locs[mask, :]
+        col = col[:, 0] * img_w + col[:, 1]
+        cols.append(col)
+
+    rows = torch.cat(rows)
+    cols = torch.cat(cols)
+    pixel_pixel_data = ((img_lab[rows, :] - img_lab[cols, :]) ** 2).sum(dim=-1)
+
+    return rows, cols, pixel_pixel_data, locs
+
+
+def get_laplacian(rows, cols, data, N, alpha=0.99):
+    if torch.cuda.is_available():
+        rows = cp.asarray(rows)
+        cols = cp.asarray(cols)
+        data = cp.asarray(data)
+        W = cp_csr_matrix(
+            (data, (rows, cols)),
+            shape=(N, N),
+        )
+
+        Wn = normalize_connection_graph_cupy(W)
+        L = cp_eye(Wn.shape[0]) - alpha * Wn
+    else:
+        W = csr_matrix(
+            (data.cpu().numpy(), (rows.cpu().numpy(), cols.cpu().numpy())),
+            shape=(N, N),
+        )
+
+        Wn = normalize_connection_graph(W)
+        L = eye(Wn.shape[0]) - alpha * Wn
+    return L
+
+
+def lposs_plus(img, preds, tau=0.01, alpha=0.95):
+    preds = preds[0, ...]
+    num_classes, h_img, w_img = preds.shape
+    preds = preds.permute((1, 2, 0))
+    preds = preds.reshape((h_img*w_img, -1))
+
+    rows, cols, pixel_pixel_data, locs = get_pixel_connections(img, neigh=6)
+    pixel_pixel_data = torch.sqrt(pixel_pixel_data)
+    pixel_pixel_data = torch.exp(-pixel_pixel_data / tau)
+    L = get_laplacian(rows, cols, pixel_pixel_data, preds.shape[0], alpha=alpha)
+
+    lp_preds = perform_lp(L, preds)
+    
+    return lp_preds.reshape((h_img, w_img, num_classes)).permute((2, 0, 1)).unsqueeze(0)

models/dino.py

@@ -0,0 +1,69 @@
+import torch
+import torch.nn as nn
+import torchvision.transforms as T
+
+NORMALIZE = T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
+
+class DINO(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.backbone = torch.hub.load('facebookresearch/dino:main', 'dino_vitb16')
+        self.hook_features = {}
+        def hook_fn_forward_qkv(module, input, output):
+            self.hook_features["qkv"] = output
+
+        self.backbone._modules["blocks"][-1]._modules["attn"]._modules[
+            "qkv"
+        ].register_forward_hook(hook_fn_forward_qkv)
+
+        self.patch_size = 16
+        self.enc_type_feats = "v"
+
+
+    @torch.no_grad()
+    def extract_feats(self, type_feats="k"):
+        """
+        DINO feature extractor. Attaches a hook on the last attention layer.
+        :param type_feats: (string) - type of features from DINO ViT
+        """
+        nh = self.backbone.blocks[-1].attn.num_heads
+        nb_im, nb_tokens, C_qkv = self.hook_features["qkv"].shape
+
+        qkv = (
+            self.hook_features["qkv"]
+                .reshape(
+                nb_im, nb_tokens, 3, nh, C_qkv // nh // 3
+            )  # 3 corresponding to |qkv|
+                .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        if type_feats == "q":
+            return q.transpose(1, 2).float()
+        elif type_feats == "k":
+            return k.transpose(1, 2).float()
+        elif type_feats == "v":
+            return v.transpose(1, 2).float()
+        else:
+            raise ValueError("Unknown features")
+
+
+    @torch.no_grad()
+    def forward(self, x):
+        x = NORMALIZE(x)
+        h_featmap = x.shape[-2] // self.patch_size
+        w_featmap = x.shape[-1] // self.patch_size
+
+        # Forward pass
+        # Encoder forward pass and get hooked intermediate values
+        _ = self.backbone(x)
+
+        # Get decoder features
+        feats = self.extract_feats(type_feats=self.enc_type_feats)
+        num_extra_tokens = 1
+
+        # B nbtokens+1 nh dim
+        feats = feats[:, num_extra_tokens:, :, :].flatten(-2, -1).permute(0, 2, 1)  # B C nbtokens
+        # B, C, nbtokens
+        feats = feats / feats.norm(dim=1, keepdim=True)  # normalize features
+
+        return feats, (h_featmap, w_featmap)

models/maskclip.py

@@ -0,0 +1,154 @@
+# ---------------------------------------------------------------------------------------------------
+# CLIP-DINOiser
+# authors: Monika Wysoczanska, Warsaw University of Technology
+
+# Copyright (c) OpenMMLab. All rights reserved.
+# Modified version of the original MaskCLIP code: https://github.com/chongzhou96/MaskCLIP/tree/master
+# ---------------------------------------------------------------------------------------------------
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import List, Tuple
+from torch import Tensor
+from open_clip import get_tokenizer,  create_model_from_pretrained
+import torchvision.transforms as T
+from .utils import imagenet_templates
+
+OPENAI_NORMALIZE = T.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
+
+class MaskClip(nn.Module):
+    def __init__(
+            self,
+            clip_model="ViT-B-16",
+            pretrained="laion2b_s34b_b88k",
+            patch_size=16,
+            img_size=(224, 224),
+            in_channels=768,
+            text_channels=512,
+        ):
+        super(MaskClip, self).__init__()
+
+        self.patch_size = patch_size
+        self.img_size = img_size
+        model, _ = create_model_from_pretrained(clip_model, pretrained=pretrained)
+        model.eval()
+        self.clip_T = OPENAI_NORMALIZE
+        self.hook_features = {}
+        self.backbone = model
+        def hook_fn_forward(module, input, output):
+            self.hook_features["v"] = output
+        self.backbone.visual.transformer.resblocks[-2].register_forward_hook(hook_fn_forward)
+        self._positional_embd = nn.Parameter(self.backbone.visual.positional_embedding.data.clone())
+        self.proj = nn.Conv2d(in_channels, text_channels, 1, bias=False)
+        self.proj.weight = nn.Parameter(model.visual.proj.t()[:, :, None, None])
+        self.tokenizer = get_tokenizer(clip_model)
+
+    @torch.no_grad()
+    def extract_feat(self, inputs: Tensor) -> Tuple[Tensor]:
+        """Extract features from images."""
+        pos_embed = self.backbone.visual.positional_embedding
+
+        B, C, H, W = inputs.shape
+        hw_shape = (H // self.patch_size, W // self.patch_size)
+        x_len, pos_len = hw_shape[0]*hw_shape[1], pos_embed.shape[0]
+
+        if x_len != pos_len:
+            if pos_len == (self.img_size[0] // self.patch_size) * (self.img_size[1] // self.patch_size) + 1:
+                pos_h = self.img_size[0] // self.patch_size
+                pos_w = self.img_size[1] // self.patch_size
+            else:
+                raise ValueError(
+                    '{}, {}'.format(x_len, pos_len))
+
+            self.backbone.visual.positional_embedding.data = self.resize_pos_embed(
+                self._positional_embd[None], hw_shape,  (pos_h, pos_w), 'bicubic')[0]
+
+        _ = self.backbone(inputs)
+        v = self.hook_features["v"]
+        v = self.extract_v(v, self.backbone.visual.transformer.resblocks[-1]).permute(1, 0, 2)
+        v = self.backbone.visual.ln_post(v)
+        # v = v[:, 1:] # was there in original code
+        v = v.permute(1, 0, 2)[:, 1:] # put this as per https://github.com/wysoczanska/clip_dinoiser/issues/10
+        v = v.reshape(B, hw_shape[0], hw_shape[1], -1).permute(0, 3, 1, 2).contiguous()
+
+        self.backbone.visual.positional_embedding.data = self._positional_embd
+        return v
+
+    @torch.no_grad()
+    def extract_v(self, x, block):
+        y = block.ln_1(x)
+        y = torch.nn.functional.linear(y, block.attn.in_proj_weight, block.attn.in_proj_bias)
+        B, N, C = y.shape
+        y = y.view(B, N, 3, C // 3).permute(2, 0, 1, 3).reshape(3 * B, N, C // 3)
+        y = F.linear(y, block.attn.out_proj.weight, block.attn.out_proj.bias)
+        q, k, v = y.tensor_split(3, dim=0)
+        v += x
+        v += block.mlp(block.ln_2(v))
+        return v
+
+
+    @staticmethod
+    def resize_pos_embed(pos_embed, input_shpae, pos_shape, mode):
+        """Resize pos_embed weights.
+
+        Resize pos_embed using bicubic interpolate method.
+        Args:
+            pos_embed (torch.Tensor): Position embedding weights.
+            input_shpae (tuple): Tuple for (downsampled input image height,
+                downsampled input image width).
+            pos_shape (tuple): The resolution of downsampled origin training
+                image.
+            mode (str): Algorithm used for upsampling:
+                ``'nearest'`` | ``'linear'`` | ``'bilinear'`` | ``'bicubic'`` |
+                ``'trilinear'``. Default: ``'nearest'``
+        Return:
+            torch.Tensor: The resized pos_embed of shape [B, L_new, C]
+        """
+        assert pos_embed.ndim == 3, 'shape of pos_embed must be [B, L, C]'
+        pos_h, pos_w = pos_shape
+        cls_token_weight = pos_embed[:, 0]
+        pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w):]
+        pos_embed_weight = pos_embed_weight.reshape(
+            1, pos_h, pos_w, pos_embed.shape[2]).permute(0, 3, 1, 2)
+        pos_embed_weight = F.interpolate(
+            pos_embed_weight, size=input_shpae, align_corners=False, mode=mode)
+        cls_token_weight = cls_token_weight.unsqueeze(1)
+        pos_embed_weight = torch.flatten(pos_embed_weight, 2).transpose(1, 2)
+        pos_embed = torch.cat((cls_token_weight, pos_embed_weight), dim=1)
+        return pos_embed
+    
+    @torch.no_grad()
+    def decode_head(self, x: Tensor) -> Tensor:
+        feat = self.proj(x)
+
+        return feat
+
+
+    @torch.no_grad()
+    def forward(self, inputs: Tensor) -> Tensor:
+        """Encode images with backbone and decode into a semantic segmentation
+        map of the same size as input."""
+        inputs = self.clip_T(inputs)
+        x = self.extract_feat(inputs)
+        feats = self.decode_head(x)
+        return feats
+    
+
+    @torch.no_grad()
+    def get_classifier(self, classnames:List[str]) -> Tensor:
+        aug_embeddings = torch.stack([self._embed_label(label) for label in classnames])
+        aug_embeddings = aug_embeddings / aug_embeddings.norm(dim=-1, keepdim=True)
+        return aug_embeddings.squeeze(1)
+
+
+    @torch.no_grad()
+    def _embed_label(self, label: str) -> Tensor:
+        """Encode label name into a single vector."""
+        all_prompts = [self.tokenizer(template.format(label)) for template in imagenet_templates]
+        all_prompts = torch.cat(all_prompts)
+        all_prompts = all_prompts.to(self.backbone.visual.positional_embedding.device)
+        out = self.backbone.encode_text(all_prompts)
+        out /= out.norm(dim=-1, keepdim=True)
+        out = out.mean(dim=0)
+        return out

models/utils.py

@@ -0,0 +1,82 @@
+imagenet_templates = [
+    'a bad photo of a {}.',
+    'a photo of many {}.',
+    'a sculpture of a {}.',
+    'a photo of the hard to see {}.',
+    'a low resolution photo of the {}.',
+    'a rendering of a {}.',
+    'graffiti of a {}.',
+    'a bad photo of the {}.',
+    'a cropped photo of the {}.',
+    'a tattoo of a {}.',
+    'the embroidered {}.',
+    'a photo of a hard to see {}.',
+    'a bright photo of a {}.',
+    'a photo of a clean {}.',
+    'a photo of a dirty {}.',
+    'a dark photo of the {}.',
+    'a drawing of a {}.',
+    'a photo of my {}.',
+    'the plastic {}.',
+    'a photo of the cool {}.',
+    'a close-up photo of a {}.',
+    'a black and white photo of the {}.',
+    'a painting of the {}.',
+    'a painting of a {}.',
+    'a pixelated photo of the {}.',
+    'a sculpture of the {}.',
+    'a bright photo of the {}.',
+    'a cropped photo of a {}.',
+    'a plastic {}.',
+    'a photo of the dirty {}.',
+    'a jpeg corrupted photo of a {}.',
+    'a blurry photo of the {}.',
+    'a photo of the {}.',
+    'a good photo of the {}.',
+    'a rendering of the {}.',
+    'a {} in a video game.',
+    'a photo of one {}.',
+    'a doodle of a {}.',
+    'a close-up photo of the {}.',
+    'a photo of a {}.',
+    'the origami {}.',
+    'the {} in a video game.',
+    'a sketch of a {}.',
+    'a doodle of the {}.',
+    'a origami {}.',
+    'a low resolution photo of a {}.',
+    'the toy {}.',
+    'a rendition of the {}.',
+    'a photo of the clean {}.',
+    'a photo of a large {}.',
+    'a rendition of a {}.',
+    'a photo of a nice {}.',
+    'a photo of a weird {}.',
+    'a blurry photo of a {}.',
+    'a cartoon {}.',
+    'art of a {}.',
+    'a sketch of the {}.',
+    'a embroidered {}.',
+    'a pixelated photo of a {}.',
+    'itap of the {}.',
+    'a jpeg corrupted photo of the {}.',
+    'a good photo of a {}.',
+    'a plushie {}.',
+    'a photo of the nice {}.',
+    'a photo of the small {}.',
+    'a photo of the weird {}.',
+    'the cartoon {}.',
+    'art of the {}.',
+    'a drawing of the {}.',
+    'a photo of the large {}.',
+    'a black and white photo of a {}.',
+    'the plushie {}.',
+    'a dark photo of a {}.',
+    'itap of a {}.',
+    'graffiti of the {}.',
+    'a toy {}.',
+    'itap of my {}.',
+    'a photo of a cool {}.',
+    'a photo of a small {}.',
+    'a tattoo of the {}.',
+]

requrements.txt

@@ -0,0 +1,47 @@
+certifi==2025.1.31
+charset-normalizer==3.4.1
+cupy-cuda12x==13.3.0
+fastrlock==0.8.3
+filelock==3.13.1
+fsspec==2024.6.1
+ftfy==6.3.1
+huggingface-hub==0.28.1
+idna==3.10
+Jinja2==3.1.4
+kornia==0.8.0
+kornia_rs==0.1.8
+MarkupSafe==2.1.5
+mpmath==1.3.0
+networkx==3.3
+numpy==2.1.2
+nvidia-cublas-cu12==12.6.4.1
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cudnn-cu12==9.5.1.17
+nvidia-cufft-cu12==11.3.0.4
+nvidia-curand-cu12==10.3.7.77
+nvidia-cusolver-cu12==11.7.1.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cusparselt-cu12==0.6.3
+nvidia-nccl-cu12==2.21.5
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-nvtx-cu12==12.6.77
+open_clip_torch==2.30.0
+packaging==24.2
+pillow==11.0.0
+PyYAML==6.0.2
+regex==2024.11.6
+requests==2.32.3
+safetensors==0.5.2
+scipy==1.15.1
+sympy==1.13.1
+timm==1.0.14
+torch==2.6.0+cu126
+torchaudio==2.6.0+cu126
+torchvision==0.21.0+cu126
+tqdm==4.67.1
+triton==3.2.0
+typing_extensions==4.12.2
+urllib3==2.3.0
+wcwidth==0.2.13