Spaces:
Runtime error
Runtime error
File size: 2,380 Bytes
2afcb7e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
from __future__ import annotations
import clip
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
class ClipSimilarity(nn.Module):
def __init__(self, name: str = "ViT-L/14"):
super().__init__()
assert name in ("RN50", "RN101", "RN50x4", "RN50x16", "RN50x64", "ViT-B/32", "ViT-B/16", "ViT-L/14", "ViT-L/14@336px") # fmt: skip
self.size = {"RN50x4": 288, "RN50x16": 384, "RN50x64": 448, "ViT-L/14@336px": 336}.get(name, 224)
self.model, _ = clip.load(name, device="cpu", download_root="./")
self.model.eval().requires_grad_(False)
self.register_buffer("mean", torch.tensor((0.48145466, 0.4578275, 0.40821073)))
self.register_buffer("std", torch.tensor((0.26862954, 0.26130258, 0.27577711)))
def encode_text(self, text: list[str]) -> torch.Tensor:
text = clip.tokenize(text, truncate=True).to(next(self.parameters()).device)
text_features = self.model.encode_text(text)
text_features = text_features / text_features.norm(dim=1, keepdim=True)
return text_features
def encode_image(self, image: torch.Tensor) -> torch.Tensor: # Input images in range [0, 1].
image = F.interpolate(image.float(), size=self.size, mode="bicubic", align_corners=False)
image = image - rearrange(self.mean, "c -> 1 c 1 1")
image = image / rearrange(self.std, "c -> 1 c 1 1")
image_features = self.model.encode_image(image)
image_features = image_features / image_features.norm(dim=1, keepdim=True)
return image_features
def forward(
self, image_0: torch.Tensor, image_1: torch.Tensor, text_0: list[str], text_1: list[str]
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
image_features_0 = self.encode_image(image_0)
image_features_1 = self.encode_image(image_1)
text_features_0 = self.encode_text(text_0)
text_features_1 = self.encode_text(text_1)
sim_0 = F.cosine_similarity(image_features_0, text_features_0)
sim_1 = F.cosine_similarity(image_features_1, text_features_1)
sim_direction = F.cosine_similarity(image_features_1 - image_features_0, text_features_1 - text_features_0)
sim_image = F.cosine_similarity(image_features_0, image_features_1)
return sim_0, sim_1, sim_direction, sim_image
|