Upload files with `vila-upload`.

Upload config.json
Upload processing_vila.py
Upload processor_config.json
Upload configuration_vila.py
Upload tokenizer_config.json
Upload generation_config.json
Upload chat_template.jinja
Upload modeling_vila.py

chat_template.jinja

@@ -0,0 +1 @@
+{% 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 %}{{ '<image>' }}{% elif content['type'] == 'video' or 'video' in content %}{{ '<video>' }}{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}{{ '<|im_end|>\n' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}

config.json

@@ -10,7 +10,6 @@
     "AutoModelForVision2Seq": "modeling_vila.VILAForConditionalGeneration"
   },
   "hidden_size": 5120,
-  "image_end_token_id": 198,
   "image_token_id": 151666,
   "mm_hidden_size": 1152,
   "mm_projector_type": "mlp_downsample_3x3_fix",
@@ -45,7 +44,7 @@
     "vocab_size": 151670
   },
   "torch_dtype": "bfloat16",
-  "transformers_version": "4.51.3",
+  "transformers_version": "4.52.3",
   "video_token_id": 151670,
   "vision_config": {
     "architectures": [

configuration_vila.py

@@ -21,7 +21,6 @@ class VILAConfig(PretrainedConfig):
     # Model configuration.
     hidden_size: int
     image_token_id: int
-    image_end_token_id: int
     mm_hidden_size: int
     mm_projector_type: str
     mm_vision_select_feature: str
@@ -30,17 +29,16 @@ class VILAConfig(PretrainedConfig):
 
     def __init__(
         self,
-        *,
         text_config: Optional[Dict[str, Any]] = None,
         vision_config: Optional[Dict[str, Any]] = None,
-        hidden_size: Optional[int] = None,
-        image_token_id: Optional[int] = None,
-        image_end_token_id: Optional[int] = None,
-        mm_hidden_size: Optional[int] = None,
-        mm_projector_type: Optional[str] = None,
-        mm_vision_select_feature: Optional[str] = None,
-        mm_vision_select_layer: Optional[int] = None,
-        video_token_id: Optional[int] = None,
+        *,
+        hidden_size: int = 1536,
+        image_token_id: int = 151649,
+        mm_hidden_size: int = 1152,
+        mm_projector_type: str = "mlp_downsample_3x3_fix",
+        mm_vision_select_feature: str = "cls_patch",
+        mm_vision_select_layer: int = -2,
+        video_token_id: int = 151650,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -48,14 +46,10 @@ class VILAConfig(PretrainedConfig):
         self.text_config = Qwen2Config(**text_config) if text_config else Qwen2Config()
         self.vision_config = SiglipVisionConfig(**vision_config) if vision_config else SiglipVisionConfig()
 
-        # By default, we use settings from NVILA-Lite.
-        self.hidden_size = hidden_size if hidden_size is not None else 1536
-        self.image_token_id = image_token_id if image_token_id is not None else 151649
-        self.image_end_token_id = image_end_token_id if image_end_token_id is not None else 198  # "\n"
-        self.mm_hidden_size = mm_hidden_size if mm_hidden_size is not None else 1152
-        self.mm_projector_type = mm_projector_type if mm_projector_type is not None else "mlp_downsample_3x3_fix"
-        self.mm_vision_select_feature = (
-            mm_vision_select_feature if mm_vision_select_feature is not None else "cls_patch"
-        )
-        self.mm_vision_select_layer = mm_vision_select_layer if mm_vision_select_layer is not None else -2
-        self.video_token_id = video_token_id if video_token_id is not None else 151650
+        self.hidden_size = hidden_size
+        self.image_token_id = image_token_id
+        self.mm_hidden_size = mm_hidden_size
+        self.mm_projector_type = mm_projector_type
+        self.mm_vision_select_feature = mm_vision_select_feature
+        self.mm_vision_select_layer = mm_vision_select_layer
+        self.video_token_id = video_token_id

generation_config.json

@@ -3,5 +3,5 @@
   "bos_token_id": 151643,
   "eos_token_id": 151645,
   "pad_token_id": 151643,
-  "transformers_version": "4.51.3"
+  "transformers_version": "4.52.3"
 }

modeling_vila.py

@@ -1,8 +1,10 @@
-from typing import List, Optional, Type
+from typing import List, Optional, Type, Union
 
 import torch
 import torch.nn as nn
-from torch import Tensor
+import torch.nn.functional as F
+from torch import LongTensor, Tensor
+from transformers.cache_utils import Cache
 from transformers.configuration_utils import PretrainedConfig
 from transformers.generation.utils import GenerationMixin
 from transformers.modeling_outputs import BaseModelOutputWithPooling, CausalLMOutputWithPast
@@ -13,68 +15,34 @@ from transformers.models.siglip.modeling_siglip import SiglipVisionModel
 from .configuration_vila import VILAConfig
 
 
-class DownSampleBlock(nn.Module):
-    @staticmethod
-    def flat_square(x: Tensor) -> Tensor:
-        n, w, h, c = x.size()
-        if w % 2 == 1:
-            x = torch.concat([x, torch.zeros((n, 1, h, c), device=x.device, dtype=x.dtype)], dim=1).contiguous()
-            n, w, h, c = x.size()
-        if h % 2 == 1:
-            x = torch.concat([x, torch.zeros((n, w, 1, c), device=x.device, dtype=x.dtype)], dim=2).contiguous()
-            n, w, h, c = x.size()
-        x = x.contiguous()
-        x = x.view(n, w, int(h / 2), int(c * 2))
-        x = x.permute(0, 2, 1, 3).contiguous()
-        x = x.view(n, int(h / 2), int(w / 2), int(c * 4))
-        x = x.permute(0, 2, 1, 3).contiguous()
-        return x
-
+class DownSample3x3BlockFix(nn.Module):
     def forward(self, x: Tensor) -> Tensor:
-        vit_embeds = x
-        h = w = int(vit_embeds.shape[1] ** 0.5)
-        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
-        vit_embeds = self.flat_square(vit_embeds)
-        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
-        return vit_embeds
+        """
+        Args:
+            x: The input tensor of shape (batch_size, sequence_length, mm_hidden_size).
 
+        Returns:
+            The output tensor of shape (batch_size, image_pad_len, mm_hidden_size * 9).
+        """
 
-class DownSample3x3BlockFix(nn.Module):
-    @staticmethod
-    def flat_square_3x3(x: Tensor) -> Tensor:
-        n, w, h, c = x.size()
-        if w % 3 != 0:
-            x = torch.concat(
-                [
-                    x,
-                    torch.zeros((n, 3 - (w % 3), h, c), device=x.device, dtype=x.dtype),
-                ],
-                dim=1,
-            ).contiguous()
-            n, w, h, c = x.size()
-        x = x.contiguous()
-        if h % 3 != 0:
-            x = torch.concat(
-                [
-                    x,
-                    torch.zeros((n, w, 3 - (h % 3), c), device=x.device, dtype=x.dtype),
-                ],
-                dim=2,
-            ).contiguous()
-            n, w, h, c = x.size()
-        x = x.view(n, w, int(h / 3), int(c * 3))
-        x = x.permute(0, 2, 1, 3).contiguous()
-        x = x.view(n, int(h / 3), int(w / 3), int(c * 9))
-        x = x.permute(0, 2, 1, 3).contiguous()
-        return x
+        batch_size, sequence_length, hidden_size = x.shape
 
-    def forward(self, x: Tensor) -> Tensor:
-        vit_embeds = x
-        h = w = int(vit_embeds.shape[1] ** 0.5)
-        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
-        vit_embeds = self.flat_square_3x3(vit_embeds)
-        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
-        return vit_embeds
+        feat_size = int(sequence_length**0.5)
+        if feat_size**2 != sequence_length:
+            raise ValueError(f"Cannot take square root: sequence_length {sequence_length} is not a perfect square")
+
+        features = x.reshape(batch_size, feat_size, feat_size, hidden_size)
+
+        pad_after = (3 - feat_size % 3) % 3
+        if pad_after > 0:
+            features = F.pad(features, (0, 0, 0, pad_after, 0, pad_after))
+            feat_size = feat_size + pad_after
+
+        features = features.reshape(batch_size, feat_size // 3, 3, feat_size // 3, 3, hidden_size)
+        features = features.permute(0, 1, 3, 2, 4, 5).contiguous()
+        features = features.reshape(batch_size, -1, 9 * hidden_size)
+
+        return features
 
 
 class MultimodalProjector(nn.Module):
@@ -88,37 +56,24 @@ class MultimodalProjector(nn.Module):
     ):
         super().__init__(*args, **kwargs)
 
-        match config.mm_projector_type:
-            case "linear":
-                self.layers = nn.Sequential(
-                    nn.Linear(config.vision_config.hidden_size, config.hidden_size),
-                )
-            case "mlp_downsample":
-                self.layers = nn.Sequential(
-                    DownSampleBlock(),
-                    nn.LayerNorm(config.mm_hidden_size * 4),
-                    nn.Linear(config.mm_hidden_size * 4, config.hidden_size),
-                    nn.GELU(),
-                    nn.Linear(config.hidden_size, config.hidden_size),
-                )
-            case "mlp_downsample_3x3_fix":
-                self.layers = nn.Sequential(
-                    DownSample3x3BlockFix(),
-                    nn.LayerNorm(config.mm_hidden_size * 9),
-                    nn.Linear(
-                        config.mm_hidden_size * 9,
-                        config.mm_hidden_size * 3,
-                    ),
-                    nn.GELU(),
-                    nn.LayerNorm(config.vision_config.hidden_size * 3),
-                    nn.Linear(config.vision_config.hidden_size * 3, config.hidden_size),
-                    nn.GELU(),
-                    nn.Linear(config.hidden_size, config.hidden_size),
-                )
-            case _:
-                raise NotImplementedError(f"mm_projector_type={config.mm_projector_type} not implemented.")
-
-        self.layers.to(dtype=config.torch_dtype)
+        if config.mm_projector_type == "mlp_downsample_3x3_fix":
+            self.layers = nn.Sequential(
+                DownSample3x3BlockFix(),
+                nn.LayerNorm(config.mm_hidden_size * 9),
+                nn.Linear(
+                    config.mm_hidden_size * 9,
+                    config.mm_hidden_size * 3,
+                ),
+                nn.GELU(),
+                nn.LayerNorm(config.vision_config.hidden_size * 3),
+                nn.Linear(config.vision_config.hidden_size * 3, config.hidden_size),
+                nn.GELU(),
+                nn.Linear(config.hidden_size, config.hidden_size),
+            )
+        else:
+            raise NotImplementedError(f"Unsupported mm_projector_type: {config.mm_projector_type}")
+
+        self.layers.type(config.torch_dtype)
 
     @property
     def device(self) -> torch.device:
@@ -129,7 +84,15 @@ class MultimodalProjector(nn.Module):
         return next(self.parameters()).dtype
 
     def forward(self, x: Tensor) -> Tensor:
-        return self.layers(x)
+        """
+        Args:
+            x: The input tensor of shape (batch_size, sequence_length, mm_hidden_size).
+
+        Returns:
+            The output tensor of shape (batch_size, image_pad_len, hidden_size).
+        """
+
+        return self.layers(x.to(device=self.device, dtype=self.dtype))
 
 
 class VILAForConditionalGeneration(PreTrainedModel, GenerationMixin):
@@ -156,9 +119,9 @@ class VILAForConditionalGeneration(PreTrainedModel, GenerationMixin):
     ):
         super().__init__(config, *args, **kwargs)
 
-        self.llm = Qwen2ForCausalLM(config.text_config, *args, **kwargs)
+        self.llm = Qwen2ForCausalLM._from_config(config.text_config, *args, **kwargs)
         self.mm_projector = MultimodalProjector(config)
-        self.vision_tower = SiglipVisionModel(config.vision_config, *args, **kwargs)
+        self.vision_tower = SiglipVisionModel._from_config(config.vision_config, *args, **kwargs)
 
         self.post_init()
 
@@ -168,36 +131,29 @@ class VILAForConditionalGeneration(PreTrainedModel, GenerationMixin):
         attention_mask: Optional[Tensor] = None,
         input_ids: Optional[Tensor] = None,
         inputs_embeds: Optional[Tensor] = None,
+        past_key_values: Optional[Cache] = None,
         pixel_values: Optional[Tensor] = None,
+        position_ids: Optional[LongTensor] = None,
+        logits_to_keep: Union[int, Tensor] = 0,
         **kwargs,
     ) -> CausalLMOutputWithPast:
-        # Vision info is only used for prefilling.
-        if kwargs.get("past_key_values", None) is not None:
-            pixel_values = None
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds.")
 
-        inputs_embeds = inputs_embeds.to(dtype=self.dtype) if inputs_embeds is not None else None
-        pixel_values = pixel_values.to(dtype=self.dtype) if pixel_values is not None else None
-
-        if inputs_embeds is None:
-            assert input_ids is not None
-
-            inputs_embeds = self._embed(input_ids, pixel_values)
-        else:
-            assert input_ids is None
-            assert pixel_values is None
+        if past_key_values is None:  # Prefill
+            if input_ids is not None:
+                inputs_embeds = self._embed(input_ids, pixel_values)
+                input_ids = None
 
         outputs = self.llm.__call__(
-            inputs_embeds=inputs_embeds.to(
-                device=self.llm.device,
-                dtype=self.llm.dtype,
-            ),
-            attention_mask=(
-                attention_mask.to(
-                    device=self.llm.device,
-                )
-                if attention_mask is not None
-                else None
+            attention_mask=(attention_mask.to(device=self.llm.device) if attention_mask is not None else None),
+            input_ids=(input_ids.to(device=self.llm.device) if input_ids is not None else None),
+            inputs_embeds=(
+                inputs_embeds.to(device=self.llm.device, dtype=self.llm.dtype) if inputs_embeds is not None else None
             ),
+            past_key_values=past_key_values,
+            position_ids=(position_ids.to(device=self.llm.device) if position_ids is not None else None),
+            logits_to_keep=logits_to_keep,
             **kwargs,
         )
 
@@ -221,8 +177,6 @@ class VILAForConditionalGeneration(PreTrainedModel, GenerationMixin):
             The embedding of the input ids and pixel values.
         """
 
-        # Video tokens should be removed during preprocessing, so there must not be any video
-        # tokens in the input ids.
         if torch.any(input_ids == self.config.video_token_id):
             raise ValueError("Video token ids should not be present in the input ids.")
 
@@ -233,56 +187,35 @@ class VILAForConditionalGeneration(PreTrainedModel, GenerationMixin):
         if pixel_values is None:
             return text_embedding
 
-        image_features: BaseModelOutputWithPooling = self.vision_tower.__call__(
-            pixel_values.to(
-                device=self.vision_tower.device,
-                dtype=self.vision_tower.dtype,
-            ),
+        vision_tower_output: BaseModelOutputWithPooling = self.vision_tower.__call__(
+            pixel_values.to(device=self.vision_tower.device, dtype=self.vision_tower.dtype),
             output_hidden_states=True,
         )
-        assert image_features.hidden_states is not None
-
-        # Select image feature.
-        selected_layer_output = image_features.hidden_states[self.config.mm_vision_select_layer]
-        match self.config.mm_vision_select_feature:
-            case "cls_patch":
-                selected_feature = selected_layer_output
-            case _:
-                raise NotImplementedError(
-                    f"mm_vision_select_feature={self.config.mm_vision_select_feature} not implemented."
-                )
 
-        # TODO: Support dynamic_s2.
+        mm_projector_input = self._vision_tower_output_to_mm_projector_input(vision_tower_output)
 
         image_embedding: Tensor = self.mm_projector.__call__(
-            selected_feature.to(
-                device=self.mm_projector.device,
-                dtype=self.mm_projector.dtype,
-            )
+            mm_projector_input.to(device=self.mm_projector.device, dtype=self.mm_projector.dtype)
         )
 
-        # Append image end token to every image embedding.
-        image_end_token_embedding: Tensor = self.llm.get_input_embeddings().__call__(
-            torch.tensor(
-                self.config.image_end_token_id,
-                device=text_embedding.device,
-                dtype=torch.long,
-            ).view(1, -1)
-        )  # Shape: (1, 1, dim_feature)
-        image_end_token_embedding = image_end_token_embedding.expand(
-            image_embedding.shape[0], 1, -1
-        )  # Shape: (n_images, 1, dim_feature)
-        image_embedding = torch.concat(
-            [
-                image_embedding.to(device=text_embedding.device),
-                image_end_token_embedding,
-            ],
-            dim=1,
+        image_embedding = image_embedding.reshape(-1, image_embedding.shape[-1])
+
+        text_embedding.masked_scatter_(
+            image_token_mask.to(device=text_embedding.device, dtype=torch.bool).unsqueeze(-1),
+            image_embedding.to(device=text_embedding.device, dtype=text_embedding.dtype).flatten(),
         )
 
-        n_images, n_feature, dim_feature = image_embedding.shape
-        image_embedding = image_embedding.view(n_images * n_feature, dim_feature)
+        return text_embedding
+
+    def _vision_tower_output_to_mm_projector_input(
+        self,
+        vision_tower_output: BaseModelOutputWithPooling,
+    ) -> Tensor:
+        assert vision_tower_output.hidden_states is not None
 
-        text_embedding[image_token_mask.to(device=text_embedding.device)] = image_embedding
+        selected_layer_hidden_states = vision_tower_output.hidden_states[self.config.mm_vision_select_layer]
 
-        return text_embedding
+        if self.config.mm_vision_select_feature == "cls_patch":
+            return selected_layer_hidden_states
+        else:
+            raise NotImplementedError(f"Unsupported mm_vision_select_feature: {self.config.mm_vision_select_feature}")

processing_vila.py

@@ -1,19 +1,22 @@
+import uuid
 from typing import List, Optional, Tuple, cast
 
 import transformers.image_transforms as image_transforms
 import transformers.image_utils as image_utils
 import transformers.utils.logging
+import transformers.video_utils as video_utils
 from PIL.Image import Image
 from torch import Tensor
 from transformers.feature_extraction_utils import BatchFeature
 from transformers.image_processing_utils import BaseImageProcessor
 from transformers.image_processing_utils_fast import BaseImageProcessorFast
-from transformers.image_utils import ImageInput, VideoInput
+from transformers.image_utils import ImageInput
 from transformers.models.siglip.image_processing_siglip import SiglipImageProcessor
 from transformers.models.siglip.image_processing_siglip_fast import SiglipImageProcessorFast
-from transformers.processing_utils import ImagesKwargs, ProcessingKwargs, ProcessorMixin, Unpack
+from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack
 from transformers.tokenization_utils import PreTrainedTokenizer
-from transformers.tokenization_utils_base import PreTrainedTokenizerBase, TextInput
+from transformers.tokenization_utils_base import BatchEncoding, PreTrainedTokenizerBase, TextInput
+from transformers.video_utils import VideoInput
 
 logger = transformers.utils.logging.get_logger(__name__)
 
@@ -41,6 +44,7 @@ class VILAProcessor(ProcessorMixin):
         "image_pad_len",
         "max_tiles",
         "min_tiles",
+        "video_max_tiles",
     ]
 
     # Attributes.
@@ -51,15 +55,17 @@ class VILAProcessor(ProcessorMixin):
     image_pad_len: int
     max_tiles: int
     min_tiles: int
+    video_max_tiles: int
 
     def __init__(
         self,
         image_processor: BaseImageProcessor,
         tokenizer: PreTrainedTokenizer,
         *,
-        image_pad_len: Optional[int] = None,
-        max_tiles: Optional[int] = None,
-        min_tiles: Optional[int] = None,
+        image_pad_len: int = 121,
+        max_tiles: int = 12,
+        min_tiles: int = 1,
+        video_max_tiles: int = 1,
         **kwargs,
     ):
         super().__init__(
@@ -68,17 +74,17 @@ class VILAProcessor(ProcessorMixin):
             **kwargs,
         )
 
-        self.image_pad_len = image_pad_len if image_pad_len is not None else 122
-        self.max_tiles = max_tiles if max_tiles is not None else 12
-        self.min_tiles = min_tiles if min_tiles is not None else 1
+        self.image_pad_len = image_pad_len
+        self.max_tiles = max_tiles
+        self.min_tiles = min_tiles
+        self.video_max_tiles = video_max_tiles
 
     def __call__(
         self,
         text: TextInput | List[TextInput],
         images: Optional[ImageInput] = None,
         videos: Optional[VideoInput] = None,
-        audio: None = None,
-        **kwargs: Unpack[VILAProcessorProcessingKwargs],
+        **kwargs: Unpack[ProcessingKwargs],
     ) -> VILAProcessorOutput:
         """Preprocesses inputs for VILA.
 
@@ -86,7 +92,6 @@ class VILAProcessor(ProcessorMixin):
             text: The text to be processed.
             images: The images to be processed.
             videos: The videos to be processed.
-            audio: Not available.
             **kwargs: Additional arguments for processing.
 
         Returns:
@@ -99,39 +104,33 @@ class VILAProcessor(ProcessorMixin):
             **kwargs,
         )
 
-        text, images, videos = self._prepare_inputs(
+        normalized_text, normalized_images, normalized_videos = self._normalize_inputs(
             text=text,
             images=images,
             videos=videos,
         )
 
-        # Process videos.
-        text, images, video_flags = self._treat_videos_as_image_seqs(
-            text=text,
-            images=images,
-            videos=videos,
+        preprocessed_text, preprocessed_media_tiles = self._preprocess_inputs(
+            text=normalized_text,
+            images=normalized_images,
+            videos=normalized_videos,
         )
 
-        # Process images.
-        image_inputs, num_cropped_images = self._process_images(
-            images=images,
-            **merged_kwargs["images_kwargs"],
-        )
-
-        # Process text.
-        text = self._pad_image_tokens_by_num_crops(
-            text,
-            num_cropped_images=num_cropped_images,
-            video_flags=video_flags,
-        )
-
-        text = self._pad_image_tokens_by_num_embeddings(text)
-
         text_inputs = self.tokenizer.__call__(
-            text,
+            preprocessed_text,
             **merged_kwargs["text_kwargs"],
         )
 
+        if len(preprocessed_media_tiles) > 0:
+            image_inputs = self.image_processor.__call__(
+                preprocessed_media_tiles,
+                **merged_kwargs["images_kwargs"],
+            )
+        else:
+            image_inputs = BatchFeature()
+
+        text_inputs = self._replace_image_tile_suffix(text_inputs)
+
         return VILAProcessorOutput(
             data={
                 **text_inputs,
@@ -139,99 +138,144 @@ class VILAProcessor(ProcessorMixin):
             }
         )
 
-    def _crop_image(
-        self,
-        image: Image,
-    ) -> List[Image]:
-        """Crops the image into multiple tiles.
+    def _find_media_token_order(self, text: List[str]) -> List[str]:
+        """Finds the order of media tokens in the text.
 
         Args:
-            image: The image to be cropped.
+            text: The text to be processed.
 
         Returns:
-            The cropped images.
+            The order of media tokens in the text. Each item is either an image token or a video
+            token.
         """
 
-        # TODO: Support more image processors.
-        if not isinstance(self.image_processor, (SiglipImageProcessor, SiglipImageProcessorFast)):
-            raise NotImplementedError
+        image_token = cast(str, self.tokenizer.image_token)
+        video_token = cast(str, self.tokenizer.video_token)
 
-        assert self.image_processor.size["height"] == self.image_processor.size["width"]
-        cropped_size = self.image_processor.size["height"]
+        return_order: List[str] = []
 
-        cropped_images: List[Image] = dynamic_preprocess(
-            image,
-            min_num=self.min_tiles,
-            max_num=self.max_tiles,
-            image_size=cropped_size,
-        )
+        for text_item in text:
+            while image_token in text_item or video_token in text_item:
+                image_pos = text_item.find(image_token)
+                video_pos = text_item.find(video_token)
+
+                if image_pos == -1 and video_pos == -1:
+                    # If no media token found, move to the next text item.
+                    break
+
+                elif image_pos == -1:
+                    # If only video token found, add it to the return order.
+                    return_order.append(video_token)
+                    text_item = text_item[video_pos + len(video_token) :]
 
-        return cropped_images
+                elif video_pos == -1:
+                    # If only image token found, add it to the return order.
+                    return_order.append(image_token)
+                    text_item = text_item[image_pos + len(image_token) :]
 
-    def _pad_image_tokens_by_num_crops(
+                else:
+                    # If both tokens found, choose the one that appears first.
+                    if image_pos < video_pos:
+                        return_order.append(image_token)
+                        text_item = text_item[image_pos + len(image_token) :]
+                    else:
+                        return_order.append(video_token)
+                        text_item = text_item[video_pos + len(video_token) :]
+
+        return return_order
+
+    def _generate_image_token_placeholder(self, text: List[str]) -> str:
+        while True:
+            placeholder = f"<|image_placeholder_{str(uuid.uuid4())}|>"
+            if all(placeholder not in text_item for text_item in text):
+                return placeholder
+
+    def _merge_media_tiles(
         self,
-        text: List[str],
-        *,
-        num_cropped_images: List[int],
-        video_flags: List[bool],
-    ) -> List[str]:
-        """Pads each \\<image> to num_cropped_images of "\\<image>\\n" for images and "\\<video>" for videos.
+        image_tiles: List[List[Image]],
+        video_tiles: List[List[List[Image]]],
+        media_token_order: List[str],
+    ) -> List[Image]:
+        """Merges the media tiles by the media token order.
 
         Args:
-            text: The text to be padded.
-            num_cropped_images: The number of cropped images for each image token.
-            video_flags: A list of flags indicating whether the num_cropped_images item is a video.
+            image_tiles: The image tiles.
+            video_tiles: The video tiles.
+            media_token_order: The order of media tokens in the text.
 
         Returns:
-            The padded text.
+            The merged media tiles.
         """
 
-        assert len(num_cropped_images) == len(
-            video_flags
-        ), "num_cropped_images and video_flags must have the same length."
+        image_token = cast(str, self.tokenizer.image_token)
+        video_token = cast(str, self.tokenizer.video_token)
 
-        image_token: str = cast(str, self.tokenizer.image_token)
+        image_tiles_idx = 0
+        video_tiles_idx = 0
 
-        return_text: List[str] = []
+        return_tiles: List[Image] = []
 
-        for text_item in text:
-            return_text_item: str = ""
+        for media_token in media_token_order:
+            if media_token == image_token:
+                return_tiles.extend(image_tiles[image_tiles_idx])
+                image_tiles_idx += 1
+            elif media_token == video_token:
+                for video_tile in video_tiles[video_tiles_idx]:
+                    return_tiles.extend(video_tile)
+                video_tiles_idx += 1
+            else:
+                raise ValueError(f"Invalid media token: {media_token}")
 
-            # Repeatedly find image_token in the text.
-            while image_token in text_item:
-                image_pos = text_item.find(image_token)
+        return return_tiles
 
-                if image_pos != -1 and len(num_cropped_images) > 0:
-                    num_crops = num_cropped_images.pop(0)
-                    video_flag = video_flags.pop(0)
+    def _normalize_inputs(
+        self,
+        text: TextInput | List[TextInput],
+        images: Optional[ImageInput],
+        videos: Optional[VideoInput],
+    ) -> Tuple[List[str], List[Image], List[List[Image]]]:
+        """Normalizes text, image, and video inputs for processing.
 
-                    return_text_item += (
-                        text_item[:image_pos] + (image_token if video_flag else (image_token + "\n")) * num_crops
-                    )
-                    text_item = text_item[image_pos + len(image_token) :]
+        This method converts various input formats into standardized lists of PIL images
+        and text strings that can be processed by the model.
 
-                else:
-                    break
+        Args:
+            text: The original input text.
+            images: The original input images.
+            videos: The original input videos.
 
-            # Must place outside the while loop.
-            if image_token in text_item:
-                raise ValueError("Too many image tokens in the text.")
+        Returns:
+            The text as a list of strings.
+            The images as a list of PIL images.
+            The videos as a list of lists of PIL images.
+        """
 
-            return_text_item += text_item
-            text_item = ""
+        prepared_text = text if isinstance(text, list) else [text]
 
-            return_text.append(return_text_item)
+        if images is not None:
+            image_list = cast(List, image_utils.make_flat_list_of_images(images))
+            prepared_images = [cast(Image, image_transforms.to_pil_image(image)) for image in image_list]
+        else:
+            prepared_images = []
 
-        if len(num_cropped_images) != 0:
-            raise ValueError("Too many images provided.")
+        if videos is not None:
+            video_list = cast(List[List], video_utils.make_batched_videos(videos))
+            prepared_videos = [
+                [cast(Image, image_transforms.to_pil_image(image)) for image in video] for video in video_list
+            ]
+        else:
+            prepared_videos = []
 
-        return return_text
+        return prepared_text, prepared_images, prepared_videos
 
-    def _pad_image_tokens_by_num_embeddings(
+    def _pad_image_tiles(
         self,
         text: List[str],
     ) -> List[str]:
-        """Pads each \\<image> to image_pad_len times of "\\<image>".
+        """Pads each media tile.
+
+        This will pad each <image> to (self.image_pad_len + 1) times. The additional one padding is
+        for the \\n token suffix.
 
         Args:
             text: The text to be padded.
@@ -240,147 +284,189 @@ class VILAProcessor(ProcessorMixin):
             The padded text.
         """
 
-        return [
-            text_item.replace(
-                cast(str, self.tokenizer.image_token), cast(str, self.tokenizer.image_token) * self.image_pad_len
-            )
-            for text_item in text
-        ]
-
-    @staticmethod
-    def _prepare_inputs(
-        text: TextInput | List[TextInput],
-        images: Optional[ImageInput],
-        videos: Optional[VideoInput],
-    ) -> Tuple[List[str], List[Image], List[List[Image]]]:
-        # Prepare text.
-        text = text if isinstance(text, list) else [text]
-
-        # Prepare images.
-        if images is not None:
-            image_list = cast(List, image_utils.make_flat_list_of_images(images))
-            images = [image_transforms.to_pil_image(image) for image in image_list]
-        else:
-            images = cast(List[Image], [])
-
-        # Prepare videos.
-        if videos is not None:
-            video_list = cast(List[List], image_utils.make_batched_videos(videos))
-            videos = [[image_transforms.to_pil_image(image) for image in video] for video in video_list]
-        else:
-            videos = cast(List[List[Image]], [])
+        image_token = cast(str, self.tokenizer.image_token)
 
-        return text, images, videos
+        return [text_item.replace(image_token, image_token * (self.image_pad_len + 1)) for text_item in text]
 
-    def _process_images(
+    def _preprocess_inputs(
         self,
+        text: List[str],
         images: List[Image],
-        **kwargs: Unpack[ImagesKwargs],
-    ) -> Tuple[BatchFeature, List[int]]:
-        cropped_images: List[Image] = []
-        num_cropped_images: List[int] = []
+        videos: List[List[Image]],
+    ) -> Tuple[List[str], List[Image]]:
+        """Preprocesses the input data for the VILA model.
 
-        for image in images:
-            single_cropped_images = self._crop_image(image)
+        This method takes a list of texts, images, and videos, and prepares them for the model.
+        It handles the interleaving of text and media, and returns the processed text and a
+        list of media tiles (images or video frames).
 
-            cropped_images.extend(single_cropped_images)
-            num_cropped_images.append(len(single_cropped_images))
+        Args:
+            text: The input text.
+            images: The input images.
+            videos: The input videos.
 
-        if len(cropped_images) == 0:
-            # The image processor may not properly handle empty image lists.
-            # This is a workaround to avoid errors.
-            return BatchFeature(), num_cropped_images
+        Returns:
+            The text ready to be tokenized.
+            The media tiles ready to be processed.
+        """
 
-        image_inputs = self.image_processor.__call__(
-            cropped_images,
-            **kwargs,
+        media_token_order = self._find_media_token_order(text)
+
+        image_token_placeholder = self._generate_image_token_placeholder(text)
+
+        preprocessed_text = text
+        preprocessed_text, preprocessed_image_tiles = self._preprocess_images(
+            preprocessed_text,
+            images,
+            image_token_placeholder=image_token_placeholder,
+        )
+        preprocessed_text, preprocessed_video_tiles = self._preprocess_videos(
+            preprocessed_text,
+            videos,
+            image_token_placeholder=image_token_placeholder,
         )
 
-        return image_inputs, num_cropped_images
+        # Convert back to the original image token.
+        image_token = cast(str, self.tokenizer.image_token)
+        preprocessed_text = [text_item.replace(image_token_placeholder, image_token) for text_item in preprocessed_text]
 
-    def _treat_videos_as_image_seqs(
-        self, text: List[str], images: List[Image], videos: List[List[Image]]
-    ) -> Tuple[List[str], List[Image], List[bool]]:
-        """Treats videos as image sequences.
+        preprocessed_text = self._pad_image_tiles(preprocessed_text)
 
-        This method will replace all video tokens in the text with #frame image tokens,
-        and insert the corresponding images into the images list.
+        preprocessed_media_tiles = self._merge_media_tiles(
+            preprocessed_image_tiles,
+            preprocessed_video_tiles,
+            media_token_order,
+        )
 
-        Args:
-            text: The text to be processed.
-            images: The images to be processed.
-            videos: The videos to be processed.
+        return preprocessed_text, preprocessed_media_tiles
 
-        Returns:
-            The processed text and images, and a list of flags indicating whether the images are from videos.
-        """
+    def _preprocess_images(
+        self,
+        text: List[str],
+        images: List[Image],
+        *,
+        image_token_placeholder: str,
+    ) -> Tuple[List[str], List[List[Image]]]:
+        single_image_token_placeholder = self._generate_image_token_placeholder(text)
 
-        image_token = cast(str, self.tokenizer.image_token)
-        video_token = cast(str, self.tokenizer.video_token)
+        preprocessed_text = text
+        preprocessed_image_tiles: List[List[Image]] = []
 
-        return_text: List[str] = []
-        return_images: List[Image] = []
-        return_video_flags: List[bool] = []
+        for image in images:
+            preprocessed_text, preprocessed_single_image_tiles = self._preprocess_single_image(
+                text,
+                image,
+                image_token_placeholder=single_image_token_placeholder,
+                is_video_frame=False,
+                use_dynamic_preprocess=(len(images) == 1),
+            )
 
-        for text_item in text:
-            return_text_item: str = ""
+            preprocessed_text = [
+                text_item.replace(
+                    single_image_token_placeholder,
+                    (image_token_placeholder + "\n") if len(images) == 1 else image_token_placeholder,
+                )
+                for text_item in preprocessed_text
+            ]
 
-            # Repeatedly find image_token or video_token in the text.
-            while image_token in text_item or video_token in text_item:
-                image_pos = text_item.find(image_token)
-                video_pos = text_item.find(video_token)
+            preprocessed_image_tiles.append(preprocessed_single_image_tiles)
 
-                # If not found, set position to the end of the text.
-                if image_pos == -1:
-                    image_pos = len(text_item)
-                if video_pos == -1:
-                    video_pos = len(text_item)
+        return preprocessed_text, preprocessed_image_tiles
 
-                if image_pos != len(text_item) and len(images) > 0 and image_pos < video_pos:
-                    # Take an image and keep the image token if:
-                    #   - an image token is found, and
-                    #   - there are images left, and
-                    #   - the image token is before the first video token.
+    def _preprocess_single_image(
+        self,
+        text: List[str],
+        image: Image,
+        *,
+        image_token_placeholder: str,
+        is_video_frame: bool,
+        use_dynamic_preprocess: bool,
+    ) -> Tuple[List[str], List[Image]]:
+        assert isinstance(self.image_processor, (SiglipImageProcessor, SiglipImageProcessorFast))
+        assert self.image_processor.size["height"] == self.image_processor.size["width"]
+        cropped_size = self.image_processor.size["height"]
 
-                    image = images.pop(0)
-                    return_images.append(image)
-                    return_video_flags.append(False)
+        if use_dynamic_preprocess:
+            if is_video_frame:
+                max_num = self.video_max_tiles
+            else:
+                max_num = self.max_tiles
+        else:
+            max_num = 1
 
-                    return_text_item += text_item[: image_pos + len(image_token)]
-                    text_item = text_item[image_pos + len(image_token) :]
+        image = image.convert("RGB")
 
-                elif video_pos != len(text_item) and len(videos) > 0 and video_pos < image_pos:
-                    # Take a video and replace the video token with #frame image tokens if:
-                    #   - a video token is found, and
-                    #   - there are videos left, and
-                    #   - the video token is before the first image token.
+        cropped_images: List[Image] = dynamic_preprocess(
+            image,
+            min_num=self.min_tiles,
+            max_num=max_num,
+            image_size=cropped_size,
+        )
+
+        image_token = cast(str, self.tokenizer.image_token)
 
-                    video = videos.pop(0)
-                    return_images.extend(video)
-                    return_video_flags.extend([True] * len(video))
+        for i in range(len(text)):
+            if image_token in text[i]:
+                text[i] = text[i].replace(image_token, image_token_placeholder * len(cropped_images))
+                break
 
-                    return_text_item += text_item[:video_pos] + image_token * len(video)
-                    text_item = text_item[video_pos + len(video_token) :]
-                else:
-                    break
+        return text, cropped_images
 
-            # Must place outside the while loop.
-            if image_token in text_item:
-                raise ValueError("Too many image tokens in the text.")
-            if video_token in text_item:
-                raise ValueError("Too many video tokens in the text.")
+    def _preprocess_videos(
+        self,
+        text: List[str],
+        videos: List[List[Image]],
+        *,
+        image_token_placeholder: str,
+    ) -> Tuple[List[str], List[List[List[Image]]]]:
+        image_token = cast(str, self.tokenizer.image_token)
+        video_token = cast(str, self.tokenizer.video_token)
 
-            return_text_item += text_item
-            text_item = ""
+        processed_text = text
+        processed_video_tiles: List[List[List[Image]]] = []
 
-            return_text.append(return_text_item)
+        for video in videos:
+            # Replace the first video token with #frame image tokens.
+            for i in range(len(processed_text)):
+                if video_token in processed_text[i]:
+                    processed_text[i] = processed_text[i].replace(video_token, image_token * len(video))
+                    break
 
-        if len(images) != 0:
-            raise ValueError("Too many images provided.")
-        if len(videos) != 0:
-            raise ValueError("Too many videos provided.")
+            processed_frame_tiles: List[List[Image]] = []
+            for frame in video:
+                processed_text, processed_single_frame_tiles = self._preprocess_single_image(
+                    processed_text,
+                    frame,
+                    image_token_placeholder=image_token_placeholder,
+                    is_video_frame=True,
+                    use_dynamic_preprocess=(self.video_max_tiles > 1),
+                )
+                processed_frame_tiles.append(processed_single_frame_tiles)
+
+            processed_video_tiles.append(processed_frame_tiles)
+
+        return processed_text, processed_video_tiles
+
+    def _replace_image_tile_suffix(self, text_inputs: BatchEncoding) -> BatchEncoding:
+        lf_token_id = cast(int, self.tokenizer.encode("\n")[0])
+        image_token_id = cast(int, self.tokenizer.image_token_id)
+
+        for i in range(len(text_inputs.input_ids)):
+            input_ids = text_inputs.input_ids[i]
+
+            idx = 0
+            while idx < len(input_ids):
+                if input_ids[idx] != image_token_id:
+                    idx += 1
+                    continue
+
+                if idx + self.image_pad_len < len(input_ids):
+                    input_ids[idx + self.image_pad_len] = lf_token_id
+                    idx += self.image_pad_len + 1
+                else:
+                    break
 
-        return return_text, return_images, return_video_flags
+        return text_inputs
 
 
 def dynamic_preprocess(image: Image, min_num: int, max_num: int, image_size: int, use_thumbnail=True) -> List[Image]:

processor_config.json

@@ -2,8 +2,9 @@
   "auto_map": {
     "AutoProcessor": "processing_vila.VILAProcessor"
   },
-  "image_pad_len": 122,
+  "image_pad_len": 121,
   "max_tiles": 12,
   "min_tiles": 1,
-  "processor_class": "VILAProcessor"
+  "processor_class": "VILAProcessor",
+  "video_max_tiles": 1
 }

tokenizer_config.json

@@ -249,7 +249,6 @@
     "AutoProcessor": "processing_vila.VILAProcessor"
   },
   "bos_token": "[BOS]",
-  "chat_template": null,
   "clean_up_tokenization_spaces": false,
   "eos_token": "<|im_end|>",
   "errors": "replace",