data/datasets_custom/validation_dataset.py

# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# coding: utf-8

import json
import os
from typing import Any, Dict, List

import numpy as np
import torch
from torch.utils.data import Dataset
import decord
from decord import VideoReader
from PIL import Image

from data.video.sampler.utils import FRAME_SAMPLER_TYPES
from data.video.sampler.frames import FrameSamplerOutput
from data.transforms import VideoTransform
from data.data_utils import (
    get_flattened_position_ids_extrapolate_video,
    len2weight,
    patchify_video_with_merge,
)
from data.system_prompt_render import render_qwenvl_prompt, expand_and_index_by_token_ids_new
from data.common import generate_system_prompt
from modeling.qwen2 import Qwen2Tokenizer
from config.config_factory import ModelArguments, DataArguments, TrainingArguments

sample_task_map = {
    't2v': 0,
    'idip': 1,
    'edit': 2,
    'refedit': 3,
}
modality_map = {
    'system_prompt': -1,
    'text': 0,
    'noise': 1,
    'ref_source': 2, # for vae
    'ref_image': 3, # for vae
    'ref_vit': 4 # for ref vit
}


class ValidationDataset(Dataset):
    def __init__(
        self,
        jsonl_path: str,
        tokenizer: Qwen2Tokenizer,
        data_args: DataArguments,
        model_args: ModelArguments,
        training_args: TrainingArguments,
        new_token_ids: Dict[str, int],
        dataset_config: None,
        local_rank: int = 0,
        world_size: int = 1,
    ):
        """
        初始化验证数据集

        Args:
            jsonl_path: JSONL文件路径
            tokenizer: 分词器
        """
        self.jsonl_path = jsonl_path
        self.tokenizer = tokenizer
        self.new_token_ids = new_token_ids

        # 读取JSONL文件
        try:
            full_data = self._read_jsonl()
        except:
            with open(jsonl_path, 'r', encoding='utf-8') as f:
                full_data = json.load(f)
            if isinstance(full_data, dict):
                # 转换为列表格式，每个元素是独立的字典
                full_data = [{"index": self.pro_index(index), "data": prompt} for index, prompt in full_data.items()]

        if world_size > 1:
            self.data = full_data[local_rank::world_size]
            print(f"Rank {local_rank}/{world_size} will process {len(self.data)} samples")
        else:
            self.data = full_data

        self.data_config = dataset_config

        self.bos_token_id = self.new_token_ids["bos_token_id"]
        self.eos_token_id = self.new_token_ids["eos_token_id"]
        self.start_of_image = self.new_token_ids["start_of_image"]
        self.end_of_image = self.new_token_ids["end_of_image"]
        self.image_token_id = self.new_token_ids["image_token_id"]

        # 视频采样
        try:
            max_duration = self.data_config.max_duration
        except:
            max_duration = 6.0

        video_frame_sampler_params = {"temporal": 4, "sample_fps": 12, "max_duration": max_duration, "assert_seconds": False, "truncate": False}

        self.frame_sampler = FRAME_SAMPLER_TYPES["multi_clips"](**video_frame_sampler_params)
        self.cpu_count = os.cpu_count() or 1

        # VideoTransform for vae: 仅在存在原始视频时才发挥作用
        if self.data_config.resolution in ["video_192p", "image_256res"]:
            resolution_vae = 256
            resolution_vit = 224
        elif self.data_config.resolution == "image_512res":
            resolution_vae = 512
            resolution_vit = 448
        elif self.data_config.resolution == "image_768res":
            resolution_vae = 768
            resolution_vit = 672
        elif self.data_config.resolution == "video_360p":
            resolution_vae = 480  # 480 for 360fps # 256 for 192p
            resolution_vit = 476  # 476 for 360fps , 224 for 192p
        elif self.data_config.resolution == "video_480p":
            resolution_vae = 640  # 480 for 360fps # 256 for 192p
            resolution_vit = 616  # 476 for 360fps , 224 for 192p
        else:
            raise ValueError(f"Unknown resolution: {self.data_config.resolution}")

        video_transform_args = {
            "resolution": resolution_vae,
            "mode": "bucket",
            "divisible_crop_size": 16,  # 32 # 16 | 32 让视频的分辨率被多少整除
            "stride_spatial": 16,  # 空间下采样倍率
            "stride_temporal": 4,  # 时间下采样倍率
            "aspect_ratios": ["21:9", "16:9", "4:3", "1:1", "3:4", "9:16"],  # 仅在 mode="bucket" 时生效
            "mean": 0.5,
            "std": 0.5,
        }
        self.transform = VideoTransform(**video_transform_args)

        # VideoTransform for vit
        vit_video_transform_args = {
            "resolution": resolution_vit,
            "mode": "bucket",
            "divisible_crop_size": 28,  # 让视频的分辨率被多少整除, qwen2.5vl需要被14整除
            "aspect_ratios": ["21:9", "16:9", "4:3", "1:1", "3:4", "9:16"],  # 仅在 mode="bucket" 时生效
            "mean": [0.48145466, 0.4578275, 0.40821073],  # Qwen2.5-VL vit 使用的mean
            "std": [0.26862954, 0.26130258, 0.27577711],
        }
        self.vit_transform = VideoTransform(**vit_video_transform_args)

        self.sample = self.set_sequence_status()

        self.frame_condition_idx = []

        if hasattr(self.data_config, 'system_prompt_type'):
            self.system_prompt_type = self.data_config.system_prompt_type
        else:
            self.system_prompt_type = 'SP0'

    def pro_index(self, index: int):
        if isinstance(index, str):
            for x in ['.mp4', '.jpg', '.png', '.jpeg']:
                index = index.replace(x, "")
        return int(index)

    def set_sequence_status(self):
        sequence_status = dict(
            curr=0,  # 指针
            sample_lens=[],
            sample_type=[],
            sample_N_target=[],
            packed_position_ids=[],
            nested_attention_masks=[],
            split_lens=[],
            attn_modes=[],
            packed_text_ids=[],
            packed_text_indexes=[],
            packed_label_ids=[],
            ce_loss_indexes=[],
            ce_loss_weights=[],
            vae_image_tensors=[],  # image
            vae_video_tensors=[],  # video
            packed_latent_position_ids=[],
            vae_latent_shapes=[],
            packed_vae_token_indexes=[],
            packed_timesteps=[],
            mse_loss_indexes=[],
            packed_vit_tokens=[],
            vit_token_seqlens=[],
            packed_vit_position_ids=[],
            packed_vit_token_indexes=[],
            vit_video_grid_thw=[],  # for vit video
            vae_video_grid_thw=[],  # for vae video
            video_grid_thw=[],  # for all video tensor
            vit_video_tensors=[],  # for vit original video tensor
            # offline 参数
            vae_video_latent=[],  # for vae video latent offline
            vae_data_mode=[],  # offline or online
            vit_data_mode=[],  # offline or online
            key_frame_mask=[],  # for key frame mask
            # sample_task for joint training
            sample_task=[],
            sample_modality=[],
        )
        return sequence_status

    def _read_jsonl(self) -> List[Dict[str, Any]]:
        """读取JSONL文件"""
        data = []
        with open(self.jsonl_path, "r", encoding="utf-8") as f:
            for line in f:
                data.append(json.loads(line.strip()))
        return data

    def __len__(self) -> int:
        return len(self.data)


    @staticmethod
    def _read_decord(video: VideoReader, frame_idx: List[int]) -> List[Image.Image]:
        # 使用 get_batch() 替换循环单帧读取，可以大幅提升性能
        frames_np = video.get_batch(frame_idx).asnumpy()
        return [Image.fromarray(frame) for frame in frames_np]

    def get_video_tensor_online(self, media_url, vision_stream, worker_id=0, element_dtype="image") -> torch.Tensor:
        self.vision_stream = vision_stream
        video_stream = media_url  # BytesIO(self.tos_cli.get_obj_by_url(media_url))

        if element_dtype == "image":
            image = Image.open(video_stream)
            if image.mode == "P":
                image = image.convert("RGBA")
            if image.mode == "RGBA":
                # 在白底上合成，去掉透明
                bg = Image.new("RGB", image.size, (255, 255, 255))
                bg.paste(image, mask=image.split()[3])  # 用 alpha 通道做掩码
                image = bg
            else:
                image = image.convert("RGB")
            video_frames = [image]
        else:  # for video
            video_reader = VideoReader(video_stream, ctx=decord.cpu(worker_id % self.cpu_count))
            total_frames = len(video_reader)

            sampler_name = self.frame_sampler.__class__.__name__
            if sampler_name == "MultiClipsFrameSampler":
                frames_info = {
                    "clip_indices": [(0, total_frames)],  # 左闭右开 默认为单个clip
                    "fps": 24,  # 默认为24
                }
            elif sampler_name == "FixedFrameSampler":
                frames_info = {
                    "start_frame": 0,
                    "end_frame": total_frames,
                    "total_frames": total_frames,
                }
            else:
                raise ValueError(f"Not verified frame sampler type: {sampler_name}")

            frames_sampler_output: FrameSamplerOutput = self.frame_sampler(frames_info)
            video_frames = self._read_decord(video_reader, frames_sampler_output.indices)

        if vision_stream == "vae_video":
            video_tensor = self.transform(video_frames)  # fix: use List input
        elif vision_stream == "vit_video":
            video_tensor = self.vit_transform(video_frames)  # fix: use List input
            if element_dtype == "image":
                video_tensor = video_tensor.repeat(1, 2, 1, 1)  # NOTE 对于单张图像，需要复制一份，因为encoder的temporal patch size = 2
            # NOTE: 视频长度必须是偶数
            if video_tensor.shape[1] % 2 == 1:
                last_frame = video_tensor[:, -1:, :, :]
                video_tensor = torch.cat([video_tensor, last_frame], dim=1)

        else:
            raise ValueError(f"Unknown vision_stream: {vision_stream}")
        return video_tensor  # , self.vision_token_count(video_tensor)

    def process_vit_video(self, video_tensor, curr: int, curr_rope_id: int, curr_split_len: int, curr_video_grid_thw: None, item_loss=0):
        if not self.data_config.text_template:
            self.sample["packed_text_ids"].append(self.start_of_image)  # 151652, <|vision_start|>
            self.sample["packed_text_indexes"].append(curr)
            curr += 1
            curr_split_len += 1

        # 在线模式下，video_tensor 为tensor, 离线模式下，video_tensor 为list [latent]
        if isinstance(video_tensor, torch.Tensor):  # online
            self.sample["vit_video_tensors"].append(video_tensor)  # CTHW 原始的视频，非latent , 仅用于validation中的可视化

            # preprocess video
            vit_tokens = patchify_video_with_merge(
                video_tensor, self.data_config.vit_patch_size, self.data_config.vit_patch_size_temporal
            )  # C T H W -> (T//2 * H//p * W//p) (p*p*2*C)
            num_video_tokens = vit_tokens.shape[0] // 4  # 实际上qwen2.5-vl还需要merge，2x2 merge成1个， hardcode for temp
            t, h, w = video_tensor.size(1), video_tensor.size(2), video_tensor.size(3)

            self.sample["packed_vit_tokens"].append(vit_tokens)
            self.sample["vit_data_mode"].append("online")

        if t is not None:
            vit_video_grid_thw = [
                t // self.data_config.vit_patch_size_temporal,
                h // self.data_config.vit_patch_size,
                w // self.data_config.vit_patch_size,
            ]  # [1, 16, 16]
        self.sample["vit_video_grid_thw"].append(vit_video_grid_thw)
        curr_video_grid_thw.append(vit_video_grid_thw)

        self.sample["vit_token_seqlens"].append(num_video_tokens)
        self.sample["packed_vit_position_ids"].append(
            torch.zeros(num_video_tokens)
        )  # TODO : 不一定是 0 ？ 对于多个vit序列会有问题

        if not self.data_config.text_template:
            self.sample["packed_vit_token_indexes"].extend(range(curr, curr + num_video_tokens))
            curr += num_video_tokens
            curr_split_len += num_video_tokens

            # NOTE dummy position_ids
            self.sample["packed_text_ids"].extend([self.image_token_id] * num_video_tokens)

            # add a <|endofimage|> token
            self.sample["packed_text_ids"].append(self.end_of_image)  # 151653, <|vision_end|>
            self.sample["packed_text_indexes"].append(curr)
            curr += 1
            curr_split_len += 1
            self.sample["packed_position_ids"].extend([curr_rope_id] * curr_split_len)
            curr_rope_id += 1

            # update sequence status
            self.sample["attn_modes"].append("full")
            self.sample["split_lens"].append(curr_split_len)

        return self.sample, curr, curr_rope_id, curr_split_len, curr_video_grid_thw, num_video_tokens

    def process_text(self, caption: str, curr: int, curr_rope_id: int, curr_split_len: int, item_loss=0):
        """处理文本，添加特殊token"""
        text_ids = self.tokenizer.encode(caption)
        shifted_text_ids = [self.bos_token_id] + text_ids  # NOTE: self.bos_token_id=151644 <|im_start|>

        self.sample["packed_text_ids"].extend(shifted_text_ids)
        self.sample["packed_text_indexes"].extend(range(curr, curr + len(shifted_text_ids)))

        # NOTE: 生成还是理解可以通过 item_loss == 1 来判定
        if item_loss == 1:
            loss_token_shift = 0  # HACK
            self.sample["ce_loss_indexes"].extend(range(curr - loss_token_shift, curr + len(shifted_text_ids)))
            self.sample["ce_loss_weights"].extend([len2weight(len(shifted_text_ids) + loss_token_shift)] * (len(shifted_text_ids) + loss_token_shift))
            self.sample["packed_label_ids"].extend(text_ids + [self.eos_token_id])  # NOTE: self.eos_token_id=151645 <|im_end|>
        curr += len(shifted_text_ids)
        curr_split_len += len(shifted_text_ids)

        # add a <|im_end|> token
        self.sample["packed_text_ids"].append(self.eos_token_id)
        self.sample["packed_text_indexes"].append(curr)
        curr += 1
        curr_split_len += 1

        # update sequence status
        self.sample["attn_modes"].append("causal")
        # if self.apply_chat_template:
        self.sample["packed_position_ids"].extend(range(curr_rope_id, curr_rope_id + curr_split_len))
        curr_rope_id += curr_split_len

        # self.sample['sample_modality'].extend([modality_map[item['type']]] * curr_split_len)

        self.sample["split_lens"].append(curr_split_len)

        return self.sample, curr, curr_rope_id, curr_split_len


    def process_vae_video(self, video_tensor, curr: int, curr_rope_id: int, curr_split_len: int, curr_video_grid_thw: None, video_sizes: list, item_loss=0):
        if not self.data_config.text_template:
            num_special_tokens = 0
            # 添加 <|startofimage|> token (视频与图像共用) TODO: 要将image和video的special token拆开嘛？
            self.sample["packed_text_ids"].append(self.start_of_image)  # self.start_of_image=151652, <|vision_start|>
            self.sample["packed_text_indexes"].append(curr)
            curr += 1
            curr_split_len += 1
            num_special_tokens += 1

        # 在线模式下，video_tensor 为tensor, 离线模式下，video_tensor 为list [latent]
        if isinstance(video_tensor, torch.Tensor):  # online
            # 预处理视频
            self.sample["vae_video_tensors"].append(video_tensor)  # CTHW 原始的视频，非latent
            # 假设 video_tensor 的形状为 (C, T, H, W)
            _, T, H, W = video_tensor.shape
            _T, _H, _W = self.data_config.vae_downsample  # NOTE: 绝对尺度的downsample，包含了patchify的！
            t = (T - 1) // _T + 1  # k*N+1 一般t维度不做patchify!! 如果t维度要做patchify，写法需要更新
            h = H // _H
            w = W // _W
            self.sample["vae_data_mode"].append("online")

            spatial_merge_size = 2  # TODO：spatial_merge_size 一定是2吗？
            vae_video_grid_thw = [
                t,
                h * spatial_merge_size,
                w * spatial_merge_size,
            ]  # 因为Qwen-VL 中的rope 处理默认存在 /spatial_merge_size 的操作（与VI处理匹配），所以对VAE 要额外进行*spatial_merge_size处理

            self.sample["vae_video_grid_thw"].append(vae_video_grid_thw)
            curr_video_grid_thw.append(vae_video_grid_thw)

            # 使用原生的 (t, h, w) latent shape
            self.sample["vae_latent_shapes"].append((t, h, w))

            # 使用3D感知的位置编码函数
            # 外插
            packed_latent_position_ids = get_flattened_position_ids_extrapolate_video(t, h, w, max_latent_size=self.data_config.max_latent_size)

            self.sample["packed_latent_position_ids"].append(packed_latent_position_ids)

            num_vid_tokens = t * h * w
            if not self.data_config.text_template:
                self.sample["packed_vae_token_indexes"].extend(range(curr, curr + num_vid_tokens))

            if item_loss == 1:
                timestep = np.random.randn()  # NOTE: 外面会sigmoid一下

                frame_condition_idx = self.frame_condition_idx
                packed_timesteps = [timestep] * num_vid_tokens

                mse_loss_indexes = list(range(curr, curr + num_vid_tokens))
                frame_condition_indexes = []
                for idx in frame_condition_idx:
                    if idx == -1:
                        idx = t - 1
                        if idx == 1:
                            continue  # 如果帧数仅两帧跳过，避免所有帧均为条件帧相同
                    frame_condition_indexes.extend(mse_loss_indexes[idx * h * w : (idx + 1) * h * w])
                    packed_timesteps[idx * h * w : (idx + 1) * h * w] = [-sys.float_info.max] * (h * w)
                if frame_condition_idx:
                    mse_loss_indexes = sorted(list(set(mse_loss_indexes) - set(frame_condition_indexes)))

                if not self.data_config.text_template:
                    self.sample["mse_loss_indexes"].extend(mse_loss_indexes)  # range(curr, curr + num_vid_tokens))
            else:
                timestep = float("-inf")
                packed_timesteps = [timestep] * num_vid_tokens

            self.sample["packed_timesteps"].extend(packed_timesteps)

            if not self.data_config.text_template:
                curr += num_vid_tokens
                curr_split_len += num_vid_tokens

                self.sample["packed_text_ids"].extend([self.image_token_id] * num_vid_tokens)

                # 添加 <|endofimage|> token
                self.sample["packed_text_ids"].append(self.end_of_image)  # self.end_of_image=151653, <|vision_end|>
                self.sample["packed_text_indexes"].append(curr)
                curr += 1
                curr_split_len += 1
                num_special_tokens += 1

                # 更新 sequence status
                if item_loss == 1:
                    self.sample["attn_modes"].append("noise")
                else:
                    self.sample["attn_modes"].append("full_noise")

                self.sample["packed_position_ids"].extend([curr_rope_id] * (num_vid_tokens + num_special_tokens))  # NOTE: 为什么rope固定？
                curr_rope_id += 1

                # update sample sequence modality
                # if item_loss == 1:
                #     self.sample['sample_modality'].extend([modality_map['noise']] * curr_split_len)
                # elif item_loss == 0 and sample_task == 'edit':
                #     self.sample['sample_modality'].extend([modality_map['ref_source']] * curr_split_len)
                # elif item_loss == 0 and sample_task == 'idip':
                #     self.sample['sample_modality'].extend([modality_map['ref_image']] * curr_split_len)

                self.sample["split_lens"].append(curr_split_len)

            video_sizes.append([T, H, W])

        return self.sample, curr, curr_rope_id, curr_split_len, curr_video_grid_thw, video_sizes, num_vid_tokens

    def process_text_template(
        self,
        text_ids,
        spans_index,
        tgt_index,
        caption_index,
        video_types: list[str],
        curr: int,
        curr_rope_id: int,
        curr_split_len: int,
        item_loss=0,
    ):
        # video_types = ['vit_video','vae_video_target','vae_video_cond'] 等信息，caption_index 即对应 search_index

        self.sample["packed_text_ids"].extend(text_ids)
        self.sample["sample_lens"] = len(text_ids)
        curr_split_idx = curr

        for video_id, span_index in enumerate(spans_index):
            vision_start, vision_end = curr_split_idx + span_index[0], curr_split_idx + span_index[-1]  # 对应第一和最后一个'<|video_pad|>' 的index
            self.sample["packed_text_indexes"].extend(range(curr, vision_start))
            if (vision_start - 1) - curr != 0:  # 确认vision前面有文本split ## HACK 相比llava 版本有修改
                curr_split_len = (vision_start - 1) - curr
                self.sample["packed_position_ids"].extend(
                    range(curr_rope_id, curr_rope_id + curr_split_len)
                )  # 注意：这里是 vision_start-1 而不是 vision_start，因为 vision_start 是 video split 起始token 的位置
                curr_rope_id += curr_split_len
                self.sample["sample_modality"].extend([modality_map["system_prompt"]] * curr_split_len)

                if caption_index != [] and caption_index[0] in range(curr, curr + curr_split_len): # NOTE： 不支持交错的文本，即文本必须连续，
                    split_len_1 = caption_index[0] - curr  # 文本前system_prompt 的长度
                    split_len_2 = len(caption_index) # 文本的长度
                    split_len_3 = curr_split_len - split_len_1 - split_len_2 # 文本后system_prompt 的长度

                    split_len_text = [split_len_1, split_len_2, split_len_3]
                    split_len_text = [x for x in split_len_text if x != 0]
                    self.sample["attn_modes"].extend(["causal"] * len(split_len_text))
                    self.sample["split_lens"].extend(split_len_text)
                else:
                    self.sample["attn_modes"].append("causal")
                    self.sample["split_lens"].append(curr_split_len)

            curr_split_len = len(span_index) + 2
            if video_types[video_id] == "vit_video":
                self.sample["packed_vit_token_indexes"].extend(range(vision_start, vision_end + 1))
                self.sample["attn_modes"].append("full")  # TODO : gen 分支也使用模版则需加上判断
                self.sample["sample_modality"].extend([modality_map["ref_vit"]] * curr_split_len)
            elif "vae_video" in video_types[video_id]:
                self.sample["packed_vae_token_indexes"].extend(range(vision_start, vision_end + 1))
                if "cond" in video_types[video_id]:
                    self.sample["attn_modes"].append("full_noise")  # TODO : gen 分支也使用模版则需加上判断
                    if self.sample_task == "edit":
                        self.sample["sample_modality"].extend([modality_map["ref_source"]] * curr_split_len)
                    elif self.sample_task == "idip":
                        self.sample["sample_modality"].extend([modality_map["ref_image"]] * curr_split_len)
                elif "target" in video_types[video_id]:
                    self.sample["mse_loss_indexes"].extend(range(vision_start, vision_end + 1))  # 目前不支持f2v
                    self.sample["attn_modes"].append("noise")  # TODO : gen 分支也使用模版则需加上判断
                    self.sample["sample_modality"].extend([modality_map["noise"]] * curr_split_len)
                else:
                    raise ValueError(f"video_types {video_types[video_id]} not supported")

            self.sample["packed_position_ids"].extend([curr_rope_id] * curr_split_len)
            # attn_modes.append("full")  # TODO : gen 分支也使用模版则需加上判断
            self.sample["split_lens"].append(len(span_index) + 2)
            curr = vision_end + 1  # 对应 '<|vision_end|>' token 的index
            curr_rope_id += 1
            self.sample["packed_text_indexes"].append(curr)
            curr += 1  # 对应下一个序列的起始token

        len_split_last = self.sample["sample_lens"] - (curr - curr_split_idx) if spans_index != [] else len(text_ids)
        if len_split_last != 0:  # 即末尾还有一段文本
            self.sample["split_lens"].append(len_split_last)
            self.sample["packed_text_indexes"].extend(range(curr, curr + len_split_last))
            self.sample["packed_position_ids"].extend(range(curr_rope_id, curr_rope_id + len_split_last))
            self.sample["attn_modes"].append("causal")
            self.sample["sample_modality"].extend([modality_map["system_prompt"]] * len_split_last)

        if item_loss == 1:  # 即代表为理解任务，需要计算ce loss
            packed_label_index = tgt_index
            self.sample["packed_label_ids"].extend(text_ids[packed_label_index[0] :])
            packed_label_index = np.asarray(packed_label_index, dtype=np.int64) + curr_split_idx
            ce_loss_indexes = (packed_label_index - 1).tolist()
            self.sample["ce_loss_indexes"].extend(ce_loss_indexes)
            self.sample["ce_loss_weights"].extend([len2weight(len(packed_label_index))] * (len(packed_label_index)))

            # 获取文本中 caption 的 index ，修改其sample_modality
        # caption_index = item.get("cap_index", [])
        if caption_index != []:
            self.sample["sample_modality"][caption_index[0] : caption_index[-1] + 1] = [modality_map["text"]] * (caption_index[-1] - caption_index[0] + 1)

        curr_split_idx += len(text_ids)
        curr = curr_split_idx
        return self.sample, curr, curr_rope_id, curr_split_len
    def process_und_template(self, system_prompt, user_prompt, answer, vit_video_tensor):
        """
        格式：
        <|im_start|>system
        {system_prompt}<|im_end|>
        <|im_start|>user
        <|vision_start|><|video_pad|><|vision_end|>{instruction_prompt}<|im_end|>
        <|im_start|>assistant
        {answer}<|im_end|>
        """
        curr = 0
        sample_lens = 0
        curr_rope_id = 0
        curr_video_grid_thw = []

        # 1. 处理第一部分的文本：
        # <|im_start|>system
        # {system_prompt}<|im_end|>
        # <|im_start|>user
        prompt_prefix = "<|im_start|>" + "system\n" + system_prompt + "<|im_end|>" + "\n" + "<|im_start|>" + "user\n"
        text_ids_prompt_prefix = self.tokenizer.encode(prompt_prefix)
        self.sample["packed_text_ids"].extend(text_ids_prompt_prefix)
        self.sample["packed_text_indexes"].extend(range(curr, curr + len(text_ids_prompt_prefix)))
        curr += len(text_ids_prompt_prefix)
        split_len_prefix = len(text_ids_prompt_prefix)

        # update sequence status
        self.sample["attn_modes"].append("causal")
        self.sample["packed_position_ids"].extend(range(curr_rope_id, curr_rope_id + split_len_prefix))
        self.sample["split_lens"].append(split_len_prefix)
        curr_rope_id += split_len_prefix

        # 2. 处理vision token部分，添加视觉tokens，在线模式下，video_tensor 为tensor, 离线模式下，video_tensor 为list [latent]
        self.sample["packed_text_ids"].append(self.start_of_image)  # 151652, <|vision_start|>
        self.sample["packed_text_indexes"].append(curr)
        curr += 1
        split_len_vision_token = 1

        if isinstance(vit_video_tensor, torch.Tensor):  # online
            self.sample["vit_video_tensors"].append(vit_video_tensor)  # CTHW 原始的视频，非latent , 仅用于validation中的可视化

            # preprocess video
            vit_tokens = patchify_video_with_merge(
                vit_video_tensor, self.data_config.vit_patch_size, self.data_config.vit_patch_size_temporal
            )  # C T H W -> (T//2 * H//p * W//p) (p*p*2*C)
            num_video_tokens = vit_tokens.shape[0] // 4  # 实际上qwen2.5-vl还需要merge，2x2 merge成1个， hardcode for temp
            t, h, w = vit_video_tensor.size(1), vit_video_tensor.size(2), vit_video_tensor.size(3)

            self.sample["packed_vit_tokens"].append(vit_tokens)
            self.sample["vit_data_mode"].append("online")

        if t is not None:
            vit_video_grid_thw = [
                t // self.data_config.vit_patch_size_temporal,
                h // self.data_config.vit_patch_size,
                w // self.data_config.vit_patch_size,
            ]  # [1, 16, 16]
        self.sample["vit_video_grid_thw"].append(vit_video_grid_thw)
        curr_video_grid_thw.append(vit_video_grid_thw)

        self.sample["vit_token_seqlens"].append(num_video_tokens)
        self.sample["packed_vit_position_ids"].append(
            torch.zeros(num_video_tokens)
        )  # TODO : 不一定是 0 ？ 对于多个vit序列会有问题

        self.sample["packed_vit_token_indexes"].extend(range(curr, curr + num_video_tokens))
        curr += num_video_tokens
        split_len_vision_token += num_video_tokens

        # dummy position_ids
        self.sample["packed_text_ids"].extend([self.image_token_id] * num_video_tokens)

        # add a <|endofimage|> token
        self.sample["packed_text_ids"].append(self.end_of_image)  # 151653, <|vision_end|>
        self.sample["packed_text_indexes"].append(curr)
        curr += 1
        split_len_vision_token += 1

        # update sequence status
        self.sample["attn_modes"].append("full")
        self.sample["packed_position_ids"].extend([curr_rope_id] * split_len_vision_token)
        self.sample["split_lens"].append(split_len_vision_token)
        curr_rope_id += 1

        # 3. 处理后半部分的文本：
        # {instruction_prompt}<|im_end|>
        # <|im_start|>assistant
        prompt_postfix = user_prompt + "<|im_end|>" + "\n" + "<|im_start|>" + "assistant"
        text_ids_prompt_postfix = self.tokenizer.encode(prompt_postfix)
        self.sample["packed_text_ids"].extend(text_ids_prompt_postfix)
        self.sample["packed_text_indexes"].extend(range(curr, curr + len(text_ids_prompt_postfix)))
        curr += len(text_ids_prompt_postfix)
        split_len_postfix = len(text_ids_prompt_postfix)

        # update sequence status
        self.sample["attn_modes"].append("causal")
        self.sample["packed_position_ids"].extend(range(curr_rope_id, curr_rope_id + split_len_postfix))
        self.sample["split_lens"].append(split_len_postfix)
        curr_rope_id += split_len_postfix

        # 4. 添加answer
        answer = "\n" + answer
        answer_ids = self.tokenizer.encode(answer)
        shifted_text_ids_answer = answer_ids + [self.eos_token_id]
        self.sample["packed_text_ids"].extend(shifted_text_ids_answer)
        self.sample["packed_text_indexes"].extend(range(curr, curr + len(shifted_text_ids_answer)))

        # item_loss == 1:
        self.sample["ce_loss_indexes"].extend(range(curr, curr + len(shifted_text_ids_answer)))
        self.sample["ce_loss_weights"].extend([len2weight(len(shifted_text_ids_answer))] * (len(shifted_text_ids_answer)))
        self.sample["packed_label_ids"].extend(shifted_text_ids_answer)  # NOTE: self.eos_token_id=151645 <|im_end|>

        curr += len(shifted_text_ids_answer)
        split_len_answer = len(shifted_text_ids_answer)

        # update sequence status
        self.sample["attn_modes"].append("causal")
        self.sample["packed_position_ids"].extend(range(curr_rope_id, curr_rope_id + split_len_answer))
        self.sample["split_lens"].append(split_len_answer)
        curr_rope_id += split_len_answer

        sample_lens = len(self.sample["packed_text_ids"])

        return sample_lens, curr_video_grid_thw

    def _finalize_sample(self, sample_lens, curr_video_grid_thw, sample_type, sample=None, additional_fields=None, video_sizes=None):
        """通用 sample 结尾处理，减少代码重复"""
        self.sample["sample_lens"] = [sample_lens]
        self.sample["video_grid_thw"] = torch.tensor([curr_video_grid_thw])
        self.sample["packed_text_ids"] = torch.tensor(self.sample["packed_text_ids"])
        self.sample["packed_text_indexes"] = torch.tensor(self.sample["packed_text_indexes"])

        self.sample["packed_vae_token_indexes"] = torch.tensor(self.sample["packed_vae_token_indexes"])
        self.sample["packed_position_ids"] = torch.tensor(self.sample["packed_position_ids"])
        self.sample["vae_video_grid_thw"] = torch.tensor(self.sample["vae_video_grid_thw"])

        self.sample["vit_video_grid_thw"] = torch.tensor(self.sample["vit_video_grid_thw"])
        self.sample["packed_vit_token_indexes"] = torch.tensor(self.sample["packed_vit_token_indexes"])

        self.sample["sample_N_target"] = torch.tensor([[1]])
        self.sample["sample_type"] = [sample_type]
        self.sample["padded_videos"] = self.sample["vae_video_tensors"]

        if "ce_loss_indexes" in self.sample and len(self.sample["ce_loss_indexes"]) > 0:
            self.sample["ce_loss_indexes"] = torch.tensor(self.sample["ce_loss_indexes"])
        # 原始代码总是处理 mse_loss_indexes，即使为空列表
        self.sample["mse_loss_indexes"] = torch.tensor(self.sample["mse_loss_indexes"])
        if video_sizes is not None:
            self.sample["video_sizes"] = torch.tensor(video_sizes)
        elif "video_sizes" in self.sample:
            self.sample["video_sizes"] = torch.tensor(self.sample["video_sizes"])
        if "sample_modality" in self.sample and len(self.sample["sample_modality"]) > 0:
            self.sample["sample_modality"] = torch.tensor(self.sample["sample_modality"])

        if sample is not None:
            for key in ["index", "category", "question", "gt"]:
                if key in sample:
                    self.sample[key] = sample[key]

        if additional_fields is not None:
            for key, value in additional_fields.items():
                self.sample[key] = value

        return self.sample

    def ti2t_sample(self, idx: int) -> Dict[str, Any]:
        """
        获取单个样本
        默认system_prompt和user_prompt中均不包含sos和eos token
        格式：
        <|im_start|>system
        {system_prompt}<|im_end|>
        <|im_start|>user
        <|vision_start|><|video_pad|><|vision_end|>{instruction_prompt}<|im_end|>
        <|im_start|>assistant
        {answer}<|im_end|>
        """
        self.sample = self.set_sequence_status()
        sample = self.data[idx]

        system_prompt = sample["system_prompt"]
        user_prompt = sample["user_prompt"]
        answer = sample["gt"]
        image_path = sample["image_path"]
        vit_image_tensor = self.get_video_tensor_online(image_path, vision_stream="vit_video", element_dtype="image")  # [C=3, T=2, H, W]

        sample_lens, curr_video_grid_thw = self.process_und_template(
            system_prompt=system_prompt,
            user_prompt=user_prompt,
            answer=answer,
            vit_video_tensor=vit_image_tensor,
        )

        self.sample["system_prompt"] = system_prompt
        self.sample["user_prompt"] = user_prompt
        self.sample["image_path"] = image_path
        self.sample["instruction"] = user_prompt

        return self._finalize_sample(
            sample_lens, curr_video_grid_thw,
            sample_type="und",
            sample=sample
        )

    def t2v_sample(self, idx: int) -> Dict[str, Any]:
        """获取单个样本"""
        _T, _H, _W = self.data_config.vae_downsample
        if self.data_config.task == "t2i":
            t = 1
            t_ = 1
            element_dtype = 'image'
        else:
            t = (self.data_config.num_frames - 1) // _T + 1  # k*N+1 一般t维度不做patchify!! 如果t维度要做patchify，写法需要更新
            t_ = self.data_config.num_frames
            element_dtype = 'video'

        self.sample = self.set_sequence_status()
        packed_text_indexes, packed_position_ids, sample_modality = [], [], []
        sample = self.data[idx]
        if "prompt_en" in sample.keys():
            user_prompt = "".join(sample["prompt_en"][0])
            # user_prompt = sample["prompt_en"][0][0] + sample["prompt_en"][0][1] # image_caption + video_caption
        else:
            user_prompt = sample["data"]

        if self.data_config.text_template:
            caption_instruction = generate_system_prompt(system_prompt_type=self.data_config.task, vision_type=element_dtype)

            text_template_user, text_template_assistant, vit_num_tokens, video_types = [], [], [], []
            if self.system_prompt_type == 'SP2':
                user_prompt = caption_instruction + " " + user_prompt # user_prompt 对应caption_q
                caption_instruction = "You are a helpful assistant. "
            elif self.system_prompt_type == 'SP1':
                # SP1: assistant
                caption_instruction = "You are a helpful assistant. " + caption_instruction

            text_template_user.append({"type": "text", "text": user_prompt})
        else:
            # 编码文本
            text_ids = self.tokenizer.encode(user_prompt)
            text_ids = [self.new_token_ids["bos_token_id"]] + text_ids + [self.new_token_ids["eos_token_id"]]
            text_split_len = len(text_ids)
            packed_text_indexes.extend(range(0, text_split_len))  # curr = 0
            packed_position_ids.extend(range(0, text_split_len))
            sample_modality.extend([modality_map['text']] * text_split_len)

        # 视频参数

        h = self.data_config.H // _H
        w = self.data_config.W // _W
        spatial_merge_size = 2  # TODO：spatial_merge_size 一定是2吗？
        # vae_video_grid_thw = torch.tensor([[t, h * spatial_merge_size, w * spatial_merge_size]])
        num_vid_tokens = t * h * w

        if self.data_config.text_template:
            text_template_assistant.append({"type":element_dtype})
        else:
            text_ids.append(self.new_token_ids["start_of_image"])
            packed_text_indexes.append(text_split_len)
            packed_vae_token_indexes = torch.tensor(range(len(text_ids), len(text_ids) + num_vid_tokens))
            text_ids.extend([self.image_token_id] * num_vid_tokens)
            text_ids.append(self.new_token_ids["end_of_image"])
            packed_text_indexes.append(len(text_ids) - 1)
            video_split_len = num_vid_tokens + 2
            packed_position_ids.extend([text_split_len] * video_split_len)
            sample_modality.extend([modality_map['noise']] * video_split_len)

        if self.data_config.text_template:
            all_token_id, spans_index, tgt_index, search_index = self.render_template(caption_instruction, text_template_assistant, text_template_user, [num_vid_tokens], search_text=user_prompt)

            # 计算
            self.sample, curr, curr_rope_id, curr_split_len = self.process_text_template(
                all_token_id,
                spans_index,
                tgt_index,
                search_index,
                video_types=['target_vae_video'],
                curr=0,
                curr_rope_id=0,
                curr_split_len=0,
                item_loss=0,
                )

        # 构造返回字典
        return {
            "packed_text_ids": torch.tensor(text_ids) if not self.data_config.text_template else torch.tensor(self.sample["packed_text_ids"]),
            "packed_text_indexes": torch.tensor(packed_text_indexes) if not self.data_config.text_template else torch.tensor(self.sample["packed_text_indexes"]),
            "packed_vae_token_indexes": packed_vae_token_indexes if not self.data_config.text_template else torch.tensor(self.sample["packed_vae_token_indexes"]),
            "vae_video_grid_thw": torch.tensor([[t, h * spatial_merge_size, w * spatial_merge_size]]),
            "video_grid_thw": torch.tensor([[[t, h * spatial_merge_size, w * spatial_merge_size]]]),
            "sample_N_target": torch.tensor([[1]]),  # 生成一个视频
            "split_lens": [text_split_len, video_split_len] if not self.data_config.text_template else self.sample["split_lens"],
            "attn_modes": ["causal", "noise"] if not self.data_config.text_template else self.sample["attn_modes"],
            "sample_lens": [text_split_len + video_split_len] if not self.data_config.text_template else [self.sample["sample_lens"]],
            "val_sample_type": ["gen"],  # 生成任务
            "padded_latent": None,
            "mse_loss_indexes": packed_vae_token_indexes if not self.data_config.text_template else torch.tensor(self.sample["mse_loss_indexes"]),
            "video_sizes": torch.tensor([[t_, self.data_config.H, self.data_config.W]]),
            "packed_position_ids": torch.tensor(packed_position_ids) if not self.data_config.text_template else torch.tensor(self.sample["packed_position_ids"]),
            "caption": user_prompt,  # 用于可视化
            "sample_type": ["gen"],  # 生成任务
            "index": sample["index"],
            "caption_cn": user_prompt,
            "original_prompt_en": sample["original_prompt_en"] if "original_prompt_en" in sample.keys() else user_prompt,  # 新增字段，用于保存的命名
            "sample_task": torch.zeros(text_split_len + video_split_len) if not self.data_config.text_template else torch.zeros(self.sample["sample_lens"]),
            "sample_modality": torch.tensor(sample_modality) if not self.data_config.text_template else torch.tensor(self.sample["sample_modality"]),
            "additional_info": sample["additional_info"] if "additional_info" in sample.keys() else None,
        }

    def tv2v_sample(self, idx: int) -> Dict[str, Any]:
        """获取单个样本 - 使用 tiv2v_sample 的通用 interleave 格式"""
        sample = self.data[idx]
        user_prompt = "Create a 2D animation based on the provided image of a maze. The blue star slides smoothly along the white path, stopping perfectly on the red flag and then acquiring a trophy. The blue star never slides or crosses into the black segments of the maze. The camera is a static, top-down view showing the entire maze."
        
        # 转换为 tiv2v 的 interleave 格式
        sample["data"] = {
            "interleave_array": [user_prompt, sample["image_path"], sample["image_path"], sample["video_path"]],
            "element_dtype_array": ["text", "image", "image", "video"],
            "istarget_in_interleave": [0, 0, 0, 1]
        }
        
        self.sample_task = 'edit'
        result = self.tiv2v_sample(idx)
        
        # 额外设置一些 tv2v 特有的字段
        result["caption"] = user_prompt
        result["caption_cn"] = user_prompt
        
        return result

    def tiv2v_sample(self, idx: int) -> Dict[str, Any]: # 构造一个统一的interleave数据处理函数
        """获取单个样本"""
        sample_modality, text_template_user, text_template_assistant, vit_num_tokens, video_types = [], [], [], [], []
        self.sample = self.set_sequence_status()
        sample_lens = 0
        sample = self.data[idx]

        index = sample["index"]
        data_sample = sample["data"] # {'interleave_array': [...], 'element_dtype_array': [...], 'istarget_in_interleave': [...]}}
        additional_info = sample["data"]["additional_info"] if "additional_info" in sample["data"] else [] #sample["data"]["additional_info"]

        interleave_array, element_dtype_array, istarget_in_interleave = data_sample["interleave_array"], data_sample["element_dtype_array"], data_sample["istarget_in_interleave"]

        curr, curr_rope_id, curr_split_len, curr_video_grid_thw, video_sizes, caption_all = 0, 0, 0, [], [], ''
        for element, element_dtype, is_target in zip(interleave_array, element_dtype_array, istarget_in_interleave):
            if element_dtype == "text":
                # 文本 序列处理
                caption_all += element
                if self.data_config.text_template:
                    text_template_user.append({"type": "text", "text": element})
                    search_text = element
                else:
                    self.sample, curr, curr_rope_id, curr_split_len = self.process_text(element, curr=curr, curr_rope_id=curr_rope_id, curr_split_len=0, item_loss=is_target)
                    sample_lens += curr_split_len
                    sample_modality.extend([modality_map['text']] * curr_split_len)
            elif element_dtype in ["image", "video"]:
                if is_target == 0: # condition 需要 vit 处理
                    vit_image_tensor = self.get_video_tensor_online(element, vision_stream="vit_video", element_dtype=element_dtype)  # [C=3, T, H, W]
                    self.sample, curr, curr_rope_id, curr_split_len, curr_video_grid_thw, num_tokens_ = self.process_vit_video(
                        vit_image_tensor, curr=curr, curr_rope_id=curr_rope_id, curr_split_len=0, curr_video_grid_thw=curr_video_grid_thw, item_loss=0
                        )
                    if self.data_config.text_template:
                        text_template_user.append({"type": element_dtype})
                        vit_num_tokens.append(num_tokens_)
                        video_types.append("vit_video")
                    else:
                        sample_lens += curr_split_len
                        sample_modality.extend([modality_map['ref_vit']] * curr_split_len)

                # vae condition/target 处理
                vae_image_tensor = self.get_video_tensor_online(element, vision_stream="vae_video", element_dtype=element_dtype)  # [C=3, T=1, H, W]
                self.sample, curr, curr_rope_id, curr_split_len, curr_video_grid_thw, video_sizes, num_tokens_ = self.process_vae_video(
                    vae_image_tensor, curr=curr, curr_rope_id=curr_rope_id, curr_split_len=0, curr_video_grid_thw=curr_video_grid_thw, video_sizes=video_sizes, item_loss=is_target
                )
                if self.data_config.text_template:
                    vit_num_tokens.append(num_tokens_)
                    if is_target == 0:
                        text_template_user.append({"type": element_dtype})
                        video_types.append("cond_vae_video")
                    else:
                        text_template_assistant.append({"type": element_dtype})
                        video_types.append("target_vae_video")
                else:
                    sample_lens += curr_split_len
                    if is_target == 0:
                        sample_modality.extend([modality_map[f'ref_{element_dtype}']] * curr_split_len)
                    else:
                        sample_modality.extend([modality_map[f'noise']] * curr_split_len)

        if self.data_config.text_template:
            if text_template_user[0]['type']=='text': # 先图像/视频后文本的处理：
                text_template_user = text_template_user[1:] + text_template_user[:1] # HACK
            caption_instruction = generate_system_prompt(system_prompt_type=self.data_config.task, vision_type=element_dtype)
            all_token_id, spans_index, tgt_index, search_index = self.render_template(caption_instruction, text_template_assistant, text_template_user, vit_num_tokens, search_text=search_text)
            # 计算
            self.sample, curr, curr_rope_id, curr_split_len = self.process_text_template(
                all_token_id,
                spans_index,
                tgt_index,
                search_index,
                video_types=video_types,
                curr=0,
                curr_rope_id=0,
                curr_split_len=0,
                item_loss=0,
                )
            sample_lens = len(all_token_id)
            sample_modality = self.sample["sample_modality"]


        additional_fields = {
            "caption": caption_all,
            "caption_cn": caption_all,
            "index": sample["index"],
            "additional_info": additional_info
        }

        if self.sample_task == 'edit':
            self.sample["sample_task"] = torch.ones(sample_lens) * sample_task_map['edit']
        elif self.sample_task == 'idip':
            self.sample["sample_task"] = torch.ones(sample_lens) * sample_task_map['idip']

        return self._finalize_sample(
            sample_lens, curr_video_grid_thw,
            sample_type="gen",
            sample=sample,
            additional_fields=additional_fields,
            video_sizes=video_sizes
        )

    def render_template(self, instruction, text_template_assistant, text_template_user, vit_num_tokens, search_text=""):
        # NOTE: 无target 文本的样本，设置 caption_a = ""
        # caption_i, caption_q, caption_a = element[0], element[1], element[2]

        # text_template_assistant.append({"type": "text", "text": caption_a}) # caption
        # if caption_q != "":
        #     text_template_user.append({"type": "text", "text": caption_q})

        messages = [
            {
                "role": "user",
                "content": text_template_user, # 原使用
            },
            {
                "role": "assistant",
                "content": text_template_assistant,
            },
        ]
        caption_all = render_qwenvl_prompt(messages, default_system=instruction, include_assistant_content=True) # NOTE: 是否添加 You are a helpful assistant.

        all_token_id, spans_index, tgt_index, search_index = expand_and_index_by_token_ids_new(
            rendered_text=caption_all.strip(), tokens=vit_num_tokens, target_text=f"assistant\n", tokenizer=self.tokenizer, search_text=search_text
        )
        assert len(all_token_id[tgt_index[0] :]) == len(tgt_index)
        return all_token_id, spans_index, tgt_index, search_index

    def x2t_sample(self, idx: int) -> Dict[str, Any]: # 构造一个统一的interleave数据处理函数
        """获取单个样本"""
        sample_modality = []
        self.sample = self.set_sequence_status()
        sample_lens = 0
        sample = self.data[idx]
        index = sample["index"]
        data_sample = sample["data"]  # {'interleave_array': [...], 'element_dtype_array': [...], 'istarget_in_interleave': [...]}}

        interleave_array, element_dtype_array, istarget_in_interleave = data_sample["interleave_array"], data_sample["element_dtype_array"], data_sample["istarget_in_interleave"]

        curr, curr_rope_id, curr_split_len, curr_video_grid_thw, video_sizes, caption_all = 0, 0, 0, [], [], ""
        if self.data_config.text_template:
            text_template_user, text_template_assistant, vit_num_tokens, video_types = [], [], [], []
        for element, element_dtype, is_target in zip(interleave_array, element_dtype_array, istarget_in_interleave):
            if element_dtype == "text":
                # 文本 序列处理
                if is_target == 1:  # 对应target 文本
                    if self.data_config.text_template:  # 即使用system_prompt
                        if isinstance(element, str):  # 即只有一条文本
                            caption_a = element
                            caption_i = generate_system_prompt(system_prompt_type="caption", vision_type=element_dtype_array[0])
                            caption_q = ""
                            element = [caption_i, caption_q, caption_a]

                        # ====================== SP1 + SP2 处理 START ======================
                        caption_i, caption_q, caption_a = element[0], element[1], element[2]
                        if self.system_prompt_type == 'SP2':
                            caption_q = caption_i + " " + caption_q
                            caption_i = "You are a helpful assistant. "
                        elif self.system_prompt_type == 'SP1':
                            # SP1: assistant
                            caption_i = "You are a helpful assistant. " + caption_i
                        element = [caption_i, caption_q, caption_a]

                        print('element',element)
                        # ====================== SP1 + SP2 处理 END ======================

                        caption_i, caption_q, caption_a = element[0], element[1], element[2]

                        text_template_assistant.append({"type": "text", "text": caption_a}) # caption
                        if caption_q != "":
                            text_template_user.append({"type": "text", "text": caption_q})

                        all_token_id, spans_index, tgt_index, search_index = self.render_template(caption_i, text_template_assistant, text_template_user, vit_num_tokens)
                        self.sample, curr, curr_rope_id, curr_split_len = self.process_text_template(
                            all_token_id,
                            spans_index,
                            tgt_index,
                            search_index,
                            video_types,
                            curr=curr,
                            curr_rope_id=curr_rope_id,
                            curr_split_len=0,
                            item_loss=is_target,
                        )
                        sample_lens += curr_split_len

                        caption_all += "\n".join(element)
                        caption_answer = element[-1]  # 传出element
                    else:
                        if isinstance(element, list):
                            element = element[-1]  # 使用 element = "" 效果是一样的，对生成理解文本无影响
                        self.sample, curr, curr_rope_id, curr_split_len = self.process_text(
                            element, curr=curr, curr_rope_id=curr_rope_id, curr_split_len=0, item_loss=is_target
                        )
                        sample_lens += curr_split_len
                        sample_modality.extend([modality_map["text"]] * curr_split_len)
                        caption_all += element
                        caption_answer = element  # NOTE unsure

            elif element_dtype in ["image", "video"]:

                vit_image_tensor = self.get_video_tensor_online(element, vision_stream="vit_video", element_dtype=element_dtype)  # [C=3, T, H, W]
                self.sample, curr, curr_rope_id, curr_split_len, curr_video_grid_thw, num_tokens_ = self.process_vit_video(
                    vit_image_tensor, curr=curr, curr_rope_id=curr_rope_id, curr_split_len=0, curr_video_grid_thw=curr_video_grid_thw, item_loss=0
                )
                sample_lens += curr_split_len
                sample_modality.extend([modality_map["ref_vit"]] * curr_split_len)
                index_video_path_name = element.split("/")[-1]

                if self.data_config.text_template:
                    text_template_user.append({"type": element_dtype})
                    vit_num_tokens.append(num_tokens_)
                    video_types.append("vit_video")

        if self.sample["sample_lens"] != []:
            sample_lens = self.sample["sample_lens"]

        if self.sample["sample_modality"] != []:
            sample_modality = self.sample["sample_modality"]
        self.sample["sample_modality"] = sample_modality
        self.sample["sample_task"] = torch.ones(self.sample["sample_lens"]) * sample_task_map["t2v"]

        additional_fields = {
            "caption": caption_all,
            "caption_cn": caption_all,
            "caption_answer": caption_answer,
            "index_item": index,
            "index": index_video_path_name,
            "additional_information": data_sample["additional_information"] if "additional_information" in data_sample.keys() else {},
            "visual_path": data_sample["interleave_array"][0],
            "question": data_sample["interleave_array"][1][1] if isinstance(data_sample["interleave_array"][1], list) and len(data_sample["interleave_array"][1]) > 1 else None,
            "answer": data_sample["interleave_array"][1][2] if isinstance(data_sample["interleave_array"][1], list) and len(data_sample["interleave_array"][1]) > 2 else None
        }

        return self._finalize_sample(
            sample_lens, curr_video_grid_thw,
            sample_type="und",
            additional_fields=additional_fields
        )

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        if self.data_config.task == "tv2v":
            return self.tv2v_sample(idx)
        elif self.data_config.task in ["t2i","t2v"]:
            return self.t2v_sample(idx)
        elif self.data_config.task == "ti2t":
            return self.ti2t_sample(idx)
        elif "tiv2v" in self.data_config.task:
            if 'edit' in self.data_config.task:
                self.sample_task = 'edit'
            elif 'idip' in self.data_config.task:
                self.sample_task = 'idip'
            return self.tiv2v_sample(idx)
        elif self.data_config.task == "video_edit":
            self.sample_task = 'edit'
            return self.tiv2v_sample(idx)
        elif self.data_config.task == "video_idip" or self.data_config.task == "video_idip_multiref":
            self.sample_task = 'idip'
            return self.tiv2v_sample(idx)
        elif self.data_config.task == "image_edit":
            self.sample_task = 'edit'
            return self.tiv2v_sample(idx)
        elif self.data_config.task == "image_idip":
            self.sample_task = 'idip'
            return self.tiv2v_sample(idx)
        elif self.data_config.task in ["x2t", "x2t_image", "x2t_video"]:
            return self.x2t_sample(idx)
        else:
            raise ValueError(f"Unknown task: {self.data_config.task}")