| import torch
|
| import torch.nn as nn
|
|
|
| from transformers import CLIPImageProcessor
|
|
|
| try:
|
| from imagebind.models import imagebind_model
|
| from imagebind.models.imagebind_model import ModalityType
|
| from imagebind.data import load_and_transform_audio_data
|
| except ImportError:
|
| pass
|
|
|
|
|
| class ImageBindWrapper(nn.Module):
|
| def __init__(self, vision_tower, select_layer, select_feature="patch", delay_load=False):
|
| super().__init__()
|
|
|
| self.is_loaded = False
|
|
|
| self.vision_tower_name = vision_tower
|
| self.select_layer = select_layer
|
| self.select_feature = select_feature
|
|
|
| if not delay_load:
|
| self.load_model()
|
|
|
| def load_model(self):
|
| self.image_processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-large-patch14")
|
| self.vision_tower = imagebind_model.imagebind_huge(pretrained=True)
|
| for p in self.vision_tower.parameters():
|
| p.requires_grad = False
|
| self.vision_tower.eval()
|
| self.is_loaded = True
|
|
|
| def train(self, mode=True):
|
| self.training = mode
|
|
|
| if self.is_loaded:
|
| self.vision_tower.eval()
|
|
|
| @torch.no_grad()
|
| def forward(self, x):
|
| if type(x) == dict:
|
| if x["audios"] is not None:
|
| inputs = {ModalityType.AUDIO: load_and_transform_audio_data(x["audios"], device=self.device).half()}
|
| embeddings = self.vision_tower(inputs)
|
| audio_embedding = embeddings[ModalityType.AUDIO]
|
| return audio_embedding.unsqueeze(1)
|
| else:
|
| inputs = {ModalityType.VISION: x.to(dtype=self.dtype)}
|
| embeddings = self.vision_tower(inputs)
|
| vision_embedding = embeddings[ModalityType.VISION]
|
| if vision_embedding.ndim == 2:
|
| return vision_embedding.unsqueeze(1)
|
| if vision_embedding.shape[1] == 257:
|
| return vision_embedding[:, 1:]
|
| raise ValueError(f"Unexpected shape: {vision_embedding.shape}")
|
|
|
| @property
|
| def dummy_feature(self):
|
| return torch.zeros(1, 1024, device=self.device, dtype=self.dtype)
|
|
|
| @property
|
| def dtype(self):
|
| return self.vision_tower.modality_preprocessors.vision.cls_token.dtype
|
|
|
| @property
|
| def device(self):
|
| return self.vision_tower.modality_preprocessors.vision.cls_token.device
|
|
|
| @property
|
| def hidden_size(self):
|
| return 1024
|
|
|
