Update cumo/model/multimodal_encoder/clip_encoder.py

cumo/model/multimodal_encoder/clip_encoder.py

@@ -48,9 +48,7 @@ class CLIPVisionTower(nn.Module):
         self.is_loaded = True
 
     def feature_select(self, image_features):
-        # Take first element if output is a tuple
-        if isinstance(image_features, tuple):
-            image_features = image_features[0]
+        #image_features = image_forward_outs.hidden_states[self.select_layer]
         if self.select_feature == 'patch':
             image_features = image_features[:, 1:]
         elif self.select_feature == 'cls_patch':
@@ -59,7 +57,6 @@ class CLIPVisionTower(nn.Module):
             raise ValueError(f'Unexpected select feature: {self.select_feature}')
         return image_features
 
-
     def split_chessboard(self, x, num_split):
         """
             x: b * c * h * w
@@ -87,13 +84,9 @@ class CLIPVisionTower(nn.Module):
         if type(images) is list:
             image_features = []
             for image in images:
-                dev = image.device if hasattr(image, "device") else torch.device("cuda" if torch.cuda.is_available() else "cpu")
-                dt = image.dtype if hasattr(image, "dtype") else torch.float16
-                print("Image shape before vision_model:", image.shape)
-                image_forward_out = self.vision_model(image.to(device=dev, dtype=dt))
+                image_forward_out = self.vision_model(image.to(device=self.device, dtype=self.dtype).unsqueeze(0), output_hidden_states=True)
                 image_feature = self.feature_select(image_forward_out).to(image.dtype)
                 image_features.append(image_feature)
-
         else:
             input_size = images.shape[3]
             img_sizes = [int(input_size * scale) for scale in self.scales]