Delete tokenizer_wrapper.py

tokenizer_wrapper.py

@@ -1,1426 +0,0 @@
-# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://github.com/Tencent-Hunyuan/HunyuanImage-3.0/blob/main/LICENSE
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-import warnings
-import random
-from typing import List, Optional, Union, Dict, Any
-from collections import defaultdict
-from copy import deepcopy
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from transformers import AutoTokenizer
-from diffusers.utils import BaseOutput
-
-
-def default(value, default_value):
-    return value if value is not None else default_value
-
-
-def ensure_list(value):
-    if value is None:
-        return []
-    if isinstance(value, (list, tuple)):
-        return list(value)
-    return [value]
-
-
-class Resolution(object):
-    def __init__(self, size, *args):
-        if isinstance(size, str):
-            if 'x' in size:
-                size = size.split('x')
-                size = (int(size[0]), int(size[1]))
-            else:
-                size = int(size)
-        if len(args) > 0:
-            size = (size, args[0])
-        if isinstance(size, int):
-            size = (size, size)
-
-        self.h = self.height = size[0]
-        self.w = self.width = size[1]
-        self.r = self.ratio = self.height / self.width
-
-    def __getitem__(self, idx):
-        if idx == 0:
-            return self.h
-        elif idx == 1:
-            return self.w
-        else:
-            raise IndexError(f'Index {idx} out of range')
-
-    def __str__(self):
-        return f'{self.h}x{self.w}'
-
-
-class ResolutionGroup(object):
-    def __init__(self, base_size=None, step=None, align=1):
-        self.align = align
-        self.base_size = base_size
-        assert base_size % align == 0, f'base_size {base_size} is not divisible by align {align}'
-        if base_size is not None and not isinstance(base_size, int):
-            raise ValueError(f'base_size must be None or int, but got {type(base_size)}')
-        if step is None:
-            step = base_size // 16
-        if step is not None and step > base_size // 2:
-            raise ValueError(f'step must be smaller than base_size // 2, but got {step} > {base_size // 2}')
-
-        self.step = step
-        self.data = self._calc_by_step()
-
-        self.ratio = np.array([x.ratio for x in self.data])
-        self.attr = ['' for _ in range(len(self.data))]
-        self.prefix_space = 0
-
-    def __len__(self):
-        return len(self.data)
-
-    def __getitem__(self, idx):
-        return self.data[idx]
-
-    def __repr__(self):
-        prefix = self.prefix_space * ' '
-        prefix_close = (self.prefix_space - 4) * ' '
-        res_str = f'ResolutionGroup(base_size={self.base_size}, step={self.step}, data='
-        attr_maxlen = max([len(x) for x in self.attr] + [5])
-        res_str += \
-            f'\n{prefix}ID: height width   ratio {" " * max(0, attr_maxlen - 4)}count  h/16 w/16    tokens\n{prefix}'
-        res_str += \
-            ('\n' + prefix).join([f'{i:2d}: ({x.h:4d}, {x.w:4d})  {self.ratio[i]:.4f}  {self.attr[i]:>{attr_maxlen}s}  '
-                                  f'({x.h // 16:3d}, {x.w // 16:3d})  {x.h // 16 * x.w // 16:6d}'
-                                  for i, x in enumerate(self.data)])
-        res_str += f'\n{prefix_close})'
-        return res_str
-
-    def _calc_by_step(self):
-        assert self.align <= self.step, f'align {self.align} must be smaller than step {self.step}'
-
-        min_height = self.base_size // 2
-        min_width = self.base_size // 2
-        max_height = self.base_size * 2
-        max_width = self.base_size * 2
-
-        resolutions = [Resolution(self.base_size, self.base_size)]
-
-        cur_height, cur_width = self.base_size, self.base_size
-        while True:
-            if cur_height >= max_height and cur_width <= min_width:
-                break
-
-            cur_height = min(cur_height + self.step, max_height)
-            cur_width = max(cur_width - self.step, min_width)
-            resolutions.append(Resolution(cur_height // self.align * self.align, cur_width // self.align * self.align))
-
-        cur_height, cur_width = self.base_size, self.base_size
-        while True:
-            if cur_height <= min_height and cur_width >= max_width:
-                break
-
-            cur_height = max(cur_height - self.step, min_height)
-            cur_width = min(cur_width + self.step, max_width)
-            resolutions.append(Resolution(cur_height // self.align * self.align, cur_width // self.align * self.align))
-
-        resolutions = sorted(resolutions, key=lambda x: x.ratio)
-
-        return resolutions
-
-    def get_target_size(self, width, height):
-        ratio = height / width
-        idx = np.argmin(np.abs(self.ratio - ratio))
-        reso = self.data[idx]
-        return reso.w, reso.h
-
-    def get_base_size_and_ratio_index(self, width, height):
-        ratio = height / width
-        idx = np.argmin(np.abs(self.ratio - ratio))
-        return self.base_size, idx
-
-
-class ImageInfo:
-    """ Class to store image information for processing and generation. """
-
-    def __init__(
-            self,
-            image_type: str = None,
-            image_tensor: torch.Tensor = None,
-            image_width: int = None,
-            image_height: int = None,
-            token_width: int = None,
-            token_height: int = None,
-            image_token_length: int = None,
-            base_size: int = None,
-            ratio_index: int = None,
-            **kwargs,
-    ):
-        self.image_type = image_type
-        self.image_tensor = image_tensor
-        self.image_width = image_width
-        self.w = image_width
-        self.image_height = image_height
-        self.h = image_height
-        self.token_width = token_width
-        self.tk_w = token_width
-        self.token_height = token_height
-        self.tk_h = token_height
-        self.image_token_length = default(
-            image_token_length,
-            token_width * token_height if token_width is not None and token_height is not None else None
-        )
-        self.base_size = base_size
-        self.ratio_index = ratio_index
-
-        self.add_timestep_token = kwargs.get("add_timestep_token", True)
-        self.add_guidance_token = kwargs.get("add_guidance_token", False)
-        self.use_front_boi_token = kwargs.get("use_front_boi_token", True)
-        self.add_image_shape_token = kwargs.get("add_image_shape_token", True)
-
-    def __getitem__(self, key: str) -> Any:
-        """Allow dictionary-like access to attributes."""
-        if hasattr(self, key):
-            return getattr(self, key)
-        raise KeyError(f"Key '{key}' not found in ImageInfo")
-
-    def __setitem__(self, key: str, value: Any) -> None:
-        """Allow dictionary-like assignment to attributes."""
-        if hasattr(self, key):
-            setattr(self, key, value)
-        else:
-            raise KeyError(f"Key '{key}' not found in ImageInfo")
-
-    def __contains__(self, key: str) -> bool:
-        """Check if the key exists in the ImageInfo object."""
-        return hasattr(self, key)
-
-    def __repr__(self):
-        return (f"ImageInfo(image_type={self.image_type}, image_tensor={self.image_tensor}, "
-                f"image_width={self.image_width}, image_height={self.image_height}, "
-                f"token_width={self.token_width}, token_height={self.token_height}, "
-                f"image_token_length={self.image_token_length}, "
-                f"base_size={self.base_size}, ratio_index={self.ratio_index}")
-
-    @property
-    def meta_info(self):
-        # Used for image sections of tkwrapper.encode_general()
-        if self.image_type in ["vae", "gen_image"]:
-            return dict(
-                token_length=self.image_token_length,
-                add_timestep_token=self.add_timestep_token,
-                add_guidance_token=self.add_guidance_token,
-                use_front_boi_token=self.use_front_boi_token,
-                add_image_shape_token=self.add_image_shape_token,
-                base_size=self.base_size,
-                ratio_idx=self.ratio_index,
-                # for rope 2d
-                token_height=self.token_height,
-                token_width=self.token_width,
-                # for bc
-                image_height=self.image_height,
-                image_width=self.image_width,
-            )
-        elif self.image_type in ["vit"]:
-            return dict(
-                token_length=self.image_token_length,
-                use_front_boi_token=self.use_front_boi_token,
-                add_image_shape_token=self.add_image_shape_token,
-                # for rope 2d
-                token_height=self.token_height,
-                token_width=self.token_width,
-                # for bc
-                image_height=self.image_height,
-                image_width=self.image_width,
-            )
-        else:
-            raise ValueError(f"Unknown image type '{self.image_type}'")
-
-    @property
-    def num_special_tokens(self):
-        if self.args is None:
-            raise ValueError("meta_info requires `args` attribute to be set.")
-        if self.image_type in ["vae", "src_image", "gen_image"]:
-            count = (
-                    2 +  # <boi> + <eoi> or <src_boi> + <src_eoi>
-                    (1 if self.add_timestep_token else 0) +
-                    (1 if self.add_guidance_token else 0) +
-                    (2 if self.add_image_shape_token else 0)
-            )
-        else:
-            raise ValueError(f"Unknown image_type: {self.image_type}")
-        return count
-
-    def copy(self, copy_image_tensor=True):
-        if copy_image_tensor and self.image_tensor is None:
-            raise ValueError("image_tensor is None, cannot copy")
-        return ImageInfo(
-            image_type=self.image_type,
-            image_tensor=self.image_tensor.clone() if copy_image_tensor else None,
-            image_width=self.image_width,
-            image_height=self.image_height,
-            token_width=self.token_width,
-            token_height=self.token_height,
-            image_token_length=self.image_token_length,
-            base_size=self.base_size,
-            ratio_index=self.ratio_index,
-        )
-
-    def zeros_(self):
-        self.image_tensor = torch.zeros_like(self.image_tensor)
-
-
-class ImageTensor(torch.Tensor):
-    # This class is just for type hinting purposes. Attribute `i` should be defined
-    # as an instance attribute of the torch.Tensor instance, like: tensor.i = ImageInfo(...)
-    i: ImageInfo
-    vision_encoder_kwargs: dict
-
-
-class JointImageInfo(object):
-    def __init__(self, vae_image_info: ImageInfo, vision_image_info: ImageInfo, vision_encoder_kwargs: dict = None):
-        self.vae_image_info = vae_image_info
-        self.vision_image_info = vision_image_info
-        self.vision_encoder_kwargs = vision_encoder_kwargs
-
-        # Define key attributes to align with ImageInfo for uniformity
-        self.image_type = "joint_image"
-        self.image_token_length = vae_image_info.image_token_length + vision_image_info.image_token_length
-
-        self.add_timestep_token = vae_image_info.add_timestep_token
-        self.use_front_boi_token = vae_image_info.use_front_boi_token
-        self.add_image_shape_token = vae_image_info.add_image_shape_token
-
-    def __repr__(self):
-        return f"JointImageInfo(vae_image={self.vae_image_info}, vision_image={self.vision_image_info})"
-
-    @property
-    def meta_info(self):
-        # Used for image sections of tkwrapper.encode_general()
-        return dict(
-            token_length=[self.vae_image_info.image_token_length, self.vision_image_info.image_token_length],
-            add_timestep_token=self.add_timestep_token,
-            use_front_boi_token=self.use_front_boi_token,
-            add_image_shape_token=self.add_image_shape_token,
-            base_size=self.vae_image_info.base_size,
-            ratio_idx=self.vae_image_info.ratio_index,
-            # for rope 2d
-            token_height=[self.vae_image_info.token_height, self.vision_image_info.token_height],
-            token_width=[self.vae_image_info.token_width, self.vision_image_info.token_width],
-            # for bc
-            image_height=[self.vae_image_info.image_height, self.vision_image_info.image_height],
-            image_width=[self.vae_image_info.image_width, self.vision_image_info.image_width],
-        )
-
-    @property
-    def num_special_tokens(self):
-        return (
-                2 +  # <boi> + <eoi>
-                (1 if self.add_timestep_token else 0) +
-                (2 if self.add_image_shape_token else 0) +
-                1   # <joint_image_sep>
-        )
-
-    def copy(self, copy_image_tensor=True):
-        if copy_image_tensor and (
-                self.vae_image_info.image_tensor is None or self.vision_image_info.image_tensor is None):
-            raise ValueError("image_tensor is None, cannot copy")
-        return JointImageInfo(
-            self.vae_image_info.copy(copy_image_tensor),
-            self.vision_image_info.copy(copy_image_tensor),
-            self.vision_encoder_kwargs,
-        )
-
-    def zeros_(self):
-        self.vae_image_info.zeros_()
-        self.vision_image_info.zeros_()
-
-
-class JointImage(object):
-    def __init__(self, vae_image: ImageTensor, vision_image: ImageTensor):
-        self.vae_image = vae_image
-        self.vision_image = vision_image
-        self.i = JointImageInfo(vae_image.i, vision_image.i)
-
-
-class TokenizerEncodeOutput(BaseOutput):
-    tokens: torch.Tensor = None
-    timestep_scatter_index: Optional[torch.Tensor] = None
-    guidance_scatter_index: Optional[torch.Tensor] = None
-    text_slices: Optional[List[slice]] = None
-    gen_image_slices: Optional[List[slice]] = None
-    joint_image_slices: Optional[List[slice]] = None
-    cond_vae_image_slices: Optional[List[slice]] = None
-    cond_vit_image_slices: Optional[List[slice]] = None
-    text_mask: Optional[torch.Tensor] = None
-    gen_image_mask: Optional[torch.Tensor] = None
-    cond_vae_image_mask: Optional[torch.Tensor] = None
-    cond_vit_image_mask: Optional[torch.Tensor] = None
-    real_pos: Optional[torch.Tensor] = None
-    all_image_slices: Optional[List[slice]] = None
-    cond_timestep_scatter_index: Optional[torch.Tensor] = None
-    gen_timestep_scatter_index: Optional[torch.Tensor] = None
-
-
-class Conversation:
-    roles: List[str] = ["User", "Assistant"]
-    sep: str = "\n\n"
-
-
-class TokenizerWrapper(object):
-    def __init__(self, tokenizer):
-        if isinstance(tokenizer, str):
-            self.tokenizer = AutoTokenizer.from_pretrained(tokenizer)
-        else:
-            self.tokenizer = tokenizer
-
-        # Define short names
-        self.bos_token_id = self.tokenizer.bos_token_id
-        self.eos_token_id = self.tokenizer.eos_token_id
-        self.pad_token_id = self.tokenizer.pad_token_id
-        self.boi_token_id = self.tokenizer.convert_tokens_to_ids("<boi>")
-        self.eoi_token_id = self.tokenizer.convert_tokens_to_ids("<eoi>")
-        self.img_token_id = self.tokenizer.convert_tokens_to_ids("<img>")
-        self.cfg_token_id = self.tokenizer.convert_tokens_to_ids("<cfg>")
-        self.end_answer_token_id = self.tokenizer.convert_tokens_to_ids("</answer>")
-        self.end_recaption_token_id = self.tokenizer.convert_tokens_to_ids("</recaption>")
-        self.ratio_token_offset = self.tokenizer.convert_tokens_to_ids("<img_ratio_0>")
-        self.special_token_map = self.tokenizer.added_tokens_encoder
-
-    def pad(self, tensor_list, dim=0, pad_val=None):
-        if pad_val is None:
-            pad_val = self.pad_token_id
-        max_len = max([t.shape[dim] for t in tensor_list])
-        padded_tensor_list = []
-        for t in tensor_list:
-            if t.shape[dim] < max_len:
-                assert pad_val is not False, "Not allowed pad."
-                t = F.pad(t, (0, max_len - t.shape[dim]), value=pad_val)
-            padded_tensor_list.append(t)
-        return padded_tensor_list
-
-    def encode(self, *args, **kwargs):
-        return self.tokenizer.encode(*args, **kwargs)
-
-    def decode(self, *args, **kwargs):
-        return self.tokenizer.decode(*args, **kwargs)
-
-    def encode_text(
-            self,
-            *texts,
-            uncond_enabled: Optional[Union[bool, List[bool]]] = None,
-            uncond_p: Optional[float] = None,
-            max_length: Optional[int] = None,
-            pad: Optional[str] = None,
-            return_lengths: bool = False,
-    ):
-        """
-        Encode text and image for AR-like model training of the text-to-image/instruction tuning tasks.
-        Support encode multiple texts at once. Each text can be separately conditioned or unconditioned
-        based on the uncond_flags and a uniform uncond_p.
-        **<bos> token is always prepended to the text tokens.**
-
-        Parameters
-        ----------
-        texts: str or List[str]
-            List of texts to be encoded.
-        uncond_enabled: bool or List[bool]
-            List of flags to indicate whether the text should be unconditioned.
-            If False, the text will never be unconditioned.
-            If True, the text will be unconditioned with uncond_p.
-        uncond_p: float
-            Probability to the unconditional text. Only works when uncond_enabled is True.
-        max_length: int
-            Maximum length of the encoded text.
-        pad: Optional[str]
-            Padding method. Can be 'left' or 'right'.
-        return_lengths: bool
-            Whether to return the length of each encoded text.
-        """
-        if pad is not None:
-            assert max_length is not None, "max_length should be provided when pad is not None."
-
-        if uncond_enabled is None:
-            uncond_enabled = [True] * len(texts)
-        elif isinstance(uncond_enabled, bool):
-            uncond_enabled = [uncond_enabled] * len(texts)
-        if len(uncond_enabled) != len(texts):
-            print(uncond_enabled, texts)
-        assert len(uncond_enabled) == len(texts), (
-            f"Length of uncond_flags should be equal to the number of texts, "
-            f"but got {len(uncond_enabled)} and {len(texts)}."
-        )
-
-        # Prepare text/uncond tokens
-        # TODO: If len(texts) > 1, such as instruction + prompt in inpainting, we need to determine how to do uncond.
-        # Now all texts will be cond or uncond at the same time.
-        do_uncond_drop = (uncond_p is not None) and (random.random() < uncond_p)
-        text_tokens, lengths = [], []
-        cum_length = 0
-        for text, uncond_flag in zip(texts, uncond_enabled):
-            # If reach the max_length and there still have unencoded texts, give a warning message and break the loop.
-            if max_length is not None and cum_length >= max_length:
-                warnings.warn(
-                    f"Text length exceeds the max_length({max_length}). The remaining texts will be ignored: "
-                    f"{text[:80]}..."
-                )
-                break
-            # Set add_special_tokens=False to avoid adding <bos> token in some LLMs.
-            if isinstance(text, str):
-                text_token = self.tokenizer.encode(text, add_special_tokens=False)
-            else:
-                text_token = text
-            if uncond_flag and do_uncond_drop:
-                text_token = [self.cfg_token_id] * len(text_token)
-            # Cutoff the text by max_length if necessary
-            if max_length is not None and (cum_length + len(text_token)) > max_length:
-                text_token = text_token[:max_length - cum_length]
-            text_tokens.extend(text_token)
-            lengths.append(len(text_token))
-            cum_length += len(text_token)
-
-        # Prepend/Append <pad> tokens if applicable
-        if pad is not None and (pad_length := max_length - len(text_tokens)) > 0:
-            if pad == 'left':
-                text_tokens = [self.pad_token_id] * pad_length + text_tokens
-            elif pad == 'right':
-                text_tokens = text_tokens + [self.pad_token_id] * pad_length
-            else:
-                raise ValueError(f"Unsupported padding method: {pad}.")
-
-        if return_lengths:
-            return text_tokens, lengths
-        return text_tokens
-
-    @staticmethod
-    def _check_key_number_matched(keys, data):
-        # Assert keys and token_source are matched
-        assert set(keys) == set(data.keys()), (
-            f"Keys in the template and token source should be matched, but got {set(keys)} and {list(data.keys())}."
-        )
-        key_counts = {k: 0 for k in keys}
-        for key in keys:
-            key_counts[key] += 1
-        for key, count in key_counts.items():
-            assert len(data[key]) == count, (
-                f"Number of `{key}` in the token source should be matched with the template, but got "
-                f"{data[key]}({len(data[key])}) and {count}."
-            )
-
-    def _add_image_meta_info_token(self, token_seq, token_count, extra_token_pos, add_timestep_token=False,
-                                   add_image_shape_token=False, base_size=None, ratio_idx=None, image_type=None,
-                                   add_guidance_token=False):
-        if add_image_shape_token:
-            token_seq.extend([
-                self.special_token_map[f"<img_size_{base_size}>"],
-                self.special_token_map[f"<img_ratio_{ratio_idx}>"]
-            ])
-            token_count += 2
-        if add_timestep_token:
-            token_seq.extend([self.special_token_map["<timestep>"]])
-            extra_token_pos['timestep'].append(token_count)
-            if image_type is not None:
-                if image_type == "gen_image":
-                    extra_token_pos['gen_timestep'].append(token_count)
-                elif image_type in ["joint_image"]:
-                    extra_token_pos['cond_timestep'].append(token_count)
-                else:
-                    raise ValueError(f"Unsupported image type: {image_type}.")
-            token_count += 1
-        if add_guidance_token:
-            token_seq.extend([self.special_token_map["<guidance>"]])
-            extra_token_pos['guidance'].append(token_count)
-            token_count += 1
-        return token_count
-
-    @staticmethod
-    def _shorten_text(text):
-        import re
-        text = re.sub(r"(<img>)+", lambda m: f"[<img>]{{{len(m.group(0)) // 5}}}", text)
-        text = re.sub(r"(<pad>)+", lambda m: f"[<pad>]{{{len(m.group(0)) // 5}}}", text)
-        return text
-
-    def encode_sequence(
-            self,
-            template: str,
-            token_source: Dict[str, List],
-            total_length=None,
-            add_timestep_token=False,
-            add_guidance_token=False,
-            last_key_only_prefix=False,
-            add_eos=True,
-            use_front_boi_token=True,
-            add_pad=True,
-            add_bos=True,
-            drop_last: Union[str, bool] = 'auto',
-            add_image_shape_token=False,
-    ):
-        """
-        Encode a sequence based on the template (e.g., `text-image` for t2i, `text-image-image` for instruction tuning)
-        and token source.
-
-        Parameters
-        ----------
-        template: str
-            Template of the sequence. E.g., "text-gen_image" means the sequence is composed of text and an image.
-            "text-text-gen_image" means the sequence is composed of two sections of text and an image.
-        token_source: Dict[str, List]
-            Token source for each key in the template, in order.
-            - text: List[Dict].
-            - gen_image: List[Dict].
-            - joint_image: List[Dict].
-        total_length: int
-            Total length of the encoded sequence, include padding tokens.
-        add_timestep_token: bool
-            Whether to add timestep token before the image tokens.
-            (Right after the <img_ratio_*><img_size_*> tokens)
-        add_guidance_token: bool
-            Whether to add guidance token before the image tokens.
-        last_key_only_prefix: bool
-            Whether to only use the modal prefix in the last key.
-        add_eos: bool or 'auto'
-            Whether to add eos token at the end of the sequence. If True, always add eos token. If 'auto',
-            add eos token only when the total_length is not reached and the last token is not <eos>.
-        use_front_boi_token: bool:
-            Whether to put the <boi> token at the front of iw, ih and timestep tokens.
-        add_pad: bool or 'auto'
-            Whether to add padding tokens to the sequence. If True and total_length is not reached, add padding tokens.
-        add_bos: bool
-            Whether to add bos token at the beginning of the sequence.
-        drop_last: bool or 'auto'
-            - If auto, drop last tokens exceeding the total_length if the total_length is provided. If cut point is
-                in the middle of the image tokens, an error will raised.
-            - If True, drop last tokens exceeding the total_length. If cut point is in the middle of the image tokens,
-                all the successive image tokens will be dropped.
-            - If False, keep the last tokens exceeding the total_length, even if the total_length is reached.
-        add_image_shape_token: bool
-            Whether to add image shape token before the image tokens. (Right before the <timestep> token)
-
-        Returns
-        -------
-        token_seq: list
-            Encoded token sequence.
-        extra_token_pos: dict
-            Positions of extra tokens.
-        """
-        if last_key_only_prefix:
-            assert add_eos is not True, "add_eos should not be True when last_key_only_prefix is True."
-        if drop_last is True and total_length is None:
-            raise ValueError("total_length should be provided when drop_last is True.")
-
-        keys = template.split('-')
-        modal_length = len(keys)
-        index_indicator = {k: 0 for k in token_source}
-        for k, v in token_source.items():
-            assert isinstance(v, (list, tuple)), (
-                f"Value of `{k}` in the token source should be a list or tuple, but got {type(v)}."
-            )
-        self._check_key_number_matched(keys, token_source)
-
-        token_seq = []
-        token_count = 0
-        extra_token_pos = defaultdict(list)
-        if add_bos:
-            token_seq.append(self.bos_token_id)
-            token_count += 1
-        # If drop_last is True, we check the token_count on the fly and exit the loop if the total_length is reached.
-        # This check is only applied to the block tokens. Block tokens mean the tokens that are unsplittable, like
-        # image tokens. Text tokens are splittable, so we don't need to check the token_count for text.
-        # If the loop is broken by drop_last, we don't add the eos token at the end because the sequence is not
-        # complete.
-        drop_last_break = False
-        for i, key in enumerate(keys):
-            source = token_source[key][index_indicator[key]]
-            if key == "text":
-                token_seq.extend(source)  # text token sequence
-                extra_token_pos["<text>_start"].append(token_count)
-                token_count += len(source)
-                extra_token_pos["<text>_end"].append(token_count - 1)
-
-            elif key == "gen_image":
-                if isinstance(source, int):
-                    source = {'length': source}
-                extra_count = 2 + (
-                    1 if source.get('timestep', add_timestep_token) else 0) + (
-                    1 if source.get('guidance', add_guidance_token) else 0) + (
-                    2 if source.get('image_shape', add_image_shape_token) else 0
-                )
-                if drop_last is True and token_count + extra_count + source['length'] > total_length:
-                    drop_last_break = True
-                    break
-                if source.get('front_boi', use_front_boi_token):
-                    token_seq.append(self.boi_token_id)
-                    extra_token_pos["boi"].append(token_count)
-                    token_count += 1
-                token_count = self._add_image_meta_info_token(
-                    token_seq=token_seq,
-                    token_count=token_count,
-                    extra_token_pos=extra_token_pos,
-                    add_timestep_token=source.get('timestep', add_timestep_token),
-                    add_guidance_token=source.get('guidance', add_guidance_token),
-                    add_image_shape_token=source.get('image_shape', add_image_shape_token),
-                    base_size=source.get('base_size'),
-                    ratio_idx=source.get('ratio_idx'),
-                    image_type=key,
-                )
-                if not source.get('front_boi', use_front_boi_token):
-                    token_seq.append(self.boi_token_id)
-                    extra_token_pos["boi"].append(token_count)
-                    token_count += 1
-                if last_key_only_prefix and i == modal_length - 1:
-                    pass  # for AR inference
-                else:
-                    token_seq.extend(
-                        [self.img_token_id] * source['length'] +  # token number
-                        [self.eoi_token_id]
-                    )
-                    extra_token_pos["<img>_start"].append(token_count)
-                    extra_token_pos["<all_img>_start"].append(token_count)
-                    token_count += source['length']
-                    extra_token_pos["<img>_end"].append(token_count - 1)
-                    extra_token_pos["<all_img>_end"].append(token_count - 1)
-                    extra_token_pos["eoi"].append(token_count)
-                    token_count += 1  # <eoi>
-
-            elif key == "joint_image":
-                assert isinstance(source['length'], list) and len(
-                    source['length']) == 2, "joint_image length should be a list of two integers"
-                extra_count = 2 + 1 + (  # boi, eoi, joint_img_sep
-                    1 if source.get('timestep', add_timestep_token) else 0) + (
-                    2 if source.get('image_shape', add_image_shape_token) else 0
-                )
-                if drop_last is True and token_count + extra_count + sum(source['length']) > total_length:
-                    drop_last_break = True
-                    break
-                if source.get('front_boi', use_front_boi_token):
-                    token_seq.append(self.boi_token_id)  # Use patched boi for Janus, otherwise useing default <boi>
-                    extra_token_pos["boi"].append(token_count)
-                    token_count += 1
-                token_count = self._add_image_meta_info_token(
-                    token_seq=token_seq,
-                    token_count=token_count,
-                    extra_token_pos=extra_token_pos,
-                    add_timestep_token=source.get('timestep', add_timestep_token),
-                    add_image_shape_token=source.get('image_shape', add_image_shape_token),
-                    base_size=source.get('base_size'),
-                    ratio_idx=source.get('ratio_idx'),
-                    image_type=key,
-                )
-                if not source.get('front_boi', use_front_boi_token):
-                    token_seq.append(self.boi_token_id)
-                    extra_token_pos["boi"].append(token_count)
-                    token_count += 1
-                if last_key_only_prefix and i == modal_length - 1:
-                    pass  # for AR inference
-                else:
-                    token_seq.extend(
-                        [self.img_token_id] * source['length'][0]
-                    )
-                    extra_token_pos["<vae_img>_start"].append(token_count)
-                    extra_token_pos["<joint_img>_start"].append(token_count)
-                    extra_token_pos["<all_img>_start"].append(token_count)
-                    token_count += source['length'][0]
-                    extra_token_pos["<vae_img>_end"].append(token_count - 1)
-                    extra_token_pos["<all_img>_end"].append(token_count - 1)
-
-                    token_seq.extend(
-                        [self.special_token_map["<joint_img_sep>"]]
-                    )
-                    extra_token_pos["joint_img_sep"].append(token_count)
-                    token_count += 1
-
-                    token_seq.extend(
-                        [self.img_token_id] * source['length'][1]
-                    )
-                    extra_token_pos["<vit_img>_start"].append(token_count)
-                    extra_token_pos["<all_img>_start"].append(token_count)
-                    token_count += source['length'][1]
-                    extra_token_pos["<vit_img>_end"].append(token_count - 1)
-                    extra_token_pos["<joint_img>_end"].append(token_count - 1)
-                    extra_token_pos["<all_img>_end"].append(token_count - 1)
-
-                    token_seq.extend(
-                        [self.eoi_token_id]
-                    )
-                    extra_token_pos["eoi"].append(token_count)
-                    token_count += 1  # <eoi>
-
-            else:
-                raise ValueError(f"Not supported key: {key}")
-            index_indicator[key] += 1
-
-        if add_eos is True and not drop_last_break:
-            # Typically used for t2i task.
-            token_seq.append(self.eos_token_id)
-            extra_token_pos["eos"].append(token_count)
-            token_count += 1
-        elif add_eos == 'auto' and not drop_last_break:
-            # Typically used for lm and mmu task.
-            if token_seq[-1] != self.eos_token_id and (total_length is None or token_count < total_length):
-                token_seq.append(self.eos_token_id)
-                extra_token_pos["eos"].append(token_count)
-                token_count += 1
-
-        if total_length:
-            # Check token count and clip sequence if necessary
-            if token_count > total_length and drop_last:
-                # Assert clip position is not in the middle of the block-wise tokens (gen_image, joint_image)
-                for start_key, end_key in [
-                    ("<img>_start", "<img>_end"), ("<joint_img>_start", "<joint_img>_end"),
-                    ("<vae_img>_start", "<vae_img>_end"), ("<vit_img>_start", "<vit_img>_end"),
-                ]:
-                    if start_key in extra_token_pos and end_key in extra_token_pos:
-                        assert all(
-                            (start > total_length or end + 1 < total_length)
-                            for start, end in zip(extra_token_pos[start_key], extra_token_pos[end_key])
-                        ), ("Clip position should not be in the middle of the image tokens.\n"
-                            f"Below is the text:\n{self._shorten_text(self.tokenizer.decode(token_seq))}")
-                token_seq = token_seq[:total_length]
-
-            # Pad the sequence if necessary
-            pad_num = max(0, total_length - len(token_seq))
-            if add_pad and pad_num:
-                token_seq.extend([self.pad_token_id] * pad_num)
-                extra_token_pos["first_pad"].append(token_count)
-
-        return token_seq, extra_token_pos
-
-    def batch_gen_infer(
-            self,
-            infer_fn,
-            prompt_list: list,
-            negative_prompt_list: list = None,
-            infer_fn_kwargs_list: List[Dict[str, int]] = None,
-            do_classifier_free_guidance=False,
-            condition_repeat_times: int = 1,
-            uncondition_repeat_times: int = 1,
-    ):
-        """
-        Batch inference for the AR-like model training of the text-to-image/instruction tuning tasks.
-
-        Parameters
-        ----------
-        infer_fn: callable
-            Inference function to encode the prompt.
-        prompt_list: list
-            List of prompts. Each element can be a single prompt or a list of prompts passed to the infer_fn.
-        negative_prompt_list: list
-            List of negative prompts. Only used when do_classifier_free_guidance is True. If None, will use <cfg>
-            token sequence as negative prompt.
-        infer_fn_kwargs_list: List[Dict[str, int]]
-            List of keyword arguments for the infer_fn.
-        do_classifier_free_guidance: bool
-            Whether to do classifier-free guidance.
-        condition_repeat_times: int
-            Support multi-condition.
-        uncondition_repeat_times: int
-            Support multi-uncondition.
-        """
-        if infer_fn_kwargs_list is None:
-            infer_fn_kwargs_list = [{} for _ in prompt_list]
-
-        # [n_output, bsz]
-        cond_results_list = None
-        uncond_results_list = None
-        output_type_list = []
-
-        for prompt_idx, (prompt, infer_fn_kwargs) in enumerate(zip(prompt_list, infer_fn_kwargs_list)):
-            if not isinstance(prompt, (list, tuple)):
-                prompt = [prompt]
-            cond_kwargs = {"uncond_p": 0.0} if do_classifier_free_guidance else {}
-            results = infer_fn(
-                *prompt,
-                **infer_fn_kwargs,
-                **cond_kwargs,
-            )
-            output_type_list.append((type(results), len(results) if isinstance(results, (list, tuple)) else 1))
-            if isinstance(results, dict):
-                raise ValueError("Make batch on dict is not supported. Please return list or tuple for infer_fn.")
-            if not isinstance(results, (list, tuple)):
-                results = (results,)
-            if cond_results_list is None:
-                cond_results_list = [[] for _ in results]
-                uncond_results_list = [[] for _ in results]
-            for i, result in enumerate(results):
-                cond_results_list[i].append(result)
-
-            if do_classifier_free_guidance:
-                if negative_prompt_list is None:
-                    uncond_kwargs = {"uncond_p": 1.0}
-                    uncond_results = infer_fn(
-                        *prompt,
-                        **infer_fn_kwargs,
-                        **uncond_kwargs,
-                    )
-                else:
-                    negative_prompt = negative_prompt_list[prompt_idx]
-                    if not isinstance(negative_prompt, (list, tuple)):
-                        negative_prompt = [negative_prompt]
-                    uncond_results = infer_fn(
-                        *negative_prompt,
-                        **infer_fn_kwargs,
-                    )
-                if isinstance(uncond_results, TokenizerEncodeOutput):
-                    uncond_results_list.append(uncond_results)
-                else:
-                    for i, result in enumerate(uncond_results):
-                        uncond_results_list[i].append(result)
-
-        assert all(output_type_list[0] == n for n in output_type_list), \
-            f"Number of outputs should be equal for all samples, but got {output_type_list}."
-        output_type, output_num = output_type_list[0]
-
-        def make_batch(batch_cond_item, batch_uncond_item):
-            # Process each output item to make batch
-            first = batch_cond_item[0]  # The first element in the batch
-            if isinstance(first, torch.Tensor):
-                stacked_item = torch.stack(self.pad(
-                    batch_cond_item * condition_repeat_times + batch_uncond_item * uncondition_repeat_times,
-                ))
-
-            elif first is None:
-                assert all(item is None for item in batch_cond_item + batch_uncond_item), \
-                    (f"The first cond item is None, but some items are not None:\n\n"
-                     f"condition: {batch_cond_item}\n\n"
-                     f"uncondition: {batch_uncond_item}")
-                stacked_item = None
-
-            elif isinstance(first, (list, tuple)):
-                # If the output item is a list or tuple, we treat it as a whole, and won't make nested batch any more.
-                stacked_item = batch_cond_item * condition_repeat_times + batch_uncond_item * uncondition_repeat_times
-
-            elif isinstance(first, TokenizerEncodeOutput):
-                stacked_item = {}
-                # Traverse not-None attributes
-                for key in list(first.keys()):
-                    merged_list = [cond_item[key] for cond_item in batch_cond_item] * condition_repeat_times + \
-                                  [uncond_item[key] for uncond_item in batch_uncond_item] * uncondition_repeat_times
-                    if isinstance(first[key], torch.Tensor):
-                        if 'mask' in key:
-                            pad_val = 0.0
-                        elif key == 'tokens':
-                            pad_val = self.special_token_map["<pad>"]
-                        else:
-                            pad_val = False  # Should not pad for other tensors
-                        stacked_item[key] = torch.stack(self.pad(merged_list, pad_val=pad_val), dim=0)
-                    elif isinstance(first[key], list):
-                        stacked_item[key] = merged_list
-                    elif first[key] is None:
-                        pass
-                    else:
-                        raise ValueError(f"Unsupported type of {key}: {type(first[key])}.")
-                stacked_item = TokenizerEncodeOutput(stacked_item)
-
-            else:
-                raise TypeError(f"Making batch on type {type(first)} is not supported.")
-
-            return stacked_item
-
-        stacked_outputs = []
-        for cond_results, uncond_results in zip(cond_results_list, uncond_results_list):
-            stacked_outputs.append(make_batch(cond_results, uncond_results))
-
-        if output_type == list:
-            return stacked_outputs
-        elif output_type == tuple:
-            return tuple(stacked_outputs)
-        elif output_num == 1:
-            return stacked_outputs[0]
-        else:
-            raise ValueError(f"Unsupported output type: {output_type}.")
-
-    @staticmethod
-    def parse_extra_token_pos(extra_token_pos, prefix, tokens, rng=None):
-        if rng is None:
-            rng = slice(None)
-        image_slices = [
-            slice(start, end + 1)
-            for start, end in zip(extra_token_pos[f'<{prefix}>_start'][rng], extra_token_pos[f'<{prefix}>_end'][rng])
-        ] if f'<{prefix}>_start' in extra_token_pos and f'<{prefix}>_end' in extra_token_pos else []
-        if image_slices:
-            image_mask = torch.zeros_like(tokens, dtype=torch.bool)
-            for image_slice in image_slices:
-                image_mask[image_slice] = True
-        else:
-            image_mask = None
-        return image_slices, image_mask
-
-    def encode_general(
-            self,
-            sections: Optional[List[Dict[str, Any]]] = None,
-            max_token_length: Optional[int] = None,
-            add_eos='auto',
-            use_text_mask=True,
-            add_pad='auto',
-            add_bos=True,
-            drop_last='auto',
-    ):
-        """
-        General encode function to encode a sequence with multiple sections of text and images.
-        Each section is a dict with a `type` key and other keys depending on the type.
-        Supported section types:
-        - text: dict with keys:
-            - text: str or List[int], text to be encoded. Either `text` or `tokens` should be provided.
-            - tokens: List[int], pre-encoded text tokens. Either `text` or `tokens` should be provided.
-            - uncond_enabled: bool, whether to enable uncondition for this text section.
-            - uncond_p: float, probability to drop the text section for uncondition.
-            - max_length: int, maximum length of the text section.
-            - ignore: bool, whether to ignore this text section in the text mask.
-            - start_offset: int, start offset of the text mask.
-            - end_offset: int, end offset of the text mask.
-        - gen_image: dict with keys:
-            - token_length: int, number of image tokens.
-            - add_timestep_token: bool, whether to add timestep token before the image tokens.
-            - add_guidance_token: bool, whether to add guidance token before the image tokens.
-            - use_front_boi_token: bool, whether to put the <boi> token at the front of size, ratio and timestep tokens.
-            - add_image_shape_token: bool, whether to add image shape token before the image tokens.
-            - base_size: int, base size of the image.
-            - ratio_idx: int, ratio index of the image.
-        - joint_image: dict with keys:
-            - token_length: List[int], number of image tokens for the two images.
-            - add_timestep_token: bool, whether to add timestep token before the image tokens.
-            - use_front_boi_token: bool, whether to put the <boi> token at the front of size, ratio and timestep tokens.
-            - add_image_shape_token: bool, whether to add image shape token before the image tokens.
-            - base_size: int, base size of the image.
-            - ratio_idx: int, ratio index of the image.
-
-        Parameters
-        ----------
-        sections: List[Dict[str, Any]]
-            List of sections to be encoded.
-        max_token_length: int
-            Maximum length of the encoded token sequence.
-        add_eos: bool or 'auto'
-            Whether to add eos token at the end of the sequence. If True, always add eos
-            token. If 'auto', add eos token only when the total_length is not reached and the last token is not <eos>.
-        use_text_mask: bool
-            Whether to generate text mask.
-        add_pad: bool or 'auto'
-            Whether to add padding tokens to the sequence. If True and total_length is not reached,
-            add padding tokens.
-        add_bos: bool
-            Whether to add bos token at the beginning of the sequence.
-        drop_last: bool or 'auto'
-            - If auto, drop last tokens exceeding the total_length if the total_length is provided.
-            If cut point is in the middle of the image tokens, an error will raised.
-            - If True, drop last tokens exceeding the total_length. If cut point is in the
-            middle of the image tokens, all the successive image tokens will be dropped.
-            - If False, keep the last tokens exceeding the total_length, even if the total_length
-            is reached.
-
-        Returns
-        -------
-        TokenizerEncodeOutput
-            Encoded token sequence and extra information.
-        """
-        if sections is None:
-            raise ValueError("sections must be provided.")
-        template = '-'.join([section['type'] for section in sections])
-
-        sections = deepcopy(sections)
-        token_source = defaultdict(list)
-        text_mask_specs = []
-        for section in sections:
-            if section['type'] == 'text':
-                text = self.encode_text(
-                    section['text'] if 'text' in section else section['tokens'],
-                    uncond_enabled=section.get('uncond_enabled'),
-                    uncond_p=section.get('uncond_p'),
-                    max_length=section.get('max_length'),
-                )
-                token_source['text'].append(text)
-                text_mask_specs.append(dict(
-                    ignore=section.get('ignore', False),
-                    start_offset=section.get('start_offset', 0),
-                    end_offset=section.get('end_offset', 0),
-                ))
-            elif section['type'] == 'gen_image':
-                token_source['gen_image'].append(dict(
-                    length=section['token_length'],
-                    timestep=section.get('add_timestep_token', False),
-                    guidance=section.get('add_guidance_token', False),
-                    front_boi=section.get('use_front_boi_token', False),
-                    image_shape=section.get('add_image_shape_token', False),
-                    base_size=section.get('base_size'),
-                    ratio_idx=section.get('ratio_idx'),
-                ))
-            elif section['type'] == 'joint_image':
-                token_source['joint_image'].append(dict(
-                    length=section['token_length'],
-                    timestep=section.get('add_timestep_token', False),
-                    front_boi=section.get('use_front_boi_token', False),
-                    image_shape=section.get('add_image_shape_token', False),
-                    base_size=section.get('base_size'),
-                    ratio_idx=section.get('ratio_idx'),
-                ))
-            else:
-                raise ValueError(f"Invalid section type: {section['type']}")
-
-        # Combine text and image tokens
-        full_token_seq, extra_token_pos = self.encode_sequence(
-            template=template,
-            token_source=dict(token_source),
-            total_length=max_token_length,
-            add_eos=add_eos,
-            add_pad=add_pad,
-            add_bos=add_bos,
-            drop_last=drop_last,
-        )
-        full_seq_token_tensor = torch.tensor(full_token_seq, dtype=torch.long)
-
-        timestep_scatter_index = torch.tensor(extra_token_pos['timestep'], dtype=torch.long) \
-            if 'timestep' in extra_token_pos else None
-        guidance_scatter_index = torch.tensor(extra_token_pos['guidance'], dtype=torch.long) \
-            if 'guidance' in extra_token_pos else None
-        cond_timestep_scatter_index = torch.tensor(extra_token_pos['cond_timestep'], dtype=torch.long) \
-            if 'cond_timestep' in extra_token_pos else None
-        gen_timestep_scatter_index = torch.tensor(extra_token_pos['gen_timestep'], dtype=torch.long) \
-            if 'gen_timestep' in extra_token_pos else None
-
-        # Gen image mask
-        gen_image_slices, gen_image_mask = self.parse_extra_token_pos(extra_token_pos, 'img', full_seq_token_tensor)
-        # Joint image
-        joint_image_slices, _ = self.parse_extra_token_pos(extra_token_pos, 'joint_img', full_seq_token_tensor)
-        # Conditional vae image
-        cond_vae_image_slices, cond_vae_image_mask = self.parse_extra_token_pos(
-            extra_token_pos, 'vae_img', full_seq_token_tensor)
-        # Conditional vit image
-        cond_vit_image_slices, cond_vit_image_mask = self.parse_extra_token_pos(
-            extra_token_pos, 'vit_img', full_seq_token_tensor)
-        # All image slices (gen_image, joint_image)
-        all_image_slices = [
-            slice(start, end + 1)
-            for start, end in zip(extra_token_pos['<all_img>_start'], extra_token_pos['<all_img>_end'])
-        ] if '<all_img>_start' in extra_token_pos and '<all_img>_end' in extra_token_pos else []
-
-        # Text mask
-        text_slices = [
-            slice(start, end + 1)
-            for start, end in zip(extra_token_pos['<text>_start'], extra_token_pos['<text>_end'])
-        ] if '<text>_start' in extra_token_pos and '<text>_end' in extra_token_pos else []
-        assert len(text_slices) <= len(text_mask_specs), \
-            (f"Number of text slices ({len(text_slices)}) should be less than or equal to "
-             f"number of text mask specs ({len(text_mask_specs)})")
-        if use_text_mask:
-            text_mask = torch.zeros_like(full_seq_token_tensor, dtype=torch.float32)
-            for text_slice, mask_spec in zip(text_slices, text_mask_specs):
-                if not mask_spec['ignore']:
-                    real_slice = slice(
-                        text_slice.start + mask_spec['start_offset'],
-                        text_slice.stop + mask_spec['end_offset']
-                    )
-                    text_mask[real_slice] = 1.0
-        else:
-            text_mask = None
-
-        # real_pos is the first position of the <pad> token
-        real_pos = torch.tensor(extra_token_pos.get('first_pad', [full_seq_token_tensor.shape[0]]), dtype=torch.long)
-
-        return TokenizerEncodeOutput(
-            tokens=full_seq_token_tensor,
-            timestep_scatter_index=timestep_scatter_index,
-            guidance_scatter_index=guidance_scatter_index,
-            text_slices=text_slices,
-            gen_image_slices=gen_image_slices,
-            joint_image_slices=joint_image_slices,
-            cond_vae_image_slices=cond_vae_image_slices,
-            cond_vit_image_slices=cond_vit_image_slices,
-            text_mask=text_mask,
-            gen_image_mask=gen_image_mask,
-            cond_vae_image_mask=cond_vae_image_mask,
-            cond_vit_image_mask=cond_vit_image_mask,
-            real_pos=real_pos,
-            all_image_slices=all_image_slices,
-            cond_timestep_scatter_index=cond_timestep_scatter_index,
-            gen_timestep_scatter_index=gen_timestep_scatter_index,
-        )
-
-    def get_cot_sections(self, cot_text, uncond_kwargs, cot_max_length=None, drop_think=False):
-        if not cot_text:  # None or empty
-            return []
-        if '<think>' in cot_text and '</think>' in cot_text:
-            before_think_sec = cot_text.split('<think>')[0]
-            after_think_sec = cot_text.split('</think>')[1]
-            think_sec = cot_text.split('<think>')[1].split('</think>')[0]
-            return self.get_cot_sections(before_think_sec, uncond_kwargs, drop_think=drop_think) + \
-                ([
-                    dict(type="text", text="<think>"),
-                    dict(type="text", text=think_sec, max_length=cot_max_length, **uncond_kwargs),
-                    dict(type="text", text="</think>")
-                ] if not drop_think else []) + \
-                self.get_cot_sections(after_think_sec, uncond_kwargs, drop_think=drop_think)
-
-        if '<recaption>' in cot_text and '</recaption>' in cot_text:
-            before_recaption_sec = cot_text.split('<recaption>')[0]
-            after_recaption_sec = cot_text.split('</recaption>')[1]
-            recaption_sec = cot_text.split('<recaption>')[1].split('</recaption>')[0]
-            return self.get_cot_sections(before_recaption_sec, uncond_kwargs, drop_think=drop_think) + \
-                [
-                    dict(type="text", text="<recaption>"),
-                    dict(type="text", text=recaption_sec, max_length=cot_max_length, **uncond_kwargs),
-                    dict(type="text", text="</recaption>")
-                ] + \
-                self.get_cot_sections(after_recaption_sec, uncond_kwargs, drop_think=drop_think)
-
-        return [
-            dict(type="text", text=cot_text, **uncond_kwargs),
-        ]
-
-    def apply_general_template(
-            self,
-            message_list,
-            max_length=None,
-            add_assistant_prefix=False,
-            answer="auto",
-            bot_task="auto",
-            sequence_template="instruct",
-            uncond_p=0.0,
-            cfg_factor=1,
-            batchify=False,
-            image_base_size=1024,
-            drop_think=False,
-    ):
-        # If cfg_factor > 1, we need to repeat the unconditioned part
-        if batchify:
-            assert isinstance(message_list[0], list), \
-                f"When batchify is True, message_list should be a list of list, but got [{type(message_list[0])}, ...]."
-            return self.batch_gen_infer(
-                infer_fn=self.apply_general_template,
-                prompt_list=[[]],
-                infer_fn_kwargs_list=[dict(
-                    message_list=message_list_i,
-                    max_length=max_length,
-                    add_assistant_prefix=add_assistant_prefix,
-                    answer=answer,
-                    bot_task=bot_task,
-                    sequence_template=sequence_template,
-                    image_base_size=image_base_size,
-                    drop_think=drop_think,
-                ) for message_list_i in message_list],
-                do_classifier_free_guidance=cfg_factor > 1,
-                condition_repeat_times=1,
-                uncondition_repeat_times=cfg_factor - 1,
-            )
-
-        conv = Conversation()
-        uncond_kwargs = dict(uncond_enabled=uncond_p == 1.0, uncond_p=uncond_p)
-
-        def process_successive_message(_message_list, _cur_message_idx, role, prefix, suffix,
-                                       answer_prefix="", answer_suffix=""):
-            _sub_sections = []
-            while _cur_message_idx < len(message_list) and _message_list[_cur_message_idx]['role'] == role:
-                message = _message_list[_cur_message_idx]
-                if message['type'] == 'text':
-                    text = message['content']
-                    if role == "system":
-                        _sub_sections.append(dict(type="text", text=text))
-                    elif role == "assistant":
-                        if ("<recaption>" in text and "</recaption>" in text) or (
-                                "<think>" in text and "</think>" in text):
-                            _sub_sections.extend(self.get_cot_sections(text, uncond_kwargs, drop_think=drop_think))
-                        else:
-                            _sub_sections.append(dict(type="text", text=text, **uncond_kwargs))
-                    else:
-                        _sub_sections.append(dict(
-                            type="text", text=f"{answer_prefix}{text}{answer_suffix}", **uncond_kwargs))
-                elif message['type'] == 'gen_image':
-                    info = message['content']
-                    assert isinstance(info, ImageInfo), f"Expected ImageInfo, but got {type(info)}"
-                    if role == "assistant":
-                        _sub_sections.append(dict(type="text", text=answer_prefix))
-                    _sub_sections.append(dict(type=message['type'], **info.meta_info))
-                    if role == "assistant":
-                        _sub_sections.append(dict(type="text", text=answer_suffix))
-                elif message['type'] == 'joint_image':
-                    info = message['content']
-                    assert isinstance(info, JointImageInfo), f"Expected JointImageInfo, but got {type(info)}"
-                    _sub_sections.append(dict(type=message['type'], **info.meta_info))
-                else:
-                    raise ValueError(f"Unknown message type: {message['type']}")
-                _cur_message_idx += 1
-            if len(_sub_sections) > 0:
-                # Add role prefix and suffix
-                _sub_sections.insert(0, dict(type='text', text=prefix))
-                _sub_sections.append(dict(type='text', text=suffix))
-            return _sub_sections, _cur_message_idx
-
-        # Define assistant prefix and suffix
-        if (answer == "auto" and sequence_template == "instruct") or answer is True:
-            answer_prefix, answer_suffix = "<answer>", "</answer>"
-        else:
-            answer_prefix, answer_suffix = "", ""
-        if sequence_template == "pretrain":
-            system_suffix = ""
-            user_prefix = ""
-            user_suffix = ""
-            bot_prefix = ""
-            bot_suffix = ""
-        else:
-            system_suffix = f"{conv.sep}"
-            user_prefix = f"{conv.roles[0]}: "
-            user_suffix = f"{conv.sep}"
-            bot_prefix = f"{conv.roles[1]}: "
-            bot_suffix = f"{conv.sep}"
-
-        # Process successive user and assistant messages
-        sections = []
-        cur_message_idx = 0
-        final_role = None
-        while cur_message_idx < len(message_list):
-            # Process successive system messages
-            sub_sections, cur_message_idx = process_successive_message(
-                message_list, cur_message_idx, role="system", prefix="", suffix=system_suffix)
-            # Add to the template and sections
-            sections.extend(sub_sections)
-            if len(sub_sections) > 0:
-                final_role = "system"
-
-            # Process successive user messages
-            sub_sections, cur_message_idx = process_successive_message(
-                message_list, cur_message_idx, role="user", prefix=user_prefix, suffix=user_suffix)
-            # Add to the template and sections
-            sections.extend(sub_sections)
-            if len(sub_sections) > 0:
-                final_role = "user"
-
-            # Process successive assistant messages
-            sub_sections, cur_message_idx = process_successive_message(
-                message_list, cur_message_idx, role="assistant", prefix=bot_prefix, suffix=bot_suffix,
-                answer_prefix=answer_prefix, answer_suffix=answer_suffix,
-            )
-            # Add to the template and sections
-            sections.extend(sub_sections)
-            if len(sub_sections) > 0:
-                final_role = "assistant"
-
-        if add_assistant_prefix:
-            if final_role == "assistant":
-                # Avoid adding prefix twice
-                _bot_prefix = ""
-                # Remove the final bot_suffix
-                if len(sections) > 0 and sections[-1]['type'] == 'text' and sections[-1]['text'] == bot_suffix:
-                    sections = sections[:-1]
-            else:
-                _bot_prefix = bot_prefix
-            # We can add special tokens for the bot lastest message according to different tasks
-            bot_response_prefix = dict(
-                auto=_bot_prefix,
-                image="",
-                think=f"{_bot_prefix}<think>",
-                recaption=f"{_bot_prefix}<recaption>",
-                img_ratio=f"{_bot_prefix}{answer_prefix}<boi><img_size_{image_base_size}>",
-            )[bot_task]
-            sections.append(dict(type='text', text=bot_response_prefix))
-
-        output = self.encode_general(
-            sections=sections,
-            use_text_mask=False,
-            add_eos=False,
-            add_pad=False,
-        )
-
-        if max_length is not None:
-            if output.tokens.shape[-1] > max_length:
-                raise ValueError(
-                    f"Encoded token length {output.tokens.shape[-1]} exceeds max_length {max_length}.\n"
-                    f"Please set a larger max_length or check the input messages:\n{message_list}"
-                )
-
-        return output, sections
-
-    def apply_chat_template(
-            self,
-            batch_prompt: Optional[List[str]] = None,
-            batch_message_list: Optional[List[List[Dict[str, Any]]]] = None,
-            mode: str = "gen_text",
-            batch_gen_image_info: Optional[List[ImageInfo]] = None,
-            batch_cond_image_info: Optional[Union[List[JointImageInfo], List[List[JointImageInfo]]]] = None,
-            batch_system_prompt: Optional[List[str]] = None,
-            batch_cot_text: Optional[List[str]] = None,
-            max_length: Optional[int] = None,
-            bot_task: str = "auto",    # auto/image/think/recaption/img_ratio
-            image_base_size: int = 1024,
-            sequence_template: str = "pretrain",
-            cfg_factor: int = 1,
-            add_assistant_prefix: Optional[bool] = None,
-            drop_think: bool = False,
-    ) -> Dict[str, Any]:
-        assert bot_task in ["image", "auto", "think", "recaption", "img_ratio"], \
-            f"bot_task should be one of ['image', 'auto', 'think', 'recaption', 'img_ratio'], but got {bot_task}."
-
-        if batch_message_list is None:
-            # Simple text-to-image or text-cot-to-image task
-            batch_size = len(batch_prompt)
-
-            # Batchify inputs
-            if not isinstance(batch_system_prompt, list):
-                batch_system_prompt = [batch_system_prompt] * batch_size
-            if not isinstance(batch_gen_image_info, list):
-                batch_gen_image_info = [batch_gen_image_info] * batch_size
-            if batch_cot_text is not None:
-                assert len(batch_cot_text) == batch_size, \
-                    (f"batch_cot_text should have the same length as batch_size ({batch_size}), "
-                     f"but got {len(batch_cot_text)}.")
-            else:
-                batch_cot_text = [None] * batch_size
-            if batch_cond_image_info is not None:
-                assert len(batch_cond_image_info) == batch_size, \
-                    (f"batch_cond_image_info should have the same length as batch_size ({batch_size}), "
-                     f"but got {len(batch_cond_image_info)}.")
-                batch_cond_image_info = [
-                    cond_image_info if isinstance(cond_image_info, list) else [cond_image_info]
-                    for cond_image_info in batch_cond_image_info
-                ]
-            else:
-                batch_cond_image_info = [[] for _ in range(batch_size)]
-
-            # Convert single round materials into standard message list
-            batch_message_list = []
-            for prompt, system_prompt, cot_text, gen_image_info, cond_image_info_list in zip(
-                    batch_prompt, batch_system_prompt, batch_cot_text, batch_gen_image_info,
-                    batch_cond_image_info,
-            ):
-                message_list = []
-                # 1. system prompt section
-                if system_prompt:
-                    message_list.append(dict(
-                        role="system", type="text", content=system_prompt, context_type="str"))
-                # 2. user inputs sections
-                #   2.1 image inputs
-                if len(cond_image_info_list) > 0:
-                    message_list.extend([
-                        dict(role="user", type="joint_image", content=cond_image_info, context_type="image_info")
-                        for cond_image_info in cond_image_info_list
-                    ])
-                #   2.2 text inputs
-                message_list.append(dict(
-                    role="user", type="text", content=prompt, context_type="str"))
-                # 3. assistant answer sections
-                if cot_text is not None:
-                    message_list.append(dict(role="assistant", type="text", content=cot_text, context_type="str"))
-                if mode == "gen_image":
-                    message_list.append(dict(
-                        role="assistant", type="gen_image", content=gen_image_info, context_type="image_info"))
-                # ---
-                batch_message_list.append(message_list)
-
-        output, sections = self.apply_general_template(
-            message_list=batch_message_list,
-            max_length=max_length,
-            add_assistant_prefix=default(add_assistant_prefix, mode != "gen_image"),
-            bot_task=bot_task,
-            sequence_template=sequence_template,
-            cfg_factor=cfg_factor,
-            batchify=True,
-            image_base_size=image_base_size,
-            drop_think=drop_think,
-        )
-        return dict(output=output, sections=sections)