Spaces:
Running on Zero
Running on Zero
| import numpy as np | |
| from torch import cuda, no_grad, relu, Tensor | |
| from torch.nn import functional as F | |
| from transformers import AutoModel, AutoProcessor | |
| class SigLip2: | |
| # MODEL_NAME = "google/siglip2-giant-opt-patch16-256" # from utils | |
| # MODEL_NAME = "google/siglip2-so400m-patch16-512" # bigger image | |
| # MODEL_NAME = "google/siglip2-so400m-patch14-384" # more patches | |
| MODEL_NAME = "google/siglip2-so400m-patch16-256" | |
| DEVICE = "cuda" if cuda.is_available() else "cpu" | |
| def scaleMinMax(data): | |
| dmin = data.min() | |
| dmax = data.max() | |
| return (data - dmin) / (dmax - dmin + 1e-10) | |
| def __init__(self, model=None): | |
| model_name = SigLip2.MODEL_NAME if model is None else model | |
| self.processor = AutoProcessor.from_pretrained(model_name) | |
| self.model = AutoModel.from_pretrained(model_name).to(self.DEVICE) | |
| self.model_grid_size = ( | |
| self.model.config.vision_config.image_size // | |
| self.model.config.vision_config.patch_size | |
| ) | |
| self.model_scale = self.model.logit_scale.exp().detach().cpu().numpy() | |
| self.model_bias = self.model.logit_bias.detach().cpu().numpy() | |
| def get_gradient_activation_map(self, img, texts, *, img_idx=0, text_idx=None): | |
| if isinstance(texts, str): | |
| texts = [texts] | |
| if isinstance(texts, list) and isinstance(texts[0], str): | |
| inputs = self.processor( | |
| text=texts, | |
| padding="max_length", max_length=64, truncation=True, | |
| return_tensors="pt" | |
| ).to(self.model.device) | |
| with no_grad(): | |
| text_embs = self.model.text_model(**inputs).pooler_output | |
| texts = text_embs.detach().cpu().numpy() | |
| if isinstance(texts, np.ndarray) and len(texts.shape) == 1: | |
| texts = texts[None, :] | |
| if (isinstance(texts, np.ndarray) or isinstance(texts, Tensor)) and len(texts.shape) == 2: | |
| texts = Tensor(texts).to(self.model.device) | |
| return self.get_gradient_activation_map_from_embeddings(img, texts, img_idx=img_idx, text_idx=text_idx) | |
| else: | |
| raise TypeError(f"Expected a 2D np.ndarray or Tensor, got {type(texts)}") | |
| def get_gradient_activation_map_from_embeddings(self, img, texts, *, img_idx=0, text_idx=None): | |
| text_idxs = range(len(texts)) if text_idx is None else [text_idx] | |
| text_activations = [] | |
| inputs = self.processor( | |
| images=img, | |
| padding="max_length", max_length=64, truncation=True, | |
| return_tensors="pt" | |
| ).to(self.model.device) | |
| outputs = self.model.vision_model(**inputs) | |
| img_embedding = F.normalize(outputs.pooler_output[img_idx], p=2, dim=-1) | |
| patch_features = outputs.last_hidden_state | |
| patch_features.retain_grad() | |
| for tidx in text_idxs: | |
| text_embedding = F.normalize(texts[tidx], p=2, dim=-1) | |
| similarity_score = (text_embedding * img_embedding).sum() | |
| self.model.zero_grad() | |
| similarity_score.backward(retain_graph=True) # need L4 GPU with 24GB (15.5GB) | |
| patch_grads = patch_features.grad | |
| patch_weights = patch_grads[img_idx].mean(dim=0) | |
| cam = (patch_weights * patch_features[img_idx]).sum(dim=-1) | |
| cam = relu(cam).detach().cpu().numpy() | |
| text_activations.append(SigLip2.scaleMinMax(cam)) | |
| cam01 = SigLip2.scaleMinMax(np.array(text_activations).sum(axis=0)) | |
| return cam01.reshape(self.model_grid_size, self.model_grid_size) | |