Upload folder using huggingface_hub

added_tokens.json

@@ -0,0 +1,295 @@
+{
+  "</box>": 151677,
+  "</img>": 151671,
+  "</quad>": 151673,
+  "</ref>": 151675,
+  "</think>": 151668,
+  "</tool_call>": 151658,
+  "</tool_response>": 151666,
+  "<FAKE_PAD_0>": 151682,
+  "<FAKE_PAD_100>": 151782,
+  "<FAKE_PAD_101>": 151783,
+  "<FAKE_PAD_102>": 151784,
+  "<FAKE_PAD_103>": 151785,
+  "<FAKE_PAD_104>": 151786,
+  "<FAKE_PAD_105>": 151787,
+  "<FAKE_PAD_106>": 151788,
+  "<FAKE_PAD_107>": 151789,
+  "<FAKE_PAD_108>": 151790,
+  "<FAKE_PAD_109>": 151791,
+  "<FAKE_PAD_10>": 151692,
+  "<FAKE_PAD_110>": 151792,
+  "<FAKE_PAD_111>": 151793,
+  "<FAKE_PAD_112>": 151794,
+  "<FAKE_PAD_113>": 151795,
+  "<FAKE_PAD_114>": 151796,
+  "<FAKE_PAD_115>": 151797,
+  "<FAKE_PAD_116>": 151798,
+  "<FAKE_PAD_117>": 151799,
+  "<FAKE_PAD_118>": 151800,
+  "<FAKE_PAD_119>": 151801,
+  "<FAKE_PAD_11>": 151693,
+  "<FAKE_PAD_120>": 151802,
+  "<FAKE_PAD_121>": 151803,
+  "<FAKE_PAD_122>": 151804,
+  "<FAKE_PAD_123>": 151805,
+  "<FAKE_PAD_124>": 151806,
+  "<FAKE_PAD_125>": 151807,
+  "<FAKE_PAD_126>": 151808,
+  "<FAKE_PAD_127>": 151809,
+  "<FAKE_PAD_128>": 151810,
+  "<FAKE_PAD_129>": 151811,
+  "<FAKE_PAD_12>": 151694,
+  "<FAKE_PAD_130>": 151812,
+  "<FAKE_PAD_131>": 151813,
+  "<FAKE_PAD_132>": 151814,
+  "<FAKE_PAD_133>": 151815,
+  "<FAKE_PAD_134>": 151816,
+  "<FAKE_PAD_135>": 151817,
+  "<FAKE_PAD_136>": 151818,
+  "<FAKE_PAD_137>": 151819,
+  "<FAKE_PAD_138>": 151820,
+  "<FAKE_PAD_139>": 151821,
+  "<FAKE_PAD_13>": 151695,
+  "<FAKE_PAD_140>": 151822,
+  "<FAKE_PAD_141>": 151823,
+  "<FAKE_PAD_142>": 151824,
+  "<FAKE_PAD_143>": 151825,
+  "<FAKE_PAD_144>": 151826,
+  "<FAKE_PAD_145>": 151827,
+  "<FAKE_PAD_146>": 151828,
+  "<FAKE_PAD_147>": 151829,
+  "<FAKE_PAD_148>": 151830,
+  "<FAKE_PAD_149>": 151831,
+  "<FAKE_PAD_14>": 151696,
+  "<FAKE_PAD_150>": 151832,
+  "<FAKE_PAD_151>": 151833,
+  "<FAKE_PAD_152>": 151834,
+  "<FAKE_PAD_153>": 151835,
+  "<FAKE_PAD_154>": 151836,
+  "<FAKE_PAD_155>": 151837,
+  "<FAKE_PAD_156>": 151838,
+  "<FAKE_PAD_157>": 151839,
+  "<FAKE_PAD_158>": 151840,
+  "<FAKE_PAD_159>": 151841,
+  "<FAKE_PAD_15>": 151697,
+  "<FAKE_PAD_160>": 151842,
+  "<FAKE_PAD_161>": 151843,
+  "<FAKE_PAD_162>": 151844,
+  "<FAKE_PAD_163>": 151845,
+  "<FAKE_PAD_164>": 151846,
+  "<FAKE_PAD_165>": 151847,
+  "<FAKE_PAD_166>": 151848,
+  "<FAKE_PAD_167>": 151849,
+  "<FAKE_PAD_168>": 151850,
+  "<FAKE_PAD_169>": 151851,
+  "<FAKE_PAD_16>": 151698,
+  "<FAKE_PAD_170>": 151852,
+  "<FAKE_PAD_171>": 151853,
+  "<FAKE_PAD_172>": 151854,
+  "<FAKE_PAD_173>": 151855,
+  "<FAKE_PAD_174>": 151856,
+  "<FAKE_PAD_175>": 151857,
+  "<FAKE_PAD_176>": 151858,
+  "<FAKE_PAD_177>": 151859,
+  "<FAKE_PAD_178>": 151860,
+  "<FAKE_PAD_179>": 151861,
+  "<FAKE_PAD_17>": 151699,
+  "<FAKE_PAD_180>": 151862,
+  "<FAKE_PAD_181>": 151863,
+  "<FAKE_PAD_182>": 151864,
+  "<FAKE_PAD_183>": 151865,
+  "<FAKE_PAD_184>": 151866,
+  "<FAKE_PAD_185>": 151867,
+  "<FAKE_PAD_186>": 151868,
+  "<FAKE_PAD_187>": 151869,
+  "<FAKE_PAD_188>": 151870,
+  "<FAKE_PAD_189>": 151871,
+  "<FAKE_PAD_18>": 151700,
+  "<FAKE_PAD_190>": 151872,
+  "<FAKE_PAD_191>": 151873,
+  "<FAKE_PAD_192>": 151874,
+  "<FAKE_PAD_193>": 151875,
+  "<FAKE_PAD_194>": 151876,
+  "<FAKE_PAD_195>": 151877,
+  "<FAKE_PAD_196>": 151878,
+  "<FAKE_PAD_197>": 151879,
+  "<FAKE_PAD_198>": 151880,
+  "<FAKE_PAD_199>": 151881,
+  "<FAKE_PAD_19>": 151701,
+  "<FAKE_PAD_1>": 151683,
+  "<FAKE_PAD_200>": 151882,
+  "<FAKE_PAD_201>": 151883,
+  "<FAKE_PAD_202>": 151884,
+  "<FAKE_PAD_203>": 151885,
+  "<FAKE_PAD_204>": 151886,
+  "<FAKE_PAD_205>": 151887,
+  "<FAKE_PAD_206>": 151888,
+  "<FAKE_PAD_207>": 151889,
+  "<FAKE_PAD_208>": 151890,
+  "<FAKE_PAD_209>": 151891,
+  "<FAKE_PAD_20>": 151702,
+  "<FAKE_PAD_210>": 151892,
+  "<FAKE_PAD_211>": 151893,
+  "<FAKE_PAD_212>": 151894,
+  "<FAKE_PAD_213>": 151895,
+  "<FAKE_PAD_214>": 151896,
+  "<FAKE_PAD_215>": 151897,
+  "<FAKE_PAD_216>": 151898,
+  "<FAKE_PAD_217>": 151899,
+  "<FAKE_PAD_218>": 151900,
+  "<FAKE_PAD_219>": 151901,
+  "<FAKE_PAD_21>": 151703,
+  "<FAKE_PAD_220>": 151902,
+  "<FAKE_PAD_221>": 151903,
+  "<FAKE_PAD_222>": 151904,
+  "<FAKE_PAD_223>": 151905,
+  "<FAKE_PAD_224>": 151906,
+  "<FAKE_PAD_225>": 151907,
+  "<FAKE_PAD_226>": 151908,
+  "<FAKE_PAD_227>": 151909,
+  "<FAKE_PAD_228>": 151910,
+  "<FAKE_PAD_229>": 151911,
+  "<FAKE_PAD_22>": 151704,
+  "<FAKE_PAD_230>": 151912,
+  "<FAKE_PAD_231>": 151913,
+  "<FAKE_PAD_232>": 151914,
+  "<FAKE_PAD_233>": 151915,
+  "<FAKE_PAD_234>": 151916,
+  "<FAKE_PAD_235>": 151917,
+  "<FAKE_PAD_236>": 151918,
+  "<FAKE_PAD_237>": 151919,
+  "<FAKE_PAD_238>": 151920,
+  "<FAKE_PAD_239>": 151921,
+  "<FAKE_PAD_23>": 151705,
+  "<FAKE_PAD_240>": 151922,
+  "<FAKE_PAD_241>": 151923,
+  "<FAKE_PAD_242>": 151924,
+  "<FAKE_PAD_243>": 151925,
+  "<FAKE_PAD_244>": 151926,
+  "<FAKE_PAD_245>": 151927,
+  "<FAKE_PAD_246>": 151928,
+  "<FAKE_PAD_247>": 151929,
+  "<FAKE_PAD_248>": 151930,
+  "<FAKE_PAD_249>": 151931,
+  "<FAKE_PAD_24>": 151706,
+  "<FAKE_PAD_250>": 151932,
+  "<FAKE_PAD_251>": 151933,
+  "<FAKE_PAD_252>": 151934,
+  "<FAKE_PAD_253>": 151935,
+  "<FAKE_PAD_25>": 151707,
+  "<FAKE_PAD_26>": 151708,
+  "<FAKE_PAD_27>": 151709,
+  "<FAKE_PAD_28>": 151710,
+  "<FAKE_PAD_29>": 151711,
+  "<FAKE_PAD_2>": 151684,
+  "<FAKE_PAD_30>": 151712,
+  "<FAKE_PAD_31>": 151713,
+  "<FAKE_PAD_32>": 151714,
+  "<FAKE_PAD_33>": 151715,
+  "<FAKE_PAD_34>": 151716,
+  "<FAKE_PAD_35>": 151717,
+  "<FAKE_PAD_36>": 151718,
+  "<FAKE_PAD_37>": 151719,
+  "<FAKE_PAD_38>": 151720,
+  "<FAKE_PAD_39>": 151721,
+  "<FAKE_PAD_3>": 151685,
+  "<FAKE_PAD_40>": 151722,
+  "<FAKE_PAD_41>": 151723,
+  "<FAKE_PAD_42>": 151724,
+  "<FAKE_PAD_43>": 151725,
+  "<FAKE_PAD_44>": 151726,
+  "<FAKE_PAD_45>": 151727,
+  "<FAKE_PAD_46>": 151728,
+  "<FAKE_PAD_47>": 151729,
+  "<FAKE_PAD_48>": 151730,
+  "<FAKE_PAD_49>": 151731,
+  "<FAKE_PAD_4>": 151686,
+  "<FAKE_PAD_50>": 151732,
+  "<FAKE_PAD_51>": 151733,
+  "<FAKE_PAD_52>": 151734,
+  "<FAKE_PAD_53>": 151735,
+  "<FAKE_PAD_54>": 151736,
+  "<FAKE_PAD_55>": 151737,
+  "<FAKE_PAD_56>": 151738,
+  "<FAKE_PAD_57>": 151739,
+  "<FAKE_PAD_58>": 151740,
+  "<FAKE_PAD_59>": 151741,
+  "<FAKE_PAD_5>": 151687,
+  "<FAKE_PAD_60>": 151742,
+  "<FAKE_PAD_61>": 151743,
+  "<FAKE_PAD_62>": 151744,
+  "<FAKE_PAD_63>": 151745,
+  "<FAKE_PAD_64>": 151746,
+  "<FAKE_PAD_65>": 151747,
+  "<FAKE_PAD_66>": 151748,
+  "<FAKE_PAD_67>": 151749,
+  "<FAKE_PAD_68>": 151750,
+  "<FAKE_PAD_69>": 151751,
+  "<FAKE_PAD_6>": 151688,
+  "<FAKE_PAD_70>": 151752,
+  "<FAKE_PAD_71>": 151753,
+  "<FAKE_PAD_72>": 151754,
+  "<FAKE_PAD_73>": 151755,
+  "<FAKE_PAD_74>": 151756,
+  "<FAKE_PAD_75>": 151757,
+  "<FAKE_PAD_76>": 151758,
+  "<FAKE_PAD_77>": 151759,
+  "<FAKE_PAD_78>": 151760,
+  "<FAKE_PAD_79>": 151761,
+  "<FAKE_PAD_7>": 151689,
+  "<FAKE_PAD_80>": 151762,
+  "<FAKE_PAD_81>": 151763,
+  "<FAKE_PAD_82>": 151764,
+  "<FAKE_PAD_83>": 151765,
+  "<FAKE_PAD_84>": 151766,
+  "<FAKE_PAD_85>": 151767,
+  "<FAKE_PAD_86>": 151768,
+  "<FAKE_PAD_87>": 151769,
+  "<FAKE_PAD_88>": 151770,
+  "<FAKE_PAD_89>": 151771,
+  "<FAKE_PAD_8>": 151690,
+  "<FAKE_PAD_90>": 151772,
+  "<FAKE_PAD_91>": 151773,
+  "<FAKE_PAD_92>": 151774,
+  "<FAKE_PAD_93>": 151775,
+  "<FAKE_PAD_94>": 151776,
+  "<FAKE_PAD_95>": 151777,
+  "<FAKE_PAD_96>": 151778,
+  "<FAKE_PAD_97>": 151779,
+  "<FAKE_PAD_98>": 151780,
+  "<FAKE_PAD_99>": 151781,
+  "<FAKE_PAD_9>": 151691,
+  "<IMG_CONTEXT>": 151669,
+  "<box>": 151676,
+  "<img>": 151670,
+  "<quad>": 151672,
+  "<ref>": 151674,
+  "<think>": 151667,
+  "<tool_call>": 151657,
+  "<tool_response>": 151665,
+  "<|action_end|>": 151679,
+  "<|action_start|>": 151678,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|interpreter|>": 151681,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|plugin|>": 151680,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

config.json

@@ -0,0 +1,99 @@
+{
+  "architectures": [
+    "NEOChatModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_neo_chat.NEOChatConfig",
+    "AutoModel": "modeling_neo_chat.NEOChatModel",
+    "AutoModelForCausalLM": "modeling_neo_chat.NEOChatModel"
+  },
+  "downsample_ratio": 0.5,
+  "eos_token_id": 151645,
+  "llm_config": {
+    "_name_or_path": null,
+    "architectures": [
+      "Qwen3ForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 151643,
+    "eos_token_id": 151645,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "intermediate_size": 12288,
+    "max_position_embeddings": 262144,
+    "max_position_embeddings_hw": 10000,
+    "max_window_layers": 42,
+    "model_type": "qwen3",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 42,
+    "num_key_value_heads": 8,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": null,
+    "rope_theta": 5000000.0,
+    "rope_theta_hw": 10000.0,
+    "sliding_window": null,
+    "torch_dtype": "bfloat16",
+    "use_cache": false,
+    "use_deepep": false,
+    "use_sliding_window": false,
+    "vocab_size": 151936,
+    "pure_llm": false
+  },
+  "model_type": "neo_chat",
+  "pad_token_id": 151643,
+  "template": "neo1_0",
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.37.2",
+  "use_backbone_lora": 0,
+  "use_llm_lora": 0,
+  "min_pixels": 65536,
+  "max_pixels": 16777216,
+  "patch_size": 16,
+  "timestep_shift": 1.0,
+  "time_schedule": "standard",
+  "time_shift_type": "exponential",
+  "base_shift": 0.5,
+  "max_shift": 1.15,
+  "base_image_seq_len": 64,
+  "max_image_seq_len": 4096,
+  "noise_scale_mode": "resolution",
+  "noise_scale_base_image_seq_len": 64,
+  "add_noise_scale_embedding": true,
+  "noise_scale_max_value": 8.0,
+  "noise_scale": 1.0,
+  "P_mean": -0.8,
+  "P_std": 0.8,
+  "t_eps": 0.05,
+  "fm_head_dim": 1536,
+  "fm_head_layers": 2,
+  "fm_head_mlp_ratio": 1,
+  "extra_num_layers_post": 0,
+  "concat_time_token_num": 0,
+  "use_pixel_head": false,
+  "use_adaLN": false,
+  "vision_config": {
+    "architectures": [
+      "NEOVisionModel"
+    ],
+    "attention_dropout": 0.0,
+    "auto_map": {
+      "AutoConfig": "configuration_neo_vit.NEOVisionConfig",
+      "AutoModel": "modeling_neo_vit.NEOVisionModel"
+    },
+    "llm_hidden_size": 4096,
+    "downsample_ratio": 0.5,
+    "hidden_size": 1024,
+    "model_type": "neo_vision",
+    "rope_theta_vision": 10000.0,
+    "max_position_embeddings_vision": 10000,
+    "num_channels": 3,
+    "patch_size": 16,
+    "torch_dtype": "bfloat16",
+    "transformers_version": "4.37.2",
+    "min_pixels": 65536,
+    "max_pixels": 16777216
+  }
+}

configuration_neo_chat.py

@@ -0,0 +1,77 @@
+import copy
+
+from transformers import Qwen3Config
+from transformers.utils import logging
+from transformers.configuration_utils import PretrainedConfig
+
+from .configuration_neo_vit import NEOVisionConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class NEOLLMConfig(Qwen3Config):
+    def __init__(self, rope_theta_hw=10000.0, max_position_embeddings_hw=10000, **kwargs):
+        super().__init__(**kwargs)
+        self.rope_theta_hw = rope_theta_hw
+        self.max_position_embeddings_hw = max_position_embeddings_hw
+
+
+class NEOChatConfig(PretrainedConfig):
+    model_type = 'neo_chat'
+    is_composition = True
+
+    def __init__(
+        self,
+        vision_config=None,
+        llm_config=None,
+        use_backbone_lora=0,
+        use_llm_lora=0,
+        downsample_ratio=0.5,
+        template=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if vision_config is None:
+            vision_config = {'architectures': ['NEOVisionModel']}
+            logger.info('vision_config is None. Initializing the NEOVisionConfig with default values.')
+
+        if llm_config is None:
+            llm_config = {'architectures': ['Qwen3ForCausalLM']}
+            logger.info('llm_config is None. Initializing the LlamaConfig config with default values (`LlamaConfig`).')
+        assert 'architectures' in llm_config, "Should specify architecture in llm_config"
+
+        if isinstance(vision_config, dict):
+            self.vision_config = NEOVisionConfig(**vision_config)
+        else:
+            self.vision_config = vision_config
+
+        if isinstance(llm_config, dict):
+            self.llm_config = NEOLLMConfig(**llm_config)
+        else:
+            self.llm_config = llm_config
+
+        self.use_backbone_lora = use_backbone_lora
+        self.use_llm_lora = use_llm_lora
+        self.downsample_ratio = downsample_ratio
+        self.template = template
+        self.tie_word_embeddings = self.llm_config.tie_word_embeddings
+
+    def to_dict(self):
+        """
+        Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`].
+
+        Returns:
+            `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
+        """
+        output = copy.deepcopy(self.__dict__)
+        output['vision_config'] = self.vision_config.to_dict()
+        output['llm_config'] = self.llm_config.to_dict()
+        output['model_type'] = self.__class__.model_type
+        output['use_backbone_lora'] = self.use_backbone_lora
+        output['use_llm_lora'] = self.use_llm_lora
+        output['downsample_ratio'] = self.downsample_ratio
+        output['template'] = self.template
+
+        return output

configuration_neo_vit.py

@@ -0,0 +1,52 @@
+import os
+from typing import Union
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class NEOVisionConfig(PretrainedConfig):
+
+    model_type = 'neo_vision'
+
+    def __init__(
+            self,
+            num_channels=3,
+            patch_size=16,
+            hidden_size=1024,
+            llm_hidden_size=2048,
+            downsample_ratio=0.5,
+            rope_theta_vision=10000.0,
+            max_position_embeddings_vision=10000,
+            min_pixels=65536,
+            max_pixels=4194304,
+            **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.hidden_size = hidden_size
+        self.llm_hidden_size = llm_hidden_size,
+        self.downsample_ratio = downsample_ratio,
+        self.rope_theta_vision = rope_theta_vision
+        self.max_position_embeddings_vision = max_position_embeddings_vision
+        self.num_channels = num_channels
+        self.patch_size = patch_size
+        self.min_pixels = min_pixels
+        self.max_pixels = max_pixels
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> 'PretrainedConfig':
+        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+
+        if 'vision_config' in config_dict:
+            config_dict = config_dict['vision_config']
+
+        if 'model_type' in config_dict and hasattr(cls, 'model_type') and config_dict['model_type'] != cls.model_type:
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f'{cls.model_type}. This is not supported for all configurations of models and can yield errors.'
+            )
+
+        return cls.from_dict(config_dict, **kwargs)

conversation.py

@@ -0,0 +1,397 @@
+"""
+Conversation prompt templates.
+
+We kindly request that you import fastchat instead of copying this file if you wish to use it.
+If you have changes in mind, please contribute back so the community can benefit collectively and continue to maintain these valuable templates.
+
+Modified from https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
+"""
+
+import dataclasses
+from enum import IntEnum, auto
+from typing import Dict, List, Tuple, Union
+
+
+class SeparatorStyle(IntEnum):
+    """Separator styles."""
+
+    ADD_COLON_SINGLE = auto()
+    ADD_COLON_TWO = auto()
+    ADD_COLON_SPACE_SINGLE = auto()
+    NO_COLON_SINGLE = auto()
+    NO_COLON_TWO = auto()
+    ADD_NEW_LINE_SINGLE = auto()
+    LLAMA2 = auto()
+    CHATGLM = auto()
+    CHATML = auto()
+    CHATINTERN = auto()
+    DOLLY = auto()
+    RWKV = auto()
+    PHOENIX = auto()
+    ROBIN = auto()
+    FALCON_CHAT = auto()
+    CHATGLM3 = auto()
+    INTERNVL_ZH = auto()
+    MPT = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that manages prompt templates and keeps all conversation history."""
+
+    # The name of this template
+    name: str
+    # The template of the system prompt
+    system_template: str = '{system_message}'
+    # The system message
+    system_message: str = ''
+    # The names of two roles
+    roles: Tuple[str] = ('USER', 'ASSISTANT')
+    # All messages. Each item is (role, message).
+    messages: List[List[str]] = ()
+    # The number of few shot examples
+    offset: int = 0
+    # The separator style and configurations
+    sep_style: SeparatorStyle = SeparatorStyle.ADD_COLON_SINGLE
+    sep: str = '\n'
+    sep2: str = None
+    # Stop criteria (the default one is EOS token)
+    stop_str: Union[str, List[str]] = None
+    # Stops generation if meeting any token in this list
+    stop_token_ids: List[int] = None
+
+    def get_prompt(self) -> str:
+        """Get the prompt for generation."""
+        if self.system_message is not None and self.system_message != '':
+            system_prompt = self.system_template.format(system_message=self.system_message)
+        else:
+            system_prompt = ''
+        
+        if self.sep_style == SeparatorStyle.ADD_COLON_SINGLE:
+            ret = '' if system_prompt == '' else system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = '' if system_prompt == '' else system_prompt + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ': ' + message + seps[i % 2]
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_SPACE_SINGLE:
+            ret = '' if system_prompt == '' else system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ': '  # must be end with a space
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_NEW_LINE_SINGLE:
+            ret = '' if system_prompt == '' else system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + message + self.sep
+                else:
+                    ret += role + '\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_SINGLE:
+            ret = system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + message + seps[i % 2]
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.RWKV:
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += (
+                        role
+                        + ': '
+                        + message.replace('\r\n', '\n').replace('\n\n', '\n')
+                    )
+                    ret += '\n\n'
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.LLAMA2:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt if system_prompt != '' else '[INST] '
+            for i, (role, message) in enumerate(self.messages):
+                tag = self.roles[i % 2]
+                if message:
+                    if i == 0:
+                        ret += message + ' '
+                    else:
+                        ret += tag + ' ' + message + seps[i % 2]
+                else:
+                    ret += tag
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATGLM:
+            # source: https://huggingface.co/THUDM/chatglm-6b/blob/1d240ba371910e9282298d4592532d7f0f3e9f3e/modeling_chatglm.py#L1302-L1308
+            # source2: https://huggingface.co/THUDM/chatglm2-6b/blob/e186c891cf64310ac66ef10a87e6635fa6c2a579/modeling_chatglm.py#L926
+            round_add_n = 1 if self.name == 'chatglm2' else 0
+            ret = '' if system_prompt == '' else system_prompt + self.sep
+
+            for i, (role, message) in enumerate(self.messages):
+                if i % 2 == 0:
+                    ret += f'[Round {i//2 + round_add_n}]{self.sep}'
+
+                if message:
+                    ret += f'{role}：{message}{self.sep}'
+                else:
+                    ret += f'{role}：'
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATML:
+            ret = '' if system_prompt == '' else system_prompt + self.sep + '\n'
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + message + self.sep + '\n'
+                else:
+                    ret += role + '\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATGLM3:
+            ret = system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + ' ' + message
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATINTERN:
+            # source: https://huggingface.co/internlm/internlm-chat-7b-8k/blob/bd546fa984b4b0b86958f56bf37f94aa75ab8831/modeling_internlm.py#L771
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                # if i % 2 == 0:
+                #     ret += "<s>"
+                if message:
+                    ret += role + ':' + message + seps[i % 2] + '\n'
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.DOLLY:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ':\n' + message + seps[i % 2]
+                    if i % 2 == 1:
+                        ret += '\n\n'
+                else:
+                    ret += role + ':\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.PHOENIX:
+            ret = system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + '<s>' + message + '</s>'
+                else:
+                    ret += role + ': ' + '<s>'
+            return ret
+        elif self.sep_style == SeparatorStyle.ROBIN:
+            ret = '' if system_prompt == '' else system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ':\n' + message + self.sep
+                else:
+                    ret += role + ':\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.FALCON_CHAT:
+            ret = '' if system_prompt == '' else system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ':'
+
+            return ret
+        elif self.sep_style == SeparatorStyle.INTERNVL_ZH:
+            seps = [self.sep, self.sep2]
+            ret = '' if system_prompt == '' else self.system_message + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ': ' + message + seps[i % 2]
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = '' if system_prompt == '' else system_prompt + self.sep
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    if i != len(self.messages) - 1 and message is not None:
+                        ret += role + self.sep
+                    else:
+                        ret += role
+            return ret
+        else:
+            raise ValueError(f'Invalid style: {self.sep_style}')
+
+    def set_system_message(self, system_message: str):
+        """Set the system message."""
+        self.system_message = system_message
+
+    def append_message(self, role: str, message: str):
+        """Append a new message."""
+        self.messages.append([role, message])
+
+    def update_last_message(self, message: str):
+        """Update the last output.
+
+        The last message is typically set to be None when constructing the prompt,
+        so we need to update it in-place after getting the response from a model.
+        """
+        self.messages[-1][1] = message
+
+    def to_gradio_chatbot(self):
+        """Convert the conversation to gradio chatbot format."""
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def to_openai_api_messages(self):
+        """Convert the conversation to OpenAI chat completion format."""
+        ret = [{'role': 'system', 'content': self.system_message}]
+
+        for i, (_, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                ret.append({'role': 'user', 'content': msg})
+            else:
+                if msg is not None:
+                    ret.append({'role': 'assistant', 'content': msg})
+        return ret
+
+    def copy(self):
+        return Conversation(
+            name=self.name,
+            system_template=self.system_template,
+            system_message=self.system_message,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            stop_str=self.stop_str,
+            stop_token_ids=self.stop_token_ids,
+        )
+
+    def dict(self):
+        return {
+            'template_name': self.name,
+            'system_message': self.system_message,
+            'roles': self.roles,
+            'messages': self.messages,
+            'offset': self.offset,
+        }
+
+
+# A global registry for all conversation templates
+conv_templates: Dict[str, Conversation] = {}
+
+
+def register_conv_template(template: Conversation, override: bool = False):
+    """Register a new conversation template."""
+    if not override:
+        assert (
+            template.name not in conv_templates
+        ), f'{template.name} has been registered.'
+
+    conv_templates[template.name] = template
+
+
+def get_conv_template(name: str) -> Conversation:
+    """Get a conversation template."""
+    return conv_templates[name].copy()
+
+
+# Both Hermes-2 and neo1_0-chat are chatml-format conversation templates. The difference
+# is that during training, the preprocessing function for the Hermes-2 template doesn't add
+# <s> at the beginning of the tokenized sequence, while the neo1_0-chat template does.
+# Therefore, they are completely equivalent during inference.
+
+register_conv_template(
+    Conversation(
+        name='Hermes-2',
+        system_template='<|im_start|>system\n{system_message}',
+        # note: The new system prompt was not used here to avoid changes in benchmark performance.
+        # system_message='我是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。',
+        system_message='你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。',
+        roles=('<|im_start|>user\n', '<|im_start|>assistant\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|im_end|>',
+        stop_str='<|endoftext|>',
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name='internlm2-chat',
+        system_template='<|im_start|>system\n{system_message}',
+        # note: The new system prompt was not used here to avoid changes in benchmark performance.
+        # system_message='我是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。',
+        system_message='你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。',
+        roles=('<|im_start|>user\n', '<|im_start|>assistant\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|im_end|>',
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name='phi3-chat',
+        system_template='<|system|>\n{system_message}',
+        # note: The new system prompt was not used here to avoid changes in benchmark performance.
+        # system_message='我是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。',
+        system_message='你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。',
+        roles=('<|user|>\n', '<|assistant|>\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|end|>',
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name='internvl2_5',
+        system_template='<|im_start|>system\n{system_message}',
+        system_message='你是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。',
+        roles=('<|im_start|>user\n', '<|im_start|>assistant\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|im_end|>\n',
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name='neo1_0',
+        system_template='<|im_start|>system\n{system_message}',
+        system_message='',
+        roles=('<|im_start|>user\n', '<|im_start|>assistant\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|im_end|>\n',
+    )
+)

model-00001-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aba7fb8fb424afb603c907326df6130c8f5ba2c7522601c71c3cb3348c3572db
+size 4937383712

model-00002-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:127b9c3cdecc1a68c76b28f4ca2228d5f29b0e19abd71bcf71f1d279cf3450bb
+size 4999870648

model-00003-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0c1c5a4c46b4121bdcd7f7c133fa3ff054f1d54b1b2b42eef6efb733eecf1595
+size 4932725608

model-00004-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:07e708038f365bf152dd81a9e0e55b34e8c0713f6f47a0aa788ef71193f0b486
+size 4999868056

model-00005-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f01f7e0ce70a92da7a2e68af86facd02cac91927a0387eee612a5a628314e618
+size 4932728352

model-00006-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:89eb05f76727185d3747b06c6bd029247f7b947b607a733fd3022007fb0e5a3c
+size 4932725624

model-00007-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d0ec946d133991b1bd2963af79548accff18b457746bccc330e30fc35d9651ef
+size 3959659416

model-00008-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c864b190def1e7c10665eb073df8c687d40eac16591123b702ec0fcc6f8eb718
+size 1409869776

modeling_fm_modules.py

@@ -0,0 +1,591 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import math
+from functools import lru_cache
+
+from torch.utils.checkpoint import checkpoint
+def modulate(x, shift, scale=None):
+    if shift is None:
+        return x * (1 + scale)
+    return x * (1 + scale) + shift
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return output * self.weight
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t: torch.Tensor, dim: int, max_period: float = 10000.0):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element. These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        half = dim // 2
+        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
+            device=t.device
+        )
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq.to(self.mlp[0].weight.dtype))
+        return t_emb
+
+class ResBlock(nn.Module):
+
+    def __init__(self, channels, mlp_ratio=1.0):
+        super().__init__()
+        self.channels = channels
+        self.intermediate_size = int(channels * mlp_ratio)
+
+        self.in_ln = nn.LayerNorm(self.channels, eps=1e-6)
+        self.mlp = nn.Sequential(
+            nn.Linear(self.channels, self.intermediate_size),
+            nn.SiLU(),
+            nn.Linear(self.intermediate_size, self.channels),
+        )
+
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(channels, 3 * channels, bias=True))
+
+    def forward(self, x, y):
+        shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(y).chunk(3, dim=-1)
+        h = modulate(self.in_ln(x), shift_mlp, scale_mlp)
+        h = self.mlp(h)
+        return x + gate_mlp * h
+
+# class FinalLayer(nn.Module):
+
+#     def __init__(self, model_channels, out_channels):
+#         super().__init__()
+#         self.norm_final = nn.LayerNorm(model_channels, elementwise_affine=False, eps=1e-6)
+#         self.linear = nn.Linear(model_channels, out_channels, bias=True)
+#         self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(model_channels, 2 * model_channels, bias=True))
+
+#     def forward(self, x, c):
+#         shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+#         x = modulate(self.norm_final(x), shift, scale)
+#         x = self.linear(x)
+#         return x
+
+# class SimpleMLPAdaLN(nn.Module):
+
+#     def __init__(self, input_dim, out_dim, dim=1536, layers=12, mlp_ratio=1.0):
+#         super().__init__()
+#         self.input_dim = input_dim
+#         self.out_dim = out_dim
+#         self.dim = dim
+#         self.layers = layers
+#         self.mlp_ratio = mlp_ratio
+
+#         self.time_embed = TimestepEmbedder(dim)
+#         self.input_proj = nn.Linear(input_dim, dim)
+
+#         res_blocks = []
+#         for _ in range(layers):
+#             res_blocks.append(ResBlock(dim, mlp_ratio))
+#         self.res_blocks = nn.ModuleList(res_blocks)
+
+#         self.final_layer = FinalLayer(dim, out_dim)
+
+#         self.grad_checkpointing = False
+
+#         self.initialize_weights()
+
+#     def initialize_weights(self):
+#         def _basic_init(module):
+#             if isinstance(module, nn.Linear):
+#                 torch.nn.init.xavier_uniform_(module.weight)
+#                 if module.bias is not None:
+#                     nn.init.constant_(module.bias, 0)
+
+#         self.apply(_basic_init)
+
+#         # Initialize timestep embedding MLP
+#         nn.init.normal_(self.time_embed.mlp[0].weight, std=0.02)
+#         nn.init.normal_(self.time_embed.mlp[2].weight, std=0.02)
+
+#         # Zero-out adaLN modulation layers
+#         for block in self.res_blocks:
+#             nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+#             nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+#         # Zero-out output layers
+#         nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+#         nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+#         nn.init.constant_(self.final_layer.linear.weight, 0)
+#         nn.init.constant_(self.final_layer.linear.bias, 0)
+
+#     def forward(self, x, t):
+#         """
+#         x.shape = (bsz, input_dim)
+#         t.shape = (bsz,)
+#         """
+
+#         x = self.input_proj(x)
+#         t = self.time_embed(t)
+
+#         y = t
+
+#         for block in self.res_blocks:
+#             if self.grad_checkpointing and self.training:
+#                 x = checkpoint(block, x, y, use_reentrant=True)
+#             else:
+#                 x = block(x, y)
+
+#         return self.final_layer(x, y)
+
+class FlowMatchingHead(nn.Module):
+
+    def __init__(self, input_dim, out_dim, dim=1536, layers=12, mlp_ratio=1.0):
+        super(FlowMatchingHead, self).__init__()
+        self.net = SimpleMLPAdaLN(input_dim=input_dim, out_dim=out_dim, dim=dim, layers=layers, mlp_ratio=mlp_ratio)
+
+    @property
+    def dtype(self):
+        return self.net.input_proj.weight.dtype
+
+    @property
+    def device(self):
+        return self.net.input_proj.weight.device
+
+    def forward(self, x, t):
+        x = self.net(x, t)
+        return x
+
+
+def precompute_freqs_cis_2d(dim: int, height: int, width:int, theta: float = 10000.0, scale=16.0):
+    # assert  H * H == end
+    # flat_patch_pos = torch.linspace(-1, 1, end) # N = end
+    x_pos = torch.linspace(0, scale, width)
+    y_pos = torch.linspace(0, scale, height)
+    y_pos, x_pos = torch.meshgrid(y_pos, x_pos, indexing="ij")
+    y_pos = y_pos.reshape(-1)
+    x_pos = x_pos.reshape(-1)
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)) # Hc/4
+    x_freqs = torch.outer(x_pos, freqs).float() # N Hc/4
+    y_freqs = torch.outer(y_pos, freqs).float() # N Hc/4
+    x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs)
+    y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs)
+    freqs_cis = torch.cat([x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1) # N,Hc/4,2
+    freqs_cis = freqs_cis.reshape(height*width, -1)
+    return freqs_cis
+
+class NerfEmbedder(nn.Module):
+    def __init__(self, in_channels, hidden_size_input, max_freqs):
+        super().__init__()
+        self.max_freqs = max_freqs
+        self.hidden_size_input = hidden_size_input
+        self.embedder = nn.Sequential(
+            nn.Linear(in_channels+max_freqs**2, hidden_size_input, bias=True),
+        )
+
+    @lru_cache
+    def fetch_pos(self, patch_size, device, dtype):
+        pos = precompute_freqs_cis_2d(self.max_freqs ** 2 * 2, patch_size, patch_size).real
+        pos = pos[None, :, :].to(device=device, dtype=dtype)
+        return pos
+
+
+    def forward(self, inputs):
+        B, P2, C = inputs.shape
+        patch_size = int(P2 ** 0.5)
+        device = inputs.device
+        dtype = inputs.dtype
+        dct = self.fetch_pos(patch_size, device, dtype)
+        dct = dct.repeat(B, 1, 1)
+        inputs = torch.cat([inputs, dct], dim=-1)
+        inputs = self.embedder(inputs)
+        return inputs
+
+class SimpleMLPAdaLN(nn.Module):
+    """
+    The MLP for Diffusion Loss.
+    :param in_channels: channels in the input Tensor.
+    :param model_channels: base channel count for the model.
+    :param out_channels: channels in the output Tensor.
+    :param z_channels: channels in the condition.
+    :param num_res_blocks: number of residual blocks per downsample.
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        model_channels,
+        out_channels,
+        z_channels,
+        num_res_blocks,
+        patch_size,
+        grad_checkpointing=False
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.num_res_blocks = num_res_blocks
+        self.grad_checkpointing = grad_checkpointing
+        self.patch_size = patch_size
+
+        self.cond_embed = nn.Linear(z_channels, patch_size**2*model_channels)
+
+        self.input_proj = nn.Linear(in_channels, model_channels)
+        
+        res_blocks = []
+        for i in range(num_res_blocks):
+            res_blocks.append(ResBlock(
+                model_channels,
+            ))
+
+        self.res_blocks = nn.ModuleList(res_blocks)
+        self.final_layer = FinalLayer(model_channels, out_channels)
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+        # Zero-out adaLN modulation layers
+        for block in self.res_blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        # Zero-out output layers
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    def forward(self, x, c):
+        """
+        Apply the model to an input batch.
+        :param x: an [N x C] Tensor of inputs.
+        :param t: a 1-D batch of timesteps.
+        :param c: conditioning from AR transformer.
+        :return: an [N x C] Tensor of outputs.
+        """
+        x = self.input_proj(x)
+        c = self.cond_embed(c)
+
+        y = c.reshape(-1, self.patch_size**2, self.model_channels)
+
+        for block in self.res_blocks:
+            x = block(x, y)
+
+        return self.final_layer(x)
+
+
+class FinalLayer(nn.Module):
+    """
+    The final layer adopted from DiT.
+    """
+    def __init__(self, model_channels, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(model_channels, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(model_channels, out_channels, bias=True)
+
+    def forward(self, x):
+        x = self.norm_final(x)
+        x = self.linear(x)
+        return x
+
+#################################################################################
+#                   Sine/Cosine Positional Embedding Functions                  #
+#################################################################################
+# https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py
+
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0, pe_interpolation=1.0):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32) / pe_interpolation
+    grid_w = np.arange(grid_size, dtype=np.float32) / pe_interpolation
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+# --------------------------------------------------------
+# Interpolate position embeddings for high-resolution
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# --------------------------------------------------------
+def interpolate_pos_embed(model_path, pe_key: str = "gen_pos_embed", new_len: int = 4096):
+    state_dict = torch.load(model_path, map_location="cpu")
+
+    pos_embed_1d = state_dict[pe_key]
+    _, ori_len, embed_dim = pos_embed_1d.shape
+
+    ori_size = int(ori_len**0.5)
+    new_size = int(new_len**0.5)
+
+    if ori_size != new_size:
+        logger.info("Position interpolate from %dx%d to %dx%d" % (ori_size, ori_size, new_size, new_size))
+        pos_embed_2d = pos_embed_1d.reshape(-1, ori_size, ori_size, embed_dim).permute(0, 3, 1, 2)
+        pos_embed_2d = torch.nn.functional.interpolate(
+            pos_embed_2d, size=(new_size, new_size), mode="bicubic", align_corners=False
+        )
+        pos_embed_1d = pos_embed_2d.permute(0, 2, 3, 1).flatten(1, 2)
+        state_dict[pe_key] = pos_embed_1d
+
+    torch.save(state_dict, model_path)
+
+class PositionEmbedding(nn.Module):
+    def __init__(self, max_num_patch_per_side, hidden_size):
+        super().__init__()
+        self.max_num_patch_per_side = max_num_patch_per_side
+        self.hidden_size = hidden_size
+        self.pos_embed = nn.Parameter(
+            torch.zeros(max_num_patch_per_side ** 2, hidden_size), 
+            requires_grad=False
+        )
+        self._init_weights()
+
+    def _init_weights(self):
+        # Initialize (and freeze) pos_embed by sin-cos embedding:
+        pos_embed = get_2d_sincos_pos_embed(self.hidden_size, self.max_num_patch_per_side)
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float())
+
+    def forward(self, position_ids):
+        return self.pos_embed[position_ids]
+
+
+class ResidualConvBlock(nn.Module):
+    def __init__(self, channels: int):
+        super().__init__()
+        self.block = nn.Sequential(
+            nn.Conv2d(channels, channels, kernel_size=3, padding=1),
+            nn.SiLU(),
+            nn.Conv2d(channels, channels, kernel_size=3, padding=1),
+        )
+        nn.init.zeros_(self.block[2].weight)
+        nn.init.zeros_(self.block[2].bias)
+
+    def forward(self, x):
+        return x + self.block(x)
+
+
+class PostConvSmoother(nn.Module):
+    def __init__(self, in_channels=3, hidden_channels=64, num_blocks=3):
+        super().__init__()
+        self.in_proj = nn.Conv2d(in_channels, hidden_channels, kernel_size=3, padding=1)
+        self.blocks = nn.Sequential(*[ResidualConvBlock(hidden_channels) for _ in range(num_blocks)])
+        self.out_proj = nn.Conv2d(hidden_channels, in_channels, kernel_size=1)
+
+        nn.init.zeros_(self.out_proj.weight)
+        nn.init.zeros_(self.out_proj.bias)
+
+    def forward(self, x):
+        h = self.in_proj(x)
+        h = self.blocks(h)
+        return x + self.out_proj(h)
+
+
+class ProgressiveConvDecoder(nn.Module):
+    def __init__(self, hidden_dim=4096, out_channels=3):
+        super().__init__()
+        
+        # self.proj = nn.Linear(hidden_dim, 1024)
+        # self.act = nn.SiLU()
+        
+        self.up_blocks = nn.ModuleList([
+            nn.Sequential(
+                nn.Upsample(scale_factor=2, mode='nearest'),
+                nn.Conv2d(hidden_dim, 512, kernel_size=3, padding=1),
+                nn.GroupNorm(32, 512),
+                nn.SiLU()
+            ),
+            nn.Sequential(
+                nn.Upsample(scale_factor=2, mode='nearest'),
+                nn.Conv2d(512, 256, kernel_size=3, padding=1),
+                nn.GroupNorm(32, 256),
+                nn.SiLU()
+            ),
+            nn.Sequential(
+                nn.Upsample(scale_factor=2, mode='nearest'),
+                nn.Conv2d(256, 64, kernel_size=3, padding=1),
+                nn.GroupNorm(32, 64),
+                nn.SiLU()
+            ),
+            nn.Sequential(
+                nn.Upsample(scale_factor=2, mode='nearest'),
+                nn.Conv2d(64, 32, kernel_size=3, padding=1),
+                nn.GroupNorm(16, 32),
+                nn.SiLU()
+            ),
+            nn.Sequential(
+                nn.Upsample(scale_factor=2, mode='nearest'),
+                nn.Conv2d(32, 16, kernel_size=3, padding=1),
+                nn.SiLU()
+            )
+        ])
+        
+        self.out_conv = nn.Conv2d(16, out_channels, kernel_size=3, padding=1)
+
+    def forward(self, x_2d):
+        # B, C, H, W = x_2d.shape
+        # x = x_2d.permute(0, 2, 3, 1).contiguous() # (B, H, W, C)
+        # x = self.proj(x)
+        # x = self.act(x)
+        # x = x.permute(0, 3, 1, 2).contiguous()    # (B, 512, H, W)
+        x = x_2d
+        for block in self.up_blocks:
+            x = block(x)
+            
+        out = self.out_conv(x)
+        return out
+
+
+class PatchDecoder_postps(nn.Module):
+    def __init__(self):
+        super().__init__()
+        # layer 1: H/32 -> H/8 (4x upscale)
+      
+        self.conv1 = nn.Conv2d(4096, 4096, kernel_size=3, padding=1)
+        self.ps1 = nn.PixelShuffle(4)
+        self.act1 = nn.GELU()
+
+        # layer 2: H/8 -> H (8x upscale)
+        self.conv2 = nn.Conv2d(256, 192, kernel_size=3, padding=1)
+        self.ps2 = nn.PixelShuffle(8)
+
+    def forward(self, x):
+        # x shape: [B, 4096, H/32, W/32]
+        x = self.ps1(self.act1(self.conv1(x))) # -> [B, 256, H/8, W/8]        
+        x = self.ps2(self.conv2(x)) # -> [B, 3, H, W]
+        return x
+
+
+class PatchDecoder_preps(nn.Module):
+    def __init__(self):
+        super().__init__()
+        # layer 1: H/32 -> H/16 (2x upscale)
+        self.ps1 = nn.PixelShuffle(2)
+        self.conv1 = nn.Conv2d(1024, 1024, kernel_size=3, padding=1)
+        self.act1 = nn.GELU()
+
+        # layer 2: H/16 -> H/8 (2x upscale)
+        self.ps2 = nn.PixelShuffle(2)
+        self.conv2 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
+        self.act2 = nn.GELU()
+
+        # layer 3: H/8 -> H (8x upscale)
+        self.ps3 = nn.PixelShuffle(8)
+        self.conv3 = nn.Conv2d(4, 3, kernel_size=3, padding=1)
+
+    def forward(self, x):
+        # x shape: [B, 4096, H/32, W/32]
+        x = self.act1(self.conv1(self.ps1((x)))) # -> [B, 256, H/16, W/16]
+        x = self.act2(self.conv2(self.ps2((x)))) # -> [B, 256, H/8, W/8]
+        x = self.conv3(self.ps3((x))) # -> [B, 3, H, W]
+        return x
+
+class PatchDecoder_preps1(nn.Module):
+    def __init__(self):
+        super().__init__()
+        # layer 1: H/32 -> H/16 (2x upscale)
+        self.ps1 = nn.PixelShuffle(2)
+        self.conv1 = nn.Conv2d(1024, 1024, kernel_size=3, padding=1)
+        self.act1 = nn.GELU()
+
+        # layer 2: H/16 -> H/8 (2x upscale)
+        self.ps2 = nn.PixelShuffle(2)
+        self.conv2 = nn.Conv2d(256, 192, kernel_size=3, padding=1)
+
+        # layer 3: H/8 -> H (8x upscale)
+        self.ps3 = nn.PixelShuffle(8)
+
+    def forward(self, x):
+        # x shape: [B, 4096, H/32, W/32]
+        x = self.act1(self.conv1(self.ps1((x)))) # -> [B, 256, H/16, W/16]
+        x = self.ps3(self.conv2(self.ps2((x)))) # -> [B, 256, H/8, W/8]
+        return x
+
+class ConvDecoder(nn.Module):
+    def __init__(self, input_dim=4096, hidden_dim=1024):
+        super().__init__()
+        # layer 1: H/32 -> H/16 (2x upscale)
+        self.ps1 = nn.PixelShuffle(2)
+        self.conv1 = nn.Conv2d(input_dim // 4, hidden_dim, kernel_size=3, padding=1)
+        self.act1 = nn.GELU()
+
+        # layer 2: H/16 -> H/8 (2x upscale)
+        self.ps2 = nn.PixelShuffle(2)
+        self.conv2 = nn.Conv2d(hidden_dim // 4, 192, kernel_size=3, padding=1)
+
+        # layer 3: H/8 -> H (8x upscale)
+        self.ps3 = nn.PixelShuffle(8)
+
+    def forward(self, x):
+        x = self.act1(self.conv1(self.ps1((x))))
+        x = self.ps3(self.conv2(self.ps2((x))))
+        return x

modeling_neo_chat.py

@@ -0,0 +1,1866 @@
+from typing import List, Optional, Tuple, Union
+import math
+import os
+import torch.utils.checkpoint
+from torch import nn
+import transformers
+from torch.nn import CrossEntropyLoss
+from transformers import GenerationConfig
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+from .configuration_neo_chat import NEOChatConfig
+from .conversation import get_conv_template
+from .modeling_neo_vit import NEOVisionModel
+from .modeling_qwen3 import Qwen3ForCausalLM, create_block_causal_mask
+from .modeling_fm_modules import PositionEmbedding, TimestepEmbedder, FlowMatchingHead, RMSNorm, NerfEmbedder, SimpleMLPAdaLN, ConvDecoder
+from .utils import load_image_native, SYSTEM_MESSAGE_FOR_GEN
+
+logger = logging.get_logger(__name__)
+
+
+def version_cmp(v1, v2, op='eq'):
+    import operator
+
+    from packaging import version
+    op_func = getattr(operator, op)
+    return op_func(version.parse(v1), version.parse(v2))
+
+def prepare_flash_kv_cache(
+    past_key_values,
+    current_len: int,
+    batch_size: int,
+):
+    """
+    Convert prefix cache from [B, H, S, D] to flash-attn friendly [B, S, H, D],
+    and preallocate full KV buffer for [prefix + current].
+
+    This is done once before denoising loop.
+    """
+    if past_key_values is None:
+        return
+
+    for layer in past_key_values.layers:
+        past_k = layer.keys
+        past_v = layer.values
+
+        if past_k is None or past_v is None:
+            layer.flash_prefix_len = 0
+            layer.flash_total_len = current_len
+            layer.flash_k_cache = None
+            layer.flash_v_cache = None
+            continue
+
+        # original cache layout assumed: [B, H, S, D]
+        past_k_flash = past_k.transpose(1, 2).contiguous()  # [B, S, H, D]
+        past_v_flash = past_v.transpose(1, 2).contiguous()  # [B, S, H, D]
+
+        prefix_len = past_k_flash.shape[1]
+        total_len = prefix_len + current_len
+
+        k_cache = torch.empty(
+            (batch_size, total_len, past_k_flash.shape[2], past_k_flash.shape[3]),
+            device=past_k_flash.device,
+            dtype=past_k_flash.dtype,
+        )
+        v_cache = torch.empty(
+            (batch_size, total_len, past_v_flash.shape[2], past_v_flash.shape[3]),
+            device=past_v_flash.device,
+            dtype=past_v_flash.dtype,
+        )
+
+        k_cache[:, :prefix_len].copy_(past_k_flash)
+        v_cache[:, :prefix_len].copy_(past_v_flash)
+
+        layer.flash_prefix_len = prefix_len
+        layer.flash_total_len = total_len
+        layer.flash_k_cache = k_cache
+        layer.flash_v_cache = v_cache
+
+def clear_flash_kv_cache(past_key_values):
+    if past_key_values is None:
+        return
+    for layer in past_key_values.layers:
+        if hasattr(layer, "flash_prefix_len"):
+            delattr(layer, "flash_prefix_len")
+        if hasattr(layer, "flash_total_len"):
+            delattr(layer, "flash_total_len")
+        if hasattr(layer, "flash_k_cache"):
+            delattr(layer, "flash_k_cache")
+        if hasattr(layer, "flash_v_cache"):
+            delattr(layer, "flash_v_cache")
+
+@torch.cuda.amp.autocast(dtype=torch.float32)
+def optimized_scale(positive_flat, negative_flat):
+
+    # Calculate dot production
+    dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)
+
+    # Squared norm of uncondition
+    squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8
+
+    # st_star = v_cond^T * v_uncond / ||v_uncond||^2
+    st_star = dot_product / squared_norm
+    
+    return st_star
+
+def build_abs_positions_from_grid_hw(grid_hw: torch.Tensor, device=None):
+    """
+    Compute patch coordinates (x, y)
+
+    Args:
+        grid_hw: (B, 2) tensor representing (H, W) per image
+    """
+    device = grid_hw.device
+    B = grid_hw.shape[0]
+
+    # Get the number of patches per image
+    H = grid_hw[:, 0]
+    W = grid_hw[:, 1]
+    N = H * W
+    N_total = N.sum()
+
+    # Create the batch index for each patch (B x patch count)
+    patch_to_sample = torch.repeat_interleave(torch.arange(B, device=device), N)  # (N_total,)
+
+    # Generate intra-image patch index (row-major order)
+    patch_id_within_image = torch.arange(N_total, device=device)
+    patch_id_within_image = patch_id_within_image - torch.cumsum(
+        torch.cat([torch.tensor([0], device=device), N[:-1]]), dim=0
+    )[patch_to_sample]
+
+    # Get H/W for each patch according to its image
+    W_per_patch = W[patch_to_sample]
+    abs_x = patch_id_within_image % W_per_patch
+    abs_y = patch_id_within_image // W_per_patch
+
+    return abs_x, abs_y
+
+
+class NEOChatModel(PreTrainedModel):
+    config_class = NEOChatConfig
+    main_input_name = 'pixel_values'
+    base_model_prefix = 'language_model'
+    _supports_flash_attn_2 = True
+    supports_gradient_checkpointing = True
+    _no_split_modules = [
+        "NEOVisionModel",
+        "Qwen3DecoderLayer",
+    ]
+
+    # support transformers 4.51.+
+    _tp_plan = ''
+
+    def __init__(self, config: NEOChatConfig, vision_model=None, language_model=None, use_flash_attn=True):
+        super().__init__(config)
+
+        assert version_cmp(transformers.__version__, '4.37.0', 'ge')
+        patch_size = config.vision_config.patch_size
+        self.patch_size = patch_size
+        self.template = config.template
+        self.downsample_ratio = config.downsample_ratio
+        config.llm_config._attn_implementation = 'eager'
+
+        if vision_model is not None:
+            self.vision_model = vision_model
+        else:
+            self.vision_model = NEOVisionModel(config.vision_config)
+            vision_model_mot_gen = NEOVisionModel(config.vision_config)
+        if language_model is not None:
+            self.language_model = language_model
+        else:
+            self.language_model = Qwen3ForCausalLM(config.llm_config)
+
+        merge_size = int(1 / self.downsample_ratio)
+        output_dim = 3*(patch_size*merge_size)**2
+        llm_hidden_size = self.config.llm_config.hidden_size
+        self.use_deep_fm_head = self.config.fm_head_layers > 2
+        self.use_pixel_head = self.config.use_pixel_head
+        if self.use_deep_fm_head:
+                fm_head = FlowMatchingHead(llm_hidden_size, output_dim, dim=self.config.fm_head_dim, layers=self.config.fm_head_layers, mlp_ratio=self.config.fm_head_mlp_ratio)
+        else:
+            fm_head = nn.Sequential(
+                    nn.Linear(llm_hidden_size, 4096, bias=True),
+                    nn.GELU(),
+                    nn.Linear(4096, output_dim, bias=True),
+                )
+
+        timestep_embedder = TimestepEmbedder(llm_hidden_size)
+        self.fm_modules = nn.ModuleDict(
+                    {   
+                        "vision_model_mot_gen": vision_model_mot_gen,
+                        "timestep_embedder": timestep_embedder,
+                        "fm_head": fm_head
+                    }
+                )
+
+        if self.use_pixel_head:
+            self.fm_modules["fm_head"] = ConvDecoder(llm_hidden_size)
+
+        self.concat_time_token_num = config.concat_time_token_num
+        self.noise_scale = config.noise_scale
+        self.noise_scale_mode = config.noise_scale_mode
+        self.noise_scale_base_image_seq_len = config.noise_scale_base_image_seq_len
+        self.add_noise_scale_embedding = config.add_noise_scale_embedding
+        self.noise_scale_max_value = config.noise_scale_max_value
+        self.time_schedule = config.time_schedule
+        self.time_shift_type = config.time_shift_type
+        self.base_shift = config.base_shift
+        self.max_shift = config.max_shift
+        self.base_image_seq_len = config.base_image_seq_len
+        self.max_image_seq_len = config.max_image_seq_len
+
+        if self.add_noise_scale_embedding:
+            noise_scale_embedder = TimestepEmbedder(llm_hidden_size)
+            self.fm_modules['noise_scale_embedder'] = noise_scale_embedder
+
+
+        self.img_context_token_id = None
+        self.img_start_token_id = 151670
+        self.last_think_content = ""
+        self.conv_template = get_conv_template(self.template)
+        self.system_message = self.conv_template.system_message
+
+    def forward(
+            self,
+            pixel_values: torch.FloatTensor,
+            input_ids: torch.LongTensor = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            image_flags: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        raise NotImplementedError('forward')
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        image_flags = image_flags.squeeze(-1)
+        input_embeds = self.language_model.get_input_embeddings()(input_ids).clone()
+
+        vit_embeds = self.extract_feature(pixel_values)
+        vit_embeds = vit_embeds[image_flags == 1]
+
+        B, N, C = input_embeds.shape
+        input_embeds = input_embeds.reshape(B * N, C)
+
+        # if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:
+        #     print(f'dynamic ViT batch size: {vit_batch_size}, images per sample: {vit_batch_size / B}, dynamic token length: {N}')
+
+        input_ids = input_ids.reshape(B * N)
+        selected = (input_ids == self.img_context_token_id)
+        try:
+            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)
+        except Exception as e:
+            vit_embeds = vit_embeds.reshape(-1, C)
+            print(f'warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, '
+                  f'vit_embeds.shape={vit_embeds.shape}')
+            n_token = min(selected.sum(), vit_embeds.size(0))
+            input_embeds[selected][:n_token] = input_embeds[selected][:n_token] * 0.0 + vit_embeds[:n_token]
+
+        input_embeds = input_embeds.reshape(B, N, C)
+
+        outputs = self.language_model(
+            inputs_embeds=input_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        logits = outputs.logits
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def extract_feature(self, pixel_values, gen_model=False, grid_hw=None):
+        if gen_model:
+            return self.fm_modules['vision_model_mot_gen'](pixel_values=pixel_values, 
+                                 output_hidden_states=False, 
+                                 return_dict=True, 
+                                 grid_hw=grid_hw).last_hidden_state
+        else:
+            return self.vision_model(pixel_values=pixel_values, 
+                                 output_hidden_states=False, 
+                                 return_dict=True, 
+                                 grid_hw=grid_hw).last_hidden_state
+
+    def batch_chat(self, tokenizer, pixel_values, questions, generation_config, num_patches_list=None,
+                   history=None, return_history=False, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>',
+                   IMG_CONTEXT_TOKEN='<IMG_CONTEXT>', verbose=False, image_counts=None):
+        raise NotImplementedError('batch_chat')
+        if history is not None or return_history:
+            print('Now multi-turn chat is not supported in batch_chat.')
+            raise NotImplementedError
+
+        if image_counts is not None:
+            num_patches_list = image_counts
+            print('Warning: `image_counts` is deprecated. Please use `num_patches_list` instead.')
+
+        img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
+        self.img_context_token_id = img_context_token_id
+
+        if verbose and pixel_values is not None:
+            image_bs = pixel_values.shape[0]
+            print(f'dynamic ViT batch size: {image_bs}')
+
+        queries = []
+        for idx, num_patches in enumerate(num_patches_list):
+            question = questions[idx]
+            if pixel_values is not None and '<image>' not in question:
+                question = '<image>\n' + question
+            template = get_conv_template(self.template)
+            template.system_message = self.system_message
+            template.append_message(template.roles[0], question)
+            template.append_message(template.roles[1], None)
+            query = template.get_prompt()
+
+            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN + IMG_END_TOKEN
+            query = query.replace('<image>', image_tokens, 1)
+            queries.append(query)
+
+        tokenizer.padding_side = 'left'
+        model_inputs = tokenizer(queries, return_tensors='pt', padding=True)
+        input_ids = model_inputs['input_ids'].to(self.device)
+        attention_mask = model_inputs['attention_mask'].to(self.device)
+        eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())
+        generation_config['eos_token_id'] = eos_token_id
+        generation_output = self.generate(
+            pixel_values=pixel_values,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            **generation_config
+        )
+        responses = tokenizer.batch_decode(generation_output, skip_special_tokens=True)
+        responses = [response.split(template.sep.strip())[0].strip() for response in responses]
+        return responses
+    
+    def patchify(self, images, patch_size, channel_first=False):
+        """
+        images: (N, 3, H, W)
+        x: (N, L, patch_size**2 *3)
+        """
+        h, w = images.shape[2] // patch_size, images.shape[3] // patch_size
+        x = images.reshape(shape=(images.shape[0], 3, h, patch_size, w, patch_size))
+
+        if channel_first:
+            x = torch.einsum('nchpwq->nhwcpq', x)
+        else:
+            x = torch.einsum('nchpwq->nhwpqc', x)
+        
+        x = x.reshape(shape=(images.shape[0], h * w, patch_size**2 * 3))
+        return x
+    
+    def unpatchify(sle, x, patch_size, h=None, w=None):
+        """
+        x: (N, L, patch_size**2 *3)
+        images: (N, 3, H, W)
+        """
+        if h is None or w is None:
+            h = w = int(x.shape[1]**.5)
+        else:
+            h = h // patch_size
+            w = w // patch_size        
+        x = x.reshape(shape=(x.shape[0], h, w, patch_size, patch_size, 3))
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        images = x.reshape(shape=(x.shape[0], 3, h * patch_size, w * patch_size))
+        return images
+    
+    def _euler_step(self, v_pred, z, t, t_next):
+        z_next = z + (t_next - t) * v_pred
+        return z_next
+
+    def _calculate_dynamic_mu(self, image_seq_len: int) -> float:
+        denom = self.max_image_seq_len - self.base_image_seq_len
+        if denom == 0:
+            return float(self.base_shift)
+        m = (self.max_shift - self.base_shift) / denom
+        b = self.base_shift - m * self.base_image_seq_len
+        return float(image_seq_len) * m + b
+
+    def _apply_time_schedule(self, t: torch.Tensor, image_seq_len: int, timestep_shift: float) -> torch.Tensor:
+        self.time_schedule = "standard"
+        sigma = 1 - t
+        if timestep_shift != 1:
+            self.time_schedule = "standard"
+        if self.time_schedule == "standard":
+            shift = timestep_shift
+            sigma = shift * sigma / (1 + (shift - 1) * sigma)
+        elif self.time_schedule == "dynamic":
+            mu = self._calculate_dynamic_mu(image_seq_len)
+            mu_t = t.new_tensor(mu)
+            if self.time_shift_type == "exponential":
+                shift = torch.exp(mu_t)
+                sigma = shift * sigma / (1 + (shift - 1) * sigma)
+            elif self.time_shift_type == "linear":
+                sigma = mu_t / (mu_t + (1 / sigma - 1))
+            else:
+                raise ValueError(f"Unsupported time_shift_type: {self.time_shift_type}")
+        else:
+            raise ValueError(f"Unsupported time_schedule: {self.time_schedule}")
+        return 1 - sigma
+
+    def _build_t2i_query(self, prompt_text, system_message=None, append_text=None):
+        template = get_conv_template(self.template)
+        template.system_message = self.system_message if system_message is None else system_message
+        template.append_message(template.roles[0], prompt_text)
+        template.append_message(template.roles[1], None)
+        if append_text is not None:
+            return template.get_prompt() + append_text
+        return template.get_prompt()
+
+    def _build_t2i_text_inputs(self, tokenizer, query: str):
+        model_inputs = tokenizer(query, return_tensors="pt")
+        input_ids = model_inputs["input_ids"].to(self.device)
+
+        t_idx = torch.arange(0, input_ids.shape[1], dtype=torch.long, device=input_ids.device)
+        h_idx = torch.zeros_like(t_idx)
+        w_idx = torch.zeros_like(t_idx)
+        indexes = torch.stack([t_idx, h_idx, w_idx], dim=0)
+
+        attention_mask = {"full_attention": create_block_causal_mask(indexes[0])}
+        return input_ids, indexes, attention_mask
+    
+    def _build_t2i_image_indexes(self, token_h, token_w, text_len, device):
+        t_image = torch.full((token_h * token_w,), text_len, dtype=torch.long, device=device)
+        idx = torch.arange(token_h * token_w, device=device, dtype=torch.long)
+        h_image = idx // token_w
+        w_image = idx % token_w
+        return torch.stack([t_image, h_image, w_image], dim=0)
+    
+    def _t2i_prefix_forward(self, input_ids, indexes, attention_mask):
+        out = self.language_model.model(
+            input_ids=input_ids,
+            indexes=indexes,
+            attention_mask=attention_mask,
+            use_cache=True,
+        )
+        return out.past_key_values, out.last_hidden_state
+
+    def _it2i_prefix_forward(self, input_imbeds, indexes, attention_mask, gen_indicators=None):
+        out = self.language_model.model(
+            inputs_embeds=input_imbeds,
+            indexes=indexes,
+            attention_mask=attention_mask,
+            use_cache=True,
+            image_gen_indicators=gen_indicators.view(1, -1) if gen_indicators is not None else None
+        )
+        return out.past_key_values, out.last_hidden_state
+
+    def _append_text_tokens_to_cache(self, cache, t_idx, input_ids):
+        if input_ids.shape[1] == 0:
+            return t_idx
+
+        device = input_ids.device
+        seq_len = input_ids.shape[1]
+        inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
+
+        t_indexes = torch.arange(t_idx + 1, t_idx + 1 + seq_len, dtype=torch.long, device=device)
+        h_indexes = torch.zeros(seq_len, dtype=torch.long, device=device)
+        w_indexes = torch.zeros(seq_len, dtype=torch.long, device=device)
+        indexes = torch.stack([t_indexes, h_indexes, w_indexes], dim=0)
+
+        past_len = cache.get_seq_length()
+        mask = torch.zeros(1, 1, seq_len, past_len + seq_len, device=device)
+        causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=device))
+        causal_mask = torch.where(causal_mask == 1, 0.0, float('-inf'))
+        mask[:, :, :, past_len:] = causal_mask
+        attention_mask_dict = {"full_attention": mask}
+
+        self.language_model(
+            inputs_embeds=inputs_embeds,
+            indexes=indexes,
+            attention_mask=attention_mask_dict,
+            past_key_values=cache,
+            use_cache=True
+        )
+        return t_idx + seq_len
+
+    def _generate_think(
+        self,
+        tokenizer,
+        prefix_outputs,
+        past_key_values,
+        t_idx,
+        IMG_START_TOKEN,
+        max_think_tokens=1024,
+    ):
+        template = get_conv_template(self.template)
+        eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())
+        think_end_token_id = tokenizer.convert_tokens_to_ids('</think>')
+        think_token_ids = []
+        next_token = torch.argmax(prefix_outputs.logits[:, -1, :], dim=-1)
+
+        for _ in range(max_think_tokens):
+            token_item = next_token.item()
+            if token_item == eos_token_id:
+                break
+            if token_item == think_end_token_id:
+                self.language_model.model.current_index = t_idx
+                outputs = self.language_model(
+                    input_ids=next_token.unsqueeze(0),
+                    past_key_values=past_key_values,
+                    use_cache=True
+                )
+                past_key_values = outputs.past_key_values
+                t_idx += 1
+                think_token_ids.append(token_item)
+                break
+
+            think_token_ids.append(token_item)
+
+            self.language_model.model.current_index = t_idx
+            outputs = self.language_model(
+                input_ids=next_token.unsqueeze(0),
+                past_key_values=past_key_values,
+                use_cache=True
+            )
+            past_key_values = outputs.past_key_values
+            t_idx += 1
+
+            next_token = torch.argmax(outputs.logits[:, -1, :], dim=-1)
+
+        append_ids = tokenizer(
+            '\n\n' + IMG_START_TOKEN,
+            return_tensors='pt',
+            add_special_tokens=False,
+        )['input_ids'].to(self.device)
+        t_idx = self._append_text_tokens_to_cache(past_key_values, t_idx, append_ids)
+
+        think_text = tokenizer.decode(think_token_ids, skip_special_tokens=False)
+
+        return past_key_values, t_idx, think_text
+    
+    def _t2i_predict_v(self, input_embeds, indexes_image, attn_mask, past_key_values, t, z, image_token_num, timestep_embeddings=None, image_size=None):
+        B, L = z.shape[0], z.shape[1]
+
+        outputs = self.language_model.model(
+            inputs_embeds=input_embeds,
+            image_gen_indicators=torch.ones((input_embeds.shape[0], input_embeds.shape[1]), dtype=torch.bool, device=input_embeds.device),
+            indexes=indexes_image,
+            attention_mask=attn_mask,
+            past_key_values=past_key_values,
+            update_cache=False,
+            use_cache=True,
+        )
+
+        if self.use_pixel_head:
+            merge_size = int(1 / self.downsample_ratio)
+            token_h = image_size[1] // (self.patch_size * merge_size)
+            token_w = image_size[0] // (self.patch_size * merge_size)
+
+            img_reshaped = outputs.last_hidden_state[:, -image_token_num:].view(B, token_h, token_w, -1)
+            img_2d = torch.einsum("b h w c -> b c h w", img_reshaped)
+            img_2d = img_2d.contiguous().view(B, -1, token_h, token_w)
+                
+            smoothed_img_2d = self.fm_modules['fm_head'](img_2d)
+                
+            smoothed_reshaped = smoothed_img_2d.view(B, 3, token_h, self.patch_size * merge_size, token_w, self.patch_size * merge_size)
+            smoothed_reshaped = torch.einsum("b c h p w q -> b h w p q c", smoothed_reshaped)
+            out_1d = smoothed_reshaped.contiguous().view(B, L, self.patch_size * merge_size * self.patch_size * merge_size * 3)
+            x_pred = out_1d
+        else:
+            if self.use_deep_fm_head:
+                x_pred = self.fm_modules["fm_head"](
+                outputs.last_hidden_state[:, -image_token_num:].view(B*L, -1), t.repeat(B*L)
+                ).view(B, L, -1)
+            else:
+                x_pred = self.fm_modules["fm_head"](
+                    outputs.last_hidden_state[:, -image_token_num:].view(B, L, -1)
+                ).view(B, L, -1)
+            
+        
+        v_pred = (x_pred - z) / (1 - t).clamp_min(self.config.t_eps)
+        return v_pred
+    
+    def _build_it2i_inputs(self, tokenizer, query, pixel_values=None, grid_hw=None):
+        model_inputs = tokenizer(query, return_tensors="pt")
+        input_ids = model_inputs["input_ids"].to(self.device)
+
+        indexes = self.get_thw_indexes(input_ids[0], grid_hw)
+
+        attention_mask = {"full_attention": create_block_causal_mask(indexes[0])}
+
+        input_embeds = self.language_model.get_input_embeddings()(input_ids)
+        B, N, C = input_embeds.shape
+        if pixel_values is not None:
+            vit_embeds = self.extract_feature(pixel_values, grid_hw=grid_hw)
+            input_embeds = input_embeds.reshape(B * N, C)
+            input_ids = input_ids.reshape(B * N)
+            selected = (input_ids == self.img_context_token_id)
+            assert selected.sum() != 0
+            input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)
+            input_embeds = input_embeds.reshape(B, N, C)
+
+        return input_embeds, indexes, attention_mask
+
+    @torch.no_grad()
+    def interleave_gen_image_only(
+            self,
+            tokenizer,
+            prompt,
+            gt_text,
+            images=None,
+            gt_images=None,
+            cfg_scale=1.0,
+            img_cfg_scale=1.0,
+            cfg_norm='none',
+            max_images=10,
+            enable_timestep_shift=True,
+            timestep_shift=1.0,
+            image_size=(256, 256),
+            num_steps=30,
+            IMG_START_TOKEN='<img>',
+            IMG_END_TOKEN='</img>',
+            IMG_CONTEXT_TOKEN='<IMG_CONTEXT>',
+            method='euler',
+            cfg_interval=(0, 1),
+            t_eps=0.02,
+            verbose=False,
+            system_message='',
+    ):
+        self.img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
+        self.img_start_token_id = tokenizer.convert_tokens_to_ids(IMG_START_TOKEN)
+        self.config.t_eps = t_eps
+
+        if isinstance(image_size, tuple):
+            image_size_list = [image_size] * max_images
+        elif isinstance(image_size, list) and isinstance(image_size[0], tuple):
+            image_size_list = image_size
+            if len(image_size) < max_images:
+                image_size_list += [image_size_list[-1]] * (max_images - len(image_size_list))
+        else:
+            assert False, "image size should be a tuple or a list of tuple"
+
+        if images is None:
+            images =[]
+
+        image_token_count = prompt.count('<image>')
+        assert len(images) >= image_token_count
+        if len(images) > image_token_count:
+            prompt = "<image>\n" * (len(images) - image_token_count) + prompt
+
+        pixel_values =[]
+        grid_hw =[]
+        for image in images:
+            cur_pixel_values, cur_grid_hw = load_image_native(image, self.patch_size, self.downsample_ratio, min_pixels=512*512, max_pixels=min(2048*2048, (4096*4096)//max(1, len(images))), upscale=False)
+            grid_hw.append(cur_grid_hw.to(self.device))
+            pixel_values.append(cur_pixel_values.to(self.device).to(torch.bfloat16))
+
+        merge_size = int(1 / self.downsample_ratio)
+        pv_tensor = torch.cat(pixel_values) if pixel_values else None
+        ghw_tensor = torch.cat(grid_hw) if grid_hw else None
+
+        # Condition Initial Cache
+        template_cond = get_conv_template(self.template)
+        template_cond.system_message = system_message
+        template_cond.append_message(template_cond.roles[0], prompt)
+        template_cond.append_message(template_cond.roles[1], None)
+        query_cond = template_cond.get_prompt() + '<think>\n\n</think>\n\n'
+
+        def replace_image_tokens(query, grid_hw_list):
+            for i in range(len(grid_hw_list)):
+                num_patch_token = int(grid_hw_list[i][0, 0] * grid_hw_list[i][0, 1] * self.downsample_ratio**2)
+                image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * num_patch_token + IMG_END_TOKEN
+                query = query.replace('<image>', image_tokens, 1)
+            return query
+
+        query_cond = replace_image_tokens(query_cond, grid_hw)
+        input_embeds_cond, indexes_cond, attention_mask_cond = self._build_it2i_inputs(tokenizer, query_cond, pv_tensor, ghw_tensor)
+        
+        outputs_cond = self.language_model(inputs_embeds=input_embeds_cond, indexes=indexes_cond, attention_mask=attention_mask_cond, use_cache=True)
+        past_key_values_cond = outputs_cond.past_key_values
+        t_index_cond = indexes_cond[0].max().item()
+
+        # Text Uncondition Cache Initial
+        question_text_uncondition = '<image>' * len(images)
+        template_tu = get_conv_template(self.template)
+        template_tu.system_message = self.system_message
+        template_tu.append_message(template_tu.roles[0], question_text_uncondition)
+        template_tu.append_message(template_tu.roles[1], None)
+        query_text_uncond = template_tu.get_prompt()
+        query_text_uncond = replace_image_tokens(query_text_uncond, grid_hw)
+
+        input_embeds_tu, indexes_tu, attention_mask_tu = self._build_it2i_inputs(tokenizer, query_text_uncond, pv_tensor, ghw_tensor)
+        outputs_tu = self.language_model(inputs_embeds=input_embeds_tu, indexes=indexes_tu, attention_mask=attention_mask_tu, use_cache=True)
+        past_key_values_tu = outputs_tu.past_key_values
+        t_index_tu = indexes_tu[0].max().item()
+
+        # Img Uncondition Cache Initial
+        query_img_uncond = self._build_t2i_query("", append_text=IMG_START_TOKEN)
+        input_embeds_iu, indexes_iu, attention_mask_iu = self._build_it2i_inputs(tokenizer, query_img_uncond)
+        outputs_iu = self.language_model(inputs_embeds=input_embeds_iu, indexes=indexes_iu, attention_mask=attention_mask_iu, use_cache=True)
+        past_key_values_iu = outputs_iu.past_key_values
+
+
+        generated_images =[]
+        img_count = 0
+        device = self.device
+
+        def append_ids_to_cache(cache, t_idx, input_ids):
+            if input_ids.shape[1] == 0:
+                return t_idx
+            seq_len = input_ids.shape[1]
+            inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
+            
+            t_indexes = torch.arange(t_idx + 1, t_idx + 1 + seq_len, dtype=torch.long, device=device)
+            h_indexes = torch.zeros(seq_len, dtype=torch.long, device=device)
+            w_indexes = torch.zeros(seq_len, dtype=torch.long, device=device)
+            indexes = torch.stack([t_indexes, h_indexes, w_indexes], dim=0)
+            
+            past_len = cache.get_seq_length()
+            mask = torch.zeros(1, 1, seq_len, past_len + seq_len, device=device)
+            causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=device))
+            causal_mask = torch.where(causal_mask == 1, 0.0, float('-inf'))
+            mask[:, :, :, past_len:] = causal_mask
+            attention_mask_dict = {"full_attention": mask}
+            
+            self.language_model(
+                inputs_embeds=inputs_embeds,
+                indexes=indexes,
+                attention_mask=attention_mask_dict,
+                past_key_values=cache,
+                use_cache=True
+            )
+            return t_idx + seq_len
+
+        def append_image_to_cache(cache, t_idx, inputs_embeds_img, N_img_tokens, abs_pos_w, abs_pos_h):
+            past_len = cache.get_seq_length()
+            tgt_len = N_img_tokens + 1
+            
+            t_indexes = torch.zeros(tgt_len, dtype=torch.long, device=device)
+            t_indexes[:N_img_tokens] = t_idx + 1
+            t_indexes[N_img_tokens] = t_idx + 2
+            
+            h_indexes = torch.zeros(tgt_len, dtype=torch.long, device=device)
+            w_indexes = torch.zeros(tgt_len, dtype=torch.long, device=device)
+            h_indexes[:N_img_tokens] = abs_pos_h
+            w_indexes[:N_img_tokens] = abs_pos_w
+            
+            indexes = torch.stack([t_indexes, h_indexes, w_indexes], dim=0)
+            
+            mask = torch.zeros(1, 1, tgt_len, past_len + tgt_len, device=device)
+            mask[0, 0, :N_img_tokens, past_len + N_img_tokens] = float('-inf')
+            attention_mask_dict = {"full_attention": mask}
+            
+            self.language_model(
+                inputs_embeds=inputs_embeds_img,
+                indexes=indexes,
+                attention_mask=attention_mask_dict,
+                past_key_values=cache,
+                use_cache=True
+            )
+            return t_idx + 2
+
+        parts = gt_text.split('<image>')
+        img_start_id_tensor = torch.tensor([[self.img_start_token_id]], device=device)
+
+        for i, part in enumerate(parts):
+            if len(part) > 0:
+                if verbose:
+                    print(part, end='', flush=True)
+                part_ids = tokenizer(part, return_tensors='pt', add_special_tokens=False)['input_ids'].to(device)
+                t_index_cond = append_ids_to_cache(past_key_values_cond, t_index_cond, part_ids)
+
+            if i < len(parts) - 1:
+                if img_count >= max_images:
+                    break
+                    
+                if verbose:
+                    print("<image>", end='', flush=True)
+
+                t_index_cond = append_ids_to_cache(past_key_values_cond, t_index_cond, img_start_id_tensor)
+                t_index_tu = append_ids_to_cache(past_key_values_tu, t_index_tu, img_start_id_tensor)
+
+                cur_image_size = image_size_list[img_count]
+                token_h = cur_image_size[1] // (self.patch_size * merge_size)
+                token_w = cur_image_size[0] // (self.patch_size * merge_size)
+
+                indexes_image_condition = self._build_t2i_image_indexes(token_h, token_w, t_index_cond + 1, device=device)
+                indexes_image_text_uncondition = self._build_t2i_image_indexes(token_h, token_w, t_index_tu + 1, device=device)
+                indexes_image_img_uncondition = self._build_t2i_image_indexes(token_h, token_w, indexes_iu[0].max() + 1, device=device)
+
+                grid_h = cur_image_size[1] // self.patch_size
+                grid_w = cur_image_size[0] // self.patch_size
+                gen_grid_hw = torch.tensor([[grid_h, grid_w]], device=device)
+
+                noise_scale = self.noise_scale
+                if self.noise_scale_mode in ("resolution", "dynamic", 'dynamic_sqrt'):
+                    noise_scale = math.sqrt((grid_h*grid_w)/(merge_size**2) / self.noise_scale_base_image_seq_len)
+                    base = float(self.noise_scale_base_image_seq_len)
+                    noise_scale = math.sqrt((grid_h*grid_w)/(merge_size**2)/base) * float(self.noise_scale)
+                    if self.noise_scale_mode == 'dynamic_sqrt':
+                        noise_scale = math.sqrt(noise_scale)
+                noise_scale = min(noise_scale, self.noise_scale_max_value)
+                image_prediction = noise_scale * torch.randn((1, 3, cur_image_size[1], cur_image_size[0]), device=device, dtype=outputs_cond.logits.dtype)
+
+                past_key_values_cond_cfg = past_key_values_cond
+                past_key_values_tu_cfg = past_key_values_tu
+                past_key_values_iu_cfg = past_key_values_iu
+
+                # attention_mask_condition = {"full_attention": torch.zeros(1, 1, token_h*token_w, past_key_values_cond.get_seq_length() + token_h*token_w, device=device)}
+                # attention_mask_text_uncondition = {"full_attention": torch.zeros(1, 1, token_h*token_w, past_key_values_tu.get_seq_length() + token_h*token_w, device=device)}
+                # attention_mask_img_uncondition = {"full_attention": torch.zeros(1, 1, token_h*token_w, past_key_values_iu.get_seq_length() + token_h*token_w, device=device)}
+                attention_mask_condition = {"full_attention": None}
+                attention_mask_text_uncondition = {"full_attention": None}
+                attention_mask_img_uncondition = {"full_attention": None}
+
+                prepare_flash_kv_cache(
+                    past_key_values_cond_cfg,
+                    current_len=token_h * token_w,
+                    batch_size=1,
+                )
+                prepare_flash_kv_cache(
+                    past_key_values_tu_cfg,
+                    current_len=token_h * token_w,
+                    batch_size=1,
+                )
+                prepare_flash_kv_cache(
+                    past_key_values_iu_cfg,
+                    current_len=token_h * token_w,
+                    batch_size=1,
+                )
+
+                timesteps = torch.linspace(0.0, 1.0, num_steps+1, device=device)
+                if enable_timestep_shift:
+                    timesteps = self._apply_time_schedule(timesteps, token_h*token_w, timestep_shift)
+
+                for step_i in range(num_steps):
+                    t = timesteps[step_i]
+                    t_next = timesteps[step_i + 1]
+
+                    z = self.patchify(image_prediction, self.patch_size * merge_size)
+                    image_input = self.patchify(image_prediction, self.patch_size, channel_first=True)
+                    image_embeds = self.extract_feature(image_input.view(1 * grid_h*grid_w, -1), gen_model=True, grid_hw=gen_grid_hw).view(1, token_h*token_w, -1)
+                    t_expanded = t.expand(token_h*token_w)
+                    timestep_embeddings = self.fm_modules['timestep_embedder'](t_expanded).view(1, token_h*token_w, -1)
+                    if self.add_noise_scale_embedding:
+                        noise_scale_tensor = torch.full_like(t_expanded, noise_scale/self.noise_scale_max_value)
+                        noise_embeddings = self.fm_modules['noise_scale_embedder'](noise_scale_tensor).view(1, token_h*token_w, -1)
+                        timestep_embeddings += noise_embeddings
+                    image_embeds = image_embeds + timestep_embeddings
+
+                    use_cfg = (t > cfg_interval[0] and t < cfg_interval[1]) or cfg_interval[0] == 0
+                    out_cond = self._t2i_predict_v(image_embeds, indexes_image_condition, attention_mask_condition, past_key_values_cond_cfg, t, z, image_token_num=token_h*token_w, timestep_embeddings=timestep_embeddings)
+                    if not use_cfg:
+                        v_pred = out_cond
+                    elif cfg_scale == 1 and img_cfg_scale == 1:
+                        v_pred = out_cond
+                    elif img_cfg_scale == 1:
+                        out_img_cond = self._t2i_predict_v(image_embeds, indexes_image_text_uncondition, attention_mask_text_uncondition, past_key_values_tu_cfg, t, z, image_token_num=token_h*token_w, timestep_embeddings=timestep_embeddings)
+                        v_pred = out_img_cond + cfg_scale * (out_cond - out_img_cond)
+                    elif cfg_scale == img_cfg_scale:
+                        out_uncond = self._t2i_predict_v(image_embeds, indexes_image_img_uncondition, attention_mask_img_uncondition, past_key_values_iu_cfg, t, z, image_token_num=token_h*token_w, timestep_embeddings=timestep_embeddings)
+                        v_pred = out_uncond + cfg_scale * (out_cond - out_uncond)
+                    else:
+                        out_img_cond = self._t2i_predict_v(image_embeds, indexes_image_text_uncondition, attention_mask_text_uncondition, past_key_values_tu_cfg, t, z, image_token_num=token_h*token_w, timestep_embeddings=timestep_embeddings)
+                        out_uncond = self._t2i_predict_v(image_embeds, indexes_image_img_uncondition, attention_mask_img_uncondition, past_key_values_iu_cfg, t, z, image_token_num=token_h*token_w, timestep_embeddings=timestep_embeddings)
+                        v_pred = (
+                            out_uncond
+                            + cfg_scale * (out_cond - out_img_cond)
+                            + img_cfg_scale * (out_img_cond - out_uncond)
+                        )
+                    if (cfg_scale > 1 or img_cfg_scale > 1) and use_cfg:
+                        if cfg_norm == 'global':
+                            norm_v_condition = torch.norm(out_cond, dim=(1,2), keepdim=True)
+                            norm_v_cfg = torch.norm(v_pred, dim=(1,2), keepdim=True)
+                            scale = (norm_v_condition / (norm_v_cfg + 1e-8)).clamp(min=0, max=1.0)
+                            v_pred = v_pred * scale
+                        elif cfg_norm == 'channel':
+                            norm_v_condition = torch.norm(out_cond, dim=-1, keepdim=True)
+                            norm_v_cfg = torch.norm(v_pred, dim=-1, keepdim=True)
+                            scale = (norm_v_condition / (norm_v_cfg + 1e-8)).clamp(min=0, max=1.0)
+                            v_pred = v_pred * scale
+
+                    z = z + (t_next - t) * v_pred
+                    image_prediction = self.unpatchify(z, self.patch_size * merge_size, cur_image_size[1], cur_image_size[0])
+
+                generated_images.append(image_prediction)
+
+                clear_flash_kv_cache(past_key_values_cond_cfg)
+                clear_flash_kv_cache(past_key_values_tu_cfg)
+                clear_flash_kv_cache(past_key_values_iu_cfg)
+
+                if gt_images is not None and img_count < len(gt_images):
+                    gt_img_pil = gt_images[img_count]
+                    gt_pixel_values, gt_grid_hw = load_image_native(gt_img_pil, self.patch_size, self.downsample_ratio, min_pixels=512*512, max_pixels=(2048*2048), upscale=False)
+                    gt_pixel_values = gt_pixel_values.to(device).to(torch.bfloat16)
+                    
+                    flatten_pixel_values = gt_pixel_values
+                    gen_grid_hw_und = gt_grid_hw
+                else:
+                    pred_img = image_prediction[0].unsqueeze(0).to(torch.bfloat16)
+                    raw_img = pred_img * 0.5 + 0.5
+                    img_mean = torch.tensor([0.485, 0.456, 0.406], dtype=raw_img.dtype, device=device).view(1, 3, 1, 1)
+                    img_std = torch.tensor([0.229, 0.224, 0.225], dtype=raw_img.dtype, device=device).view(1, 3, 1, 1)
+                    und_img = (raw_img - img_mean) / img_std
+                    
+                    c, h, w = und_img[0].shape
+                    ps = self.patch_size
+                    p_grid_h = h // ps
+                    p_grid_w = w // ps
+                    flatten_pixel_values = (
+                        und_img[0].view(c, p_grid_h, ps, p_grid_w, ps)
+                        .permute(1, 3, 0, 2, 4)
+                        .reshape(p_grid_h * p_grid_w, c * ps ** 2)
+                    )
+                    gen_grid_hw_und = torch.tensor([[p_grid_h, p_grid_w]], device=device)
+
+                vit_embeds = self.extract_feature(flatten_pixel_values, grid_hw=gen_grid_hw_und[:1]).unsqueeze(0)
+                
+                img_end_id = tokenizer.convert_tokens_to_ids(IMG_END_TOKEN)
+                img_end_embed = self.language_model.get_input_embeddings()(torch.tensor([[img_end_id]], device=device))
+                inputs_embeds_img = torch.cat([vit_embeds, img_end_embed], dim=1) # (1, N + 1, C)
+                
+                N_img_tokens = vit_embeds.shape[1]
+                abs_pos_w, abs_pos_h = build_abs_positions_from_grid_hw(gen_grid_hw_und[:1] // int(1 / self.downsample_ratio), device=device)
+
+                t_index_cond = append_image_to_cache(past_key_values_cond, t_index_cond, inputs_embeds_img, N_img_tokens, abs_pos_w, abs_pos_h)
+                t_index_tu = append_image_to_cache(past_key_values_tu, t_index_tu, inputs_embeds_img, N_img_tokens, abs_pos_w, abs_pos_h)
+
+                img_count += 1
+
+        return generated_images
+
+    @torch.no_grad()
+    def interleave_gen(
+            self,
+            tokenizer,
+            prompt,
+            images=None,
+            generation_config=None,
+            cfg_scale=1.0,
+            img_cfg_scale=1.0,
+            cfg_norm='none',
+            max_images=10,
+            enable_timestep_shift=True,
+            timestep_shift=1.0,
+            image_size=(256, 256),
+            num_steps=30,
+            IMG_START_TOKEN='<img>',
+            IMG_END_TOKEN='</img>',
+            IMG_CONTEXT_TOKEN='<IMG_CONTEXT>',
+            method='euler',
+            cfg_interval=(0, 1),
+            t_eps=0.02,
+            verbose=False,
+            system_message='',
+            think_mode=False
+    ):
+        self.img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
+        self.img_start_token_id = tokenizer.convert_tokens_to_ids(IMG_START_TOKEN)
+        self.config.t_eps = t_eps
+
+        if isinstance(image_size, tuple):
+            image_size_list = [image_size] * max_images
+        elif isinstance(image_size, list) and isinstance(image_size[0], tuple):
+            image_size_list = image_size
+            if len(image_size) < max_images:
+                image_size_list += [image_size_list[-1]] * (max_images - len(image_size_list))
+        else:
+            assert False, "image size should be a tuple or a list of tuple"
+
+        if generation_config and hasattr(generation_config, 'max_new_tokens') and generation_config.max_new_tokens is not None:
+            max_new_tokens = generation_config.max_new_tokens
+        else:
+            max_new_tokens = 1024
+
+        current_generated_tokens = 0
+
+        if images is None:
+            images = []
+
+        template = get_conv_template(self.template)
+        template.system_message = self.system_message
+        eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())
+
+        image_token_count = prompt.count('<image>')
+        assert len(images) >= image_token_count
+        if len(images) > image_token_count:
+            prompt = "<image>\n" * (len(images) - image_token_count) + prompt
+
+        pixel_values =[]
+        grid_hw =[]
+        for image in images:
+            cur_pixel_values, cur_grid_hw = load_image_native(image, self.patch_size, self.downsample_ratio, min_pixels=512*512, max_pixels=min(2048*2048, (4096*4096)//max(1, len(images))), upscale=False)
+            grid_hw.append(cur_grid_hw.to(self.device))
+            pixel_values.append(cur_pixel_values.to(self.device).to(torch.bfloat16))
+
+        merge_size = int(1 / self.downsample_ratio)
+        pv_tensor = torch.cat(pixel_values) if pixel_values else None
+        ghw_tensor = torch.cat(grid_hw) if grid_hw else None
+
+        # Condition
+        template_cond = get_conv_template(self.template)
+        template_cond.system_message = system_message
+        template_cond.append_message(template_cond.roles[0], prompt)
+        template_cond.append_message(template_cond.roles[1], None)
+        query_cond = template_cond.get_prompt()
+
+        if not think_mode:
+            query_cond = query_cond + '<think>\n\n</think>\n\n'
+        else:
+            query_cond = query_cond + '<think>\n'
+
+        def replace_image_tokens(query, grid_hw_list):
+            for i in range(len(grid_hw_list)):
+                num_patch_token = int(grid_hw_list[i][0, 0] * grid_hw_list[i][0, 1] * self.downsample_ratio**2)
+                image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * num_patch_token + IMG_END_TOKEN
+                query = query.replace('<image>', image_tokens, 1)
+            return query
+
+        query_cond = replace_image_tokens(query_cond, grid_hw)
+        input_embeds_cond, indexes_cond, attention_mask_cond = self._build_it2i_inputs(tokenizer, query_cond, pv_tensor, ghw_tensor)
+        
+        outputs_cond = self.language_model(inputs_embeds=input_embeds_cond, indexes=indexes_cond, attention_mask=attention_mask_cond, use_cache=True)
+        past_key_values_cond = outputs_cond.past_key_values
+        t_index_cond = indexes_cond[0].max().item()
+
+        # Initialize Text Uncondition Cache
+        question_text_uncondition = '<image>' * len(images)
+        template_tu = get_conv_template(self.template)
+        template_tu.system_message = self.system_message
+        template_tu.append_message(template_tu.roles[0], question_text_uncondition)
+        template_tu.append_message(template_tu.roles[1], None)
+        query_text_uncond = template_tu.get_prompt()
+        query_text_uncond = replace_image_tokens(query_text_uncond, grid_hw)
+
+        input_embeds_tu, indexes_tu, attention_mask_tu = self._build_it2i_inputs(tokenizer, query_text_uncond, pv_tensor, ghw_tensor)
+        outputs_tu = self.language_model(inputs_embeds=input_embeds_tu, indexes=indexes_tu, attention_mask=attention_mask_tu, use_cache=True)
+        past_key_values_tu = outputs_tu.past_key_values
+        t_index_tu = indexes_tu[0].max().item()
+
+        # Initialize Img (ALL) Uncondition Cache
+        query_img_uncond = self._build_t2i_query("", append_text=IMG_START_TOKEN)
+        input_embeds_iu, indexes_iu, attention_mask_iu = self._build_it2i_inputs(tokenizer, query_img_uncond)
+        outputs_iu = self.language_model(inputs_embeds=input_embeds_iu, indexes=indexes_iu, attention_mask=attention_mask_iu, use_cache=True)
+        past_key_values_iu = outputs_iu.past_key_values
+
+
+        generated_text = ""
+        generated_images =[]
+        max_images = 10
+        img_count = 0
+
+        next_token = torch.argmax(outputs_cond.logits[:, -1, :], dim=-1)
+
+        while True:
+            # text generation
+            gen_tokens = []
+            hit_max_tokens = False
+            while True:
+                token_item = next_token.item()
+                if token_item == eos_token_id or token_item == self.img_start_token_id:
+                    break
+                gen_tokens.append(token_item)
+                current_generated_tokens += 1
+
+                self.language_model.model.current_index = t_index_cond
+                outputs_cond = self.language_model(
+                    input_ids=next_token.unsqueeze(0),
+                    past_key_values=past_key_values_cond,
+                    use_cache=True
+                )
+                past_key_values_cond = outputs_cond.past_key_values
+                t_index_cond += 1 
+                next_token = torch.argmax(outputs_cond.logits[:, -1, :], dim=-1)
+
+                if current_generated_tokens >= max_new_tokens:
+                    hit_max_tokens = True
+                    break
+
+            if len(gen_tokens) > 0:
+                chunk_text = tokenizer.decode(gen_tokens, skip_special_tokens=True)
+                generated_text += chunk_text
+                if verbose:
+                    print(chunk_text, end='', flush=True)
+
+            if next_token.item() == eos_token_id or hit_max_tokens:
+                break
+
+            if next_token.item() == self.img_start_token_id:
+                if img_count >= max_images:
+                    break
+
+                generated_text += "<image>"
+                if verbose:
+                    print("<image>", end='', flush=True)
+
+                # Add the img_start_token for condition and text_uncondition branch
+                self.language_model.model.current_index = t_index_cond
+                outputs_cond = self.language_model(input_ids=next_token.unsqueeze(0), past_key_values=past_key_values_cond, use_cache=True)
+                past_key_values_cond = outputs_cond.past_key_values
+                t_index_cond += 1
+
+                self.language_model.model.current_index = t_index_tu
+                outputs_tu = self.language_model(input_ids=next_token.unsqueeze(0), past_key_values=past_key_values_tu, use_cache=True)
+                past_key_values_tu = outputs_tu.past_key_values
+                t_index_tu += 1
+
+                image_size = image_size_list[img_count]
+                # Image Generation
+                token_h = image_size[1] // (self.patch_size * merge_size)
+                token_w = image_size[0] // (self.patch_size * merge_size)
+                device = self.device
+
+                indexes_image_condition = self._build_t2i_image_indexes(token_h, token_w, t_index_cond + 1, device=device)
+                indexes_image_text_uncondition = self._build_t2i_image_indexes(token_h, token_w, t_index_tu + 1, device=device)
+                indexes_image_img_uncondition = self._build_t2i_image_indexes(token_h, token_w, indexes_iu[0].max() + 1, device=device)
+
+                grid_h = image_size[1] // self.patch_size
+                grid_w = image_size[0] // self.patch_size
+                gen_grid_hw = torch.tensor([[grid_h, grid_w]], device=device)
+
+                noise_scale = self.noise_scale
+                if self.noise_scale_mode in ("resolution", "dynamic", 'dynamic_sqrt'):
+                    noise_scale = math.sqrt((grid_h*grid_w)/(merge_size**2) / self.noise_scale_base_image_seq_len)
+                    base = float(self.noise_scale_base_image_seq_len)
+                    noise_scale = math.sqrt((grid_h*grid_w)/(merge_size**2)/base) * float(self.noise_scale)
+                    if self.noise_scale_mode == 'dynamic_sqrt':
+                        noise_scale = math.sqrt(noise_scale)
+                noise_scale = min(noise_scale, self.noise_scale_max_value)
+                image_prediction = noise_scale * torch.randn((1, 3, image_size[1], image_size[0]), device=device, dtype=outputs_cond.logits.dtype)
+
+                past_key_values_cond_cfg = past_key_values_cond
+                past_key_values_tu_cfg = past_key_values_tu
+                past_key_values_iu_cfg = past_key_values_iu
+
+                attention_mask_condition = {"full_attention": None}
+                attention_mask_text_uncondition = {"full_attention": None}
+                attention_mask_img_uncondition = {"full_attention": None}
+
+                prepare_flash_kv_cache(
+                    past_key_values_cond_cfg,
+                    current_len=token_h * token_w,
+                    batch_size=1,
+                )
+                prepare_flash_kv_cache(
+                    past_key_values_tu_cfg,
+                    current_len=token_h * token_w,
+                    batch_size=1,
+                )
+                prepare_flash_kv_cache(
+                    past_key_values_iu_cfg,
+                    current_len=token_h * token_w,
+                    batch_size=1,
+                )
+
+                timesteps = torch.linspace(0.0, 1.0, num_steps+1, device=device)
+                if enable_timestep_shift:
+                    timesteps = self._apply_time_schedule(timesteps, token_h*token_w, timestep_shift)
+
+                for step_i in range(num_steps):
+                    t = timesteps[step_i]
+                    t_next = timesteps[step_i + 1]
+
+                    z = self.patchify(image_prediction, self.patch_size * merge_size)
+                    image_input = self.patchify(image_prediction, self.patch_size, channel_first=True)
+                    image_embeds = self.extract_feature(image_input.view(1 * grid_h*grid_w, -1), gen_model=True, grid_hw=gen_grid_hw).view(1, token_h*token_w, -1)
+                    t_expanded = t.expand(token_h*token_w)
+                    timestep_embeddings = self.fm_modules['timestep_embedder'](t_expanded).view(1, token_h*token_w, -1)
+                    if self.add_noise_scale_embedding:
+                        noise_scale_tensor = torch.full_like(t_expanded, noise_scale/self.noise_scale_max_value)
+                        noise_embeddings = self.fm_modules['noise_scale_embedder'](noise_scale_tensor).view(1, token_h*token_w, -1)
+                        timestep_embeddings += noise_embeddings
+                    image_embeds = image_embeds + timestep_embeddings
+
+                    use_cfg = (t > cfg_interval[0] and t < cfg_interval[1]) or cfg_interval[0] == 0
+                    out_cond = self._t2i_predict_v(image_embeds, indexes_image_condition, attention_mask_condition, past_key_values_cond_cfg, t, z, image_token_num=token_h*token_w, timestep_embeddings=timestep_embeddings)
+                    if not use_cfg:
+                        v_pred = out_cond
+                    elif cfg_scale == 1 and img_cfg_scale == 1:
+                        v_pred = out_cond
+                    elif img_cfg_scale == 1:
+                        out_img_cond = self._t2i_predict_v(image_embeds, indexes_image_text_uncondition, attention_mask_text_uncondition, past_key_values_tu_cfg, t, z, image_token_num=token_h*token_w, timestep_embeddings=timestep_embeddings)
+                        v_pred = out_img_cond + cfg_scale * (out_cond - out_img_cond)
+                    elif cfg_scale == img_cfg_scale:
+                        out_uncond = self._t2i_predict_v(image_embeds, indexes_image_img_uncondition, attention_mask_img_uncondition, past_key_values_iu_cfg, t, z, image_token_num=token_h*token_w, timestep_embeddings=timestep_embeddings)
+                        v_pred = out_uncond + cfg_scale * (out_cond - out_uncond)
+                    else:
+                        out_img_cond = self._t2i_predict_v(image_embeds, indexes_image_text_uncondition, attention_mask_text_uncondition, past_key_values_tu_cfg, t, z, image_token_num=token_h*token_w, timestep_embeddings=timestep_embeddings)
+                        out_uncond = self._t2i_predict_v(image_embeds, indexes_image_img_uncondition, attention_mask_img_uncondition, past_key_values_iu_cfg, t, z, image_token_num=token_h*token_w, timestep_embeddings=timestep_embeddings)
+                        v_pred = (
+                            out_uncond
+                            + cfg_scale * (out_cond - out_img_cond)
+                            + img_cfg_scale * (out_img_cond - out_uncond)
+                        )
+                    if (cfg_scale > 1 or img_cfg_scale > 1 and use_cfg):
+                        if cfg_norm == 'global':
+                            norm_v_condition = torch.norm(out_cond, dim=(1,2), keepdim=True)
+                            norm_v_cfg = torch.norm(v_pred, dim=(1,2), keepdim=True)
+                            scale = (norm_v_condition / (norm_v_cfg + 1e-8)).clamp(min=0, max=1.0)
+                            v_pred = v_pred * scale
+                        elif cfg_norm == 'channel':
+                            norm_v_condition = torch.norm(out_cond, dim=-1, keepdim=True)
+                            norm_v_cfg = torch.norm(v_pred, dim=-1, keepdim=True)
+                            scale = (norm_v_condition / (norm_v_cfg + 1e-8)).clamp(min=0, max=1.0)
+                            v_pred = v_pred * scale
+
+                    z = z + (t_next - t) * v_pred
+                    image_prediction = self.unpatchify(z, self.patch_size * merge_size, image_size[1], image_size[0])
+
+                generated_images.append(image_prediction)
+
+                clear_flash_kv_cache(past_key_values_cond_cfg)
+                clear_flash_kv_cache(past_key_values_tu_cfg)
+                clear_flash_kv_cache(past_key_values_iu_cfg)
+
+                img_count += 1
+
+                # re-encode the generated image using the und-branch
+                pred_img = image_prediction[0].unsqueeze(0).to(torch.bfloat16)
+                # re-normalize the image
+                raw_img = pred_img * 0.5 + 0.5
+                img_mean = torch.tensor([0.485, 0.456, 0.406], dtype=raw_img.dtype, device=device).view(1, 3, 1, 1)
+                img_std = torch.tensor([0.229, 0.224, 0.225], dtype=raw_img.dtype, device=device).view(1, 3, 1, 1)
+                und_img = (raw_img - img_mean) / img_std
+                c, h, w = und_img[0].shape
+                ps = self.patch_size
+                p_grid_h = h // ps
+                p_grid_w = w // ps
+                flatten_pixel_values = (
+                    und_img[0].view(c, p_grid_h, ps, p_grid_w, ps)
+                    .permute(1, 3, 0, 2, 4)  # [grid_h, grid_w, c, patch_size, patch_size]
+                    .reshape(p_grid_h * p_grid_w, c * ps ** 2)
+                )
+                vit_embeds = self.extract_feature(flatten_pixel_values, grid_hw=gen_grid_hw[:1]).unsqueeze(0)
+                
+                img_end_id = tokenizer.convert_tokens_to_ids(IMG_END_TOKEN)
+                img_end_embed = self.language_model.get_input_embeddings()(torch.tensor([[img_end_id]], device=device))
+                inputs_embeds_img = torch.cat([vit_embeds, img_end_embed], dim=1) # (1, N + 1, C)
+                
+                N_img_tokens = vit_embeds.shape[1]
+                abs_pos_w, abs_pos_h = build_abs_positions_from_grid_hw(gen_grid_hw[:1] // int(1 / self.downsample_ratio), device=device)
+
+                def append_image_to_cache(cache, t_idx):
+                    past_len = cache.get_seq_length()
+                    tgt_len = N_img_tokens + 1
+                    
+                    t_indexes = torch.zeros(tgt_len, dtype=torch.long, device=device)
+                    t_indexes[:N_img_tokens] = t_idx + 1
+                    t_indexes[N_img_tokens] = t_idx + 2
+                    
+                    h_indexes = torch.zeros(tgt_len, dtype=torch.long, device=device)
+                    w_indexes = torch.zeros(tgt_len, dtype=torch.long, device=device)
+                    h_indexes[:N_img_tokens] = abs_pos_h
+                    w_indexes[:N_img_tokens] = abs_pos_w
+                    
+                    indexes = torch.stack([t_indexes, h_indexes, w_indexes], dim=0)
+                    
+                    mask = torch.zeros(1, 1, tgt_len, past_len + tgt_len, device=device)
+                    mask[0, 0, :N_img_tokens, past_len + N_img_tokens] = float('-inf')
+                    attention_mask_dict = {"full_attention": mask}
+                    
+                    outputs = self.language_model(
+                        inputs_embeds=inputs_embeds_img,
+                        indexes=indexes,
+                        attention_mask=attention_mask_dict,
+                        past_key_values=cache,
+                        use_cache=True
+                    )
+                    return outputs, t_idx + 2
+
+                outputs_cond, t_index_cond = append_image_to_cache(past_key_values_cond, t_index_cond)
+                outputs_tu, t_index_tu = append_image_to_cache(past_key_values_tu, t_index_tu)
+
+                next_token = torch.argmax(outputs_cond.logits[:, -1, :], dim=-1)
+
+        return generated_text, generated_images
+
+    @torch.no_grad()
+    def it2i_generate(self, tokenizer, prompt, images, cfg_scale=1, img_cfg_scale=1, cfg_norm='none', enable_timestep_shift=True, timestep_shift=1, image_size=(256, 256), num_steps=30, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>', method='euler', cfg_interval=(0, 1), batch_size=1, t_eps=0.02, think_mode=False):
+        assert cfg_norm in ['none', 'global', 'channel']
+
+        self.img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
+        self.config.t_eps = t_eps
+
+        image_token_count = prompt.count('<image>')
+        assert len(images) >= image_token_count
+        if len(images) > image_token_count:
+            prompt = "<image>\n" * (len(images) - image_token_count) + prompt
+
+        pixel_values = []
+        grid_hw = []
+        for image in images:
+            cur_pixel_values, cur_grid_hw = load_image_native(
+                image,
+                self.patch_size,
+                self.downsample_ratio,
+                min_pixels=512 * 512,
+                max_pixels=min(2048*2048, (4096 * 4096) // len(images)),
+                upscale=False,
+            )
+            cur_grid_hw = cur_grid_hw.to(self.device)
+            cur_pixel_values = cur_pixel_values.to(self.device).to(torch.bfloat16)
+            pixel_values.append(cur_pixel_values)
+            grid_hw.append(cur_grid_hw)
+        pixel_values = torch.cat(pixel_values)
+        grid_hw = torch.cat(grid_hw)
+
+        merge_size = int(1 / self.downsample_ratio)
+        question_condition = f"{prompt}"
+        question_img_condition = '<image>' * len(images)
+        question_uncondition = ""
+        think_text = ""
+
+        think_content = '<think>\n' if think_mode else '<think>\n\n</think>\n\n' + IMG_START_TOKEN
+        query_condition = self._build_t2i_query(question_condition, system_message=SYSTEM_MESSAGE_FOR_GEN, append_text=think_content)
+        query_img_condition = self._build_t2i_query(question_img_condition, append_text=IMG_START_TOKEN)
+        query_uncondition = self._build_t2i_query(question_uncondition, append_text=IMG_START_TOKEN)
+
+        for i in range(grid_hw.shape[0]):
+            num_patch_token = int(grid_hw[i, 0] * grid_hw[i, 1] * self.downsample_ratio**2)
+            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * num_patch_token + IMG_END_TOKEN
+            query_condition = query_condition.replace('<image>', image_tokens, 1)
+            query_img_condition = query_img_condition.replace('<image>', image_tokens, 1)
+
+        input_embeds_condition, indexes_condition, attention_mask_condition_prefix = self._build_it2i_inputs(
+            tokenizer, query_condition, pixel_values, grid_hw
+        )
+        input_embeds_img_condition, indexes_img_condition, attention_mask_img_condition_prefix = self._build_it2i_inputs(
+            tokenizer, query_img_condition, pixel_values, grid_hw
+        )
+        input_embeds_uncondition, indexes_uncondition, attention_mask_uncondition_prefix = self._build_it2i_inputs(
+            tokenizer, query_uncondition
+        )
+
+        token_h = image_size[1] // (self.patch_size * merge_size)
+        token_w = image_size[0] // (self.patch_size * merge_size)
+
+        indexes_image_condition = self._build_t2i_image_indexes(
+            token_h, token_w, indexes_condition[0].max() + 1, device=input_embeds_condition.device
+        )
+        indexes_image_img_condition = self._build_t2i_image_indexes(
+            token_h, token_w, indexes_img_condition[0].max() + 1, device=input_embeds_img_condition.device
+        )
+        indexes_image_uncondition = self._build_t2i_image_indexes(
+            token_h, token_w, indexes_uncondition[0].max() + 1, device=input_embeds_uncondition.device
+        )
+
+        if think_mode:
+            outputs_condition = self.language_model(
+                inputs_embeds=input_embeds_condition,
+                indexes=indexes_condition,
+                attention_mask=attention_mask_condition_prefix,
+                use_cache=True,
+                output_hidden_states=True,
+            )
+            past_key_values_condition = outputs_condition.past_key_values
+            hidden_states_condition = outputs_condition.hidden_states[-1]
+            t_index_condition = indexes_condition[0].max().item()
+            past_key_values_condition, t_index_condition, think_text = self._generate_think(
+                tokenizer,
+                outputs_condition,
+                past_key_values_condition,
+                t_index_condition,
+                IMG_START_TOKEN,
+            )
+            indexes_image_condition = self._build_t2i_image_indexes(
+                token_h, token_w, t_index_condition + 1, device=input_embeds_condition.device
+            )
+        else:
+            past_key_values_condition, hidden_states_condition = self._it2i_prefix_forward(
+                input_embeds_condition, indexes_condition, attention_mask_condition_prefix
+            )
+        past_key_values_img_condition, hidden_states_img_condition = self._it2i_prefix_forward(
+            input_embeds_img_condition, indexes_img_condition, attention_mask_img_condition_prefix
+        )
+        past_key_values_uncondition, hidden_states_uncondition = self._it2i_prefix_forward(
+            input_embeds_uncondition, indexes_uncondition, attention_mask_uncondition_prefix
+        )
+
+        for layer_idx in range(len(past_key_values_condition.layers)):
+            past_key_values_condition.layers[layer_idx].keys = past_key_values_condition.layers[layer_idx].keys.expand(
+                batch_size, *past_key_values_condition.layers[layer_idx].keys.shape[1:]
+            )
+            past_key_values_condition.layers[layer_idx].values = past_key_values_condition.layers[layer_idx].values.expand(
+                batch_size, *past_key_values_condition.layers[layer_idx].values.shape[1:]
+            )
+            past_key_values_img_condition.layers[layer_idx].keys = past_key_values_img_condition.layers[layer_idx].keys.expand(
+                batch_size, *past_key_values_img_condition.layers[layer_idx].keys.shape[1:]
+            )
+            past_key_values_img_condition.layers[layer_idx].values = past_key_values_img_condition.layers[layer_idx].values.expand(
+                batch_size, *past_key_values_img_condition.layers[layer_idx].values.shape[1:]
+            )
+            past_key_values_uncondition.layers[layer_idx].keys = past_key_values_uncondition.layers[layer_idx].keys.expand(
+                batch_size, *past_key_values_uncondition.layers[layer_idx].keys.shape[1:]
+            )
+            past_key_values_uncondition.layers[layer_idx].values = past_key_values_uncondition.layers[layer_idx].values.expand(
+                batch_size, *past_key_values_uncondition.layers[layer_idx].values.shape[1:]
+            )
+
+        prepare_flash_kv_cache(
+            past_key_values_condition,
+            current_len=token_h * token_w,
+            batch_size=batch_size,
+        )
+        prepare_flash_kv_cache(
+            past_key_values_img_condition,
+            current_len=token_h * token_w,
+            batch_size=batch_size,
+        )
+        prepare_flash_kv_cache(
+            past_key_values_uncondition,
+            current_len=token_h * token_w,
+            batch_size=batch_size,
+        )
+
+        device = hidden_states_condition.device
+        dtype = hidden_states_condition.dtype
+
+        grid_h = image_size[1] // self.patch_size
+        grid_w = image_size[0] // self.patch_size
+        grid_hw = torch.tensor([[grid_h, grid_w]] * batch_size, device=device)
+
+        noise_scale = self.noise_scale
+        if self.noise_scale_mode in ("resolution", "dynamic", "dynamic_sqrt"):
+            noise_scale = math.sqrt((grid_h * grid_w) / (merge_size**2) / self.noise_scale_base_image_seq_len)
+            base = float(self.noise_scale_base_image_seq_len)
+            scale = math.sqrt((grid_h * grid_w) / (merge_size**2) / base)
+            noise_scale = scale * float(self.noise_scale)
+            if self.noise_scale_mode == 'dynamic_sqrt':
+                noise_scale = math.sqrt(noise_scale)
+        noise_scale = min(noise_scale, self.noise_scale_max_value)
+        image_prediction = noise_scale * torch.randn(
+            (batch_size, 3, image_size[1], image_size[0]), device=device, dtype=dtype
+        )
+
+        attention_mask_condition = {"full_attention": None}
+        attention_mask_img_condition = {"full_attention": None}
+        attention_mask_uncondition = {"full_attention": None}
+
+        timesteps = torch.linspace(0.0, 1.0, num_steps + 1, device=device)
+        if enable_timestep_shift:
+            timesteps = self._apply_time_schedule(timesteps, token_h * token_w, timestep_shift)
+
+        for step_i in range(num_steps):
+            t = timesteps[step_i]
+            t_next = timesteps[step_i + 1]
+            use_cfg = (t > cfg_interval[0] and t < cfg_interval[1]) or cfg_interval[0] == 0
+
+            z = self.patchify(image_prediction, self.patch_size * merge_size)
+            image_input = self.patchify(image_prediction, self.patch_size, channel_first=True)
+            image_embeds = self.extract_feature(
+                image_input.view(batch_size * grid_h * grid_w, -1),
+                gen_model=True,
+                grid_hw=grid_hw,
+            ).view(batch_size, token_h * token_w, -1)
+            t_expanded = t.expand(batch_size * token_h * token_w)
+            timestep_embeddings = self.fm_modules['timestep_embedder'](t_expanded).view(batch_size, token_h * token_w, -1)
+            if self.add_noise_scale_embedding:
+                noise_scale_tensor = torch.full_like(t_expanded, noise_scale / self.noise_scale_max_value)
+                noise_embeddings = self.fm_modules['noise_scale_embedder'](noise_scale_tensor).view(batch_size, token_h * token_w, -1)
+                timestep_embeddings += noise_embeddings
+            image_embeds = image_embeds + timestep_embeddings
+
+            out_cond = self._t2i_predict_v(
+                image_embeds,
+                indexes_image_condition,
+                attention_mask_condition,
+                past_key_values_condition,
+                t,
+                z,
+                image_token_num=token_h * token_w,
+                timestep_embeddings=timestep_embeddings,
+                image_size=image_size,
+            )
+
+            if not use_cfg:
+                v_pred = out_cond
+            elif cfg_scale == 1 and img_cfg_scale == 1:
+                v_pred = out_cond
+            elif img_cfg_scale == 1:
+                out_img_cond = self._t2i_predict_v(
+                    image_embeds,
+                    indexes_image_img_condition,
+                    attention_mask_img_condition,
+                    past_key_values_img_condition,
+                    t,
+                    z,
+                    image_token_num=token_h * token_w,
+                    timestep_embeddings=timestep_embeddings,
+                    image_size=image_size,
+                )
+                v_pred = out_img_cond + cfg_scale * (out_cond - out_img_cond)
+            elif cfg_scale == img_cfg_scale:
+                out_uncond = self._t2i_predict_v(
+                    image_embeds,
+                    indexes_image_uncondition,
+                    attention_mask_uncondition,
+                    past_key_values_uncondition,
+                    t,
+                    z,
+                    image_token_num=token_h * token_w,
+                    timestep_embeddings=timestep_embeddings,
+                    image_size=image_size,
+                )
+                v_pred = out_uncond + cfg_scale * (out_cond - out_uncond)
+            else:
+                out_img_cond = self._t2i_predict_v(
+                    image_embeds,
+                    indexes_image_img_condition,
+                    attention_mask_img_condition,
+                    past_key_values_img_condition,
+                    t,
+                    z,
+                    image_token_num=token_h * token_w,
+                    timestep_embeddings=timestep_embeddings,
+                    image_size=image_size,
+                )
+                out_uncond = self._t2i_predict_v(
+                    image_embeds,
+                    indexes_image_uncondition,
+                    attention_mask_uncondition,
+                    past_key_values_uncondition,
+                    t,
+                    z,
+                    image_token_num=token_h * token_w,
+                    timestep_embeddings=timestep_embeddings,
+                    image_size=image_size,
+                )
+                v_pred = (
+                    out_uncond
+                    + cfg_scale * (out_cond - out_img_cond)
+                    + img_cfg_scale * (out_img_cond - out_uncond)
+                )
+            if (cfg_scale > 1 or img_cfg_scale > 1) and use_cfg:
+                if cfg_norm == 'global':
+                    norm_v_condition = torch.norm(out_cond, dim=(1, 2), keepdim=True)
+                    norm_v_cfg = torch.norm(v_pred, dim=(1, 2), keepdim=True)
+                    scale = (norm_v_condition / (norm_v_cfg + 1e-8)).clamp(min=0, max=1.0)
+                    v_pred = v_pred * scale
+                elif cfg_norm == 'channel':
+                    norm_v_condition = torch.norm(out_cond, dim=-1, keepdim=True)
+                    norm_v_cfg = torch.norm(v_pred, dim=-1, keepdim=True)
+                    scale = (norm_v_condition / (norm_v_cfg + 1e-8)).clamp(min=0, max=1.0)
+                    v_pred = v_pred * scale
+
+            z = z + (t_next - t) * v_pred
+            image_prediction = self.unpatchify(z, self.patch_size * merge_size, image_size[1], image_size[0])
+
+        clear_flash_kv_cache(past_key_values_condition)
+        clear_flash_kv_cache(past_key_values_img_condition)
+        clear_flash_kv_cache(past_key_values_uncondition)
+
+        self.last_think_content = think_text
+        if think_mode:
+            return image_prediction, think_text
+        return image_prediction
+
+    @torch.no_grad()
+    def t2i_generate(self, tokenizer, prompt, cfg_scale=1, timestep_shift=1, enable_timestep_shift=True, cfg_norm='none', image_size=(256, 256), num_steps=30, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>', method='euler', cfg_interval=(0, 1), batch_size=1, t_eps=0.02, think_mode=False):
+        assert self.concat_time_token_num == 0
+        assert cfg_norm in ['cfg_zero_star', 'global', 'none', 'channel']
+        merge_size = int(1 / self.downsample_ratio)
+
+        self.config.t_eps = t_eps
+        # question_condition = f"Please generate an image based on the following description: {prompt}"
+        question_condition = f"{prompt}"
+        # question_condition += f"\nThe resolution of the image should be {image_size}"
+
+        question_uncondition = f""
+        # question_uncondition += f"\nThe resolution of the image should be {image_size}"
+        think_text = ""
+
+        think_content = '<think>\n' if think_mode else '<think>\n\n</think>\n\n' + IMG_START_TOKEN
+        query_condition = self._build_t2i_query(question_condition, system_message=SYSTEM_MESSAGE_FOR_GEN, append_text=think_content)
+        query_uncondition = self._build_t2i_query(question_uncondition, append_text=IMG_START_TOKEN)
+
+        input_ids_condition, indexes_condition, attention_mask_condition_prefix = self._build_t2i_text_inputs(tokenizer, query_condition)
+        input_ids_uncondition, indexes_uncondition, attention_mask_uncondition_prefix = self._build_t2i_text_inputs(tokenizer, query_uncondition)
+       
+        token_h = image_size[1] // (self.patch_size * merge_size)
+        token_w = image_size[0] // (self.patch_size * merge_size)
+
+        indexes_image_condition = self._build_t2i_image_indexes(token_h, token_w, indexes_condition.shape[1], device=input_ids_condition.device)
+        indexes_image_uncondition = self._build_t2i_image_indexes(token_h, token_w, indexes_uncondition.shape[1], device=input_ids_uncondition.device)
+
+        if think_mode:
+            outputs_condition = self.language_model(
+                input_ids=input_ids_condition,
+                indexes=indexes_condition,
+                attention_mask=attention_mask_condition_prefix,
+                use_cache=True,
+                output_hidden_states=True,
+            )
+            past_key_values_condition = outputs_condition.past_key_values
+            hidden_states_condition = outputs_condition.hidden_states[-1]
+            t_index_condition = indexes_condition[0].max().item()
+            past_key_values_condition, t_index_condition, think_text = self._generate_think(
+                tokenizer,
+                outputs_condition,
+                past_key_values_condition,
+                t_index_condition,
+                IMG_START_TOKEN,
+            )
+            indexes_image_condition = self._build_t2i_image_indexes(
+                token_h, token_w, t_index_condition + 1, device=input_ids_condition.device
+            )
+        else:
+            past_key_values_condition, hidden_states_condition = self._t2i_prefix_forward(input_ids_condition, indexes_condition, attention_mask_condition_prefix)
+        past_key_values_uncondition, hidden_states_uncondition = self._t2i_prefix_forward(input_ids_uncondition, indexes_uncondition, attention_mask_uncondition_prefix)
+
+        for layer_idx in range(len(past_key_values_condition.layers)):
+            past_key_values_condition.layers[layer_idx].keys = past_key_values_condition.layers[layer_idx].keys.expand(batch_size, *past_key_values_condition.layers[layer_idx].keys.shape[1:])
+            past_key_values_condition.layers[layer_idx].values = past_key_values_condition.layers[layer_idx].values.expand(batch_size, *past_key_values_condition.layers[layer_idx].values.shape[1:])
+            past_key_values_uncondition.layers[layer_idx].keys = past_key_values_uncondition.layers[layer_idx].keys.expand(batch_size, *past_key_values_uncondition.layers[layer_idx].keys.shape[1:])
+            past_key_values_uncondition.layers[layer_idx].values = past_key_values_uncondition.layers[layer_idx].values.expand(batch_size, *past_key_values_uncondition.layers[layer_idx].values.shape[1:])
+
+        # prepare flash cache once
+        prepare_flash_kv_cache(
+            past_key_values_condition,
+            current_len=token_h * token_w,
+            batch_size=batch_size,
+        )
+        prepare_flash_kv_cache(
+            past_key_values_uncondition,
+            current_len=token_h * token_w,
+            batch_size=batch_size,
+        )
+
+        device = hidden_states_condition.device
+        dtype = hidden_states_condition.dtype
+
+        # init noise image tokens
+        grid_h = image_size[1] // self.patch_size
+        grid_w = image_size[0] // self.patch_size
+        grid_hw = torch.tensor([[grid_h, grid_w]]*batch_size, device=device)
+
+        noise_scale = self.noise_scale
+        if self.noise_scale_mode in ("resolution", "dynamic", 'dynamic_sqrt'):
+            noise_scale = math.sqrt((grid_h*grid_w)/(merge_size**2) / self.noise_scale_base_image_seq_len)
+            base = float(self.noise_scale_base_image_seq_len)
+            scale = math.sqrt((grid_h*grid_w)/(merge_size**2)/base)
+            noise_scale = scale * float(self.noise_scale)
+            if self.noise_scale_mode == 'dynamic_sqrt':
+                noise_scale = math.sqrt(noise_scale)
+        noise_scale = min(noise_scale, self.noise_scale_max_value)
+        image_prediction = noise_scale * torch.randn((batch_size, 3, image_size[1], image_size[0]), device=device, dtype=dtype)
+
+        # attention_mask_condition = {"full_attention": torch.zeros(batch_size, 1, token_h*token_w, input_ids_condition.shape[1]+token_h*token_w, device=device)}
+        # attention_mask_uncondition = {"full_attention": torch.zeros(batch_size, 1, token_h*token_w, input_ids_uncondition.shape[1]+token_h*token_w, device=device)}
+
+        attention_mask_condition = {"full_attention": None}
+        attention_mask_uncondition = {"full_attention": None}
+
+        timesteps = torch.linspace(0.0, 1.0, num_steps+1, device=device)
+        if enable_timestep_shift:
+            timesteps = self._apply_time_schedule(timesteps, token_h*token_w, timestep_shift)
+
+        for step_i in range(num_steps):
+            t = timesteps[step_i]
+            t_next = timesteps[step_i + 1]
+
+            z = self.patchify(image_prediction, self.patch_size * merge_size)
+            image_input = self.patchify(image_prediction, self.patch_size, channel_first=True)
+            image_embeds = self.extract_feature(image_input.view(batch_size * grid_h*grid_w, -1), gen_model=True, grid_hw=grid_hw).view(batch_size, token_h*token_w, -1)
+            t_expanded = t.expand(batch_size*token_h*token_w)
+            timestep_embeddings = self.fm_modules['timestep_embedder'](t_expanded).view(batch_size, token_h*token_w, -1)
+            if self.add_noise_scale_embedding:
+                noise_scale_tensor = torch.full_like(t_expanded, noise_scale / self.noise_scale_max_value)
+                noise_embeddings = self.fm_modules['noise_scale_embedder'](noise_scale_tensor).view(batch_size, token_h*token_w, -1)
+                timestep_embeddings += noise_embeddings
+            image_embeds = image_embeds + timestep_embeddings
+
+            v_pred_condition = self._t2i_predict_v(image_embeds, indexes_image_condition, attention_mask_condition, past_key_values_condition, t, z, image_token_num=token_h*token_w, timestep_embeddings=timestep_embeddings, image_size=image_size)
+            
+            if t >= cfg_interval[0] and t <= cfg_interval[1] and cfg_scale > 1:
+                v_pred_uncondition = self._t2i_predict_v(image_embeds, indexes_image_uncondition, attention_mask_uncondition, past_key_values_uncondition, t, z, image_token_num=token_h*token_w, timestep_embeddings=timestep_embeddings, image_size=image_size)
+                if cfg_norm == 'cfg_zero_star':
+                    positive_flat = v_pred_condition.view(batch_size, -1)  
+                    negative_flat = v_pred_uncondition.view(batch_size, -1)  
+
+                    alpha = optimized_scale(positive_flat,negative_flat)
+                    alpha = alpha.view(batch_size, *([1] * (len(v_pred_condition.shape) - 1)))
+                    alpha = alpha.to(positive_flat.dtype)
+
+                    if (step_i <= 0):
+                        v_pred = v_pred_condition*0.
+                    else:
+                        v_pred = v_pred_uncondition * alpha + cfg_scale * (v_pred_condition - v_pred_uncondition * alpha)
+                else: 
+                    v_pred = v_pred_uncondition + cfg_scale * (v_pred_condition - v_pred_uncondition)
+                    if cfg_norm == 'global':
+                        norm_v_condition = torch.norm(v_pred_condition, dim=(1,2), keepdim=True)
+                        norm_v_cfg = torch.norm(v_pred, dim=(1,2), keepdim=True)
+                        scale = (norm_v_condition / (norm_v_cfg + 1e-8)).clamp(min=0, max=1.0)
+                        v_pred = v_pred * scale
+                    elif cfg_norm == 'channel':
+                        norm_v_condition = torch.norm(v_pred_condition, dim=-1, keepdim=True)
+                        norm_v_cfg = torch.norm(v_pred, dim=-1, keepdim=True)
+                        scale = (norm_v_condition / (norm_v_cfg + 1e-8)).clamp(min=0, max=1.0)
+                        v_pred = v_pred * scale
+            else:
+                v_pred = v_pred_condition
+
+            z = z + (t_next - t) * v_pred
+
+            image_prediction = self.unpatchify(z, self.patch_size * merge_size, image_size[1], image_size[0])
+
+        clear_flash_kv_cache(past_key_values_condition)
+        clear_flash_kv_cache(past_key_values_uncondition)
+
+        self.last_think_content = think_text
+        if think_mode:
+            return image_prediction, think_text
+        return image_prediction
+
+    def chat(self, tokenizer, pixel_values, question, generation_config, history=None, return_history=False, grid_hw=None, 
+             IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>', verbose=False):
+
+        if history is None and pixel_values is not None and '<image>' not in question:
+            question = '<image>\n' + question
+
+        img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
+        self.img_context_token_id = img_context_token_id
+        self.img_start_token_id = tokenizer.convert_tokens_to_ids(IMG_START_TOKEN)
+
+        template = get_conv_template(self.template)
+        template.system_message = self.system_message
+        eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())
+
+        history = [] if history is None else history
+        for (old_question, old_answer) in history:
+            template.append_message(template.roles[0], old_question)
+            template.append_message(template.roles[1], old_answer)
+        template.append_message(template.roles[0], question)
+        template.append_message(template.roles[1], None)
+        query = template.get_prompt()
+
+        if verbose and pixel_values is not None:
+            print(f'dynamic image size: {grid_hw[0] * self.patch_size}')
+
+        for i in range(grid_hw.shape[0]):
+            num_patch_token = int(grid_hw[i, 0] * grid_hw[i, 1] * self.downsample_ratio**2)
+            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * num_patch_token + IMG_END_TOKEN
+            query = query.replace('<image>', image_tokens, 1)
+
+        model_inputs = tokenizer(query, return_tensors='pt')
+        input_ids = model_inputs['input_ids'].to(self.device)
+        attention_mask = model_inputs['attention_mask'].to(self.device)
+        generation_config['eos_token_id'] = eos_token_id
+        generation_output = self.generate(
+            pixel_values=pixel_values,
+            input_ids=input_ids,
+            grid_hw=grid_hw,
+            attention_mask=attention_mask,
+            **generation_config
+        )
+        response = tokenizer.batch_decode(generation_output, skip_special_tokens=True)[0]
+        response = response.split(template.sep.strip())[0].strip()
+        history.append((question, response))
+        if return_history:
+            return response, history
+        else:
+            query_to_print = query.replace(IMG_CONTEXT_TOKEN, '')
+            query_to_print = query_to_print.replace(f'{IMG_START_TOKEN}{IMG_END_TOKEN}', '<image>')
+            if verbose:
+                print(query_to_print, response)
+            return response
+
+    @torch.no_grad()
+    def generate(
+            self,
+            pixel_values: Optional[torch.FloatTensor] = None,
+            input_ids: Optional[torch.FloatTensor] = None,
+            grid_hw: Optional[torch.LongTensor] = None,
+            attention_mask: Optional[torch.LongTensor] = None,
+            visual_features: Optional[torch.FloatTensor] = None,
+            generation_config: Optional[GenerationConfig] = None,
+            output_hidden_states: Optional[bool] = None,
+            **generate_kwargs,
+    ) -> torch.LongTensor:
+        assert input_ids.shape[0] == 1
+        assert self.img_context_token_id is not None
+        indexes = self.get_thw_indexes(input_ids[0], grid_hw)
+        if pixel_values is not None:
+            if visual_features is not None:
+                vit_embeds = visual_features
+            else:
+                vit_embeds = self.extract_feature(pixel_values, grid_hw=grid_hw)
+        
+            input_embeds = self.language_model.get_input_embeddings()(input_ids)
+            B, N, C = input_embeds.shape
+            input_embeds = input_embeds.reshape(B * N, C)
+
+            input_ids = input_ids.reshape(B * N)
+            selected = (input_ids == self.img_context_token_id)
+            assert selected.sum() != 0
+            input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)
+
+            input_embeds = input_embeds.reshape(B, N, C)
+        else:
+            input_embeds = self.language_model.get_input_embeddings()(input_ids)
+
+        outputs = self.language_model.generate(
+            inputs_embeds=input_embeds,
+            indexes=indexes,
+            attention_mask=attention_mask,
+            generation_config=generation_config,
+            output_hidden_states=output_hidden_states,
+            use_cache=True,
+            **generate_kwargs,
+        )
+
+        return outputs
+
+    @property
+    def lm_head(self):
+        return self.language_model.get_output_embeddings()
+
+    def get_output_embeddings(self):
+        return self.language_model.get_output_embeddings()
+
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        return self.language_model.set_input_embeddings(value)
+
+    def set_output_embeddings(self, value):
+        return self.language_model.set_output_embeddings(value)
+    
+    def get_thw_indexes(self, input_ids, grid_hw=None):
+        img_start_shift = torch.cat([torch.zeros(1, dtype=torch.long).to(input_ids.device), 
+                                     (input_ids == self.img_start_token_id).long()], dim=0)[:-1]
+        not_img_token = (input_ids != self.img_context_token_id).long()
+        t_indexes = ((img_start_shift + not_img_token).cumsum(0) - 1)
+        h_indexes = torch.zeros_like(t_indexes).to(t_indexes.device)
+        w_indexes = torch.zeros_like(t_indexes).to(t_indexes.device)
+
+        if grid_hw is not None:
+            selected = (input_ids == self.img_context_token_id)
+            if selected.long().sum() > 0:
+                abs_pos_w, abs_pos_h = build_abs_positions_from_grid_hw(
+                    grid_hw // int(1 / self.downsample_ratio), device=t_indexes.device)
+                h_indexes[selected] = abs_pos_h.to(t_indexes.device, t_indexes.dtype)
+                w_indexes[selected] = abs_pos_w.to(t_indexes.device, t_indexes.dtype)
+        return torch.stack([t_indexes, h_indexes, w_indexes], dim=0)

modeling_neo_vit.py

@@ -0,0 +1,235 @@
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.modeling_utils import PreTrainedModel
+
+from .configuration_neo_vit import NEOVisionConfig
+
+
+def precompute_rope_freqs_sincos(
+    dim: int, max_position: int, base: float = 10000.0, device=None
+):
+    """预计算 RoPE 的 cos 和 sin 值 (1D)。"""
+    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=device).float() / dim))
+    t = torch.arange(max_position, device=device).type_as(inv_freq)
+    freqs = torch.outer(t, inv_freq)
+    return torch.cos(freqs), torch.sin(freqs)
+
+
+def build_abs_positions_from_grid_hw(grid_hw: torch.Tensor, device=None):
+    """
+    Compute patch coordinates (x, y)
+
+    Args:
+        grid_hw: (B, 2) tensor representing (H, W) per image
+    """
+    device = grid_hw.device
+    B = grid_hw.shape[0]
+
+    # Get the number of patches per image
+    H = grid_hw[:, 0]
+    W = grid_hw[:, 1]
+    N = H * W
+    N_total = N.sum()
+
+    # Create the batch index for each patch (B x patch count)
+    patch_to_sample = torch.repeat_interleave(torch.arange(B, device=device), N)  # (N_total,)
+
+    # Generate intra-image patch index (row-major order)
+    patch_id_within_image = torch.arange(N_total, device=device)
+    patch_id_within_image = patch_id_within_image - torch.cumsum(
+        torch.cat([torch.tensor([0], device=device), N[:-1]]), dim=0
+    )[patch_to_sample]
+
+    # Get H/W for each patch according to its image
+    W_per_patch = W[patch_to_sample]
+    abs_x = patch_id_within_image % W_per_patch
+    abs_y = patch_id_within_image // W_per_patch
+
+    return abs_x, abs_y
+
+
+def apply_rotary_emb_1d(
+    x: torch.Tensor,
+    cos_cached: torch.Tensor,
+    sin_cached: torch.Tensor,
+    positions: torch.Tensor,
+):
+    """对输入张量的一部分应用1D RoPE。"""
+    # x: (..., seq_len, dim_part)
+    # positions: (..., seq_len)
+    # cos_cached: (max_pos, dim_part / 2)
+    
+    cos = cos_cached[positions] # Shape: (positions.shape, dim_part / 2)
+    sin = sin_cached[positions] # Shape: (positions.shape, dim_part / 2)
+    
+    x1 = x[..., 0::2]
+    x2 = x[..., 1::2]
+    
+    rotated_x1 = x1 * cos - x2 * sin
+    rotated_x2 = x1 * sin + x2 * cos
+    
+    x_rotated = torch.empty_like(x)
+    x_rotated[..., 0::2] = rotated_x1
+    x_rotated[..., 1::2] = rotated_x2
+    return x_rotated
+
+
+def apply_2d_rotary_pos_emb(
+    x: torch.Tensor,
+    cos_cached_x: torch.Tensor,
+    sin_cached_x: torch.Tensor,
+    cos_cached_y: torch.Tensor,
+    sin_cached_y: torch.Tensor,
+    abs_positions_x: torch.Tensor,
+    abs_positions_y: torch.Tensor
+):
+    """应用2D RoPE到输入张量x。"""
+    dim = x.shape[-1]
+    dim_half = dim // 2
+
+    # 假设我们将embedding的前半部分用于一个方向的RoPE，后半部分用于另一个方向
+    # 例如，前一半给X坐标，后一半给Y坐标 (或者反过来，但要保持一致)
+    x_part_1 = x[..., :dim_half]
+    x_part_2 = x[..., dim_half:]
+
+    # 将与 abs_positions_x 相关的旋转应用于 x_part_1
+    rotated_part_1 = apply_rotary_emb_1d(
+        x_part_1, cos_cached_x, sin_cached_x, abs_positions_x
+    )
+    # 将与 abs_positions_y 相关的旋转应用于 x_part_2
+    rotated_part_2 = apply_rotary_emb_1d(
+        x_part_2, cos_cached_y, sin_cached_y, abs_positions_y
+    )
+    
+    # 将它们重新拼接起来。确保顺序与你分割时一致。
+    return torch.cat((rotated_part_1, rotated_part_2), dim=-1)
+
+
+class NEOVisionEmbeddings(nn.Module):
+    """
+    Embedding Module for Vision.
+    """
+
+    def __init__(self, config: NEOVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.llm_embed_dim = config.llm_hidden_size[0]
+        self.downsample_factor = int(1 / config.downsample_ratio[0])
+        self.patch_size = config.patch_size
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=config.num_channels, out_channels=self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size
+        )
+        self.dense_embedding = nn.Conv2d(
+            in_channels=self.embed_dim, out_channels=self.llm_embed_dim, kernel_size=self.downsample_factor, stride=self.downsample_factor
+        )
+        self.gelu = nn.GELU()
+
+        self.rope_dim_part = self.embed_dim // 2
+        cos_x, sin_x = precompute_rope_freqs_sincos(
+            self.rope_dim_part, config.max_position_embeddings_vision, base=config.rope_theta_vision, device=None
+        )
+        cos_y, sin_y = precompute_rope_freqs_sincos(
+            self.rope_dim_part, config.max_position_embeddings_vision, base=config.rope_theta_vision, device=None
+        )
+        
+        self.register_buffer("cos_cached_x", cos_x, persistent=False)
+        self.register_buffer("sin_cached_x", sin_x, persistent=False)
+        self.register_buffer("cos_cached_y", cos_y, persistent=False)
+        self.register_buffer("sin_cached_y", sin_y, persistent=False)
+
+    def _apply_2d_rotary_pos_emb(self, patch_embeds, grid_hw):
+        """
+        Apply 2D Rotary Position Embedding to the patch embeddings.
+        """
+        abs_pos_x, abs_pos_y = build_abs_positions_from_grid_hw(grid_hw, device=patch_embeds.device)
+        embeddings = apply_2d_rotary_pos_emb(
+            patch_embeds.to(torch.float32), # RoPE calculations are often more stable in float32
+            self.cos_cached_x, self.sin_cached_x,
+            self.cos_cached_y, self.sin_cached_y,
+            abs_pos_x,
+            abs_pos_y
+        ).to(self.patch_embedding.weight.dtype)
+        return embeddings
+        
+    def forward(self, pixel_values: torch.FloatTensor, grid_hw=None) -> torch.Tensor:
+        
+        pixel_values = pixel_values.view(  # 
+            -1,
+            3,
+            self.patch_size,
+            self.patch_size,
+        )   #  [28072, 768] -> [28072, 3, 16, 16]
+        patch_embeds = self.gelu(self.patch_embedding(pixel_values)).view(-1, self.embed_dim)
+        self.cos_cached_x = self.cos_cached_x.to(patch_embeds.device)
+        self.sin_cached_x = self.sin_cached_x.to(patch_embeds.device)
+        self.cos_cached_y = self.cos_cached_y.to(patch_embeds.device)
+        self.sin_cached_y = self.sin_cached_y.to(patch_embeds.device)
+        patch_embeds = self._apply_2d_rotary_pos_emb(patch_embeds, grid_hw) # [28072, 1024]
+        assert (grid_hw[:,0] * grid_hw[:,1]).sum() == patch_embeds.shape[0]
+
+        patches_list = []
+        cur_position = 0
+        for i in range(grid_hw.shape[0]):
+            h, w = grid_hw[i]
+            patches_per_img = patch_embeds[cur_position : cur_position + h * w].view(h, w, -1).unsqueeze(0)
+            patches_per_img = self.dense_embedding(patches_per_img.permute(0, 3, 1, 2))
+            patches_per_img = patches_per_img.permute(0, 2, 3, 1)
+            patches_list.append(patches_per_img.view(-1, patches_per_img.shape[-1]))
+            cur_position += h * w
+
+        embeddings = torch.cat(patches_list, dim=0)  # (N_total // downsample_factor**2, C)
+
+        assert cur_position == patch_embeds.shape[0]
+        assert embeddings.shape[0] == int(patch_embeds.shape[0] / self.downsample_factor**2)
+
+        return embeddings
+
+
+class NEOVisionModel(PreTrainedModel):
+    main_input_name = 'pixel_values'
+    _supports_flash_attn_2 = True
+    supports_gradient_checkpointing = True
+    config_class = NEOVisionConfig
+    # support transformers 4.51.+
+    _tp_plan = ''
+
+    def __init__(self, config: NEOVisionConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = NEOVisionEmbeddings(config)
+
+    def forward(
+            self,
+            pixel_values: Optional[torch.FloatTensor] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+            pixel_embeds: Optional[torch.FloatTensor] = None,
+            grid_hw: Optional[torch.Tensor] = None
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if pixel_values is None and pixel_embeds is None:
+            raise ValueError('You have to specify pixel_values or pixel_embeds')
+
+        if pixel_embeds is not None:
+            hidden_states = pixel_embeds
+        else:
+            assert pixel_values.dim() == 2, f"pixel_values must be 2D for native resolution, got: {pixel_values.dim()}"
+            hidden_states = self.embeddings(pixel_values, grid_hw=grid_hw)
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=hidden_states,
+            pooler_output=None,
+            hidden_states=None,
+            attentions=None,
+        )

modeling_qwen3.py

@@ -0,0 +1,1117 @@
+from typing import Callable, Optional, Union
+
+import torch
+import torch._dynamo
+from torch import nn
+
+import copy
+import math
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.integrations import use_kernel_forward_from_hub
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_layers import (
+    GenericForQuestionAnswering,
+    GenericForSequenceClassification,
+    GenericForTokenClassification,
+    GradientCheckpointingLayer,
+)
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
+from transformers.utils.deprecation import deprecate_kwarg
+from transformers.utils.generic import check_model_inputs
+from transformers import Qwen3Config
+
+from flash_attn import flash_attn_func
+
+
+def create_block_causal_mask(index: torch.Tensor):
+    """
+    index: (L)
+    return: (1, 1, L, L) block-wise causal attention mask
+    """
+    L = index.size(0)
+    idx_i = index.unsqueeze(1).expand(L, L)
+    idx_j = index.unsqueeze(0).expand(L, L)
+
+    arange = torch.arange(L, device=index.device)
+    mask = (idx_j == idx_i) | (arange.unsqueeze(0) <= arange.unsqueeze(1))
+
+    return torch.where(mask[None, None, :, :] > 0, torch.tensor(0.0), torch.tensor(float('-inf')))
+
+
+def visualize_mask(mask: torch.Tensor, i: int = 0, j: int = 12):
+    """
+    mask: (1,1, L, L)
+    """
+    submask = torch.where(mask[0, 0, :, :] == 0, torch.tensor(1.0), torch.tensor(0.0))
+    submask = mask[i:j, i:j].int().cpu().numpy()
+    for row in submask:
+        print(" ".join(map(str, row)))
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Qwen3RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        Qwen3RMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class Qwen3MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+class Qwen3RotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, config: Qwen3Config, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        base_rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        def _rope_init_fn_keep_freq_range(cfg: Qwen3Config, dev=None):
+            inv_freq, attention_scaling = base_rope_init_fn(cfg, dev)
+
+            cfg2 = copy.deepcopy(cfg)
+            head_dim = getattr(cfg2, "head_dim", None)
+            if head_dim is None:
+                head_dim = cfg2.hidden_size // cfg2.num_attention_heads
+                setattr(cfg2, "head_dim", head_dim)
+            cfg2.head_dim = int(head_dim) * 2
+
+            inv_freq_full, _ = base_rope_init_fn(cfg2, dev)
+            inv_freq = inv_freq_full[::2]
+
+            return inv_freq, attention_scaling
+
+        self.rope_init_fn = _rope_init_fn_keep_freq_range
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+class Qwen3Attention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: Qwen3Config, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.q_proj_mot_gen = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj_mot_gen = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.v_proj_mot_gen = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+        self.o_proj_mot_gen = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+
+        self.q_norm = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps)  # unlike olmo, only on the head dim!
+        self.q_norm_mot_gen = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps)
+        self.q_norm_hw = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps)
+        self.q_norm_hw_mot_gen = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps)
+
+        self.k_norm = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps)  # thus post q_norm does not need reshape
+        self.k_norm_mot_gen = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps)
+        self.k_norm_hw = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps)  # thus post q_norm does not need reshape
+        self.k_norm_hw_mot_gen = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps)
+
+        self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None
+
+        t_config = copy.deepcopy(config)
+        t_config.head_dim = config.head_dim // 2
+        self.rotary_emb = Qwen3RotaryEmbedding(config=t_config)
+
+        hw_config = copy.deepcopy(config)
+        hw_config.head_dim = config.head_dim // 4
+        hw_config.rope_theta = config.rope_theta_hw
+        hw_config.max_position_embeddings = config.max_position_embeddings_hw
+        self.rotary_emb_hw = Qwen3RotaryEmbedding(config=hw_config)
+    
+    def forward_und(
+        self,
+        hidden_states: torch.Tensor,
+        indexes: Optional[torch.LongTensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        assert self.config._attn_implementation == "eager"
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape)
+        query_states_t, query_states_hw = query_states.chunk(2, dim=-1)
+        query_states_t = self.q_norm(query_states_t).transpose(1, 2)
+        query_states_hw = self.q_norm_hw(query_states_hw).transpose(1, 2)
+        query_states_h, query_states_w = query_states_hw.chunk(2, dim=-1)
+
+        key_states = self.k_proj(hidden_states).view(hidden_shape)
+        key_states_t, key_states_hw = key_states.chunk(2, dim=-1)
+        key_states_t = self.k_norm(key_states_t).transpose(1, 2)
+        key_states_hw = self.k_norm_hw(key_states_hw).transpose(1, 2)
+        key_states_h, key_states_w = key_states_hw.chunk(2, dim=-1)
+
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos_t, sin_t = self.rotary_emb(hidden_states, indexes[0].unsqueeze(0))
+        query_states_t, key_states_t = apply_rotary_pos_emb(query_states_t, key_states_t, cos_t, sin_t)
+
+        cos_h, sin_h = self.rotary_emb_hw(hidden_states, indexes[1].unsqueeze(0))
+        query_states_h, key_states_h = apply_rotary_pos_emb(query_states_h, key_states_h, cos_h, sin_h)
+
+        cos_w, sin_w = self.rotary_emb_hw(hidden_states, indexes[2].unsqueeze(0))
+        query_states_w, key_states_w = apply_rotary_pos_emb(query_states_w, key_states_w, cos_w, sin_w)
+
+        query_states = torch.cat([query_states_t, query_states_h, query_states_w], dim=-1)
+        key_states = torch.cat([key_states_t, key_states_h, key_states_w], dim=-1)
+
+
+        if past_key_values is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            # cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            # key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+            update_cache = kwargs.get("update_cache", True)
+            if update_cache:
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs=None)
+            else:
+                # only use the past key values but do not append the current one
+                layer = past_key_values.layers[self.layer_idx]
+                past_k, past_v = layer.keys, layer.values
+
+                if past_k is not None:
+                    key_states   = torch.cat([past_k, key_states], dim=2)   # concat on seq_len
+                    value_states = torch.cat([past_v, value_states], dim=2)
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,  # diff with Llama
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+    # def forward_gen(
+    #     self,
+    #     hidden_states: torch.Tensor,
+    #     indexes: Optional[torch.LongTensor],
+    #     attention_mask: Optional[torch.Tensor],
+    #     past_key_values: Optional[Cache] = None,
+    #     cache_position: Optional[torch.LongTensor] = None,
+    #     **kwargs: Unpack[FlashAttentionKwargs],
+    # ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+    #     assert self.config._attn_implementation == "eager"
+    #     input_shape = hidden_states.shape[:-1]
+    #     hidden_shape = (*input_shape, -1, self.head_dim)
+
+    #     query_states = self.q_proj_mot_gen(hidden_states).view(hidden_shape)
+    #     query_states_t, query_states_hw = query_states.chunk(2, dim=-1)
+    #     query_states_t = self.q_norm_mot_gen(query_states_t).transpose(1, 2)
+    #     query_states_hw = self.q_norm_hw_mot_gen(query_states_hw).transpose(1, 2)
+    #     query_states_h, query_states_w = query_states_hw.chunk(2, dim=-1)
+
+    #     key_states = self.k_proj_mot_gen(hidden_states).view(hidden_shape)
+    #     key_states_t, key_states_hw = key_states.chunk(2, dim=-1)
+    #     key_states_t = self.k_norm_mot_gen(key_states_t).transpose(1, 2)
+    #     key_states_hw = self.k_norm_hw_mot_gen(key_states_hw).transpose(1, 2)
+    #     key_states_h, key_states_w = key_states_hw.chunk(2, dim=-1)
+
+    #     value_states = self.v_proj_mot_gen(hidden_states).view(hidden_shape).transpose(1, 2)
+
+    #     cos_t, sin_t = self.rotary_emb(hidden_states, indexes[0].unsqueeze(0))
+    #     query_states_t, key_states_t = apply_rotary_pos_emb(query_states_t, key_states_t, cos_t, sin_t)
+
+    #     cos_h, sin_h = self.rotary_emb_hw(hidden_states, indexes[1].unsqueeze(0))
+    #     query_states_h, key_states_h = apply_rotary_pos_emb(query_states_h, key_states_h, cos_h, sin_h)
+
+    #     cos_w, sin_w = self.rotary_emb_hw(hidden_states, indexes[2].unsqueeze(0))
+    #     query_states_w, key_states_w = apply_rotary_pos_emb(query_states_w, key_states_w, cos_w, sin_w)
+
+    #     query_states = torch.cat([query_states_t, query_states_h, query_states_w], dim=-1)
+    #     key_states = torch.cat([key_states_t, key_states_h, key_states_w], dim=-1)
+
+
+    #     if past_key_values is not None:
+    #         # sin and cos are specific to RoPE models; cache_position needed for the static cache
+    #         # cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+    #         # key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+    #         update_cache = kwargs.get("update_cache", True)
+    #         if update_cache:
+    #             key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs=None)
+    #         else:
+    #             # only use the past key values but do not append the current one
+    #             layer = past_key_values.layers[self.layer_idx]
+    #             past_k, past_v = layer.keys, layer.values
+
+    #             if past_k is not None:
+    #                 key_states   = torch.cat([past_k, key_states], dim=2)   # concat on seq_len
+    #                 value_states = torch.cat([past_v, value_states], dim=2)
+
+    #     attention_interface: Callable = eager_attention_forward
+    #     if self.config._attn_implementation != "eager":
+    #         attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+    #     attn_output, attn_weights = attention_interface(
+    #         self,
+    #         query_states,
+    #         key_states,
+    #         value_states,
+    #         attention_mask,
+    #         dropout=0.0 if not self.training else self.attention_dropout,
+    #         scaling=self.scaling,
+    #         sliding_window=self.sliding_window,  # diff with Llama
+    #         **kwargs,
+    #     )
+
+    #     attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+    #     attn_output = self.o_proj_mot_gen(attn_output)
+    #     return attn_output, attn_weights
+
+    def forward_gen(
+        self,
+        hidden_states: torch.Tensor,
+        indexes: Optional[torch.LongTensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        # -----------------------------
+        # Build q / k / v for current tokens
+        # Internal layout before flash:
+        #   q/k/v: [B, H, S, D]
+        # Flash layout:
+        #   q/k/v: [B, S, H, D]
+        # -----------------------------
+        query_states = self.q_proj_mot_gen(hidden_states).view(hidden_shape)
+        query_states_t, query_states_hw = query_states.chunk(2, dim=-1)
+        query_states_t = self.q_norm_mot_gen(query_states_t).transpose(1, 2)   # [B,H,S,D/2]
+        query_states_hw = self.q_norm_hw_mot_gen(query_states_hw).transpose(1, 2)
+        query_states_h, query_states_w = query_states_hw.chunk(2, dim=-1)
+
+        key_states = self.k_proj_mot_gen(hidden_states).view(hidden_shape)
+        key_states_t, key_states_hw = key_states.chunk(2, dim=-1)
+        key_states_t = self.k_norm_mot_gen(key_states_t).transpose(1, 2)       # [B,H,S,D/2]
+        key_states_hw = self.k_norm_hw_mot_gen(key_states_hw).transpose(1, 2)
+        key_states_h, key_states_w = key_states_hw.chunk(2, dim=-1)
+
+        value_states = self.v_proj_mot_gen(hidden_states).view(hidden_shape).transpose(1, 2)  # [B,H,S,D]
+
+        # RoPE
+        cos_t, sin_t = self.rotary_emb(hidden_states, indexes[0].unsqueeze(0))
+        query_states_t, key_states_t = apply_rotary_pos_emb(query_states_t, key_states_t, cos_t, sin_t)
+
+        cos_h, sin_h = self.rotary_emb_hw(hidden_states, indexes[1].unsqueeze(0))
+        query_states_h, key_states_h = apply_rotary_pos_emb(query_states_h, key_states_h, cos_h, sin_h)
+
+        cos_w, sin_w = self.rotary_emb_hw(hidden_states, indexes[2].unsqueeze(0))
+        query_states_w, key_states_w = apply_rotary_pos_emb(query_states_w, key_states_w, cos_w, sin_w)
+
+        # concat along head_dim
+        # query/key current layout: [B, H, S, D]
+        query_states = torch.cat([query_states_t, query_states_h, query_states_w], dim=-1)
+        key_states = torch.cat([key_states_t, key_states_h, key_states_w], dim=-1)
+
+        update_cache = kwargs.get("update_cache", True)
+
+        # ------------------------------------------------------------------
+        # Flash path:
+        # Only use when there is no explicit dense mask.
+        # This is exactly the t2i denoising use case:
+        #   current image tokens attend to [prefix + current image tokens]
+        #   fully bidirectional inside current block => causal=False
+        # ------------------------------------------------------------------
+        if attention_mask is None:
+            # Convert current q/k/v to flash layout [B, S, H, D]
+            q = query_states.transpose(1, 2).contiguous()
+            k_cur = key_states.transpose(1, 2).contiguous()
+            v_cur = value_states.transpose(1, 2).contiguous()
+
+            if past_key_values is not None:
+                if update_cache:
+                    # Rare path, keep compatibility.
+                    # past_key_values.update expects [B,H,S,D]
+                    key_states, value_states = past_key_values.update(
+                        key_states, value_states, self.layer_idx, cache_kwargs=None
+                    )
+                    k = key_states.transpose(1, 2).contiguous()
+                    v = value_states.transpose(1, 2).contiguous()
+                else:
+                    # Optimized path:
+                    # use preallocated flash_k_cache / flash_v_cache
+                    layer = past_key_values.layers[self.layer_idx]
+
+                    if (
+                        hasattr(layer, "flash_k_cache")
+                        and layer.flash_k_cache is not None
+                        and hasattr(layer, "flash_v_cache")
+                        and layer.flash_v_cache is not None
+                    ):
+                        prefix_len = layer.flash_prefix_len
+                        cur_len = k_cur.shape[1]
+
+                        # overwrite current segment in-place
+                        layer.flash_k_cache[:, prefix_len:prefix_len + cur_len].copy_(k_cur)
+                        layer.flash_v_cache[:, prefix_len:prefix_len + cur_len].copy_(v_cur)
+
+                        k = layer.flash_k_cache[:, :prefix_len + cur_len]
+                        v = layer.flash_v_cache[:, :prefix_len + cur_len]
+                    else:
+                        # fallback if user forgot to prepare flash cache
+                        layer = past_key_values.layers[self.layer_idx]
+                        past_k, past_v = layer.keys, layer.values
+
+                        if past_k is not None:
+                            past_k = past_k.transpose(1, 2).contiguous()
+                            past_v = past_v.transpose(1, 2).contiguous()
+                            k = torch.cat([past_k, k_cur], dim=1)
+                            v = torch.cat([past_v, v_cur], dim=1)
+                        else:
+                            k = k_cur
+                            v = v_cur
+            else:
+                k = k_cur
+                v = v_cur
+
+            # sanity checks
+            assert q.ndim == 4 and k.ndim == 4 and v.ndim == 4
+            assert q.shape[0] == k.shape[0] == v.shape[0], (q.shape, k.shape, v.shape)
+            assert k.shape[1] == v.shape[1], (k.shape, v.shape)
+            assert k.shape[2] == v.shape[2], (k.shape, v.shape)
+            assert q.shape[3] == k.shape[3] == v.shape[3], (q.shape, k.shape, v.shape)
+
+            attn_output = flash_attn_func(
+                q,
+                k,
+                v,
+                dropout_p=0.0 if not self.training else self.attention_dropout,
+                softmax_scale=self.scaling,
+                causal=False,
+            )  # [B, S_q, H_q, D]
+
+            attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+            attn_output = self.o_proj_mot_gen(attn_output)
+            return attn_output, None
+
+        # ------------------------------------------------------------------
+        # Original eager fallback path
+        # ------------------------------------------------------------------
+        if past_key_values is not None:
+            if update_cache:
+                key_states, value_states = past_key_values.update(
+                    key_states, value_states, self.layer_idx, cache_kwargs=None
+                )
+            else:
+                layer = past_key_values.layers[self.layer_idx]
+                past_k, past_v = layer.keys, layer.values
+                if past_k is not None:
+                    key_states = torch.cat([past_k, key_states], dim=2)
+                    value_states = torch.cat([past_v, value_states], dim=2)
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj_mot_gen(attn_output)
+        return attn_output, attn_weights
+
+    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        image_gen_indicators: torch.Tensor,
+        exist_non_image_gen_tokens: bool,
+        exist_image_gen_tokens: bool,
+        indexes: Optional[torch.LongTensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if exist_non_image_gen_tokens and not exist_image_gen_tokens:
+            return self.forward_und(hidden_states, indexes, attention_mask, past_key_values, cache_position, **kwargs)
+        if not exist_non_image_gen_tokens and exist_image_gen_tokens:
+            return self.forward_gen(hidden_states, indexes, attention_mask, past_key_values, cache_position, **kwargs)
+        
+        assert self.config._attn_implementation == "eager"
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = hidden_states.new_zeros((*input_shape, self.config.num_attention_heads*self.head_dim))
+        if exist_non_image_gen_tokens:
+            query_states[~image_gen_indicators] = self.q_proj(hidden_states[~image_gen_indicators])
+        if exist_image_gen_tokens:
+            query_states[image_gen_indicators] = self.q_proj_mot_gen(hidden_states[image_gen_indicators])
+        query_states_t, query_states_hw = query_states.chunk(2, dim=-1)
+
+        _query_states_hw = query_states_hw.new_zeros(query_states_hw.shape)
+        if exist_non_image_gen_tokens:
+            _query_states_hw[~image_gen_indicators] = self.q_norm_hw(query_states_hw[~image_gen_indicators])
+        if exist_image_gen_tokens:
+            _query_states_hw[image_gen_indicators] = self.q_norm_hw_mot_gen(query_states_h[image_gen_indicators])
+        query_states_hw = _query_states_hw.transpose(1, 2)
+        query_states_h, query_states_w = query_states_hw.chunk(2, dim=-1)
+
+        key_states = hidden_states.new_zeros((*input_shape, self.config.num_key_value_heads*self.head_dim))
+        if exist_non_image_gen_tokens:
+            key_states[~image_gen_indicators] = self.k_proj(hidden_states[~image_gen_indicators])
+        if exist_image_gen_tokens:
+            key_states[image_gen_indicators] = self.k_proj_mot_gen(hidden_states[image_gen_indicators])
+        key_states_t, key_states_hw = key_states.chunk(2, dim=-1)
+
+        _key_states_hw = key_states_hw.new_zeros(key_states_hw.shape)
+        if exist_non_image_gen_tokens:
+            _key_states_hw[~image_gen_indicators] = self.k_norm_hw(key_states_hw[~image_gen_indicators])
+        if exist_image_gen_tokens:
+            _key_states_hw[image_gen_indicators] = self.k_norm_hw_mot_gen(key_states_h[image_gen_indicators])
+        key_states_hw = _key_states_hw.transpose(1, 2)
+        key_states_h, key_states_w = key_states_hw.chunk(2, dim=-1)
+
+        value_states = hidden_states.new_zeros((*input_shape, self.config.num_key_value_heads*self.head_dim))
+        if exist_non_image_gen_tokens:
+            value_states[~image_gen_indicators] = self.v_proj(hidden_states[~image_gen_indicators])
+        if exist_image_gen_tokens:
+            value_states[image_gen_indicators] = self.v_proj_mot_gen(hidden_states[image_gen_indicators])
+        value_states = value_states.view(hidden_shape).transpose(1, 2)
+
+        cos_t, sin_t = self.rotary_emb(hidden_states, indexes[0].unsqueeze(0))
+        query_states_t, key_states_t = apply_rotary_pos_emb(query_states_t, key_states_t, cos_t, sin_t)
+
+        cos_h, sin_h = self.rotary_emb_hw(hidden_states, indexes[1].unsqueeze(0))
+        query_states_h, key_states_h = apply_rotary_pos_emb(query_states_h, key_states_h, cos_h, sin_h)
+
+        cos_w, sin_w = self.rotary_emb_hw(hidden_states, indexes[2].unsqueeze(0))
+        query_states_w, key_states_w = apply_rotary_pos_emb(query_states_w, key_states_w, cos_w, sin_w)
+
+        query_states = torch.cat([query_states_t, query_states_h, query_states_w], dim=-1)
+        key_states = torch.cat([key_states_t, key_states_h, key_states_w], dim=-1)
+
+
+        if past_key_values is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            # cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            # key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+            update_cache = kwargs.get("update_cache", True)
+            if update_cache:
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs=None)
+            else:
+                # only use the past key values but do not append the current one
+                layer = past_key_values.layers[self.layer_idx]
+                past_k, past_v = layer.keys, layer.values
+
+                if past_k is not None:
+                    key_states   = torch.cat([past_k, key_states], dim=2)   # concat on seq_len
+                    value_states = torch.cat([past_v, value_states], dim=2)
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,  # diff with Llama
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+
+        _attn_output = attn_output.new_zeros((*input_shape, self.config.hidden_size))
+        if exist_non_image_gen_tokens:
+            _attn_output[~image_gen_indicators] = self.o_proj(attn_output[~image_gen_indicators])
+        if exist_image_gen_tokens:
+            _attn_output[image_gen_indicators] = self.o_proj_mot_gen(attn_output[image_gen_indicators])
+
+        attn_output = _attn_output
+        return attn_output, attn_weights
+
+
+class Qwen3DecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: Qwen3Config, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = Qwen3Attention(config=config, layer_idx=layer_idx)
+
+        self.mlp = Qwen3MLP(config)
+        self.mlp_mot_gen = Qwen3MLP(config)
+        self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.input_layernorm_mot_gen = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm_mot_gen = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.attention_type = config.layer_types[layer_idx]
+
+    def forward_und(
+        self,
+        hidden_states: torch.Tensor,
+        image_gen_indicators: torch.Tensor,
+        exist_non_image_gen_tokens: bool,
+        exist_image_gen_tokens: bool,
+        indexes: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        # Self Attention
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            image_gen_indicators=image_gen_indicators,
+            exist_non_image_gen_tokens=exist_non_image_gen_tokens,
+            exist_image_gen_tokens=exist_image_gen_tokens,
+            indexes=indexes,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+    def forward_gen(
+        self,
+        hidden_states: torch.Tensor,
+        image_gen_indicators: torch.Tensor,
+        exist_non_image_gen_tokens: bool,
+        exist_image_gen_tokens: bool,
+        indexes: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm_mot_gen(hidden_states)
+        # Self Attention
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            image_gen_indicators=image_gen_indicators,
+            exist_non_image_gen_tokens=exist_non_image_gen_tokens,
+            exist_image_gen_tokens=exist_image_gen_tokens,
+            indexes=indexes,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm_mot_gen(hidden_states)
+        hidden_states = self.mlp_mot_gen(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        image_gen_indicators: torch.Tensor,
+        exist_non_image_gen_tokens: bool,
+        exist_image_gen_tokens: bool,
+        indexes: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        if exist_non_image_gen_tokens and not exist_image_gen_tokens:
+            return self.forward_und(hidden_states, image_gen_indicators, exist_non_image_gen_tokens, exist_image_gen_tokens, indexes, attention_mask, position_ids, past_key_values, use_cache, cache_position, **kwargs)
+        if not exist_non_image_gen_tokens and exist_image_gen_tokens:
+            return self.forward_gen(hidden_states, image_gen_indicators, exist_non_image_gen_tokens, exist_image_gen_tokens, indexes, attention_mask, position_ids, past_key_values, use_cache, cache_position, **kwargs)
+
+        residual = hidden_states
+
+        _hidden_states = hidden_states.new_zeros(hidden_states.shape)
+        if exist_non_image_gen_tokens:
+            _hidden_states[~image_gen_indicators] = self.input_layernorm(hidden_states[~image_gen_indicators])
+        if exist_image_gen_tokens:
+            _hidden_states[image_gen_indicators] = self.input_layernorm_mot_gen(hidden_states[image_gen_indicators])
+        hidden_states = _hidden_states
+
+        # Self Attention
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            image_gen_indicators=image_gen_indicators,
+            exist_non_image_gen_tokens=exist_non_image_gen_tokens,
+            exist_image_gen_tokens=exist_image_gen_tokens,
+            indexes=indexes,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+
+        _hidden_states = hidden_states.new_zeros(hidden_states.shape)
+        if exist_non_image_gen_tokens:
+            _hidden_states[~image_gen_indicators] = self.mlp(self.post_attention_layernorm(hidden_states[~image_gen_indicators]))
+
+        if exist_image_gen_tokens:
+            _hidden_states[image_gen_indicators] = self.mlp_mot_gen(self.post_attention_layernorm_mot_gen(hidden_states[image_gen_indicators]))
+
+        hidden_states = _hidden_states
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+@auto_docstring
+class Qwen3PreTrainedModel(PreTrainedModel):
+    config: Qwen3Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Qwen3DecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": Qwen3DecoderLayer,
+        "attentions": Qwen3Attention,
+    }
+
+
+@auto_docstring
+class Qwen3Model(Qwen3PreTrainedModel):
+    def __init__(self, config: Qwen3Config):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [Qwen3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.norm_mot_gen = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        
+        self.gradient_checkpointing = False
+        self.has_sliding_layers = "sliding_attention" in self.config.layer_types
+        self.current_index = -1
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @check_model_inputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        image_gen_indicators: Optional[torch.Tensor] = None,
+        indexes: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        
+        # assert position_ids is not None
+        # assert cache_position is not None
+        # assert past_key_values is not None 
+
+        if image_gen_indicators is None:
+            exist_non_image_gen_tokens = True
+            exist_image_gen_tokens = False
+        else:
+            exist_non_image_gen_tokens = (~image_gen_indicators).any()
+            exist_image_gen_tokens = image_gen_indicators.any()
+        
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        # It may already have been prepared by e.g. `generate`
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            # Prepare mask arguments
+            if input_ids is not None:
+                mask_kwargs = {
+                    "config": self.config,
+                    "input_embeds": inputs_embeds,
+                    "attention_mask": attention_mask,
+                    "cache_position": cache_position,
+                    "past_key_values": past_key_values,
+                    "position_ids": position_ids,
+                }
+                # Create the masks
+                causal_mask_mapping = {
+                    "full_attention": create_causal_mask(**mask_kwargs),
+                }
+                self.current_index += 1
+                indexes = torch.LongTensor([[self.current_index], [0], [0]]).to(input_ids.device)
+            else:
+                causal_mask_mapping = {
+                    "full_attention": create_block_causal_mask(indexes[0]),
+                }
+                self.current_index = indexes[0].max()
+        else:
+            self.current_index = indexes[0].max()
+            # raise NotImplementedError('not isinstance(causal_mask_mapping := attention_mask, dict)')
+
+            # The sliding window alternating layers are not always activated depending on the config
+            # if self.has_sliding_layers:
+            #     causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
+
+        hidden_states = inputs_embeds
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            hidden_states = decoder_layer(
+                hidden_states,
+                image_gen_indicators=image_gen_indicators,
+                exist_non_image_gen_tokens=exist_non_image_gen_tokens,
+                exist_image_gen_tokens=exist_image_gen_tokens,
+                indexes=indexes,
+                attention_mask=causal_mask_mapping[decoder_layer.attention_type],
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                **kwargs,
+            )
+        if not exist_image_gen_tokens:
+            hidden_states = self.norm(hidden_states)
+        elif not exist_non_image_gen_tokens:
+            hidden_states = self.norm_mot_gen(hidden_states)
+        else:
+            _hidden_states = hidden_states.new_zeros(hidden_states.shape)
+            _hidden_states[~image_gen_indicators] = self.norm(hidden_states[~image_gen_indicators])
+            _hidden_states[image_gen_indicators] = self.norm_mot_gen(hidden_states[image_gen_indicators])
+            hidden_states = _hidden_states
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+
+@auto_docstring
+class Qwen3ForCausalLM(Qwen3PreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = Qwen3Model(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        indexes: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CausalLMOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, Qwen3ForCausalLM
+
+        >>> model = Qwen3ForCausalLM.from_pretrained("Qwen/Qwen3-8B")
+        >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            indexes=indexes,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class Qwen3ForSequenceClassification(GenericForSequenceClassification, Qwen3PreTrainedModel):
+    pass
+
+
+class Qwen3ForTokenClassification(GenericForTokenClassification, Qwen3PreTrainedModel):
+    pass
+
+
+class Qwen3ForQuestionAnswering(GenericForQuestionAnswering, Qwen3PreTrainedModel):
+    base_model_prefix = "transformer"  # For BC, where `transformer` was used instead of `model`
+
+
+__all__ = [
+    "Qwen3ForCausalLM",
+    "Qwen3ForQuestionAnswering",
+    "Qwen3PreTrainedModel",
+    "Qwen3Model",
+    "Qwen3ForSequenceClassification",
+    "Qwen3ForTokenClassification",
+]

special_tokens_map.json

@@ -0,0 +1,1900 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>",
+    {
+      "content": "<IMG_CONTEXT>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<img>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "</img>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<quad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "</quad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<ref>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "</ref>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<box>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "</box>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|action_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|action_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|plugin|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|interpreter|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_0>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_1>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_2>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_3>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_4>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_5>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_6>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_7>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_8>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_9>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_10>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_11>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_12>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_13>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_14>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_15>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_16>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_17>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_18>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_19>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_20>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_21>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_22>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_23>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_24>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_25>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_26>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_27>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_28>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_29>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_30>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_31>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_32>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_33>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_34>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_35>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_36>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_37>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_38>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_39>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_40>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_41>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_42>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_43>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_44>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_45>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_46>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_47>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_48>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_49>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_50>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_51>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_52>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_53>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_54>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_55>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_56>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_57>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_58>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_59>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_60>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_61>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_62>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_63>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_64>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_65>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_66>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_67>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_68>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_69>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_70>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_71>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_72>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_73>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_74>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_75>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_76>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_77>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_78>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_79>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_80>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_81>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_82>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_83>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_84>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_85>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_86>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_87>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_88>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_89>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_90>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_91>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_92>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_93>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_94>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_95>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_96>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_97>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_98>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_99>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_100>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_101>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_102>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_103>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_104>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_105>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_106>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_107>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_108>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_109>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_110>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_111>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_112>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_113>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_114>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_115>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_116>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_117>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_118>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_119>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_120>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_121>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_122>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_123>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_124>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_125>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_126>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_127>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_128>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_129>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_130>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_131>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_132>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_133>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_134>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_135>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_136>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_137>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_138>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_139>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_140>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_141>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_142>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_143>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_144>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_145>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_146>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_147>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_148>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_149>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_150>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_151>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_152>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_153>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_154>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_155>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_156>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_157>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_158>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_159>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_160>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_161>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_162>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_163>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_164>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_165>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_166>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_167>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_168>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_169>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_170>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_171>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_172>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_173>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_174>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_175>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_176>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_177>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_178>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_179>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_180>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_181>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_182>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_183>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_184>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_185>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_186>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_187>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_188>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_189>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_190>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_191>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_192>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_193>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_194>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_195>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_196>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_197>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_198>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_199>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_200>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_201>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_202>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_203>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_204>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_205>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_206>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_207>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_208>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_209>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_210>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_211>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_212>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_213>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_214>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_215>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_216>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_217>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_218>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_219>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_220>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_221>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_222>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_223>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_224>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_225>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_226>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_227>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_228>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_229>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_230>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_231>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_232>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_233>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_234>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_235>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_236>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_237>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_238>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_239>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_240>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_241>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_242>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_243>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_244>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_245>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_246>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_247>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_248>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_249>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_250>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_251>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_252>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<FAKE_PAD_253>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    }
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,2643 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151669": {
+      "content": "<IMG_CONTEXT>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151670": {
+      "content": "<img>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151671": {
+      "content": "</img>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151672": {
+      "content": "<quad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151673": {
+      "content": "</quad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151674": {
+      "content": "<ref>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151675": {
+      "content": "</ref>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151676": {
+      "content": "<box>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151677": {
+      "content": "</box>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151678": {
+      "content": "<|action_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151679": {
+      "content": "<|action_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151680": {
+      "content": "<|plugin|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151681": {
+      "content": "<|interpreter|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151682": {
+      "content": "<FAKE_PAD_0>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151683": {
+      "content": "<FAKE_PAD_1>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151684": {
+      "content": "<FAKE_PAD_2>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151685": {
+      "content": "<FAKE_PAD_3>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151686": {
+      "content": "<FAKE_PAD_4>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151687": {
+      "content": "<FAKE_PAD_5>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151688": {
+      "content": "<FAKE_PAD_6>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151689": {
+      "content": "<FAKE_PAD_7>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151690": {
+      "content": "<FAKE_PAD_8>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151691": {
+      "content": "<FAKE_PAD_9>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151692": {
+      "content": "<FAKE_PAD_10>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151693": {
+      "content": "<FAKE_PAD_11>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151694": {
+      "content": "<FAKE_PAD_12>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151695": {
+      "content": "<FAKE_PAD_13>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151696": {
+      "content": "<FAKE_PAD_14>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151697": {
+      "content": "<FAKE_PAD_15>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151698": {
+      "content": "<FAKE_PAD_16>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151699": {
+      "content": "<FAKE_PAD_17>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151700": {
+      "content": "<FAKE_PAD_18>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151701": {
+      "content": "<FAKE_PAD_19>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151702": {
+      "content": "<FAKE_PAD_20>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151703": {
+      "content": "<FAKE_PAD_21>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151704": {
+      "content": "<FAKE_PAD_22>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151705": {
+      "content": "<FAKE_PAD_23>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151706": {
+      "content": "<FAKE_PAD_24>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151707": {
+      "content": "<FAKE_PAD_25>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151708": {
+      "content": "<FAKE_PAD_26>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151709": {
+      "content": "<FAKE_PAD_27>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151710": {
+      "content": "<FAKE_PAD_28>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151711": {
+      "content": "<FAKE_PAD_29>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151712": {
+      "content": "<FAKE_PAD_30>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151713": {
+      "content": "<FAKE_PAD_31>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151714": {
+      "content": "<FAKE_PAD_32>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151715": {
+      "content": "<FAKE_PAD_33>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151716": {
+      "content": "<FAKE_PAD_34>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151717": {
+      "content": "<FAKE_PAD_35>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151718": {
+      "content": "<FAKE_PAD_36>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151719": {
+      "content": "<FAKE_PAD_37>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151720": {
+      "content": "<FAKE_PAD_38>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151721": {
+      "content": "<FAKE_PAD_39>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151722": {
+      "content": "<FAKE_PAD_40>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151723": {
+      "content": "<FAKE_PAD_41>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151724": {
+      "content": "<FAKE_PAD_42>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151725": {
+      "content": "<FAKE_PAD_43>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151726": {
+      "content": "<FAKE_PAD_44>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151727": {
+      "content": "<FAKE_PAD_45>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151728": {
+      "content": "<FAKE_PAD_46>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151729": {
+      "content": "<FAKE_PAD_47>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151730": {
+      "content": "<FAKE_PAD_48>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151731": {
+      "content": "<FAKE_PAD_49>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151732": {
+      "content": "<FAKE_PAD_50>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151733": {
+      "content": "<FAKE_PAD_51>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151734": {
+      "content": "<FAKE_PAD_52>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151735": {
+      "content": "<FAKE_PAD_53>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151736": {
+      "content": "<FAKE_PAD_54>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151737": {
+      "content": "<FAKE_PAD_55>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151738": {
+      "content": "<FAKE_PAD_56>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151739": {
+      "content": "<FAKE_PAD_57>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151740": {
+      "content": "<FAKE_PAD_58>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151741": {
+      "content": "<FAKE_PAD_59>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151742": {
+      "content": "<FAKE_PAD_60>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151743": {
+      "content": "<FAKE_PAD_61>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151744": {
+      "content": "<FAKE_PAD_62>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151745": {
+      "content": "<FAKE_PAD_63>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151746": {
+      "content": "<FAKE_PAD_64>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151747": {
+      "content": "<FAKE_PAD_65>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151748": {
+      "content": "<FAKE_PAD_66>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151749": {
+      "content": "<FAKE_PAD_67>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151750": {
+      "content": "<FAKE_PAD_68>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151751": {
+      "content": "<FAKE_PAD_69>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151752": {
+      "content": "<FAKE_PAD_70>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151753": {
+      "content": "<FAKE_PAD_71>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151754": {
+      "content": "<FAKE_PAD_72>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151755": {
+      "content": "<FAKE_PAD_73>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151756": {
+      "content": "<FAKE_PAD_74>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151757": {
+      "content": "<FAKE_PAD_75>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151758": {
+      "content": "<FAKE_PAD_76>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151759": {
+      "content": "<FAKE_PAD_77>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151760": {
+      "content": "<FAKE_PAD_78>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151761": {
+      "content": "<FAKE_PAD_79>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151762": {
+      "content": "<FAKE_PAD_80>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151763": {
+      "content": "<FAKE_PAD_81>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151764": {
+      "content": "<FAKE_PAD_82>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151765": {
+      "content": "<FAKE_PAD_83>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151766": {
+      "content": "<FAKE_PAD_84>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151767": {
+      "content": "<FAKE_PAD_85>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151768": {
+      "content": "<FAKE_PAD_86>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151769": {
+      "content": "<FAKE_PAD_87>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151770": {
+      "content": "<FAKE_PAD_88>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151771": {
+      "content": "<FAKE_PAD_89>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151772": {
+      "content": "<FAKE_PAD_90>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151773": {
+      "content": "<FAKE_PAD_91>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151774": {
+      "content": "<FAKE_PAD_92>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151775": {
+      "content": "<FAKE_PAD_93>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151776": {
+      "content": "<FAKE_PAD_94>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151777": {
+      "content": "<FAKE_PAD_95>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151778": {
+      "content": "<FAKE_PAD_96>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151779": {
+      "content": "<FAKE_PAD_97>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151780": {
+      "content": "<FAKE_PAD_98>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151781": {
+      "content": "<FAKE_PAD_99>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151782": {
+      "content": "<FAKE_PAD_100>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151783": {
+      "content": "<FAKE_PAD_101>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151784": {
+      "content": "<FAKE_PAD_102>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151785": {
+      "content": "<FAKE_PAD_103>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151786": {
+      "content": "<FAKE_PAD_104>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151787": {
+      "content": "<FAKE_PAD_105>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151788": {
+      "content": "<FAKE_PAD_106>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151789": {
+      "content": "<FAKE_PAD_107>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151790": {
+      "content": "<FAKE_PAD_108>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151791": {
+      "content": "<FAKE_PAD_109>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151792": {
+      "content": "<FAKE_PAD_110>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151793": {
+      "content": "<FAKE_PAD_111>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151794": {
+      "content": "<FAKE_PAD_112>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151795": {
+      "content": "<FAKE_PAD_113>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151796": {
+      "content": "<FAKE_PAD_114>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151797": {
+      "content": "<FAKE_PAD_115>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151798": {
+      "content": "<FAKE_PAD_116>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151799": {
+      "content": "<FAKE_PAD_117>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151800": {
+      "content": "<FAKE_PAD_118>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151801": {
+      "content": "<FAKE_PAD_119>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151802": {
+      "content": "<FAKE_PAD_120>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151803": {
+      "content": "<FAKE_PAD_121>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151804": {
+      "content": "<FAKE_PAD_122>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151805": {
+      "content": "<FAKE_PAD_123>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151806": {
+      "content": "<FAKE_PAD_124>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151807": {
+      "content": "<FAKE_PAD_125>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151808": {
+      "content": "<FAKE_PAD_126>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151809": {
+      "content": "<FAKE_PAD_127>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151810": {
+      "content": "<FAKE_PAD_128>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151811": {
+      "content": "<FAKE_PAD_129>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151812": {
+      "content": "<FAKE_PAD_130>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151813": {
+      "content": "<FAKE_PAD_131>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151814": {
+      "content": "<FAKE_PAD_132>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151815": {
+      "content": "<FAKE_PAD_133>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151816": {
+      "content": "<FAKE_PAD_134>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151817": {
+      "content": "<FAKE_PAD_135>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151818": {
+      "content": "<FAKE_PAD_136>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151819": {
+      "content": "<FAKE_PAD_137>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151820": {
+      "content": "<FAKE_PAD_138>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151821": {
+      "content": "<FAKE_PAD_139>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151822": {
+      "content": "<FAKE_PAD_140>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151823": {
+      "content": "<FAKE_PAD_141>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151824": {
+      "content": "<FAKE_PAD_142>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151825": {
+      "content": "<FAKE_PAD_143>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151826": {
+      "content": "<FAKE_PAD_144>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151827": {
+      "content": "<FAKE_PAD_145>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151828": {
+      "content": "<FAKE_PAD_146>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151829": {
+      "content": "<FAKE_PAD_147>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151830": {
+      "content": "<FAKE_PAD_148>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151831": {
+      "content": "<FAKE_PAD_149>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151832": {
+      "content": "<FAKE_PAD_150>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151833": {
+      "content": "<FAKE_PAD_151>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151834": {
+      "content": "<FAKE_PAD_152>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151835": {
+      "content": "<FAKE_PAD_153>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151836": {
+      "content": "<FAKE_PAD_154>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151837": {
+      "content": "<FAKE_PAD_155>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151838": {
+      "content": "<FAKE_PAD_156>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151839": {
+      "content": "<FAKE_PAD_157>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151840": {
+      "content": "<FAKE_PAD_158>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151841": {
+      "content": "<FAKE_PAD_159>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151842": {
+      "content": "<FAKE_PAD_160>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151843": {
+      "content": "<FAKE_PAD_161>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151844": {
+      "content": "<FAKE_PAD_162>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151845": {
+      "content": "<FAKE_PAD_163>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151846": {
+      "content": "<FAKE_PAD_164>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151847": {
+      "content": "<FAKE_PAD_165>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151848": {
+      "content": "<FAKE_PAD_166>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151849": {
+      "content": "<FAKE_PAD_167>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151850": {
+      "content": "<FAKE_PAD_168>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151851": {
+      "content": "<FAKE_PAD_169>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151852": {
+      "content": "<FAKE_PAD_170>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151853": {
+      "content": "<FAKE_PAD_171>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151854": {
+      "content": "<FAKE_PAD_172>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151855": {
+      "content": "<FAKE_PAD_173>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151856": {
+      "content": "<FAKE_PAD_174>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151857": {
+      "content": "<FAKE_PAD_175>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151858": {
+      "content": "<FAKE_PAD_176>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151859": {
+      "content": "<FAKE_PAD_177>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151860": {
+      "content": "<FAKE_PAD_178>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151861": {
+      "content": "<FAKE_PAD_179>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151862": {
+      "content": "<FAKE_PAD_180>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151863": {
+      "content": "<FAKE_PAD_181>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151864": {
+      "content": "<FAKE_PAD_182>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151865": {
+      "content": "<FAKE_PAD_183>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151866": {
+      "content": "<FAKE_PAD_184>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151867": {
+      "content": "<FAKE_PAD_185>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151868": {
+      "content": "<FAKE_PAD_186>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151869": {
+      "content": "<FAKE_PAD_187>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151870": {
+      "content": "<FAKE_PAD_188>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151871": {
+      "content": "<FAKE_PAD_189>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151872": {
+      "content": "<FAKE_PAD_190>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151873": {
+      "content": "<FAKE_PAD_191>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151874": {
+      "content": "<FAKE_PAD_192>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151875": {
+      "content": "<FAKE_PAD_193>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151876": {
+      "content": "<FAKE_PAD_194>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151877": {
+      "content": "<FAKE_PAD_195>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151878": {
+      "content": "<FAKE_PAD_196>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151879": {
+      "content": "<FAKE_PAD_197>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151880": {
+      "content": "<FAKE_PAD_198>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151881": {
+      "content": "<FAKE_PAD_199>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151882": {
+      "content": "<FAKE_PAD_200>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151883": {
+      "content": "<FAKE_PAD_201>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151884": {
+      "content": "<FAKE_PAD_202>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151885": {
+      "content": "<FAKE_PAD_203>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151886": {
+      "content": "<FAKE_PAD_204>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151887": {
+      "content": "<FAKE_PAD_205>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151888": {
+      "content": "<FAKE_PAD_206>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151889": {
+      "content": "<FAKE_PAD_207>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151890": {
+      "content": "<FAKE_PAD_208>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151891": {
+      "content": "<FAKE_PAD_209>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151892": {
+      "content": "<FAKE_PAD_210>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151893": {
+      "content": "<FAKE_PAD_211>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151894": {
+      "content": "<FAKE_PAD_212>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151895": {
+      "content": "<FAKE_PAD_213>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151896": {
+      "content": "<FAKE_PAD_214>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151897": {
+      "content": "<FAKE_PAD_215>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151898": {
+      "content": "<FAKE_PAD_216>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151899": {
+      "content": "<FAKE_PAD_217>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151900": {
+      "content": "<FAKE_PAD_218>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151901": {
+      "content": "<FAKE_PAD_219>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151902": {
+      "content": "<FAKE_PAD_220>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151903": {
+      "content": "<FAKE_PAD_221>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151904": {
+      "content": "<FAKE_PAD_222>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151905": {
+      "content": "<FAKE_PAD_223>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151906": {
+      "content": "<FAKE_PAD_224>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151907": {
+      "content": "<FAKE_PAD_225>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151908": {
+      "content": "<FAKE_PAD_226>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151909": {
+      "content": "<FAKE_PAD_227>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151910": {
+      "content": "<FAKE_PAD_228>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151911": {
+      "content": "<FAKE_PAD_229>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151912": {
+      "content": "<FAKE_PAD_230>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151913": {
+      "content": "<FAKE_PAD_231>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151914": {
+      "content": "<FAKE_PAD_232>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151915": {
+      "content": "<FAKE_PAD_233>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151916": {
+      "content": "<FAKE_PAD_234>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151917": {
+      "content": "<FAKE_PAD_235>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151918": {
+      "content": "<FAKE_PAD_236>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151919": {
+      "content": "<FAKE_PAD_237>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151920": {
+      "content": "<FAKE_PAD_238>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151921": {
+      "content": "<FAKE_PAD_239>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151922": {
+      "content": "<FAKE_PAD_240>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151923": {
+      "content": "<FAKE_PAD_241>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151924": {
+      "content": "<FAKE_PAD_242>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151925": {
+      "content": "<FAKE_PAD_243>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151926": {
+      "content": "<FAKE_PAD_244>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151927": {
+      "content": "<FAKE_PAD_245>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151928": {
+      "content": "<FAKE_PAD_246>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151929": {
+      "content": "<FAKE_PAD_247>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151930": {
+      "content": "<FAKE_PAD_248>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151931": {
+      "content": "<FAKE_PAD_249>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151932": {
+      "content": "<FAKE_PAD_250>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151933": {
+      "content": "<FAKE_PAD_251>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151934": {
+      "content": "<FAKE_PAD_252>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151935": {
+      "content": "<FAKE_PAD_253>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>",
+    "<IMG_CONTEXT>",
+    "<img>",
+    "</img>",
+    "<quad>",
+    "</quad>",
+    "<ref>",
+    "</ref>",
+    "<box>",
+    "</box>",
+    "<|action_start|>",
+    "<|action_end|>",
+    "<|plugin|>",
+    "<|interpreter|>",
+    "<FAKE_PAD_0>",
+    "<FAKE_PAD_1>",
+    "<FAKE_PAD_2>",
+    "<FAKE_PAD_3>",
+    "<FAKE_PAD_4>",
+    "<FAKE_PAD_5>",
+    "<FAKE_PAD_6>",
+    "<FAKE_PAD_7>",
+    "<FAKE_PAD_8>",
+    "<FAKE_PAD_9>",
+    "<FAKE_PAD_10>",
+    "<FAKE_PAD_11>",
+    "<FAKE_PAD_12>",
+    "<FAKE_PAD_13>",
+    "<FAKE_PAD_14>",
+    "<FAKE_PAD_15>",
+    "<FAKE_PAD_16>",
+    "<FAKE_PAD_17>",
+    "<FAKE_PAD_18>",
+    "<FAKE_PAD_19>",
+    "<FAKE_PAD_20>",
+    "<FAKE_PAD_21>",
+    "<FAKE_PAD_22>",
+    "<FAKE_PAD_23>",
+    "<FAKE_PAD_24>",
+    "<FAKE_PAD_25>",
+    "<FAKE_PAD_26>",
+    "<FAKE_PAD_27>",
+    "<FAKE_PAD_28>",
+    "<FAKE_PAD_29>",
+    "<FAKE_PAD_30>",
+    "<FAKE_PAD_31>",
+    "<FAKE_PAD_32>",
+    "<FAKE_PAD_33>",
+    "<FAKE_PAD_34>",
+    "<FAKE_PAD_35>",
+    "<FAKE_PAD_36>",
+    "<FAKE_PAD_37>",
+    "<FAKE_PAD_38>",
+    "<FAKE_PAD_39>",
+    "<FAKE_PAD_40>",
+    "<FAKE_PAD_41>",
+    "<FAKE_PAD_42>",
+    "<FAKE_PAD_43>",
+    "<FAKE_PAD_44>",
+    "<FAKE_PAD_45>",
+    "<FAKE_PAD_46>",
+    "<FAKE_PAD_47>",
+    "<FAKE_PAD_48>",
+    "<FAKE_PAD_49>",
+    "<FAKE_PAD_50>",
+    "<FAKE_PAD_51>",
+    "<FAKE_PAD_52>",
+    "<FAKE_PAD_53>",
+    "<FAKE_PAD_54>",
+    "<FAKE_PAD_55>",
+    "<FAKE_PAD_56>",
+    "<FAKE_PAD_57>",
+    "<FAKE_PAD_58>",
+    "<FAKE_PAD_59>",
+    "<FAKE_PAD_60>",
+    "<FAKE_PAD_61>",
+    "<FAKE_PAD_62>",
+    "<FAKE_PAD_63>",
+    "<FAKE_PAD_64>",
+    "<FAKE_PAD_65>",
+    "<FAKE_PAD_66>",
+    "<FAKE_PAD_67>",
+    "<FAKE_PAD_68>",
+    "<FAKE_PAD_69>",
+    "<FAKE_PAD_70>",
+    "<FAKE_PAD_71>",
+    "<FAKE_PAD_72>",
+    "<FAKE_PAD_73>",
+    "<FAKE_PAD_74>",
+    "<FAKE_PAD_75>",
+    "<FAKE_PAD_76>",
+    "<FAKE_PAD_77>",
+    "<FAKE_PAD_78>",
+    "<FAKE_PAD_79>",
+    "<FAKE_PAD_80>",
+    "<FAKE_PAD_81>",
+    "<FAKE_PAD_82>",
+    "<FAKE_PAD_83>",
+    "<FAKE_PAD_84>",
+    "<FAKE_PAD_85>",
+    "<FAKE_PAD_86>",
+    "<FAKE_PAD_87>",
+    "<FAKE_PAD_88>",
+    "<FAKE_PAD_89>",
+    "<FAKE_PAD_90>",
+    "<FAKE_PAD_91>",
+    "<FAKE_PAD_92>",
+    "<FAKE_PAD_93>",
+    "<FAKE_PAD_94>",
+    "<FAKE_PAD_95>",
+    "<FAKE_PAD_96>",
+    "<FAKE_PAD_97>",
+    "<FAKE_PAD_98>",
+    "<FAKE_PAD_99>",
+    "<FAKE_PAD_100>",
+    "<FAKE_PAD_101>",
+    "<FAKE_PAD_102>",
+    "<FAKE_PAD_103>",
+    "<FAKE_PAD_104>",
+    "<FAKE_PAD_105>",
+    "<FAKE_PAD_106>",
+    "<FAKE_PAD_107>",
+    "<FAKE_PAD_108>",
+    "<FAKE_PAD_109>",
+    "<FAKE_PAD_110>",
+    "<FAKE_PAD_111>",
+    "<FAKE_PAD_112>",
+    "<FAKE_PAD_113>",
+    "<FAKE_PAD_114>",
+    "<FAKE_PAD_115>",
+    "<FAKE_PAD_116>",
+    "<FAKE_PAD_117>",
+    "<FAKE_PAD_118>",
+    "<FAKE_PAD_119>",
+    "<FAKE_PAD_120>",
+    "<FAKE_PAD_121>",
+    "<FAKE_PAD_122>",
+    "<FAKE_PAD_123>",
+    "<FAKE_PAD_124>",
+    "<FAKE_PAD_125>",
+    "<FAKE_PAD_126>",
+    "<FAKE_PAD_127>",
+    "<FAKE_PAD_128>",
+    "<FAKE_PAD_129>",
+    "<FAKE_PAD_130>",
+    "<FAKE_PAD_131>",
+    "<FAKE_PAD_132>",
+    "<FAKE_PAD_133>",
+    "<FAKE_PAD_134>",
+    "<FAKE_PAD_135>",
+    "<FAKE_PAD_136>",
+    "<FAKE_PAD_137>",
+    "<FAKE_PAD_138>",
+    "<FAKE_PAD_139>",
+    "<FAKE_PAD_140>",
+    "<FAKE_PAD_141>",
+    "<FAKE_PAD_142>",
+    "<FAKE_PAD_143>",
+    "<FAKE_PAD_144>",
+    "<FAKE_PAD_145>",
+    "<FAKE_PAD_146>",
+    "<FAKE_PAD_147>",
+    "<FAKE_PAD_148>",
+    "<FAKE_PAD_149>",
+    "<FAKE_PAD_150>",
+    "<FAKE_PAD_151>",
+    "<FAKE_PAD_152>",
+    "<FAKE_PAD_153>",
+    "<FAKE_PAD_154>",
+    "<FAKE_PAD_155>",
+    "<FAKE_PAD_156>",
+    "<FAKE_PAD_157>",
+    "<FAKE_PAD_158>",
+    "<FAKE_PAD_159>",
+    "<FAKE_PAD_160>",
+    "<FAKE_PAD_161>",
+    "<FAKE_PAD_162>",
+    "<FAKE_PAD_163>",
+    "<FAKE_PAD_164>",
+    "<FAKE_PAD_165>",
+    "<FAKE_PAD_166>",
+    "<FAKE_PAD_167>",
+    "<FAKE_PAD_168>",
+    "<FAKE_PAD_169>",
+    "<FAKE_PAD_170>",
+    "<FAKE_PAD_171>",
+    "<FAKE_PAD_172>",
+    "<FAKE_PAD_173>",
+    "<FAKE_PAD_174>",
+    "<FAKE_PAD_175>",
+    "<FAKE_PAD_176>",
+    "<FAKE_PAD_177>",
+    "<FAKE_PAD_178>",
+    "<FAKE_PAD_179>",
+    "<FAKE_PAD_180>",
+    "<FAKE_PAD_181>",
+    "<FAKE_PAD_182>",
+    "<FAKE_PAD_183>",
+    "<FAKE_PAD_184>",
+    "<FAKE_PAD_185>",
+    "<FAKE_PAD_186>",
+    "<FAKE_PAD_187>",
+    "<FAKE_PAD_188>",
+    "<FAKE_PAD_189>",
+    "<FAKE_PAD_190>",
+    "<FAKE_PAD_191>",
+    "<FAKE_PAD_192>",
+    "<FAKE_PAD_193>",
+    "<FAKE_PAD_194>",
+    "<FAKE_PAD_195>",
+    "<FAKE_PAD_196>",
+    "<FAKE_PAD_197>",
+    "<FAKE_PAD_198>",
+    "<FAKE_PAD_199>",
+    "<FAKE_PAD_200>",
+    "<FAKE_PAD_201>",
+    "<FAKE_PAD_202>",
+    "<FAKE_PAD_203>",
+    "<FAKE_PAD_204>",
+    "<FAKE_PAD_205>",
+    "<FAKE_PAD_206>",
+    "<FAKE_PAD_207>",
+    "<FAKE_PAD_208>",
+    "<FAKE_PAD_209>",
+    "<FAKE_PAD_210>",
+    "<FAKE_PAD_211>",
+    "<FAKE_PAD_212>",
+    "<FAKE_PAD_213>",
+    "<FAKE_PAD_214>",
+    "<FAKE_PAD_215>",
+    "<FAKE_PAD_216>",
+    "<FAKE_PAD_217>",
+    "<FAKE_PAD_218>",
+    "<FAKE_PAD_219>",
+    "<FAKE_PAD_220>",
+    "<FAKE_PAD_221>",
+    "<FAKE_PAD_222>",
+    "<FAKE_PAD_223>",
+    "<FAKE_PAD_224>",
+    "<FAKE_PAD_225>",
+    "<FAKE_PAD_226>",
+    "<FAKE_PAD_227>",
+    "<FAKE_PAD_228>",
+    "<FAKE_PAD_229>",
+    "<FAKE_PAD_230>",
+    "<FAKE_PAD_231>",
+    "<FAKE_PAD_232>",
+    "<FAKE_PAD_233>",
+    "<FAKE_PAD_234>",
+    "<FAKE_PAD_235>",
+    "<FAKE_PAD_236>",
+    "<FAKE_PAD_237>",
+    "<FAKE_PAD_238>",
+    "<FAKE_PAD_239>",
+    "<FAKE_PAD_240>",
+    "<FAKE_PAD_241>",
+    "<FAKE_PAD_242>",
+    "<FAKE_PAD_243>",
+    "<FAKE_PAD_244>",
+    "<FAKE_PAD_245>",
+    "<FAKE_PAD_246>",
+    "<FAKE_PAD_247>",
+    "<FAKE_PAD_248>",
+    "<FAKE_PAD_249>",
+    "<FAKE_PAD_250>",
+    "<FAKE_PAD_251>",
+    "<FAKE_PAD_252>",
+    "<FAKE_PAD_253>"
+  ],
+  "bos_token": null,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set content = message.content %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is defined and message.reasoning_content is not none %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- else %}\n            {%- if '</think>' in message.content %}\n                {%- set content = message.content.split('</think>')[-1].lstrip('\\n') %}\n                {%- set reasoning_content = message.content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n            {%- endif %}\n        {%- endif %}\n        {%- if loop.index0 > ns.last_query_index %}\n            {%- if loop.last or (not loop.last and reasoning_content) %}\n                {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n            {%- else %}\n                {{- '<|im_start|>' + message.role + '\\n' + content }}\n            {%- endif %}\n        {%- else %}\n            {{- '<|im_start|>' + message.role + '\\n' + content }}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "model_max_length": 12288,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

utils.py

@@ -0,0 +1,146 @@
+import re
+import math
+import torch
+import string
+import numpy as np
+import pandas as pd
+from PIL import Image
+import torch.distributed as dist
+import torchvision.transforms as T
+from transformers import AutoModel, AutoTokenizer
+
+SYSTEM_MESSAGE_FOR_GEN = "You are an image generation and editing assistant that accurately understands and executes user intent.\n\nYou support two modes:\n\n1. Think Mode:\nIf the task requires reasoning, you MUST start with a <think></think> block. Put all reasoning inside the block using plain text. DO NOT include any image tags. Keep it reasonable and directly useful for producing the final image.\n\n2. Non-Think Mode:\nIf no reasoning is needed, directly produce the final image.\n\nTask Types:\n\nA. Text-to-Image Generation:\n- Generate a high-quality image based on the user's description.\n- Ensure visual clarity, semantic consistency, and completeness.\n- DO NOT introduce elements that contradict or override the user's intent.\n\nB. Image Editing:\n- Use the provided image(s) as input or reference for modification or transformation.\n- The result can be an edited image or a new image based on the reference(s).\n- Preserve all unspecified attributes unless explicitly changed.\n\nGeneral Rules:\n- For any visible text in the image, follow the language specified for the rendered text in the user's description, not the language of the prompt. If no language is specified, use the user's input language."
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+def round_by_factor(number: int, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+
+def ceil_by_factor(number: int, factor: int) -> int:
+    """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+    return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: int, factor: int) -> int:
+    """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+    return math.floor(number / factor) * factor
+
+
+# copy from https://github.com/QwenLM/Qwen2.5-VL/blob/main/qwen-vl-utils/src/qwen_vl_utils/vision_process.py#L60
+def smart_resize(
+    height: int, width: int, factor: int = 32, min_pixels: int = 65536, max_pixels: int = 4194304
+) -> tuple[int, int]:
+    """
+    Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+    3. The aspect ratio of the image is maintained as closely as possible.
+    """
+    if max(height, width) / min(height, width) > 200:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than {200}, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = max(factor, floor_by_factor(height / beta, factor))
+        w_bar = max(factor, floor_by_factor(width / beta, factor))
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(height * beta, factor)
+        w_bar = ceil_by_factor(width * beta, factor)
+    return h_bar, w_bar
+
+
+def dynamic_preprocess_native_resolution(
+    image, size_factor=32, min_pixels=65536, max_pixels=4194304, **kwargs
+):
+    width, height = image.size
+    resized_height, resized_width = smart_resize(
+        height,
+        width,
+        factor=size_factor,
+        min_pixels=min_pixels,
+        max_pixels=max_pixels,
+    )
+    image = image.resize((resized_width, resized_height))
+
+    return image
+
+def preprocess_pixel_values(pixel_values, patch_size=16):
+    c, h, w = pixel_values.shape
+    grid_h = h // patch_size
+    grid_w = w // patch_size
+
+    flatten_pixel_values = (
+        pixel_values.view(c, grid_h, patch_size, grid_w, patch_size)
+        .permute(1, 3, 0, 2, 4)  # [grid_h, grid_w, c, patch_size, patch_size]
+        .reshape(grid_h * grid_w, c * patch_size ** 2)
+    )
+
+    grid_hw = torch.tensor([[grid_h, grid_w]]).to(device=pixel_values.device)
+
+    return flatten_pixel_values, grid_hw
+
+def get_contrasting_background(image):
+    """
+    Calculate the color (white or black) that is different from the average foreground color
+    to use as the background color
+    """
+    image_np = np.array(image)
+    if (image_np[:, :, 3] == 0).any():
+        non_transparent_pixels = image_np[:, :, :3][image_np[:, :, 3] > 0]
+        if non_transparent_pixels.size == 0:
+            return None
+        pixel_mean = non_transparent_pixels.mean()
+        contrasting_color = (0, 0, 0) if pixel_mean > 382.5 else (255, 255, 255)
+        return contrasting_color
+    else:
+        return None
+
+def load_image_native(
+    image, patch_size=16, downsample_ratio=0.5, min_pixels=65536, max_pixels=4194304, upscale=False
+):
+    """
+    Load and preprocess an image file, converting it to RGB mode,
+    resizing, normalizing, and optionally adding a thumbnail version.
+    """
+    if not isinstance(image, Image.Image):
+        image = Image.open(image)
+    if image.mode == "RGBA":
+        bg_color = get_contrasting_background(image)
+        if bg_color:
+            background = Image.new("RGB", image.size, bg_color)
+            background.paste(image, mask=image.split()[3])
+            image = background.convert("RGB")
+        else:
+            image = image.convert("RGB")
+    else:
+        image = image.convert("RGB")
+    
+    if upscale:
+        image = image.resize((image.width * 2, image.height * 2), Image.BILINEAR)
+
+    transform = T.Compose(
+        [
+            T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),
+            T.ToTensor(),
+            T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+        ]
+    )
+
+    new_image = dynamic_preprocess_native_resolution(
+        image, size_factor=int(patch_size // downsample_ratio), min_pixels=min_pixels, max_pixels=max_pixels
+    )
+    pixel_values, grid_hw = preprocess_pixel_values(transform(new_image).to(torch.float32), patch_size=patch_size)
+
+    # print(f'Transfer image_size from ({image.height, image.width}) to ({new_image.height, new_image.width})')
+
+    return pixel_values, grid_hw