feat: can generate now but not precise

chat_template.json

@@ -1,3 +1,3 @@
 {
-  "chat_template": "{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n{% endif %}<|im_start|>{{ message['role'] }}\n{% if message['content'] is string %}{{ message['content'] }}<|im_end|>\n{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<image>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>\n{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
+  "chat_template": "{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system\nYou are a helpful assistant<|im_end|>\n{% endif %}<|im_start|>{{ message['role'] }}\n{% if message['content'] is string %}{{ message['content'] }}<|im_end|>\n{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<image>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>\n{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
 }

modeling_vila.py

@@ -1,13 +1,14 @@
-from typing import Optional, Tuple, Type, override
+from typing import Dict, List, Optional, Tuple, Type, cast, override
 
-from configuration_vila import VILAConfig
-from torch import LongTensor, Tensor
+import torch
+from torch import Tensor
 from transformers.configuration_utils import PretrainedConfig
 from transformers.generation.utils import GenerationMixin
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
 
 from .auto_model import VILAForCausalLM
+from .configuration_vila import VILAConfig
 
 
 class VILAForConditionalGeneration(PreTrainedModel, GenerationMixin):
@@ -16,6 +17,7 @@ class VILAForConditionalGeneration(PreTrainedModel, GenerationMixin):
     is_parallelizable: bool = True
     main_input_name: str = "input_ids"
 
+    config: PretrainedConfig
     model: VILAForCausalLM
 
     def __init__(
@@ -27,18 +29,24 @@ class VILAForConditionalGeneration(PreTrainedModel, GenerationMixin):
 
     def forward(
         self,
-        input_ids: LongTensor,
-        attention_mask: Tensor,
-        pixel_values: Tensor,
         *,
-        return_dict: Optional[bool] = None,
+        attention_mask: Optional[Tensor] = None,
+        input_ids: Optional[Tensor] = None,
+        inputs_embeds: Optional[Tensor] = None,
+        pixel_values: Optional[Tensor] = None,
         **kwargs,
-    ) -> Tuple | CausalLMOutputWithPast:
+    ) -> CausalLMOutputWithPast:
+        if inputs_embeds is None:
+            assert input_ids is not None
 
-        outputs = self.model.forward(
-            input_ids=input_ids,
+            inputs_embeds, _ = self._embed(input_ids, pixel_values, attention_mask)
+        else:
+            assert input_ids is None
+            assert pixel_values is None
+
+        outputs = self.model.llm.forward(
+            inputs_embeds=inputs_embeds,
             attention_mask=attention_mask,
-            return_dict=return_dict,
             **kwargs,
         )
 
@@ -58,4 +66,62 @@ class VILAForConditionalGeneration(PreTrainedModel, GenerationMixin):
     @override
     def save_pretrained(self, *args, **kwargs) -> None:
         self.model.save_pretrained(*args, **kwargs)
-        self.model.save_pretrained(*args, **kwargs)
+
+    def _embed(
+        self,
+        input_ids: Tensor,
+        pixel_values: Optional[Tensor],
+        attention_mask: Optional[Tensor],
+    ) -> Tuple[Tensor, Tensor]:
+        """Gets the embedding of the input ids and pixel values.
+
+        Args:
+            input_ids: The input ids.
+            pixel_values: The pixel values.
+            attention_mask: The attention mask.
+
+        Returns:
+            A tuple of the embedding of the input ids and attention mask.
+        """
+
+        image_token_ids_map = cast(Dict[str, int], self.model.tokenizer.media_token_ids)
+        image_token_ids = list(image_token_ids_map.values())
+        image_token_idx = torch.isin(
+            input_ids,
+            torch.tensor(image_token_ids).to(input_ids.device),
+        )
+        image_token_count = image_token_idx.sum()
+
+        images = list(pixel_values) if pixel_values is not None else []
+
+        if image_token_count < len(images):
+            images = images[:image_token_count]
+
+        media = (
+            {
+                "image": images,
+            }
+            if image_token_count > 0
+            else {}
+        )
+        media_config = (
+            {
+                "image": {},
+            }
+            if image_token_count > 0
+            else {}
+        )
+
+        outputs = self.model._embed(
+            input_ids,
+            media,
+            media_config,
+            labels=None,
+            attention_mask=(
+                attention_mask.to(dtype=torch.bool)
+                if attention_mask is not None
+                else None
+            ),
+        )
+
+        return outputs[0], outputs[2]

processing_vila.py

@@ -15,8 +15,8 @@ from transformers.tokenization_utils import PreTrainedTokenizer
 from transformers.tokenization_utils_base import TextInput
 
 from . import mm_utils
+from .constants import DEFAULT_IMAGE_TOKEN
 
-_IMAGE_TOKEN = "<image>"
 _PLACEHOLDER_TOKEN = "<|placeholder|>"
 
 
@@ -124,20 +124,20 @@ class VILAProcessor(ProcessorMixin):
         idx_image_splice = 0
 
         for i in range(len(text)):
-            while _IMAGE_TOKEN in text[i]:
+            while DEFAULT_IMAGE_TOKEN in text[i]:
                 if idx_image_splice >= len(num_image_splices):
                     raise ValueError(
-                        f"Too many {_IMAGE_TOKEN} tokens in text. "
+                        f"Too many {DEFAULT_IMAGE_TOKEN} tokens in text. "
                         f"Expected {len(num_image_splices)} tokens, "
                         f"but found {idx_image_splice} tokens."
                     )
 
                 text[i] = text[i].replace(
-                    _IMAGE_TOKEN,
+                    DEFAULT_IMAGE_TOKEN,
                     _PLACEHOLDER_TOKEN * num_image_splices[idx_image_splice],
                 )
                 idx_image_splice += 1
-            text[i] = text[i].replace(_PLACEHOLDER_TOKEN, _IMAGE_TOKEN)
+            text[i] = text[i].replace(_PLACEHOLDER_TOKEN, f"{DEFAULT_IMAGE_TOKEN}\n")
 
         return text