File size: 3,067 Bytes
85efb5b |
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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import torch
import torch.nn.functional as F
from transformers import SiglipImageProcessor, SiglipVisionConfig, SiglipVisionModel
from .base_encoder import BaseVisionTower, ProcessorWrapper
class SiglipVisionTower(BaseVisionTower):
def __init__(self, vision_tower_name, args, delay_load=False):
super(SiglipVisionTower, self).__init__(vision_tower_name, args, delay_load)
model_path = "./checkpoints/siglip-so400m-patch14-384"
base_model_name, res, interp = model_path, 384, 576
self.vision_tower_name = base_model_name
self._image_size = res if res is not None else 512
self._interp_size = interp
if not self.delay_load:
self.load_model()
elif self.unfreeze_mm_vision_tower:
self.load_model()
else:
self._hidden_size = 1152
def load_model(self, device_map=None):
self.vision_model = "siglip"
# clip_model, processor = create_model_from_pretrained(self.vision_tower_name)
self.vision_tower = SiglipVisionModel.from_pretrained(self.vision_tower_name)
# self.vision_tower = clip_model.visual.trunk
self.vision_tower.output_tokens = True
self._hidden_size = self.vision_tower.config.hidden_size
self._image_size = self.vision_tower.config.image_size
self._patch_size = self.vision_tower.config.patch_size
self.image_processor = SiglipImageProcessor.from_pretrained(
self.vision_tower_name
)
self.vision_tower.requires_grad_(self.unfreeze_mm_vision_tower)
self.is_loaded = True
def interpolate(self, image_features):
if self._interp_size is None:
return image_features
b, num_tokens, dim = image_features.shape
if num_tokens != self.num_patches:
target_h = target_w = int(self._interp_size**0.5)
h = w = int(num_tokens**0.5)
image_features = image_features.view(b, h, w, dim)
image_features = image_features.permute(0, 3, 1, 2).contiguous()
image_features = F.interpolate(
image_features.to(torch.float32),
size=(target_h, target_w),
mode="bilinear",
align_corners=False,
).to(image_features.dtype)
# Permute the dimensions back to (b, target_h, target_w, dim)
image_features = image_features.permute(0, 2, 3, 1).contiguous()
# Flatten the spatial dimensions (target_h, target_w) into a single dimension
image_features = image_features.flatten(1, 2)
return image_features
def _forward(self, images, interpolate_token=576):
with torch.set_grad_enabled(self.unfreeze_mm_vision_tower):
image_features = self.vision_tower.forward(
images.to(device=self.device, dtype=self.dtype),
output_hidden_states=True,
).hidden_states[-1]
interp_features = self.interpolate(image_features)
return interp_features
|