| 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 |
|
|
