Ajout du module Ovis

ovis/__init__.py

@@ -0,0 +1,3 @@
+import os
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"

ovis/model/__init__.py

@@ -0,0 +1,2 @@
+from .visual_tokenizer.clip_visual_tokenizer import ClipVisualTokenizerConfig, ClipVisualTokenizer
+from .visual_tokenizer.siglip_visual_tokenizer import SiglipVisualTokenizerConfig, SiglipVisualTokenizer

ovis/model/configuration_ovis.py

@@ -0,0 +1,41 @@
+from typing import Union, Optional
+
+from transformers import PretrainedConfig, AutoConfig
+
+
+class OvisConfig(PretrainedConfig):
+    model_type = "ovis"
+
+    def __init__(
+        self,
+        llm_config: Optional[Union[PretrainedConfig, dict]] = None,
+        visual_tokenizer_config: Optional[Union[PretrainedConfig, dict]] = None,
+        multimodal_max_length=8192,
+        hidden_size=None,
+        conversation_formatter_class=None,
+        llm_attn_implementation=None,
+        disable_tie_weight=False,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        if llm_config is not None:
+            assert isinstance(llm_config, (PretrainedConfig, dict)), \
+                f"expect `llm_config` to be instance of PretrainedConfig or dict, but got {type(llm_config)} type"
+            if not isinstance(llm_config, PretrainedConfig):
+                model_type = llm_config['model_type']
+                llm_config.pop('model_type')
+                llm_config = AutoConfig.for_model(model_type, **llm_config)
+        self.llm_config = llm_config
+        if visual_tokenizer_config is not None:
+            assert isinstance(visual_tokenizer_config, (PretrainedConfig, dict)), \
+                f"expect `visual_tokenizer_config` to be instance of PretrainedConfig or dict, but got {type(visual_tokenizer_config)} type"
+            if not isinstance(visual_tokenizer_config, PretrainedConfig):
+                model_type = visual_tokenizer_config['model_type']
+                visual_tokenizer_config.pop('model_type')
+                visual_tokenizer_config = AutoConfig.for_model(model_type, **visual_tokenizer_config)
+        self.visual_tokenizer_config = visual_tokenizer_config
+        self.multimodal_max_length = multimodal_max_length
+        self.hidden_size = hidden_size
+        self.conversation_formatter_class = conversation_formatter_class
+        self.llm_attn_implementation = llm_attn_implementation
+        self.disable_tie_weight = disable_tie_weight

ovis/model/conversation_formatter.py

@@ -0,0 +1,233 @@
+from abc import ABC, abstractmethod
+from typing import List, Dict
+
+from ovis.util.constants import IMAGE_TOKEN_ID, IGNORE_ID, IMAGE_TOKEN
+
+
+class ConversationFormatter(ABC):
+    support_tokenizer_types = None
+
+    def __init__(self, tokenizer):
+        tokenizer_type = type(tokenizer).__name__
+        assert tokenizer_type in self.support_tokenizer_types, \
+            f'Invalid tokenizer type, expected one from `{self.support_tokenizer_types}`, but got `{tokenizer_type}`'
+        self.tokenizer = tokenizer
+        self.image_token = IMAGE_TOKEN
+        self.image_token_id = IMAGE_TOKEN_ID
+        self.ignore_id = IGNORE_ID
+
+    def _tokenize_with_image_symbol(self, text):
+        text_chunks = [self.tokenizer(chunk, add_special_tokens=False).input_ids for chunk in
+                       text.split(self.image_token)]
+        token_ids = []
+        num_chuck = len(text_chunks)
+        for i, chunk in enumerate(text_chunks):
+            token_ids.extend(chunk)
+            if i < num_chuck - 1:
+                token_ids.append(self.image_token_id)
+        return token_ids
+
+    @abstractmethod
+    def format(self, conversations: List[Dict], generation_preface=None):
+        pass
+
+    @abstractmethod
+    def format_query(self, query, generation_preface=""):
+        pass
+
+
+class QwenConversationFormatter(ConversationFormatter):
+    support_tokenizer_types = ['QWenTokenizer', 'Qwen2TokenizerFast']
+
+    def __init__(self, tokenizer):
+        super().__init__(tokenizer)
+        self.from2role = {
+            "system": "<|im_start|>system\n",
+            "human": "<|im_start|>user\n",
+            "gpt": "<|im_start|>assistant\n",
+        }
+        self.gpt_token_num = None
+        self.im_end = "<|im_end|>\n"
+        self.default_system_prompt = "You are a helpful assistant."
+
+    def format(self, conversations: List[Dict], generation_preface=None):
+        if self.gpt_token_num is None:
+            self.gpt_token_num = len(self.tokenizer(self.from2role["gpt"], add_special_tokens=False).input_ids)
+
+        if conversations[0]["from"] != "system":
+            conversations.insert(0, {
+                "from": "system",
+                "value": self.default_system_prompt
+            })
+
+        if generation_preface is not None:
+            conversations.append({
+                "from": "gpt",
+                "value": generation_preface
+            })
+
+        prompt = ""
+        input_ids = []
+        labels = []
+        num_conversation = len(conversations)
+        for i, conversation in enumerate(conversations):
+            frm = conversation["from"]
+            role = self.from2role[frm]
+            message = conversation["value"]
+            text = role + message
+            if i < num_conversation - 1 or generation_preface is None:
+                text += self.im_end
+            prompt += text
+            token_ids = self._tokenize_with_image_symbol(text)
+            input_ids.extend(token_ids)
+            label_ids = [self.ignore_id] * len(token_ids)
+            if frm == "gpt" and generation_preface is None:
+                # learning `\n` following `im_end` is meaningless, so the last `\n` token is ignored in label
+                label_ids[self.gpt_token_num:-1] = token_ids[self.gpt_token_num:-1]
+            labels.extend(label_ids)
+
+        assert self._tokenize_with_image_symbol(prompt) == input_ids
+        assert len(input_ids) == len(labels)
+
+        return prompt, input_ids, labels
+
+    def format_query(self, query, generation_preface=""):
+        prompt, input_ids, _ = self.format([{
+            "from": "human",
+            "value": query
+        }], generation_preface=generation_preface)
+
+        return prompt, input_ids
+
+
+class Llama3ConversationFormatter(ConversationFormatter):
+    support_tokenizer_types = ['PreTrainedTokenizerFast']
+
+    def __init__(self, tokenizer):
+        super().__init__(tokenizer)
+        self.from2role = {
+            "system": "<|start_header_id|>system<|end_header_id|>\n\n",
+            "human": "<|start_header_id|>user<|end_header_id|>\n\n",
+            "gpt": "<|start_header_id|>assistant<|end_header_id|>\n\n",
+        }
+        self.gpt_token_num = None
+        self.im_end = "<|eot_id|>"
+        self.default_system_prompt = "You are a helpful and honest multimodal assistant."
+        self.bos_token = "<|begin_of_text|>"
+        self.bos_token_ids = None
+
+    def format(self, conversations: List[Dict], generation_preface=None):
+        if self.gpt_token_num is None:
+            self.gpt_token_num = len(self.tokenizer(self.from2role["gpt"], add_special_tokens=False).input_ids)
+
+        if self.bos_token_ids is None:
+            self.bos_token_ids = self.tokenizer(self.bos_token, add_special_tokens=False).input_ids
+
+        if conversations[0]["from"] != "system":
+            conversations.insert(0, {
+                "from": "system",
+                "value": self.default_system_prompt
+            })
+
+        if generation_preface is not None:
+            conversations.append({
+                "from": "gpt",
+                "value": generation_preface
+            })
+
+        prompt = "" + self.bos_token
+        input_ids = [] + self.bos_token_ids
+        labels = [] + [IGNORE_ID] * len(input_ids)
+        num_conversation = len(conversations)
+        for i, conversation in enumerate(conversations):
+            frm = conversation["from"]
+            role = self.from2role[frm]
+            message = conversation["value"].strip()
+            text = role + message
+            if i < num_conversation - 1 or generation_preface is None:
+                text += self.im_end
+            prompt += text
+            token_ids = self._tokenize_with_image_symbol(text)
+            input_ids.extend(token_ids)
+            label_ids = [self.ignore_id] * len(token_ids)
+            if frm == "gpt":
+                label_ids[self.gpt_token_num:] = token_ids[self.gpt_token_num:]
+            labels.extend(label_ids)
+
+        assert self._tokenize_with_image_symbol(prompt) == input_ids
+        assert len(input_ids) == len(labels)
+
+        return prompt, input_ids, labels
+
+    def format_query(self, query, generation_preface=""):
+        prompt, input_ids, _ = self.format([{
+            "from": "human",
+            "value": query
+        }], generation_preface=generation_preface)
+
+        return prompt, input_ids
+
+
+class GemmaConversationFormatter(ConversationFormatter):
+    support_tokenizer_types = ['GemmaTokenizer', 'GemmaTokenizerFast']
+
+    def __init__(self, tokenizer):
+        super().__init__(tokenizer)
+        # Gemma does not support system prompt
+        self.from2role = {
+            "human": "<start_of_turn>user\n",
+            "gpt": "<start_of_turn>model\n",
+        }
+        self.gpt_token_num = None
+        self.im_end = "<end_of_turn>\n"
+        self.bos_token = "<bos>"
+        self.bos_token_ids = None
+
+    def format(self, conversations: List[Dict], generation_preface=None):
+        if self.gpt_token_num is None:
+            self.gpt_token_num = len(self.tokenizer(self.from2role["gpt"], add_special_tokens=False).input_ids)
+
+        if self.bos_token_ids is None:
+            self.bos_token_ids = self.tokenizer(self.bos_token, add_special_tokens=False).input_ids
+
+        if conversations[0]["from"] == "system":
+            raise ValueError("Gemma does not support system prompt")
+
+        if generation_preface is not None:
+            conversations.append({
+                "from": "gpt",
+                "value": generation_preface
+            })
+
+        prompt = "" + self.bos_token
+        input_ids = [] + self.bos_token_ids
+        labels = [] + [IGNORE_ID] * len(input_ids)
+        num_conversation = len(conversations)
+        for i, conversation in enumerate(conversations):
+            frm = conversation["from"]
+            role = self.from2role[frm]
+            message = conversation["value"].strip()
+            text = role + message
+            if i < num_conversation - 1 or generation_preface is None:
+                text += self.im_end
+            prompt += text
+            token_ids = self._tokenize_with_image_symbol(text)
+            input_ids.extend(token_ids)
+            label_ids = [self.ignore_id] * len(token_ids)
+            if frm == "gpt":
+                # learning `\n` following `im_end` is meaningless, so the last `\n` token is ignored in label
+                label_ids[self.gpt_token_num:-1] = token_ids[self.gpt_token_num:-1]
+            labels.extend(label_ids)
+
+        assert self._tokenize_with_image_symbol(prompt) == input_ids
+        assert len(input_ids) == len(labels)
+
+        return prompt, input_ids, labels
+
+    def format_query(self, query, generation_preface=""):
+        prompt, input_ids, _ = self.format([{
+            "from": "human",
+            "value": query
+        }], generation_preface=generation_preface)
+
+        return prompt, input_ids

ovis/model/modeling_ovis.py

@@ -0,0 +1,414 @@
+import logging
+import os
+from datetime import datetime
+from importlib import import_module
+from typing import List, Union, Callable, Optional, Dict
+
+import PIL.Image
+import deepspeed
+import torch
+from torch import Tensor
+from torch.nn import init
+from transformers import PreTrainedModel, AutoConfig, AutoModel, AutoTokenizer, AutoModelForCausalLM
+from transformers.cache_utils import HybridCache
+from transformers.generation.utils import GenerateOutput
+from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled, deepspeed_config
+
+from ovis.model.configuration_ovis import OvisConfig
+from ovis.model.conversation_formatter import ConversationFormatter
+from ovis.util.constants import IGNORE_ID, BEGIN_LINE, END_LINE, IMAGE_ATOM_ID, IMAGE_INDICATOR_IDS, \
+    IMAGE_TOKEN_ID
+from ovis.util.utils import rank0_print
+
+
+class VisualEmbedding(torch.nn.Embedding):
+    def forward(self, visual_tokens: Tensor) -> Tensor:
+        if visual_tokens.dtype in [torch.int8, torch.int16, torch.int32, torch.int64, torch.long]:
+            return super().forward(visual_tokens)
+        return torch.matmul(visual_tokens, self.weight)
+
+    def reset_parameters(self, mean=0., std=1.) -> None:
+        init.normal_(self.weight, mean=mean, std=std)
+        self._fill_padding_idx_with_zero()
+
+
+class OvisPreTrainedModel(PreTrainedModel):
+    config_class = OvisConfig
+    base_model_prefix = "ovis"
+
+
+class Ovis(OvisPreTrainedModel):
+
+    def __init__(self, config: OvisConfig, *inputs, **kwargs):
+        super().__init__(config, *inputs, **kwargs)
+        if kwargs.get('train_from_scratch'):
+            self.llm = kwargs['llm']
+            self.generation_config = self.llm.generation_config
+            self.config.llm_config = self.llm.config
+            self.config.hidden_size = self.llm.config.hidden_size  # for deepspeed auto configuration
+            self.text_tokenizer = kwargs['text_tokenizer']
+            self.visual_tokenizer = kwargs['visual_tokenizer']
+            self.config.visual_tokenizer_config = self.visual_tokenizer.config
+        else:
+            attn_kwargs = dict()
+            if self.config.llm_attn_implementation:
+                attn_kwargs['attn_implementation'] = self.config.llm_attn_implementation
+            self.llm = AutoModelForCausalLM.from_config(self.config.llm_config, **attn_kwargs)
+            assert self.config.hidden_size == self.llm.config.hidden_size, "hidden size mismatch"
+            self.text_tokenizer = AutoTokenizer.from_pretrained(self.config.name_or_path)
+            self.visual_tokenizer = AutoModel.from_config(self.config.visual_tokenizer_config,
+                                                          image_processor_name_or_path=self.config.name_or_path)
+
+        # initialize vte
+        if is_deepspeed_zero3_enabled():
+            with deepspeed.zero.Init(config_dict_or_path=deepspeed_config()):
+                self.vte = VisualEmbedding(self.config.visual_tokenizer_config.vocab_size, self.config.hidden_size)
+        else:
+            self.vte = VisualEmbedding(self.config.visual_tokenizer_config.vocab_size, self.config.hidden_size,
+                                       device=self.visual_tokenizer.device, dtype=self.visual_tokenizer.dtype)
+
+        def _merge_modules(modules_list: tuple):
+            merged_modules = []
+            for modules in modules_list:
+                merged_modules.extend(modules if modules else [])
+            return merged_modules
+
+        self._no_split_modules = _merge_modules((self.llm._no_split_modules, self.visual_tokenizer._no_split_modules))
+        self._skip_keys_device_placement = self.llm._skip_keys_device_placement
+        self._keep_in_fp32_modules = _merge_modules(
+            (self.llm._keep_in_fp32_modules, self.visual_tokenizer._keep_in_fp32_modules))
+        self.is_parallelizable = all((self.llm.is_parallelizable, self.visual_tokenizer.is_parallelizable))
+        self.supports_gradient_checkpointing = all(
+            (self.llm.supports_gradient_checkpointing, self.visual_tokenizer.supports_gradient_checkpointing))
+        self._supports_flash_attn_2 = all(
+            (self.llm._supports_flash_attn_2, self.visual_tokenizer._supports_flash_attn_2))
+        self._supports_sdpa = all((self.llm._supports_sdpa, self.visual_tokenizer._supports_sdpa))
+
+    def get_text_tokenizer(self):
+        return self.text_tokenizer
+
+    def get_visual_tokenizer(self):
+        return self.visual_tokenizer
+
+    def tie_weights(self):
+        if not self.config.disable_tie_weight:
+            self.get_llm().tie_weights()
+
+    def re_init_vte(self, mean, std):
+        vte = self.get_vte()
+        rank0_print(BEGIN_LINE)
+        rank0_print(f'[{datetime.now()}] Before re-initialization of vte: ')
+        with deepspeed.zero.GatheredParameters([vte.weight]):
+            rank0_print(f'vte.weight: {vte.weight}')
+        with deepspeed.zero.GatheredParameters([vte.weight], modifier_rank=0):
+            if not is_deepspeed_zero3_enabled() or deepspeed.comm.get_rank() == 0:
+                vte.reset_parameters(mean, std)
+        rank0_print(f'[{datetime.now()}] After re-initialization of vte:')
+        with deepspeed.zero.GatheredParameters([vte.weight]):
+            rank0_print(f'vte.weight: {vte.weight}')
+        rank0_print(END_LINE)
+
+    def get_monitor_tensors(self):
+        monitor_tensors = dict(
+            wte=self.get_wte().weight,
+            lm_head=self.get_lm_head().weight,
+            vte=self.get_vte().weight
+        )
+        monitor_tensors.update(
+            {f'visual_tokenizer_{k}': v for k, v in self.get_visual_tokenizer().get_monitor_tensors().items()})
+        return monitor_tensors
+
+    def get_lm_head(self):
+        return self.get_llm().get_output_embeddings()
+
+    def get_llm(self):
+        return self.llm
+
+    def get_vte(self):
+        return self.vte
+
+    def get_wte(self):
+        return self.llm.get_input_embeddings()
+
+    def get_conversation_formatter(self) -> ConversationFormatter:
+        if getattr(self, 'conversation_formatter', None) is None:
+            self.conversation_formatter = getattr(import_module(".conversation_formatter", __package__),
+                                                  self.config.conversation_formatter_class)(self.text_tokenizer)
+        return self.conversation_formatter
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        labels: Optional[torch.Tensor],
+        pixel_values: List[Optional[torch.Tensor]],
+        **kwargs
+    ):
+        assert self.training, "`forward` can only be used in training. For inference, use `generate`."
+        _, inputs_embeds, labels, attention_mask = self.merge_multimodal(
+            text_input_ids=input_ids,
+            text_attention_masks=attention_mask,
+            text_labels=labels,
+            pixel_values=pixel_values
+        )
+        return self.llm(inputs_embeds=inputs_embeds, labels=labels, attention_mask=attention_mask, **kwargs)
+
+    def merge_multimodal(
+        self,
+        text_input_ids: torch.Tensor,
+        text_attention_masks: torch.Tensor,
+        text_labels: Optional[torch.Tensor],
+        pixel_values: List[Optional[torch.Tensor]]
+    ):
+        input_device = text_input_ids.device
+        visual_vocab_szie = self.get_visual_tokenizer().config.vocab_size
+        visual_indicator_embeds = self.get_vte()(
+            torch.tensor(
+                list(range(visual_vocab_szie - 5, visual_vocab_szie)),
+                dtype=torch.long,
+                device=self.get_visual_tokenizer().device
+            )
+        ).to(device=input_device)
+
+        if self.training:
+            # When training, to be compatible with deepspeed zero, each sample has to include pixel_value tensor.
+            # For text-only sample, one can simply use a full zero tensor as pixel_value, which will be ignored
+            # (see below in this function); so, the gradient will not be affected.
+            num_images = [x.shape[0] for x in pixel_values]
+            visual_tokens = self.visual_tokenizer(torch.cat([x for x in pixel_values], dim=0))
+            visual_embeds = torch.split(self.get_vte()(visual_tokens).to(dtype=self.dtype, device=input_device),
+                                        split_size_or_sections=num_images, dim=0)
+            visual_input_ids = torch.split(torch.argmax(visual_tokens, dim=-1).to(device=input_device),
+                                           split_size_or_sections=num_images, dim=0)
+            visual_labels = [torch.full(x.shape, IGNORE_ID, dtype=torch.long, device=input_device) for x in
+                             visual_input_ids]
+        else:
+            # When inference, sample can include only text with `None` pixel_value
+            num_images = [x.shape[0] if x is not None else 0 for x in pixel_values]
+            if sum(num_images) > 0:
+                visual_tokens = self.visual_tokenizer(torch.cat([x for x in pixel_values if x is not None], dim=0))
+                visual_embeds = torch.split(self.get_vte()(visual_tokens).to(dtype=self.dtype, device=input_device),
+                                            split_size_or_sections=num_images, dim=0)
+                visual_input_ids = torch.split(torch.argmax(visual_tokens, dim=-1).to(device=input_device),
+                                               split_size_or_sections=num_images, dim=0)
+                visual_labels = [torch.full(x.shape, IGNORE_ID, dtype=torch.long, device=input_device) for x in
+                                 visual_input_ids]
+            else:
+                # just placeholders
+                visual_embeds = [None] * len(num_images)
+                visual_input_ids = [None] * len(num_images)
+                visual_labels = [None] * len(num_images)
+            # just placeholders
+            text_labels = torch.full(text_input_ids.shape, IGNORE_ID, dtype=torch.long, device=input_device)
+
+        input_embeds = []
+        attention_masks = []
+        labels = []
+        for text_input_id, text_label, text_attention_mask, visual_embed, visual_input_id, visual_label in zip(
+                text_input_ids, text_labels, text_attention_masks, visual_embeds, visual_input_ids, visual_labels
+        ):
+            placeholder_token_mask = torch.lt(text_input_id, 0)
+            text_embed = self.get_wte()(torch.masked_fill(text_input_id, placeholder_token_mask, 0))
+            for i, indicator_id in enumerate(IMAGE_INDICATOR_IDS):
+                text_embed[text_input_id == indicator_id] = visual_indicator_embeds[i]
+            image_atom_positions = torch.where(torch.eq(text_input_id, IMAGE_ATOM_ID))[0].tolist()
+            if len(image_atom_positions) > 0:
+                input_embed_parts = []
+                attention_mask_parts = []
+                label_parts = []
+                prev_image_atom_position = -1
+                for index, image_atom_position in enumerate(image_atom_positions):
+                    input_embed_parts.append(
+                        text_embed[prev_image_atom_position + 1:image_atom_position, :])
+                    label_parts.append(
+                        text_label[prev_image_atom_position + 1:image_atom_position])
+                    attention_mask_parts.append(
+                        text_attention_mask[prev_image_atom_position + 1:image_atom_position])
+                    input_embed_parts.append(visual_embed[index])
+                    attention_mask_parts.append(
+                        torch.ones_like(visual_label[index], dtype=torch.bool))
+                    label_parts.append(visual_label[index])
+                    prev_image_atom_position = image_atom_position
+                if prev_image_atom_position + 1 < text_input_id.shape[0]:
+                    input_embed_parts.append(
+                        text_embed[prev_image_atom_position + 1:, :])
+                    attention_mask_parts.append(
+                        text_attention_mask[prev_image_atom_position + 1:])
+                    label_parts.append(
+                        text_label[prev_image_atom_position + 1:])
+                input_embed = torch.cat(input_embed_parts, dim=0)
+                attention_mask = torch.cat(attention_mask_parts, dim=0)
+                label = torch.cat(label_parts, dim=0)
+            else:
+                input_embed = text_embed
+                attention_mask = text_attention_mask
+                label = text_label
+                if self.training:
+                    # Make visual_embed & visual_indicator_embeds involved in the backward graph,
+                    # to be compatible with deepspeed zero and ddp.
+                    input_embed += torch.sum(visual_embed * 0.0) + torch.sum(visual_indicator_embeds * 0.0)
+            input_embeds.append(input_embed)
+            attention_masks.append(attention_mask)
+            labels.append(label)
+
+        if self.training:  # padding to self.config.multimodal_max_length for increased training speed
+            padding_size = max(0, self.config.multimodal_max_length - len(input_embeds[0]))
+            input_embeds[0] = torch.nn.ConstantPad2d((0, 0, 0, padding_size), 0.0)(input_embeds[0])
+            attention_masks[0] = torch.nn.ConstantPad1d((0, padding_size), False)(attention_masks[0])
+            labels[0] = torch.nn.ConstantPad1d((0, padding_size), IGNORE_ID)(labels[0])
+        batch_input_embeds = torch.nn.utils.rnn.pad_sequence(input_embeds, batch_first=True, padding_value=0.0)[:,
+                             :self.config.multimodal_max_length, :]
+        batch_attention_mask = torch.nn.utils.rnn.pad_sequence(attention_masks, batch_first=True, padding_value=False)[
+                               :,
+                               :self.config.multimodal_max_length]
+        batch_labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_ID)[:,
+                       :self.config.multimodal_max_length]
+
+        return visual_input_ids, batch_input_embeds, batch_labels, batch_attention_mask
+
+    def preprocess_inputs(
+        self,
+        text_or_conversations: Union[List[Dict], str],
+        images: Optional[List[PIL.Image.Image]],
+        max_partition=9,
+        generation_preface='',
+        return_labels=False,
+        propagate_exception=True
+    ):
+        # convert text to conversations
+        if isinstance(text_or_conversations, str):
+            conversations = [{
+                "from": "human",
+                "value": text_or_conversations
+            }]
+        elif isinstance(text_or_conversations, list):
+            conversations = text_or_conversations
+        else:
+            raise ValueError(f'Invalid type of `text_or_conversations`, expected `List[Dict]` or `str`,'
+                             f' but got {type(text_or_conversations)}')
+
+        # format conversations
+        prompt, raw_input_ids, raw_labels = self.get_conversation_formatter().format(
+            conversations, generation_preface=generation_preface)
+
+        # place image placeholders
+        input_ids = []
+        labels = []
+        pixel_values = []
+        invalidate_label = False
+        image_token_indices = [i for i, v in enumerate(raw_input_ids) if v == IMAGE_TOKEN_ID]
+        last_image_token_index = -1
+        for i in range(len(image_token_indices)):
+            head = 0 if i == 0 else image_token_indices[i - 1] + 1
+            tail = image_token_indices[i]
+            last_image_token_index = tail
+            input_ids.extend(raw_input_ids[head:tail])
+            labels.extend(raw_labels[head:tail])
+            try:
+                image = images[i]
+                raw_pixel_values, image_placeholders = self.visual_tokenizer.preprocess_image(
+                    image, max_partition=max_partition)
+            except Exception as e:
+                if propagate_exception:
+                    raise e
+                logging.exception(e)
+                invalidate_label = True
+                raw_pixel_values, image_placeholders = self.visual_tokenizer.mock_input()
+            input_ids.extend(image_placeholders)
+            labels.extend([IGNORE_ID] * len(image_placeholders))
+            pixel_values.append(raw_pixel_values)
+        input_ids.extend(raw_input_ids[last_image_token_index + 1:])
+        labels.extend(raw_labels[last_image_token_index + 1:])
+
+        # return tensors
+        input_ids = torch.tensor(input_ids, dtype=torch.long)
+        labels = torch.tensor([IGNORE_ID] * len(labels) if invalidate_label else labels, dtype=torch.long)
+        pixel_values = torch.cat(pixel_values, dim=0) if len(pixel_values) > 0 else None
+
+        if return_labels:
+            return prompt, input_ids, pixel_values, labels
+        else:
+            return prompt, input_ids, pixel_values
+
+    def save_pretrained(
+        self,
+        save_directory: Union[str, os.PathLike],
+        is_main_process: bool = True,
+        state_dict: Optional[dict] = None,
+        save_function: Callable = torch.save,
+        push_to_hub: bool = False,
+        max_shard_size: Union[int, str] = "5GB",
+        safe_serialization: bool = True,
+        variant: Optional[str] = None,
+        token: Optional[Union[str, bool]] = None,
+        save_peft_format: bool = True,
+        **kwargs
+    ):
+        super().save_pretrained(save_directory,
+                                is_main_process=is_main_process,
+                                state_dict=state_dict,
+                                save_function=save_function,
+                                safe_serialization=safe_serialization)
+        self.get_text_tokenizer().save_pretrained(save_directory)
+        self.get_visual_tokenizer().get_image_processor().save_pretrained(save_directory)
+
+        # uncomment the following will additionally save a separate visual tokenizer
+        # visual_tokenizer_directory = os.path.join(save_directory, 'visual_tokenizer')
+        # self.get_visual_tokenizer().save_pretrained(visual_tokenizer_directory,
+        #                                             is_main_process=is_main_process,
+        #                                             state_dict=None,
+        #                                             save_function=save_function,
+        #                                             safe_serialization=safe_serialization)
+        # self.get_visual_tokenizer().get_image_processor().save_pretrained(visual_tokenizer_directory)
+
+    def _get_hybrid_cache_for_llm(self, max_batch_size: int, max_cache_len: int):
+        cache_cls = HybridCache
+        llm = self.get_llm()
+
+        need_new_cache = (
+            not hasattr(llm, "_cache")
+            or (not isinstance(llm._cache, cache_cls))
+            or llm._cache.max_batch_size != max_batch_size
+            or llm._cache.max_cache_len < max_cache_len
+        )
+
+        if need_new_cache:
+            if hasattr(llm.config, "_pre_quantization_dtype"):
+                cache_dtype = llm.config._pre_quantization_dtype
+            else:
+                cache_dtype = llm.dtype
+            llm._cache = cache_cls(
+                config=llm.config,
+                max_batch_size=max_batch_size,
+                max_cache_len=max_cache_len,
+                device=llm.device,
+                dtype=cache_dtype,
+            )
+        else:
+            llm._cache.reset()
+        return llm._cache
+
+    # TODO: support batch generation
+    def generate(
+        self,
+        inputs: Optional[torch.Tensor] = None,
+        **kwargs
+    ) -> Union[GenerateOutput, torch.LongTensor]:
+        assert inputs.shape[0] == 1, 'Currently, only support `batch_size=1`'
+        _, inputs_embeds, labels, attention_mask = self.merge_multimodal(
+            text_input_ids=inputs,
+            text_attention_masks=kwargs.pop('attention_mask'),
+            text_labels=None,
+            pixel_values=kwargs.pop('pixel_values')
+        )
+        if getattr(self.generation_config, 'cache_implementation') == 'hybrid':  # mainly for Gemma2
+            kwargs['past_key_values'] = self._get_hybrid_cache_for_llm(
+                getattr(kwargs, "num_beams", 1), kwargs['max_new_tokens'] + inputs_embeds.shape[-2])
+            self.get_llm()._supports_cache_class = True
+            kwargs['cache_implementation'] = None
+
+        return self.llm.generate(inputs=None, inputs_embeds=inputs_embeds, attention_mask=attention_mask, **kwargs)
+
+
+AutoConfig.register("ovis", OvisConfig)
+AutoModelForCausalLM.register(OvisConfig, Ovis)

ovis/model/visual_tokenizer/base_visual_tokenizer.py

@@ -0,0 +1,264 @@
+from typing import Union, Optional
+
+import PIL.Image
+import torch
+from torch.nn.functional import softmax, gumbel_softmax, pad
+from transformers import PretrainedConfig, PreTrainedModel, AutoImageProcessor, AutoModel, AutoConfig
+from ovis.util.constants import IMAGE_INDICATOR_IDS, IMAGE_ATOM_ID
+
+
+class BaseVisualTokenizerConfig(PretrainedConfig):
+    def __init__(
+        self,
+        vocab_size=16384,
+        tokenize_function="softmax",
+        tau=1.0,
+        depths=None,
+        drop_cls_token=False,
+        backbone_config: Optional[Union[PretrainedConfig, dict]] = None,
+        hidden_stride: int = 1,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.tokenize_function = tokenize_function
+        self.tau = tau
+        if isinstance(depths, str):
+            depths = [int(x) for x in depths.split('|')]
+        self.depths = depths
+        self.backbone_kwargs = {}
+        self.drop_cls_token = drop_cls_token
+        if backbone_config is not None:
+            assert isinstance(backbone_config, (PretrainedConfig, dict)), \
+                f"expect `backbone_config` to be instance of PretrainedConfig or dict, but got {type(backbone_config)} type"
+            if not isinstance(backbone_config, PretrainedConfig):
+                model_type = backbone_config['model_type']
+                backbone_config.pop('model_type')
+                backbone_config = AutoConfig.for_model(model_type, **backbone_config)
+        self.backbone_config = backbone_config
+        self.hidden_stride = hidden_stride
+
+
+class BaseVisualTokenizer(PreTrainedModel):
+    base_model_prefix = "backbone"
+    main_input_name = None
+    _image_processor_class = None
+    _image_processor_kwargs = {}
+    _backbone_class = None
+    _backbone_name_or_path = None
+
+    def __init__(self, config: BaseVisualTokenizerConfig, *inputs, **kwargs):
+        super().__init__(config, *inputs, **kwargs)
+        if kwargs.get('train_from_scratch'):
+            self.image_processor = self._image_processor_class.from_pretrained(self._backbone_name_or_path,
+                                                                               **self._image_processor_kwargs)
+            self.backbone = self._backbone_class.from_pretrained(self._backbone_name_or_path,
+                                                                 **self.config.backbone_kwargs)
+            self.config.backbone_config = self.backbone.config
+        else:
+            self.image_processor = AutoImageProcessor.from_pretrained(kwargs['image_processor_name_or_path'])
+            self.backbone = AutoModel.from_config(self.config.backbone_config)
+        head_dim = self.config.vocab_size - len(IMAGE_INDICATOR_IDS)  # reserved tokens for IMAGE_INDICATORS
+        self.head = torch.nn.Sequential(
+            torch.nn.Linear(
+                self.backbone.config.hidden_size * self.config.hidden_stride * self.config.hidden_stride, head_dim,
+                bias=False
+            ),
+            torch.nn.LayerNorm(head_dim)
+        )
+
+        assert all((self.image_processor.do_resize,
+                    not getattr(self.image_processor, 'do_center_crop', False),
+                    self.image_processor.do_rescale,
+                    self.image_processor.do_normalize
+                    )), f"image_processor `{self.image_processor}` is not supported currently"
+
+    def get_backbone(self):
+        return self.backbone
+
+    def get_monitor_tensors(self):
+        raise NotImplementedError
+
+    def get_image_processor(self):
+        return self.image_processor
+
+    def mock_input(self):
+        height, width = self.get_image_size()
+        return torch.zeros(1, 3, height, width), self.construct_image_placeholders((1, 1))
+
+    def get_head(self):
+        return self.head
+
+    def get_image_size(self):
+        raise NotImplementedError
+
+    @staticmethod
+    def construct_image_placeholders(grid):
+        image_placeholders = [IMAGE_INDICATOR_IDS[0], IMAGE_ATOM_ID, IMAGE_INDICATOR_IDS[1]]
+        if grid[0] * grid[1] > 1:
+            for r in range(grid[0]):
+                for c in range(grid[1]):
+                    image_placeholders.append(IMAGE_ATOM_ID)
+                    if c < grid[1] - 1:
+                        image_placeholders.append(IMAGE_INDICATOR_IDS[2])
+                if r < grid[0] - 1:
+                    image_placeholders.append(IMAGE_INDICATOR_IDS[3])
+        image_placeholders.append(IMAGE_INDICATOR_IDS[4])
+        return image_placeholders
+
+    def preprocess_image(self, image: PIL.Image.Image, max_partition=9, covering_threshold=0.9, convert_to_rgb=True):
+        def _preprocess(img: PIL.Image.Image, side):
+            # first resize and preprocess
+            w, h = img.size
+            if w == h:
+                new_width = new_height = side
+            elif w > h:
+                new_width = side
+                new_height = int(h / w * new_width)
+            else:
+                new_height = side
+                new_width = int(w / h * new_height)
+            new_size = dict(height=new_height, width=new_width)
+            pixel_values = self.image_processor.preprocess(img, size=new_size, return_tensors='pt')['pixel_values']
+
+            # then pad to square
+            square_values = torch.zeros([1, 3, side, side], dtype=pixel_values.dtype, device=pixel_values.device)
+            new_height, new_width = pixel_values.shape[2:]
+            if new_height == new_width:
+                square_values[:, :, :, :] = pixel_values
+            elif new_height > new_width:
+                from_index = (side - new_width) // 2
+                square_values[:, :, :, from_index:from_index + new_width] = pixel_values
+            else:
+                from_index = (side - new_height) // 2
+                square_values[:, :, from_index:from_index + new_height, :] = pixel_values
+
+            return square_values
+
+        def _partition(img, grid):
+            w, h = img.size
+            row_height = h // grid[0]
+            col_width = w // grid[1]
+
+            partition = []
+            for row in range(grid[0]):
+                for col in range(grid[1]):
+                    left = col * col_width
+                    upper = row * row_height
+                    right = w if col == grid[1] - 1 else (col + 1) * col_width
+                    lower = h if row == grid[0] - 1 else (row + 1) * row_height
+                    partition.append((left, upper, right, lower))
+
+            return partition
+
+        def _covering_area(left, upper, right, lower, side):
+            w = right - left
+            h = lower - upper
+            w, h = max(w, h), min(w, h)
+            if w > side:
+                h = h / w * side
+                w = side
+            return w * h
+
+        def _get_best_grid(img, side):
+            img_area = img.size[0] * img.size[1]
+
+            candidate_grids = []
+            for i in range(1, max_partition + 1):
+                for j in range(1, max_partition + 1):
+                    if i * j <= max_partition:
+                        candidate_grids.append((i, j))
+
+            all_grids = []
+            good_grids = []
+            for grid in candidate_grids:
+                partition = _partition(img, grid)
+                covering_ratio = sum([_covering_area(*p, side) for p in partition]) / img_area
+                assert covering_ratio <= 1.0
+                all_grids.append((grid, covering_ratio))
+                if covering_ratio > covering_threshold:
+                    good_grids.append((grid, covering_ratio))
+
+            if len(good_grids) > 0:
+                # pick the good partition with minimum #sub_images and break the tie using covering_ratio
+                return sorted(good_grids, key=lambda x: (x[0][0] * x[0][1], -x[1]))[0][0]
+            else:
+                # pick the partition with maximum covering_ratio and break the tie using #sub_images
+                return sorted(all_grids, key=lambda x: (-x[1], x[0][0] * x[0][1]))[0][0]
+
+        if convert_to_rgb and image.mode != 'RGB':
+            image = image.convert('RGB')
+
+        sides = self.get_image_size()
+        if sides[0] != sides[1]:
+            raise ValueError('get_image_size() returns non-square size')
+        side = sides[0]
+        grid = _get_best_grid(image, side)
+        partition = _partition(image, grid)
+        crops = [image.crop(p) for p in partition]
+        if len(crops) > 1:
+            crops.insert(0, image)
+        pixel_values = torch.cat([_preprocess(crop, side) for crop in crops], dim=0)
+        image_placeholders = self.construct_image_placeholders(grid)
+        return pixel_values, image_placeholders
+
+    def get_backbone_layer(self, index):
+        return self.backbone.vision_model.encoder.layers[index]
+
+    def tokenize(self, logits):
+        def st_argmax(y_soft, dim):  # straight-through softmax
+            index = y_soft.max(dim, keepdim=True)[1]
+            y_hard = torch.zeros_like(y_soft, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0)
+            ret = y_hard - y_soft.detach() + y_soft
+            return ret
+
+        if self.config.tokenize_function == 'softmax':
+            tokens = softmax(logits, dim=-1)
+        elif self.config.tokenize_function == 'gumbel_argmax':
+            tokens = gumbel_softmax(logits, tau=self.config.tau, hard=True)
+        elif self.config.tokenize_function == 'st_argmax':
+            tokens = st_argmax(logits, dim=-1)
+        else:
+            raise ValueError(
+                f'Invalid `max_type`, expected softmax or gumbel_argmax or st_argmax, but got {self.config.tokenize_function}')
+        return tokens
+
+    def encode(self, pixel_values):
+        output = self.backbone(pixel_values, output_hidden_states=True, return_dict=True)
+        features = output.hidden_states[-1]
+        if self.config.drop_cls_token:
+            features = features[:, 1:, :]
+
+        # merge number of `hidden_stride * hidden_stride` hidden states together to reduce token sequence length
+        # e.g., for hidden_stride=3, this leads to a token length reduction: 729 -> 81 for siglip
+        if self.config.hidden_stride > 1:
+            n, l, d = features.shape  # this `d` maybe different from the above `d
+            sqrt_l = int(l ** 0.5)
+            assert sqrt_l ** 2 == l, "The token sequence length should be a perfect square."
+            features = features.reshape(n, sqrt_l, sqrt_l, d)
+            pl = (self.config.hidden_stride - (sqrt_l % self.config.hidden_stride)) % self.config.hidden_stride
+            features = pad(features, (0, 0, 0, pl, 0, pl), "constant", 0)
+            sqrt_l += pl
+            features = features.reshape(n, sqrt_l // self.config.hidden_stride, self.config.hidden_stride,
+                                        sqrt_l // self.config.hidden_stride, self.config.hidden_stride, d)
+            features = features.permute(0, 1, 3, 2, 4, 5)  # [n, sqrt_l/hs, sqrt_l/hs, hs, hs, d]
+            features = features.flatten(3)  # [n, sqrt_l/hs, sqrt_l/hs, hs*hs*d]
+            features = features.reshape(
+                n, -1, self.config.hidden_stride * self.config.hidden_stride * d)
+
+        return features
+
+    def forward(self, pixel_values) -> torch.Tensor:  # [BatchSize, ImageShape] -> [BatchSize, #Token, VocabSize]
+        features = self.encode(pixel_values)
+        logits = self.head(features)
+        tokens = self.tokenize(logits)
+        # tokens' shape is [BatchSize, #Token, VocabSize-5], so padding with [BatchSize, #Token, 5], after
+        # which, tokens' shape should become [BatchSize, #Token, VocabSize]
+        batch_size, token_len, _ = tokens.shape
+        padding_tensor = torch.zeros(size=(batch_size, token_len, len(IMAGE_INDICATOR_IDS)),
+                                     dtype=tokens.dtype,
+                                     device=tokens.device,
+                                     layout=tokens.layout,
+                                     requires_grad=False)
+        tokens = torch.cat((tokens, padding_tensor), dim=2)
+        return tokens

ovis/model/visual_tokenizer/clip_visual_tokenizer.py

@@ -0,0 +1,41 @@
+from transformers import AutoConfig, AutoModel
+from transformers import CLIPVisionModel, CLIPImageProcessor
+from .base_visual_tokenizer import BaseVisualTokenizerConfig, BaseVisualTokenizer
+
+MODEL_TYPE = "clip_visual_tokenizer"
+
+
+class ClipVisualTokenizerConfig(BaseVisualTokenizerConfig):
+    model_type = MODEL_TYPE
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        if self.depths:
+            assert len(self.depths) == 1
+            self.backbone_kwargs['num_hidden_layers'] = self.depths[0]
+
+
+class ClipVisualTokenizer(BaseVisualTokenizer):
+    config_class = ClipVisualTokenizerConfig
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["CLIPEncoderLayer"]
+    _image_processor_class = CLIPImageProcessor
+    _image_processor_kwargs = dict(do_center_crop=False)
+    _backbone_class = CLIPVisionModel
+    _backbone_name_or_path = "openai/clip-vit-large-patch14-336"
+
+    def get_monitor_tensors(self):
+        return dict(
+            backbone_bottom=self.backbone.vision_model.encoder.layers[0].self_attn.k_proj.weight,
+            backbone_top=self.backbone.vision_model.encoder.layers[-1].self_attn.out_proj.weight,
+            head=self.head[0].weight
+        )
+
+    def get_image_size(self):
+        height = self.image_processor.crop_size["height"]
+        width = self.image_processor.crop_size["width"]
+        return height, width
+
+
+AutoConfig.register(MODEL_TYPE, ClipVisualTokenizerConfig)
+AutoModel.register(ClipVisualTokenizerConfig, ClipVisualTokenizer)

ovis/model/visual_tokenizer/siglip_visual_tokenizer.py

@@ -0,0 +1,43 @@
+from transformers import AutoConfig, AutoModel
+from transformers import SiglipVisionModel, SiglipImageProcessor
+from .base_visual_tokenizer import BaseVisualTokenizerConfig, BaseVisualTokenizer
+
+MODEL_TYPE = "siglip_visual_tokenizer"
+
+
+class SiglipVisualTokenizerConfig(BaseVisualTokenizerConfig):
+    model_type = MODEL_TYPE
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        if self.drop_cls_token:
+            self.drop_cls_token = False
+        if self.depths:
+            assert len(self.depths) == 1
+            self.backbone_kwargs['num_hidden_layers'] = self.depths[0]
+
+
+class SiglipVisualTokenizer(BaseVisualTokenizer):
+    config_class = SiglipVisualTokenizerConfig
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["SiglipVisionTransformer"]
+    _image_processor_class = SiglipImageProcessor
+    _image_processor_kwargs = {}
+    _backbone_class = SiglipVisionModel
+    _backbone_name_or_path = "google/siglip-so400m-patch14-384"
+
+    def get_monitor_tensors(self):
+        return dict(
+            backbone_bottom=self.backbone.vision_model.encoder.layers[0].self_attn.k_proj.weight,
+            backbone_top=self.backbone.vision_model.encoder.layers[-1].self_attn.out_proj.weight,
+            head=self.head[0].weight
+        )
+
+    def get_image_size(self):
+        height = self.image_processor.size["height"]
+        width = self.image_processor.size["width"]
+        return height, width
+
+
+AutoConfig.register(MODEL_TYPE, SiglipVisualTokenizerConfig)
+AutoModel.register(SiglipVisualTokenizerConfig, SiglipVisualTokenizer)

ovis/serve/runner.py

@@ -0,0 +1,105 @@
+from dataclasses import field, dataclass
+from typing import Optional, Union, List
+
+import torch
+from PIL import Image
+
+from ovis.model.modeling_ovis import Ovis
+from ovis.util.constants import IMAGE_TOKEN
+
+
+@dataclass
+class RunnerArguments:
+    model_path: str
+    max_new_tokens: int = field(default=512)
+    do_sample: bool = field(default=False)
+    top_p: Optional[float] = field(default=None)
+    top_k: Optional[int] = field(default=None)
+    temperature: Optional[float] = field(default=None)
+    max_partition: int = field(default=9)
+
+
+class OvisRunner:
+    def __init__(self, args: RunnerArguments):
+        self.model_path = args.model_path
+        self.dtype = torch.bfloat16
+        self.device = torch.cuda.current_device()
+        self.dtype = torch.bfloat16
+        self.model = Ovis.from_pretrained(self.model_path, torch_dtype=self.dtype, multimodal_max_length=8192)
+        self.model = self.model.eval().to(device=self.device)
+        self.eos_token_id = self.model.generation_config.eos_token_id
+        self.text_tokenizer = self.model.get_text_tokenizer()
+        self.pad_token_id = self.text_tokenizer.pad_token_id
+        self.visual_tokenizer = self.model.get_visual_tokenizer()
+        self.conversation_formatter = self.model.get_conversation_formatter()
+        self.image_placeholder = IMAGE_TOKEN
+        self.max_partition = args.max_partition
+        self.gen_kwargs = dict(
+            max_new_tokens=args.max_new_tokens,
+            do_sample=args.do_sample,
+            top_p=args.top_p,
+            top_k=args.top_k,
+            temperature=args.temperature,
+            repetition_penalty=None,
+            eos_token_id=self.eos_token_id,
+            pad_token_id=self.pad_token_id,
+            use_cache=True
+        )
+
+    def preprocess(self, inputs: List[Union[Image.Image, str]]):
+        # for single image and single text inputs, ensure image ahead
+        if len(inputs) == 2 and isinstance(inputs[0], str) and isinstance(inputs[1], Image.Image):
+            inputs = reversed(inputs)
+
+        # build query
+        query = ''
+        images = []
+        for data in inputs:
+            if isinstance(data, Image.Image):
+                query += self.image_placeholder + '\n'
+                images.append(data)
+            elif isinstance(data, str):
+                query += data.replace(self.image_placeholder, '')
+            elif data is not None:
+                raise RuntimeError(f'Invalid input type, expected `PIL.Image.Image` or `str`, but got {type(data)}')
+
+        # format conversation
+        prompt, input_ids, pixel_values = self.model.preprocess_inputs(
+            query, images, max_partition=self.max_partition)
+        attention_mask = torch.ne(input_ids, self.text_tokenizer.pad_token_id)
+        input_ids = input_ids.unsqueeze(0).to(device=self.device)
+        attention_mask = attention_mask.unsqueeze(0).to(device=self.device)
+        if pixel_values is not None:
+            pixel_values = [pixel_values.to(device=self.device, dtype=self.dtype)]
+        else:
+            pixel_values = [None]
+
+        return prompt, input_ids, attention_mask, pixel_values
+
+    def run(self, inputs: List[Union[Image.Image, str]]):
+        prompt, input_ids, attention_mask, pixel_values = self.preprocess(inputs)
+        output_ids = self.model.generate(
+            input_ids,
+            pixel_values=pixel_values,
+            attention_mask=attention_mask,
+            **self.gen_kwargs
+        )
+        output = self.text_tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        input_token_len = input_ids.shape[1]
+        output_token_len = output_ids.shape[1]
+        response = dict(
+            prompt=prompt,
+            output=output,
+            prompt_tokens=input_token_len,
+            total_tokens=input_token_len + output_token_len
+        )
+        return response
+
+
+if __name__ == '__main__':
+    runner_args = RunnerArguments(model_path='<model_path>')
+    runner = OvisRunner(runner_args)
+    image = Image.open('<image_path>')
+    text = '<prompt>'
+    response = runner.run([image, text])
+    print(response['output'])

ovis/serve/server.py

@@ -0,0 +1,41 @@
+import argparse
+import os.path
+
+import gradio as gr
+from gradio.components import Textbox, Image
+
+from ovis.serve.runner import RunnerArguments, OvisRunner
+
+
+class Server:
+    def __init__(self, runner: OvisRunner):
+        self.runner = runner
+
+    def __call__(self, image, text):
+        response = self.runner.run([image, text])
+        output = response["output"]
+        return output
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Ovis Server')
+    parser.add_argument('--model_path', type=str, required=True)
+    parser.add_argument('--flagging_dir', type=str, default=os.path.expanduser('~/ovis-flagged'))
+    parser.add_argument('--max_partition', type=int, default=9)
+    parser.add_argument('--port', type=int, required=True)
+    args = parser.parse_args()
+
+    os.makedirs(args.flagging_dir, exist_ok=True)
+    runner_args = RunnerArguments(
+        model_path=args.model_path,
+        max_partition=args.max_partition
+    )
+    demo = gr.Interface(
+        fn=Server(OvisRunner(runner_args)),
+        inputs=[Image(type='pil', label='image'),
+                Textbox(placeholder='Enter your text here...', label='prompt')],
+        outputs=gr.Markdown(),
+        title=args.model_path.split('/')[-1],
+        flagging_dir=args.flagging_dir
+    )
+    demo.launch(server_port=args.port)

ovis/train/arguments.py

@@ -0,0 +1,48 @@
+from dataclasses import dataclass, field
+from typing import Optional
+
+import transformers
+
+
+@dataclass
+class ModelArguments:
+    llm_name_or_path: Optional[str] = field(default=None)
+    visual_tokenizer_type: str = field(default=None)
+    visual_vocab_size: int = field(default=8192)
+    visual_drop_cls_token: bool = field(default=False)
+    visual_tokenize_function: str = field(default='softmax')
+    visual_tau: float = field(default=1.0)
+    visual_depths: Optional[str] = field(default=None)
+    visual_hidden_stride: int = field(default=1)
+    multimodal_max_length: int = field(default=2048)
+    conversation_formatter_class: str = field(default=None)
+    pad_token_id: Optional[int] = field(default=None)
+    llm_attn_implementation: Optional[str] = field(default=None)
+    disable_tie_weight: bool = field(default=False)
+
+
+@dataclass
+class TrainingArguments(transformers.TrainingArguments):
+    dataset_names: Optional[str] = field(default=None)  # a|b|c
+    dataset_info: Optional[str] = field(default='dataset_info_v1_6')
+    ovis_pretrained_path: Optional[str] = field(default=None)
+    visual_tokenizer_pretrained_path: Optional[str] = field(default=None)
+    caption_template: Optional[str] = field(default=None)
+    stage: Optional[int] = field(default=None)
+    train_modules: Optional[str] = field(default=None)
+    cache_dir: Optional[str] = field(default=None)
+    optim: str = field(default="adamw_torch")
+    visual_max_tau: float = field(default=5.0)
+    visual_min_tau: float = field(default=0.05)
+    save_safetensors: bool = field(default=True)
+    monitor_step: int = field(default=100)
+    vte_re_init: bool = field(default=False)
+    text_max_length: int = field(default=1024)
+    max_partitions: str = field(default="9|1|1")
+
+    def __post_init__(self):
+        if self.gradient_checkpointing:
+            self.gradient_checkpointing_kwargs = {"use_reentrant": False}
+        if self.stage < 3:
+            self.save_safetensors = False
+        super().__post_init__()

ovis/train/callback.py

@@ -0,0 +1,37 @@
+import deepspeed
+import torch
+from transformers import TrainerCallback, TrainingArguments, TrainerState, TrainerControl
+
+from ovis.util.constants import END_LINE, BEGIN_LINE
+from ovis.util.utils import rank0_print
+
+
+class TuneTauCallback(TrainerCallback):
+    def on_step_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        visual_tokenizer = kwargs['model'].get_visual_tokenizer()
+        current_step = state.global_step
+        max_step = state.max_steps
+        ratio = current_step / max_step
+        visual_tokenizer.config.tau = args.visual_max_tau - (args.visual_max_tau - args.visual_min_tau) * ratio
+
+
+class MonitorCallback(TrainerCallback):
+    def _monitoring(self, model, step):
+        with torch.no_grad():
+            with deepspeed.zero.GatheredParameters(model.get_monitor_tensors().values()):
+                for k, v in model.get_monitor_tensors().items():
+                    rank0_print(BEGIN_LINE)
+                    rank0_print(f'{k} @ step {step} with sum: {v.sum().item()} and content: ')
+                    rank0_print(v)
+                    rank0_print(END_LINE)
+
+    def on_step_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        model = kwargs['model']
+        step = state.global_step
+        if step % args.monitor_step == 0 or step == 10:  # monitor at step 10 for fast check
+            self._monitoring(model, step)
+
+    def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        model = kwargs['model']
+        step = state.global_step
+        self._monitoring(model, step)

ovis/train/dataset/caption_dataset.py

@@ -0,0 +1,67 @@
+import logging
+from datetime import datetime
+from typing import Dict
+
+import pandas
+import torch
+
+from ovis.train.dataset.multimodal_dataset import MultimodalDataset
+from ovis.util.constants import IMAGE_TOKEN, IGNORE_ID
+from ovis.util.utils import rank0_print
+
+
+class CaptionDataset(MultimodalDataset):
+
+    def load(self):
+        rank0_print(f"[{datetime.now()}] Loading dataset {self.name} from {self.meta_file} begin")
+        samples = pandas.read_parquet(self.meta_file, engine='pyarrow')
+        rank0_print(f"[{datetime.now()}] Loading dataset {self.name} end")
+        return samples
+
+    def __getitem__(self, i: int) -> Dict[str, torch.Tensor]:
+        sample = self.samples.iloc[i]
+        text = sample['caption']
+        image_path = sample['image_path']
+
+        # read and preprocess image
+        pixel_values, image_placeholders = self.visual_tokenizer.mock_input()
+        valid_image = False
+        image, e = self.read_image(image_path)
+        if image is None:
+            logging.warning(
+                f'reading image failed with index: {i}, image path: {image_path}, and exception: {e}')
+        else:
+            try:
+                pixel_values, image_placeholders = self.visual_tokenizer.preprocess_image(
+                    image, max_partition=self.max_partitions[0])
+                valid_image = True
+            except Exception as e:
+                logging.warning(
+                    f'preprocessing image failed with index: {i}, image path: {image_path}, and exception: {e}')
+
+        # preprocess text
+        if text is None:
+            logging.warning(f'text is `None`, index: {i}')
+            text = ""
+        if not valid_image:
+            logging.warning(f'image is not valid, so set text as empty, index: {i}, image path: {image_path}')
+            text = ""
+        text = text.replace(IMAGE_TOKEN, '').strip()
+        head, tail = self.caption_template.split(IMAGE_TOKEN)
+        head_ids = self.text_tokenizer(head, add_special_tokens=False).input_ids
+        tail_ids = self.text_tokenizer(tail, add_special_tokens=False).input_ids
+        text_ids = self.text_tokenizer(text, add_special_tokens=False).input_ids
+        input_ids = head_ids + image_placeholders + tail_ids + text_ids
+        labels = [IGNORE_ID] * (len(input_ids) - len(text_ids)) + text_ids
+
+        input_ids = input_ids[:self.text_max_length]
+        labels = labels[:self.text_max_length]
+
+        input_ids = torch.tensor(input_ids, dtype=torch.long)
+        labels = torch.tensor(labels, dtype=torch.long)
+
+        return dict(
+            pixel_values=pixel_values,
+            input_ids=input_ids,
+            labels=labels
+        )

ovis/train/dataset/conversation_dataset.py

@@ -0,0 +1,67 @@
+import copy
+import json
+import logging
+from datetime import datetime
+from typing import Dict
+
+import torch
+
+from ovis.train.dataset.multimodal_dataset import MultimodalDataset
+from ovis.util.utils import rank0_print
+
+
+class ConversationDataset(MultimodalDataset):
+    def load(self):
+        rank0_print(f"[{datetime.now()}] Loading dataset {self.name} from {self.meta_file} begin")
+        with open(self.meta_file, 'r', encoding='utf-8') as f:
+            samples = json.load(f)
+        rank0_print(f'#samples: {len(samples)}')
+        rank0_print(f'sample: {samples[0]}')
+        rank0_print(f"[{datetime.now()}] Loading dataset {self.name} end")
+        return samples
+
+    def __getitem__(self, i: int) -> Dict[str, torch.Tensor]:
+        sample = self.samples[i]
+        conversations = copy.deepcopy(sample["conversations"])
+
+        images = None
+        max_partition = None
+        if 'image' in sample:
+            image_paths = sample['image']
+            if isinstance(image_paths, str):
+                image_paths = [image_paths]
+            images = []
+            for image_path in image_paths:
+                image, e = self.read_image(image_path)
+                if image is None:
+                    logging.warning(
+                        f'reading image failed with index: {i}, image path: {image_path}, and exception: {e}')
+                    images = None
+                    break
+                images.append(image)
+        elif 'video' in sample:
+            raise RuntimeError('video is to be supported')
+
+        if images:
+            max_partition = self.max_partitions[0] if len(images) == 1 else self.max_partitions[1]
+
+        prompt, input_ids, pixel_values, labels = self.model.preprocess_inputs(
+            conversations,
+            images,
+            max_partition=max_partition,
+            generation_preface=None,
+            return_labels=True,
+            propagate_exception=False
+        )
+
+        if pixel_values is None:
+            pixel_values, _ = self.visual_tokenizer.mock_input()
+
+        input_ids = input_ids[:self.text_max_length]
+        labels = labels[:self.text_max_length]
+
+        return dict(
+            pixel_values=pixel_values,
+            input_ids=input_ids,
+            labels=labels
+        )

ovis/train/dataset/multimodal_dataset.py

@@ -0,0 +1,72 @@
+import logging
+import os
+from typing import Dict, Sequence, Union, List
+
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+from transformers import PreTrainedTokenizer
+
+from ovis.model.modeling_ovis import Ovis
+from ovis.train.arguments import TrainingArguments
+from ovis.util.constants import IGNORE_ID
+
+
+class MultimodalDataset(Dataset):
+    def __init__(self, name: str, info: Dict, model: Ovis, training_args: TrainingArguments):
+        self.name = name
+        self.meta_file = info['meta_file']
+        self.image_dir = info['image_dir']
+        self.caption_template = info.get('caption_template', None)
+        self.text_tokenizer = model.get_text_tokenizer()
+        self.visual_tokenizer = model.get_visual_tokenizer()
+        self.image_height, self.image_width = self.visual_tokenizer.get_image_size()
+        self.model = model
+        self.text_max_length = training_args.text_max_length
+        self.max_partitions = [int(m.strip()) for m in training_args.max_partitions.split('|')]
+        self.samples = self.load()
+
+    def load(self):
+        raise NotImplementedError
+
+    def __getitem__(self, i: int) -> Dict[str, torch.Tensor]:
+        raise NotImplementedError
+
+    def __len__(self):
+        return len(self.samples)
+
+    def read_image(self, path):
+        try:
+            full_path = os.path.join(self.image_dir, path)
+            image = Image.open(full_path).convert('RGB')
+            return image, None
+        except Exception as e:
+            return None, e
+
+
+class DataCollatorForMultimodalDataset:
+    def __init__(self, text_tokenizer: PreTrainedTokenizer):
+        self.text_tokenizer = text_tokenizer
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, Union[torch.Tensor, List[torch.Tensor]]]:
+        pixel_values, input_ids, labels = tuple([instance[key] for instance in instances]
+                                                for key in ("pixel_values", "input_ids", "labels"))
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids,
+            batch_first=True,
+            padding_value=self.text_tokenizer.pad_token_id)
+        attention_mask = torch.ne(input_ids, self.text_tokenizer.pad_token_id)
+        labels = torch.nn.utils.rnn.pad_sequence(
+            labels,
+            batch_first=True,
+            padding_value=IGNORE_ID)
+        num_valid_label = torch.not_equal(labels, IGNORE_ID).sum().item()
+        if num_valid_label == 0:
+            logging.warning(
+                f'[DataCollatorForMultimodalDataset] All labels in a batch are ignored, which may lead to training instability\n{input_ids=}\n{attention_mask=}\n{labels=}')
+        return dict(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+            pixel_values=pixel_values
+        )

ovis/train/train.py

@@ -0,0 +1,206 @@
+import json
+import os
+import pathlib
+
+import deepspeed
+import torch
+import transformers
+from deepspeed import get_accelerator
+from torch.utils.data import ConcatDataset
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel, AutoConfig
+from transformers import Trainer
+from transformers.integrations.deepspeed import unset_hf_deepspeed_config, set_hf_deepspeed_config
+
+from callback import TuneTauCallback, MonitorCallback
+from ovis.model.configuration_ovis import OvisConfig
+from ovis.model.modeling_ovis import Ovis
+from ovis.train.arguments import ModelArguments, TrainingArguments
+from ovis.train.dataset.caption_dataset import CaptionDataset
+from ovis.train.dataset.conversation_dataset import ConversationDataset
+from ovis.train.dataset.multimodal_dataset import DataCollatorForMultimodalDataset
+from ovis.util.constants import BEGIN_LINE, END_LINE
+from ovis.util.utils import smart_unit, rank0_print
+
+
+def train():
+    # parse args
+    parser = transformers.HfArgumentParser(
+        (ModelArguments, TrainingArguments))
+    model_args, training_args = parser.parse_args_into_dataclasses()
+
+    # save args to checkpoint dir
+    with training_args.main_process_first(local=False):
+        if training_args.process_index == 0:
+            def args2dict(args):
+                return {k: str(v) for k, v in args.__dict__.items()}
+
+            args_log = json.dumps(dict(
+                model_args=args2dict(model_args),
+                training_args=args2dict(training_args)
+            ), ensure_ascii=False, indent=2)
+            print(args_log)
+            os.makedirs(training_args.output_dir, exist_ok=True)
+            with open(os.path.join(training_args.output_dir, 'model_training_args.json'), 'w',
+                      encoding='utf-8') as f:
+                f.write(args_log + '\n')
+
+    # construct or load ovis model
+    if not training_args.ovis_pretrained_path:  # construct model (S1)
+        # 1. construct ovis config
+        ovis_config = OvisConfig(
+            multimodal_max_length=model_args.multimodal_max_length,
+            conversation_formatter_class=model_args.conversation_formatter_class,
+            llm_attn_implementation=model_args.llm_attn_implementation
+        )
+        # 2. load pretrained llm and text tokenizer
+        attn_kwargs = dict()
+        if model_args.llm_attn_implementation:
+            attn_kwargs['attn_implementation'] = model_args.llm_attn_implementation
+        llm = AutoModelForCausalLM.from_pretrained(model_args.llm_name_or_path, **attn_kwargs)
+        text_tokenizer = AutoTokenizer.from_pretrained(model_args.llm_name_or_path)
+        if text_tokenizer.pad_token_id is None and model_args.pad_token_id is not None:
+            text_tokenizer.pad_token_id = model_args.pad_token_id
+        # 3. construct visual tokenizer
+        # deepspeed zero.Init with bfloat16 fail for visual_tokenizer, so temporarily disable zero.Init here
+        unset_hf_deepspeed_config()
+        if training_args.visual_tokenizer_pretrained_path is not None:
+            visual_tokenizer = AutoModel.from_pretrained(
+                training_args.visual_tokenizer_pretrained_path,
+                image_processor_name_or_path=training_args.visual_tokenizer_pretrained_path
+            )
+        else:
+            visual_tokenizer_config = AutoConfig.for_model(
+                model_type=model_args.visual_tokenizer_type + "_visual_tokenizer",
+                vocab_size=model_args.visual_vocab_size,
+                tokenize_function=model_args.visual_tokenize_function,
+                tau=model_args.visual_tau,
+                depths=model_args.visual_depths,
+                drop_cls_token=model_args.visual_drop_cls_token,
+                hidden_stride=model_args.visual_hidden_stride,
+            )
+            visual_tokenizer = AutoModel.from_config(visual_tokenizer_config, train_from_scratch=True)
+        visual_tokenizer = visual_tokenizer.to(
+            device=torch.device(get_accelerator().device_name(os.getenv("LOCAL_RANK"))))
+        if getattr(training_args, 'hf_deepspeed_config', None) is not None:
+            set_hf_deepspeed_config(training_args.hf_deepspeed_config)
+        # 4. construct ovis model
+        model = Ovis(ovis_config, llm=llm, text_tokenizer=text_tokenizer, visual_tokenizer=visual_tokenizer,
+                     train_from_scratch=True)
+    else:  # load pretrained ovis model
+        model, loading_info = Ovis.from_pretrained(training_args.ovis_pretrained_path,
+                                                   multimodal_max_length=model_args.multimodal_max_length,
+                                                   output_loading_info=True)
+        rank0_print(BEGIN_LINE)
+        rank0_print(f'Loading info of Ovis:\n{loading_info}')
+        rank0_print(END_LINE)
+        training_args.vte_re_init = False
+
+    model.get_llm().config.use_cache = False
+    model.config.use_cache = False
+    text_tokenizer = model.get_text_tokenizer()
+
+    rank0_print(BEGIN_LINE)
+    rank0_print(f'model.config:\n{model.config}')
+    rank0_print(END_LINE)
+
+    # maybe re-init vte
+    if training_args.vte_re_init:
+        with deepspeed.zero.GatheredParameters([model.get_wte().weight]):
+            mean = model.get_wte().weight.mean().item()
+            std = model.get_wte().weight.std().item()
+        rank0_print(f'Statistics of embedding table of LLM: {mean=}, {std=}')
+        model.re_init_vte(mean, std)
+
+    # select train modules
+    model.requires_grad_(False)
+    for module in training_args.train_modules.split('|'):
+        if module == 'all':
+            model.requires_grad_(True)
+        elif module == 'llm':
+            model.get_llm().requires_grad_(True)
+        elif module == 'visual_tokenizer':
+            model.get_visual_tokenizer().requires_grad_(True)
+        elif module == 'visual_tokenizer.backbone':
+            model.get_visual_tokenizer().get_backbone().requires_grad_(True)
+        elif module.startswith('visual_tokenizer.backbone.layer.'):
+            layer_index = int(module[len('visual_tokenizer.backbone.layer.'):])
+            layer = model.get_visual_tokenizer().get_backbone_layer(layer_index)
+            layer.requires_grad_(True)
+        elif module == 'visual_tokenizer.head':
+            model.get_visual_tokenizer().get_head().requires_grad_(True)
+        elif module == 'vte':
+            model.get_vte().requires_grad_(True)
+        else:
+            raise ValueError(f'Invalid train module name: {module}')
+
+    rank0_print(BEGIN_LINE)
+    rank0_print('Parameters to train:')
+    for name, param in model.named_parameters():
+        if param.requires_grad:
+            rank0_print(name)
+    rank0_print(f'LLM\'s attn implementation: {model.get_llm().config._attn_implementation}')
+    rank0_print(END_LINE)
+
+    # construct data module
+    datasets = []
+    dataset_info_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                                     f'dataset/{training_args.dataset_info}.json')
+    with open(dataset_info_path, 'r', encoding='utf-8') as f:
+        dataset_info = json.load(f)
+    for name in training_args.dataset_names.split('|'):
+        info = dataset_info[name]
+        data_format = info['data_format']
+        if data_format == 'caption':
+            dataset = CaptionDataset(name, info, model, training_args)
+        elif data_format == 'conversation':
+            dataset = ConversationDataset(name, info, model, training_args)
+        else:
+            raise ValueError(f'Invalid data format `{data_format}` for dataset `{name}`')
+        datasets.append(dataset)
+    data_module = dict(
+        train_dataset=ConcatDataset(datasets),
+        data_collator=DataCollatorForMultimodalDataset(text_tokenizer)
+    )
+
+    # train
+    train_callbacks = [MonitorCallback]
+    if model_args.visual_tokenize_function == 'gumbel_argmax':
+        train_callbacks.append(TuneTauCallback)
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        callbacks=train_callbacks,
+        **data_module
+    )
+    rank0_print(BEGIN_LINE)
+    rank0_print('Dataset sample tensor:')
+    rank0_print(data_module['train_dataset'][0])
+    rank0_print(END_LINE)
+    rank0_print(BEGIN_LINE)
+    rank0_print('Dataset sample input_ids decoding:')
+    rank0_print(text_tokenizer.decode([x for x in data_module['train_dataset'][0]['input_ids'] if x >= 0]))
+    rank0_print(END_LINE)
+    rank0_print(BEGIN_LINE)
+    rank0_print('Dataset sample labels decoding:')
+    rank0_print(text_tokenizer.decode([x for x in data_module['train_dataset'][0]['labels'] if x >= 0]))
+    rank0_print(END_LINE)
+    rank0_print(BEGIN_LINE)
+    rank0_print(f'#param of model: {smart_unit(model.num_parameters())}')
+    rank0_print(f'#param of llm: {smart_unit(model.get_llm().num_parameters())}')
+    rank0_print(f'#param of visual_tokenizer: {smart_unit(model.get_visual_tokenizer().num_parameters())}')
+    rank0_print(f'#param of vte: {smart_unit(model.get_vte().weight.numel())}')
+    rank0_print(END_LINE)
+    if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")):
+        trainer.train(resume_from_checkpoint=True)
+    else:
+        trainer.train()
+    trainer.save_state()
+
+    # save model
+    model.get_llm().config.use_cache = True
+    model.config.use_cache = True
+    trainer.save_model()
+
+
+if __name__ == '__main__':
+    train()

ovis/util/constants.py

@@ -0,0 +1,11 @@
+# Model Constants
+IGNORE_ID = -100
+IMAGE_TOKEN_ID = -200
+IMAGE_TOKEN = "<image>"
+
+IMAGE_ATOM_ID = -300
+IMAGE_INDICATOR_IDS = [-301, -302, -303, -304, -305]
+
+# Log & Print
+BEGIN_LINE = '========================************========================'
+END_LINE = '------------------------------------------------------------'

ovis/util/utils.py

@@ -0,0 +1,26 @@
+import os
+from importlib import import_module
+
+
+def rank0_print(*args):
+    if int(os.getenv("LOCAL_PROCESS_RANK", os.getenv("LOCAL_RANK", 0))) == 0:
+        print(*args)
+
+
+def smart_unit(num):
+    if num / 1.0e9 >= 1:
+        return f'{num / 1.0e9:.2f}B'
+    else:
+        return f'{num / 1.0e6:.2f}M'
+
+
+def import_class_from_string(full_class_string):
+    # Split the path to get separate module and class names
+    module_path, _, class_name = full_class_string.rpartition('.')
+
+    # Import the module using the module path
+    module = import_module(module_path)
+
+    # Get the class from the imported module
+    cls = getattr(module, class_name)
+    return cls