Update modeling_vlm.py

modeling_vlm.py

@@ -3,7 +3,7 @@ from typing import Any, override
 import torch
 import torch.nn as nn
 from torch import FloatTensor, LongTensor, Tensor
-from transformers import AutoModel, PreTrainedModel, LlamaForCausalLM, LlamaModel
+from transformers import AutoModel, LlamaForCausalLM, LlamaModel, PreTrainedModel
 
 from .configuration_vlm import VLMConfig
 from .connectors import Connector, connector_map
@@ -11,6 +11,7 @@ from .connectors import Connector, connector_map
 
 class VLM(LlamaModel):
     config_class = VLMConfig
+
     @override
     def __init__(self, config):
         super().__init__(config)
@@ -41,6 +42,7 @@ class VLM(LlamaModel):
 
 class VLMForCausalLM(LlamaForCausalLM):
     config_class = VLMConfig
+
     @override
     def __init__(self, config):
         super().__init__(config)
@@ -132,7 +134,7 @@ class VLMForCausalLM(LlamaForCausalLM):
     ):
         images = kwargs.pop("images", None)
         image_sizes = kwargs.pop("image_sizes", None)
-        inputs = super(self.__class__, self).prepare_inputs_for_generation(
+        inputs = super().prepare_inputs_for_generation(
             input_ids,
             past_key_values=past_key_values,
             inputs_embeds=inputs_embeds,
@@ -145,8 +147,31 @@ class VLMForCausalLM(LlamaForCausalLM):
             inputs["image_sizes"] = image_sizes
         return inputs
 
-    def encode_images(self: Any, images: Tensor) -> tuple[Tensor, ...]:
-        image_features = self.model.vision_model(images)
+    def encode_images(self: Any, images: list[Tensor] | Tensor) -> list[Tensor] | Tensor:
+        if type(images) is list:
+            image_features: list[Tensor] | Tensor = []
+            for image in images:
+                outputs = self.model.vision_model(
+                    image.unsqueeze(0),
+                    output_hidden_states=True,
+                )
+                hidden_states: Tensor = outputs.hidden_states[self.output_layer].to(image.dtype)
+                if not self.config.vision_config.use_cls_token:
+                    image_features.append(hidden_states[:, 1:])
+                else:
+                    image_features.append(hidden_states)
+        else:
+            outputs = self.model.vision_model(
+                images,
+                output_hidden_states=True,
+            )
+            hidden_states = outputs.hidden_states[self.config.vision_config.output_layer].to(
+                images.dtype
+            )
+            if not self.config.vision_config.use_cls_token:
+                image_features = hidden_states[:, 1:]
+            else:
+                image_features = hidden_states
         image_features = self.model.connector(image_features)
 
         return image_features
@@ -398,4 +423,4 @@ class VLMForCausalLM(LlamaForCausalLM):
         return None, position_ids, attention_mask, past_key_values, new_input_embeds, new_labels
 
 
-AutoModel.register(VLMConfig, VLMForCausalLM)
+AutoModel.register(VLMConfig, VLMForCausalLM)