preprocessing/1/model.py

# Copyright 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import base64
import io
import json
import os
from typing import List

import numpy as np
import requests
import triton_python_backend_utils as pb_utils
from PIL import Image
from transformers import AutoProcessor, AutoTokenizer, T5Tokenizer


class TritonPythonModel:
    """Your Python model must use the same class name. Every Python model
    that is created must have "TritonPythonModel" as the class name.
    """

    def initialize(self, args):
        """`initialize` is called only once when the model is being loaded.
        Implementing `initialize` function is optional. This function allows
        the model to initialize any state associated with this model.
        Parameters
        ----------
        args : dict
          Both keys and values are strings. The dictionary keys and values are:
          * model_config: A JSON string containing the model configuration
          * model_instance_kind: A string containing model instance kind
          * model_instance_device_id: A string containing model instance device ID
          * model_repository: Model repository path
          * model_version: Model version
          * model_name: Model name
        """
        # Parse model configs
        model_config = json.loads(args['model_config'])
        tokenizer_dir = model_config['parameters']['tokenizer_dir'][
            'string_value']

        add_special_tokens = model_config['parameters'].get(
            'add_special_tokens')
        visual_model_path = model_config['parameters']['visual_model_path'][
            'string_value']
        max_num_images = model_config['parameters'].get('max_num_images')

        if max_num_images is not None:
            max_num_images_str = max_num_images['string_value']
            if max_num_images_str.isdigit():
                self.max_num_images = int(max_num_images_str)
            else:
                print(
                    f"[TensorRT-LLM][WARNING] 'max_num_images' parameter is not set correctly (value is {max_num_images_str}). Will be set to None"
                )
                self.max_num_images = None
        else:
            print(
                f"[TensorRT-LLM][WARNING] Don't setup 'max_num_images'. Set it as None by default."
            )
            self.max_num_images = None
        if visual_model_path == "${visual_model_path}" or visual_model_path == "":
            visual_model_path = None

        if add_special_tokens is not None:
            add_special_tokens_str = add_special_tokens['string_value'].lower()
            if add_special_tokens_str in [
                    'true', 'false', '1', '0', 't', 'f', 'y', 'n', 'yes', 'no'
            ]:
                self.add_special_tokens = add_special_tokens_str in [
                    'true', '1', 't', 'y', 'yes'
                ]
            else:
                print(
                    f"[TensorRT-LLM][WARNING] Don't setup 'add_special_tokens' correctly (set value is {add_special_tokens['string_value']}). Set it as True by default."
                )
                self.add_special_tokens = True
        else:
            print(
                f"[TensorRT-LLM][WARNING] Don't setup 'add_special_tokens'. Set it as True by default."
            )
            self.add_special_tokens = True

        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir,
                                                       legacy=False,
                                                       padding_side='left',
                                                       trust_remote_code=True)

        if isinstance(self.tokenizer, T5Tokenizer):
            self.tokenizer_bos_id = self.tokenizer.sp_model.bos_id()

        if not self.tokenizer.pad_token:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        self.tokenizer_end_id = self.tokenizer.encode(
            self.tokenizer.eos_token, add_special_tokens=False)[0]
        self.tokenizer_pad_id = self.tokenizer.encode(
            self.tokenizer.pad_token, add_special_tokens=False)[0]
        self.vocab_size = self.tokenizer.vocab_size

        self.is_multimodal = False
        self.model_type = None
        self.vision_preprocessor = None

        if visual_model_path is not None:
            self.is_multimodal = True
            visual_model_path = os.path.join(visual_model_path, 'config.json')
            with open(visual_model_path, 'r') as f:
                visual_model_config = json.load(f)
            self.model_type = visual_model_config['builder_config'][
                'model_type']

            assert self.model_type in [
                'llava', 'blip2-opt', 'vila', 'mllama', 'llava_onevision'
            ], f"[TensorRT-LLM][ERROR] Currently supported multi-modal models are llava, blip2-opt, vila, mllama and llava_onevision. Got {self.model_type}."

            assert self.model_type != 'llava_onevison' or self.max_num_images is None or self.max_num_images <= 1, f"LLaVA-OneVsion is not support multi image inference currently."

            llm_model_path = model_config['parameters']['gpt_model_path'][
                'string_value']
            llm_model_path = os.path.join(llm_model_path, 'config.json')
            with open(llm_model_path, 'r') as f:
                llm_model_config = json.load(f)
            self.vocab_size = int(
                llm_model_config["pretrained_config"]["vocab_size"])
            self._setup_ptable_shape(llm_model_config)

            if self.model_type == 'mllama' or self.model_type == 'llava_onevision':
                self.vision_preprocessor = VisionPreProcessor(
                    self.model_type,
                    AutoProcessor.from_pretrained(tokenizer_dir), model_config)

        # Parse model output configs and convert Triton types to numpy types
        output_names = [
            "INPUT_ID", "DECODER_INPUT_ID", "REQUEST_INPUT_LEN",
            "REQUEST_DECODER_INPUT_LEN", "BAD_WORDS_IDS", "STOP_WORDS_IDS",
            "OUT_END_ID", "OUT_PAD_ID", "OUT_PROMPT_TABLE_EXTRA_IDS",
            "PIXEL_VALUES", "IMAGE_SIZES"
        ]
        input_names = ["EMBEDDING_BIAS_WORDS", "EMBEDDING_BIAS_WEIGHTS"]
        for input_name in input_names:
            setattr(
                self,
                input_name.lower() + "_dtype",
                pb_utils.triton_string_to_numpy(
                    pb_utils.get_input_config_by_name(
                        model_config, input_name)['data_type']))

        for output_name in output_names:
            setattr(
                self,
                output_name.lower() + "_dtype",
                pb_utils.triton_string_to_numpy(
                    pb_utils.get_output_config_by_name(
                        model_config, output_name)['data_type']))

    def _setup_ptable_shape(self, llm_model_config):
        max_prompt_embedding_table_size = llm_model_config['build_config'][
            'max_prompt_embedding_table_size']
        max_batch_size = llm_model_config['build_config']['max_batch_size']

        num_visual_features = max_prompt_embedding_table_size // max_batch_size
        hidden_size = llm_model_config['pretrained_config']['hidden_size']
        if self.max_num_images is not None:
            num_visual_features = num_visual_features // self.max_num_images

        self.ptable_shape = (-1, num_visual_features, hidden_size)

    def execute(self, requests):
        """`execute` must be implemented in every Python model. `execute`
        function receives a list of pb_utils.InferenceRequest as the only
        argument. This function is called when an inference is requested
        for this model. Depending on the batching configuration (e.g. Dynamic
        Batching) used, `requests` may contain multiple requests. Every
        Python model, must create one pb_utils.InferenceResponse for every
        pb_utils.InferenceRequest in `requests`. If there is an error, you can
        set the error argument when creating a pb_utils.InferenceResponse.
        Parameters
        ----------
        requests : list
          A list of pb_utils.InferenceRequest
        Returns
        -------
        list
          A list of pb_utils.InferenceResponse. The length of this list must
          be the same as `requests`
        """

        responses = []

        # Every Python backend must iterate over everyone of the requests
        # and create a pb_utils.InferenceResponse for each of them.
        for idx, request in enumerate(requests):
            # Get input tensors
            query = pb_utils.get_input_tensor_by_name(request,
                                                      'QUERY').as_numpy()
            batch_size = query.shape[0]

            decoder_query = pb_utils.get_input_tensor_by_name(
                request, 'DECODER_QUERY')
            if decoder_query is not None:
                decoder_query = decoder_query.as_numpy()

            request_output_len = pb_utils.get_input_tensor_by_name(
                request, 'REQUEST_OUTPUT_LEN').as_numpy()

            bad_words_dict = pb_utils.get_input_tensor_by_name(
                request, 'BAD_WORDS_DICT')
            if bad_words_dict is not None:
                bad_words_dict = bad_words_dict.as_numpy()

            stop_words_dict = pb_utils.get_input_tensor_by_name(
                request, 'STOP_WORDS_DICT')
            if stop_words_dict is not None:
                stop_words_dict = stop_words_dict.as_numpy()

            embedding_bias_words = pb_utils.get_input_tensor_by_name(
                request, 'EMBEDDING_BIAS_WORDS')
            if embedding_bias_words is not None:
                embedding_bias_words = embedding_bias_words.as_numpy()

            embedding_bias_weights = pb_utils.get_input_tensor_by_name(
                request, 'EMBEDDING_BIAS_WEIGHTS')
            if embedding_bias_weights is not None:
                embedding_bias_weights = embedding_bias_weights.as_numpy()

            # Take the end_id from the input tensors
            # If not specified, use tokenizer to get end_id
            end_id = pb_utils.get_input_tensor_by_name(request, 'END_ID')
            if end_id is not None:
                end_id = end_id.as_numpy()
            else:
                end_id = [[self.tokenizer_end_id]] * batch_size

            # Take the pad_id from the input tensors
            # If not specified, use tokenizer to get pad_id
            pad_id = pb_utils.get_input_tensor_by_name(request, 'PAD_ID')
            if pad_id is not None:
                pad_id = pad_id.as_numpy()
            else:
                pad_id = [[self.tokenizer_pad_id]] * batch_size

            # Take the extra_id from the input tensors
            # Extra id is used in kv cache reuse for p-tuning
            prompt_table_extra_id = pb_utils.get_input_tensor_by_name(
                request, 'PROMPT_TABLE_EXTRA_ID')
            if prompt_table_extra_id is not None:
                prompt_table_extra_id = prompt_table_extra_id.as_numpy()
                assert prompt_table_extra_id.shape[
                    0] == batch_size, "Prompt table extra id must have the same batch size as Query"
                assert prompt_table_extra_id.shape[
                    1] == 1, "Multiple IDs cannot be provided for a single image"

            # Preprocessing vision input passed as a url or bytes tensor
            img_urls = pb_utils.get_input_tensor_by_name(request, 'IMAGE_URL')
            image_bytes = pb_utils.get_input_tensor_by_name(
                request, 'IMAGE_BYTES')
            video_bytes = pb_utils.get_input_tensor_by_name(
                request, 'VIDEO_BYTES')
            vision_processed_tensors = []
            visual_tokens = []
            if self.is_multimodal and (img_urls or image_bytes or video_bytes):
                assert self.vision_preprocessor != None, "Vision preprocessor for preparing images before encoding is None"
                processed_tensors = {}
                if self.model_type == 'mllama':
                    processed_tensors = self.vision_preprocessor.mllama_process(
                        queries=query.astype(str).tolist(),
                        img_urls=img_urls,
                        image_bytes=image_bytes,
                    )
                elif self.model_type == 'llava_onevision':
                    if video_bytes is None:
                        processed_tensors, visual_tokens = self.vision_preprocessor.llava_onevision_process_image(
                            queries=query.astype(str).tolist(),
                            img_urls=img_urls,
                            image_bytes=image_bytes,
                        )
                    else:
                        processed_tensors, visual_tokens = self.vision_preprocessor.llava_onevision_process_video(
                            queries=query.astype(str).tolist(),
                            video_bytes=video_bytes,
                        )
                else:
                    raise ValueError(
                        "Unsupported model type for IMAGE_BYTES or IMAGE_URL inputs"
                    )
                vision_processed_tensors = [
                    pb_utils.Tensor.from_dlpack(k, v)
                    for k, v in processed_tensors.items()
                ]
            else:
                assert self.model_type != "llava_onevision", "Image processing requires IMAGE_BYTES or IMAGE_URL to be provided"

            # Preprocessing input data.
            # For the LLaVA_OneVision model, num_visual_features is not a fixed value
            input_id, request_input_len = self._create_request(
                query, visual_tokens)
            if decoder_query is not None:
                decoder_input_id, request_decoder_input_len = self._create_request(
                    decoder_query)
            else:
                decoder_input_id = pad_id * np.ones((batch_size, 1), np.int32)
                request_decoder_input_len = 1 * np.ones(
                    (batch_size, 1), np.int32)

            bad_words = self._to_word_list_format(bad_words_dict, batch_size)
            stop_words = self._to_word_list_format(stop_words_dict, batch_size)

            embedding_bias = self._get_embedding_bias(
                embedding_bias_words, embedding_bias_weights,
                self.embedding_bias_weights_dtype, batch_size)

            if prompt_table_extra_id is not None:
                prompt_table_extra_ids = np.zeros_like(input_id)
                for i in range(batch_size):
                    prompt_table_extra_ids[i] = np.where(
                        input_id[i] >= self.vocab_size,
                        prompt_table_extra_id[i], 0)

            # Create output tensors. You need pb_utils.Tensor
            # objects to create pb_utils.InferenceResponse.
            input_id_tensor = pb_utils.Tensor(
                'INPUT_ID', input_id.astype(self.input_id_dtype))
            request_input_len_tensor = pb_utils.Tensor(
                'REQUEST_INPUT_LEN',
                request_input_len.astype(self.request_input_len_dtype))
            decoder_input_id_tensor = pb_utils.Tensor(
                'DECODER_INPUT_ID',
                decoder_input_id.astype(self.decoder_input_id_dtype))
            request_decoder_input_len_tensor = pb_utils.Tensor(
                'REQUEST_DECODER_INPUT_LEN',
                request_decoder_input_len.astype(
                    self.request_decoder_input_len_dtype))
            request_output_len_tensor = pb_utils.Tensor(
                'REQUEST_OUTPUT_LEN', request_output_len)
            bad_words_ids_tensor = pb_utils.Tensor('BAD_WORDS_IDS', bad_words)
            stop_words_ids_tensor = pb_utils.Tensor('STOP_WORDS_IDS',
                                                    stop_words)
            embedding_bias_tensor = pb_utils.Tensor('EMBEDDING_BIAS',
                                                    embedding_bias)
            end_id_tensor = pb_utils.Tensor('OUT_END_ID',
                                            np.array(end_id, dtype=np.int32))
            pad_id_tensor = pb_utils.Tensor('OUT_PAD_ID',
                                            np.array(pad_id, dtype=np.int32))

            if prompt_table_extra_id is not None:
                prompt_table_extra_ids_tensor = pb_utils.Tensor(
                    'OUT_PROMPT_TABLE_EXTRA_IDS',
                    np.array(prompt_table_extra_ids,
                             dtype=self.out_prompt_table_extra_ids_dtype))
                inference_response = pb_utils.InferenceResponse(output_tensors=[
                    input_id_tensor, decoder_input_id_tensor,
                    bad_words_ids_tensor, stop_words_ids_tensor,
                    request_input_len_tensor, request_decoder_input_len_tensor,
                    request_output_len_tensor, embedding_bias_tensor,
                    end_id_tensor, pad_id_tensor, prompt_table_extra_ids_tensor
                ] + vision_processed_tensors)
            else:
                inference_response = pb_utils.InferenceResponse(
                    output_tensors=[
                        input_id_tensor, decoder_input_id_tensor,
                        bad_words_ids_tensor, stop_words_ids_tensor,
                        request_input_len_tensor,
                        request_decoder_input_len_tensor,
                        request_output_len_tensor, embedding_bias_tensor,
                        end_id_tensor, pad_id_tensor
                    ] + vision_processed_tensors)
            responses.append(inference_response)

        # You should return a list of pb_utils.InferenceResponse. Length
        # of this list must match the length of `requests` list.
        return responses

    def finalize(self):
        """`finalize` is called only once when the model is being unloaded.
        Implementing `finalize` function is optional. This function allows
        the model to perform any necessary clean ups before exit.
        """
        print('Cleaning up...')

    def _split_prompt_by_images(self,
                                concatenated_ids,
                                image_token_index=-200):
        """
        Splits tokenized prompts by image placeholders for each sample in the batch.

        Args:
            concatenated_ids (np.ndarray): A batch of concatenated token IDs, where image placeholders are indicated by `image_token_index`.

        Returns:
            List[List[np.ndarray]]: A list containing lists of token ID arrays for each prompt segment, per batch sample.
        """
        batch_splits = []
        for batch in concatenated_ids:
            zero_indices = np.where(batch == image_token_index)[0]
            start_idx = 0
            splits = []
            for idx in zero_indices:
                if start_idx != idx:
                    splits.append(batch[start_idx:idx].reshape(1, -1))
                start_idx = idx + 1
            if start_idx < len(batch):
                splits.append(batch[start_idx:].reshape(1, -1))

            splits = [split for split in splits if split.size > 0]
            batch_splits.append(splits)

        return batch_splits

    def _setup_fake_prompts(self, batch_size, batch_split_prompts):
        """
        Replaces image placeholders with unique fake prompt IDs for multi-image inputs.

        Args:
            batch_size (int): The number of samples in the batch.
            batch_split_prompts (List[List[np.ndarray]]): Tokenized prompt segments for each batch sample.

        Returns:
            np.ndarray: An array of input IDs with image placeholders replaced by fake prompt IDs.
        """

        num_visual_features = self.ptable_shape[1]
        input_ids_list = []

        for batch_idx in range(batch_size):
            splits = batch_split_prompts[batch_idx]
            sample_input_ids = [splits[0]]
            sample_fake_prompt_counter = self.vocab_size

            for split_idx in range(len(splits) - 1):
                fake_prompt_id = np.arange(
                    sample_fake_prompt_counter,
                    sample_fake_prompt_counter + num_visual_features)
                sample_fake_prompt_counter += num_visual_features
                fake_prompt_id = np.expand_dims(fake_prompt_id, axis=0)
                sample_input_ids.append(fake_prompt_id)
                sample_input_ids.append(splits[split_idx + 1])

            sample_input_ids = np.concatenate(sample_input_ids, axis=1)
            input_ids_list.append(sample_input_ids)

        # Pad the input_ids to the same length for bs > 1
        max_seq_len = max(
            [sample_input_ids.shape[1] for sample_input_ids in input_ids_list])
        input_ids_padded = []
        for sample_input_ids in input_ids_list:
            seq_len = sample_input_ids.shape[1]
            pad_width = max_seq_len - seq_len
            if pad_width > 0:
                sample_input_ids_padded = np.pad(
                    sample_input_ids, ((0, 0), (0, pad_width)),
                    'constant',
                    constant_values=self.tokenizer_pad_id)
            else:
                sample_input_ids_padded = sample_input_ids
            input_ids_padded.append(sample_input_ids_padded)

        input_ids = np.stack(input_ids_padded)
        input_ids = input_ids.reshape(batch_size, -1).astype(np.int32)

        return input_ids

    def _process_multi_image_inputs(self, query, image_token_index=-200):
        """
        Processes input queries that contain multiple images by tokenizing the input strings and inserting image_token_index between the parts.

        Args:
            query (np.ndarray): Batch of input strings.

        Returns:
            List[np.ndarray]: List of tokenized input IDs for each sample.
        """
        start_ids = []
        for s in query:
            parts = s[0].decode().split('<image>')
            num_images = len(parts) - 1
            if num_images > self.max_num_images:
                raise ValueError(
                    f"The number of images in the request ({num_images}) exceeds the maximum allowed ({self.max_num_images})."
                )
            tokenized_parts = [
                self.tokenizer.encode(part, add_special_tokens=False)
                for part in parts
            ]

            # Insert `image_token_index` between the parts to represent <image>
            final_ids = []
            for i, part in enumerate(tokenized_parts):
                final_ids.extend(part)
                if i < len(tokenized_parts) - 1:
                    final_ids.append(image_token_index)

            start_ids.append(np.array(final_ids).astype(int))

        return start_ids

    def _create_request(self, query, visual_tokens=None):
        """
            query : batch string (2D numpy array)
        """
        if isinstance(self.tokenizer, T5Tokenizer):
            start_ids = [
                np.array([self.tokenizer_bos_id] + self.tokenizer.encode(
                    s[0].decode(), add_special_tokens=self.add_special_tokens)
                         ).astype(int) for s in query
            ]
        else:
            if self.is_multimodal and self.max_num_images and self.max_num_images > 1:
                start_ids = self._process_multi_image_inputs(query)

            else:
                start_ids = [
                    np.array(
                        self.tokenizer.encode(s[0].decode(),
                                              add_special_tokens=self.
                                              add_special_tokens)).astype(int)
                    for s in query
                ]

        if self.is_multimodal:
            if 'blip2' in self.model_type or 'mllama' == self.model_type:
                pre_prompt = None
                post_prompt = None
            elif 'llava' == self.model_type:
                pre_prompt = "USER:\n"
                post_prompt = " ASSISTANT:"
            elif 'vila' == self.model_type:
                pre_prompt = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: "
                post_prompt = " ASSISTANT:"
            elif 'llava_onevision' == self.model_type:
                pre_prompt = "<|im_start|>user "
                post_prompt = "<|im_end|><|im_start|>assistant\n"

            pre_prompt_id = np.array(
                self.tokenizer.encode(
                    pre_prompt,
                    add_special_tokens=self.add_special_tokens,
                    padding=True)) if pre_prompt is not None else np.array(
                        [], dtype=int)

            post_prompt_id = np.array(
                self.tokenizer.encode(
                    post_prompt,
                    add_special_tokens=self.add_special_tokens,
                    padding=True)) if post_prompt is not None else np.array(
                        [], dtype=int)

            if self.max_num_images and self.max_num_images > 1:
                concatenated_ids = [
                    np.concatenate((pre_prompt_id, ids, post_prompt_id),
                                   axis=0) for ids in start_ids
                ]
                batch_split_prompts = self._split_prompt_by_images(
                    concatenated_ids)
                start_ids = self._setup_fake_prompts(query.shape[0],
                                                     batch_split_prompts)
            elif self.model_type == 'llava_onevision':
                fake_prompt_ids = []
                extra_id = np.array(
                    self.tokenizer.encode(
                        '\n',
                        add_special_tokens=self.add_special_tokens,
                        padding=True))
                for tokens in visual_tokens:
                    prompt_id = np.arange(self.vocab_size,
                                          self.vocab_size + tokens)
                    fake_prompt_ids.append(prompt_id)
                start_ids = [
                    np.concatenate((pre_prompt_id, prompt_id, extra_id, ids,
                                    post_prompt_id),
                                   axis=0)
                    for prompt_id, ids in zip(fake_prompt_ids, start_ids)
                ]
            else:
                fake_prompt_id = np.arange(
                    self.vocab_size, self.vocab_size + self.ptable_shape[1])
                start_ids = [
                    np.concatenate(
                        (pre_prompt_id, fake_prompt_id, ids, post_prompt_id),
                        axis=0) for ids in start_ids
                ]

        start_lengths = np.array([[len(ids)] for ids in start_ids]).astype(int)

        max_len = 0
        for seq in start_ids:
            max_len = max(max_len, seq.shape[0])
        start_ids = np.stack([
            np.pad(seq, (0, max_len - seq.shape[0]),
                   'constant',
                   constant_values=(0, self.tokenizer_pad_id))
            for seq in start_ids
        ])

        return start_ids, start_lengths

    def _to_word_list_format(self, word_lists: List[List[str | bytes]],
                             batch_size):
        '''
        word_lists format:
            len(word_lists) == batch_size
            word_lists[i] means the words associated to batch item i. A "word" may actually be any string. Like "lorem" or "lorem ipsum".
        '''
        assert self.tokenizer != None, "need to set tokenizer"

        if word_lists is None:
            # Return an empty array of shape (1,2,0)
            return np.empty([batch_size, 2, 0], dtype="int32")

        flat_ids = []
        offsets = []
        for word_list in word_lists:
            item_flat_ids = []
            item_offsets = []

            for word in word_list:
                if isinstance(word, bytes):
                    word = word.decode()

                ids = self.tokenizer.encode(word, add_special_tokens=False)
                if len(ids) == 0:
                    continue

                item_flat_ids += ids
                item_offsets.append(len(ids))

            flat_ids.append(np.array(item_flat_ids))
            offsets.append(np.cumsum(np.array(item_offsets)))

        pad_to = max(1, max(len(ids) for ids in flat_ids))

        for i, (ids, offs) in enumerate(zip(flat_ids, offsets)):
            flat_ids[i] = np.pad(ids, (0, pad_to - len(ids)),
                                 constant_values=0)
            offsets[i] = np.pad(offs, (0, pad_to - len(offs)),
                                constant_values=-1)

        return np.array([flat_ids, offsets], dtype="int32").transpose(
            (1, 0, 2))

    def _get_embedding_bias(self, embedding_bias_words, embedding_bias_weights,
                            bias_dtype, batch_size):

        assert self.tokenizer != None, "need to set tokenizer"

        if embedding_bias_words is None or embedding_bias_weights is None:
            return np.empty([batch_size, 0],
                            dtype=self.embedding_bias_weights_dtype)

        batch_embedding_bias = []
        for words, weights in zip(embedding_bias_words,
                                  embedding_bias_weights):

            vocab_size = len(self.tokenizer.vocab)
            embedding_bias = [0.] * vocab_size

            assert len(words) == len(
                weights
            ), "Embedding bias words must have same dimension as embedding bias weights"

            for word, weight in zip(words, weights):
                if isinstance(word, bytes):
                    word = word.decode()
                ids = self.tokenizer.encode(word)

                if len(ids) == 0:
                    continue

                for id in ids:
                    embedding_bias[id] += weight

            batch_embedding_bias.append(np.array(embedding_bias))

        return np.array(batch_embedding_bias, dtype=bias_dtype)


class VisionPreProcessor:
    """ A class that can load images from url requests, and process them via a vision model processor,
    in preparation for the vision encoder.
    """

    def __init__(self,
                 vision_model_type,
                 vision_model_processor,
                 preprocessor_model_config={}):
        # import libraries that are only relevant for multimodal models
        import torch
        from torch.utils.dlpack import from_dlpack

        # NOTE: Due to the behavior of MPI initialization, it is recommended to avoid using import tensorrt_llm
        #       except for the specific modules tensorrt_llm and multimodal_encoders.
        #       As a result, the function str_dtype_to_torch has been copied directly from tensorrt_llm._utils.
        _str_to_torch_dtype_dict = dict(
            bfloat16=torch.bfloat16,
            float16=torch.float16,
            float32=torch.float32,
            int64=torch.int64,
            int32=torch.int32,
            int8=torch.int8,
            bool=torch.bool,
            fp8=torch.float8_e4m3fn,
        )

        def str_dtype_to_torch(dtype):
            ret = _str_to_torch_dtype_dict.get(dtype)
            assert ret is not None, f'Unsupported dtype: {dtype}'
            return ret

        self.load_images_tensor = lambda tensor: tensor if not hasattr(
            tensor, 'to_dlpack') else from_dlpack(tensor.to_dlpack())

        # extract expected output tensor dtype
        self.output_str_dtypes = {}
        for properties in preprocessor_model_config.get('output', []):
            dtype = properties['data_type']
            self.output_str_dtypes[properties['name']] = np.dtype(
                pb_utils.triton_string_to_numpy(dtype)).name

        # create method for converting output tensors batch to the expected type
        self.convert_tensor_list_to_tensor = lambda tensor_list: torch.concat(
            [
                torch.from_numpy(x) if isinstance(x, np.ndarray) else x
                for x in tensor_list
            ],
            dim=0)
        self.convert_tensor_to_str_dtype = lambda tensor, dtype: tensor.to(
            str_dtype_to_torch(dtype))

        # create model-specific processor
        self.vision_model_processor = vision_model_processor
        self.vision_model_type = vision_model_type

    def load_images_from_urls(self, img_urls):
        images = []
        for img_url in img_urls:
            img_url = img_url.decode()
            if img_url.startswith("data:image/jpeg;base64,"):
                image_base64 = img_url.split(",")[1]
                # Decode the base64 string
                image_data = base64.b64decode(image_base64)
                # Create a BytesIO object from the decoded data
                image_buffer = io.BytesIO(image_data)
                images.append(Image.open(image_buffer))
            else:
                images.append(
                    Image.open(requests.get(img_url, stream=True).raw))
        return images

    def mllama_process(self, queries, img_urls=None, image_bytes=None):
        vision_processed_tensors = {}
        if img_urls is not None or image_bytes is not None:
            if img_urls is not None:
                # download and read images
                images = [
                    self.load_images_from_urls(urls)
                    for urls in img_urls.as_numpy()
                ]
            else:
                images = [
                    img for img_list in self.load_images_tensor(image_bytes)
                    for img in img_list
                ]

            batch_size = len(images)

            preprocessor_outputs = {}
            possible_output_names = [
                'PIXEL_VALUES', 'ASPECT_RATIO_IDS', 'ASPECT_RATIO_MASK',
                'CROSS_ATTENTION_MASK'
            ]
            for batch_id in range(batch_size):
                # Preprocess images and query
                processed_vision_data = self.vision_model_processor(
                    images=images[batch_id],
                    text=queries[batch_id],
                    return_tensors="pt")

                # Reshape pixel_values to [num_images, *HWC/CHW]
                val = processed_vision_data["pixel_values"]

                val = val.reshape(1, -1, *(val.shape[-3:]))
                processed_vision_data["pixel_values"] = val
                # Create vision output tensors
                for key in possible_output_names:
                    val = processed_vision_data.get(key.lower())
                    if val is not None:
                        if key not in preprocessor_outputs:
                            preprocessor_outputs[key] = []
                        preprocessor_outputs[key].append(val)

            for key, tensor_list in preprocessor_outputs.items():
                val = self.convert_tensor_list_to_tensor(tensor_list)
                if key in self.output_str_dtypes:
                    val = self.convert_tensor_to_str_dtype(
                        val, self.output_str_dtypes[key])
                vision_processed_tensors[key] = val
        return vision_processed_tensors

    def llava_onevision_process_image(self,
                                      queries,
                                      img_urls=None,
                                      image_bytes=None):

        import torch
        vision_processed_tensors = {}
        if img_urls is not None:
            # download and read images
            images = [
                self.load_images_from_urls(urls)
                for urls in img_urls.as_numpy()
            ]
        else:
            images = [
                img for img_list in self.load_images_tensor(image_bytes)
                for img in img_list
            ]

        batch_size = len(images)
        assert len(
            queries
        ) == batch_size, f"Image must have the same batch size as Query."
        preprocessor_outputs = {}
        possible_output_names = ['PIXEL_VALUES', 'IMAGE_SIZES']
        visual_tokens = []
        for batch_id in range(batch_size):
            # Preprocess images and query
            processed_vision_data = self.vision_model_processor(
                images=images[batch_id], text='<image>', return_tensors="pt")
            visual_tokens.append(processed_vision_data['input_ids'].shape[1])

            # Create vision output tensors
            for key in possible_output_names:
                val = processed_vision_data.get(key.lower())
                if val is not None:
                    if key not in preprocessor_outputs:
                        preprocessor_outputs[key] = []
                    preprocessor_outputs[key].append(val)

        max_patch = max(x.shape[1]
                        for x in preprocessor_outputs['PIXEL_VALUES'])
        preprocessor_outputs['PIXEL_VALUES'] = [
            torch.nn.functional.pad(
                image, (0, 0, 0, 0, 0, 0, 0, max_patch - image.shape[1], 0, 0),
                mode='constant')
            for image in preprocessor_outputs['PIXEL_VALUES']
        ]
        for key, tensor_list in preprocessor_outputs.items():
            val = self.convert_tensor_list_to_tensor(tensor_list)
            if key in self.output_str_dtypes:
                val = self.convert_tensor_to_str_dtype(
                    val, self.output_str_dtypes[key])
            vision_processed_tensors[key] = val
        return vision_processed_tensors, visual_tokens

    def llava_onevision_process_video(self, queries, video_bytes=None):
        import torch
        vision_processed_tensors = {}
        videos = [video for video in self.load_images_tensor(video_bytes)]

        batch_size = len(videos)
        assert len(
            queries
        ) == batch_size, f"Video must have the same batch size as Query."
        preprocessor_outputs = {}
        preprocessor_outputs['PIXEL_VALUES'] = []
        preprocessor_outputs['IS_VIDEO_INPUT'] = []
        visual_tokens = []
        for batch_id in range(len(queries)):
            processed_vision_data = self.vision_model_processor(
                videos=list(videos[batch_id]),
                text='<video>',
                return_tensors="pt")
            visual_tokens.append(processed_vision_data['input_ids'].shape[1])
            preprocessor_outputs['PIXEL_VALUES'].append(
                processed_vision_data['pixel_values_videos'])
            preprocessor_outputs['IS_VIDEO_INPUT'].append(
                torch.ones((1, 1), dtype=torch.bool))

        for key, tensor_list in preprocessor_outputs.items():
            val = self.convert_tensor_list_to_tensor(tensor_list)
            if key in self.output_str_dtypes:
                val = self.convert_tensor_to_str_dtype(
                    val, self.output_str_dtypes[key])
            vision_processed_tensors[key] = val
        return vision_processed_tensors, visual_tokens