Initial commit

llava/model/multimodal_encoder/imagebind.py

@@ -0,0 +1,73 @@
+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