wehub-resource-sync
1232 lines
57 KiB
Python
1232 lines
57 KiB
Python
# 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 sys
|
|
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,
|
|
'ref_image': 3,
|
|
'ref_vit': 4
|
|
}
|
|
|
|
|
|
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,
|
|
):
|
|
"""
|
|
Initialize the validation dataset.
|
|
|
|
Args:
|
|
jsonl_path: Path to the JSONL file.
|
|
tokenizer: Tokenizer instance.
|
|
"""
|
|
self.jsonl_path = jsonl_path
|
|
self.tokenizer = tokenizer
|
|
self.new_token_ids = new_token_ids
|
|
|
|
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
|
|
|
|
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
|
|
resolution_vit = 476
|
|
elif self.data_config.resolution == "video_480p":
|
|
resolution_vae = 640
|
|
resolution_vit = 616
|
|
else:
|
|
raise ValueError(f"Unknown resolution: {self.data_config.resolution}")
|
|
|
|
video_transform_args = {
|
|
"resolution": resolution_vae,
|
|
"mode": "bucket",
|
|
"divisible_crop_size": 16,
|
|
"stride_spatial": 16,
|
|
"stride_temporal": 4,
|
|
"aspect_ratios": ["21:9", "16:9", "4:3", "1:1", "3:4", "9:16"],
|
|
"mean": 0.5,
|
|
"std": 0.5,
|
|
}
|
|
self.transform = VideoTransform(**video_transform_args)
|
|
|
|
vit_video_transform_args = {
|
|
"resolution": resolution_vit,
|
|
"mode": "bucket",
|
|
"divisible_crop_size": 28,
|
|
"aspect_ratios": ["21:9", "16:9", "4:3", "1:1", "3:4", "9:16"],
|
|
"mean": [0.48145466, 0.4578275, 0.40821073],
|
|
"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=[],
|
|
vae_video_tensors=[],
|
|
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=[],
|
|
vae_video_grid_thw=[],
|
|
video_grid_thw=[],
|
|
vit_video_tensors=[],
|
|
vae_video_latent=[],
|
|
vae_data_mode=[],
|
|
vit_data_mode=[],
|
|
sample_task=[],
|
|
sample_modality=[],
|
|
save_fps=12,
|
|
)
|
|
return sequence_status
|
|
|
|
def _read_jsonl(self) -> List[Dict[str, Any]]:
|
|
"""Read the JSONL file."""
|
|
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 _maybe_enhance_t2v_prompt(self, prompt: str) -> str:
|
|
if self.data_config.task != "t2v":
|
|
return prompt
|
|
if not getattr(self.data_config, "enhance_prompt", False):
|
|
return prompt
|
|
|
|
from common.utils.caption_rewrite import has_rewrite_api_key, rewrite_prompt
|
|
|
|
if not has_rewrite_api_key():
|
|
return prompt
|
|
|
|
try:
|
|
enhanced_prompt = rewrite_prompt(prompt)
|
|
except Exception as exc:
|
|
print(f"[enhance_prompt][t2v][warning] prompt rewrite failed, use original prompt. error={exc}")
|
|
return prompt
|
|
print(f"[enhance_prompt][t2v][original] {prompt}")
|
|
print(f"[enhance_prompt][t2v][rewritten] {enhanced_prompt}")
|
|
return enhanced_prompt
|
|
|
|
def _maybe_enhance_i2v_prompt(self, prompt: str, image_path: str) -> str:
|
|
if "i2v" not in self.data_config.task:
|
|
return prompt
|
|
if not getattr(self.data_config, "enhance_prompt", False):
|
|
return prompt
|
|
|
|
from common.utils.caption_rewrite import has_rewrite_api_key, rewrite_i2v_prompt
|
|
|
|
if not has_rewrite_api_key():
|
|
return prompt
|
|
|
|
try:
|
|
enhanced_prompt = rewrite_i2v_prompt(prompt, image_path=image_path)
|
|
except Exception as exc:
|
|
print(f"[enhance_prompt][i2v][warning] prompt rewrite failed, use original prompt. error={exc}")
|
|
return prompt
|
|
print(f"[enhance_prompt][i2v][image] {image_path}")
|
|
print(f"[enhance_prompt][i2v][original] {prompt}")
|
|
print(f"[enhance_prompt][i2v][rewritten] {enhanced_prompt}")
|
|
return enhanced_prompt
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.data)
|
|
|
|
|
|
@staticmethod
|
|
def _read_decord(video: VideoReader, frame_idx: List[int]) -> List[Image.Image]:
|
|
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
|
|
|
|
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])
|
|
image = bg
|
|
else:
|
|
image = image.convert("RGB")
|
|
video_frames = [image]
|
|
else:
|
|
video_reader = VideoReader(video_stream, ctx=decord.cpu(worker_id % self.cpu_count))
|
|
total_frames = len(video_reader)
|
|
|
|
try:
|
|
fps = int(round(float(video_reader.get_avg_fps())))
|
|
except Exception:
|
|
fps = 24
|
|
frames_info = {
|
|
"clip_indices": [(0, total_frames)],
|
|
"fps": fps,
|
|
}
|
|
|
|
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)
|
|
elif vision_stream == "vit_video":
|
|
video_tensor = self.vit_transform(video_frames)
|
|
if element_dtype == "image":
|
|
video_tensor = video_tensor.repeat(1, 2, 1, 1)
|
|
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
|
|
|
|
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)
|
|
self.sample["packed_text_indexes"].append(curr)
|
|
curr += 1
|
|
curr_split_len += 1
|
|
|
|
if isinstance(video_tensor, torch.Tensor):
|
|
self.sample["vit_video_tensors"].append(video_tensor)
|
|
vit_tokens = patchify_video_with_merge(
|
|
video_tensor, self.data_config.vit_patch_size, self.data_config.vit_patch_size_temporal
|
|
)
|
|
num_video_tokens = vit_tokens.shape[0] // 4
|
|
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,
|
|
]
|
|
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)
|
|
)
|
|
|
|
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
|
|
|
|
self.sample["packed_text_ids"].extend([self.image_token_id] * num_video_tokens)
|
|
self.sample["packed_text_ids"].append(self.end_of_image)
|
|
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
|
|
|
|
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):
|
|
"""Process text and append special tokens."""
|
|
text_ids = self.tokenizer.encode(caption)
|
|
shifted_text_ids = [self.bos_token_id] + text_ids
|
|
self.sample["packed_text_ids"].extend(shifted_text_ids)
|
|
self.sample["packed_text_indexes"].extend(range(curr, curr + len(shifted_text_ids)))
|
|
|
|
if item_loss == 1:
|
|
loss_token_shift = 0
|
|
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])
|
|
curr += len(shifted_text_ids)
|
|
curr_split_len += len(shifted_text_ids)
|
|
|
|
# Append the <|im_end|> 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
|
|
self.sample["attn_modes"].append("causal")
|
|
self.sample["packed_position_ids"].extend(range(curr_rope_id, curr_rope_id + curr_split_len))
|
|
curr_rope_id += 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
|
|
self.sample["packed_text_ids"].append(self.start_of_image)
|
|
self.sample["packed_text_indexes"].append(curr)
|
|
curr += 1
|
|
curr_split_len += 1
|
|
num_special_tokens += 1
|
|
|
|
if isinstance(video_tensor, torch.Tensor):
|
|
self.sample["vae_video_tensors"].append(video_tensor)
|
|
_, T, H, W = video_tensor.shape
|
|
_T, _H, _W = self.data_config.vae_downsample
|
|
t = (T - 1) // _T + 1
|
|
h = H // _H
|
|
w = W // _W
|
|
self.sample["vae_data_mode"].append("online")
|
|
|
|
spatial_merge_size = 2
|
|
vae_video_grid_thw = [
|
|
t,
|
|
h * spatial_merge_size,
|
|
w * spatial_merge_size,
|
|
]
|
|
|
|
self.sample["vae_video_grid_thw"].append(vae_video_grid_thw)
|
|
curr_video_grid_thw.append(vae_video_grid_thw)
|
|
self.sample["vae_latent_shapes"].append((t, h, w))
|
|
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()
|
|
|
|
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)
|
|
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)
|
|
|
|
# Append the <|endofimage|> image end token.
|
|
self.sample["packed_text_ids"].append(self.end_of_image)
|
|
self.sample["packed_text_indexes"].append(curr)
|
|
curr += 1
|
|
curr_split_len += 1
|
|
num_special_tokens += 1
|
|
|
|
# Update sequence state.
|
|
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))
|
|
curr_rope_id += 1
|
|
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,
|
|
):
|
|
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]
|
|
self.sample["packed_text_indexes"].extend(range(curr, vision_start))
|
|
if (vision_start - 1) - curr != 0:
|
|
curr_split_len = (vision_start - 1) - curr
|
|
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["system_prompt"]] * curr_split_len)
|
|
|
|
if caption_index != [] and caption_index[0] in range(curr, curr + curr_split_len):
|
|
split_len_1 = caption_index[0] - curr
|
|
split_len_2 = len(caption_index)
|
|
split_len_3 = curr_split_len - split_len_1 - split_len_2
|
|
|
|
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")
|
|
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")
|
|
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))
|
|
self.sample["attn_modes"].append("noise")
|
|
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)
|
|
self.sample["split_lens"].append(len(span_index) + 2)
|
|
curr = vision_end + 1
|
|
curr_rope_id += 1
|
|
self.sample["packed_text_indexes"].append(curr)
|
|
curr += 1
|
|
|
|
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:
|
|
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)))
|
|
|
|
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):
|
|
curr = 0
|
|
sample_lens = 0
|
|
curr_rope_id = 0
|
|
curr_video_grid_thw = []
|
|
|
|
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 state.
|
|
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
|
|
|
|
self.sample["packed_text_ids"].append(self.start_of_image)
|
|
self.sample["packed_text_indexes"].append(curr)
|
|
curr += 1
|
|
split_len_vision_token = 1
|
|
|
|
if isinstance(vit_video_tensor, torch.Tensor):
|
|
self.sample["vit_video_tensors"].append(vit_video_tensor)
|
|
|
|
# Preprocess the video.
|
|
vit_tokens = patchify_video_with_merge(
|
|
vit_video_tensor, self.data_config.vit_patch_size, self.data_config.vit_patch_size_temporal
|
|
)
|
|
num_video_tokens = vit_tokens.shape[0] // 4
|
|
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,
|
|
]
|
|
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)
|
|
)
|
|
|
|
self.sample["packed_vit_token_indexes"].extend(range(curr, curr + num_video_tokens))
|
|
curr += num_video_tokens
|
|
split_len_vision_token += num_video_tokens
|
|
|
|
# Fill placeholder position_ids.
|
|
self.sample["packed_text_ids"].extend([self.image_token_id] * num_video_tokens)
|
|
|
|
# Append the <|endofimage|> image end token.
|
|
self.sample["packed_text_ids"].append(self.end_of_image)
|
|
self.sample["packed_text_indexes"].append(curr)
|
|
curr += 1
|
|
split_len_vision_token += 1
|
|
|
|
# Update sequence state.
|
|
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
|
|
|
|
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)
|
|
|
|
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
|
|
|
|
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)))
|
|
|
|
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)
|
|
|
|
curr += len(shifted_text_ids_answer)
|
|
split_len_answer = len(shifted_text_ids_answer)
|
|
|
|
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):
|
|
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"])
|
|
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]:
|
|
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")
|
|
|
|
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]:
|
|
"""Get a single sample."""
|
|
thw_video, thw_downsample = self.get_thw()
|
|
t, h, w = thw_downsample
|
|
num_vid_tokens = t * h * w
|
|
spatial_merge_size = 2
|
|
|
|
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])
|
|
else:
|
|
user_prompt = sample["data"]
|
|
user_prompt = self._maybe_enhance_t2v_prompt(user_prompt)
|
|
|
|
if self.data_config.text_template:
|
|
caption_instruction = generate_system_prompt(system_prompt_type=self.data_config.task, vision_type=self.data_config.target_modality)
|
|
|
|
text_template_user, text_template_assistant, vit_num_tokens, video_types = [], [], [], []
|
|
if self.system_prompt_type == 'SP2':
|
|
user_prompt = caption_instruction + " " + user_prompt
|
|
caption_instruction = "You are a helpful assistant. "
|
|
elif self.system_prompt_type == 'SP1':
|
|
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))
|
|
packed_position_ids.extend(range(0, text_split_len))
|
|
sample_modality.extend([modality_map['text']] * text_split_len)
|
|
|
|
if self.data_config.text_template:
|
|
text_template_assistant.append({"type":self.data_config.target_modality})
|
|
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([thw_video]),
|
|
"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 get_thw(self):
|
|
_T, _H, _W = self.data_config.vae_downsample
|
|
if self.data_config.target_modality == "image":
|
|
t = 1
|
|
t_ = 1
|
|
elif self.data_config.target_modality == "video":
|
|
t = (self.data_config.num_frames - 1) // _T + 1
|
|
t_ = self.data_config.num_frames
|
|
|
|
h = self.data_config.H // _H
|
|
w = self.data_config.W // _W
|
|
return [t_, self.data_config.H, self.data_config.W], [t, h, w] # Original video size and downsampled size.
|
|
|
|
|
|
def gen_timesteps(self, t, h, w, curr, num_vid_tokens):
|
|
timestep = np.random.randn()
|
|
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:
|
|
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)))
|
|
return packed_timesteps, mse_loss_indexes
|
|
|
|
def tiv2v_sample(self, idx: int) -> Dict[str, Any]:
|
|
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"]
|
|
additional_info = sample["data"]["additional_info"] if "additional_info" in sample["data"] else []
|
|
|
|
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:
|
|
vit_image_tensor = self.get_video_tensor_online(element, vision_stream="vit_video", element_dtype=element_dtype)
|
|
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)
|
|
|
|
# Process VAE conditioning input.
|
|
vae_image_tensor = self.get_video_tensor_online(element, vision_stream="vae_video", element_dtype=element_dtype)
|
|
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
|
|
)
|
|
vit_num_tokens.append(num_tokens_)
|
|
text_template_user.append({"type": element_dtype})
|
|
video_types.append("cond_vae_video")
|
|
|
|
if self.sample_task == 'edit':
|
|
self.data_config.num_frames, self.data_config.H, self.data_config.W = vae_image_tensor.shape[1], vae_image_tensor.shape[2], vae_image_tensor.shape[3]
|
|
|
|
# Process the VAE target input.
|
|
thw_video, thw_downsample = self.get_thw()
|
|
video_sizes.append(thw_video)
|
|
t, h, w = thw_downsample
|
|
num_vid_tokens = t * h * w
|
|
self.sample["vae_data_mode"].append("online")
|
|
spatial_merge_size = 2
|
|
vae_video_grid_thw = [
|
|
t,
|
|
h * spatial_merge_size,
|
|
w * spatial_merge_size,
|
|
]
|
|
curr_video_grid_thw.append(vae_video_grid_thw)
|
|
self.sample["vae_video_grid_thw"].append(vae_video_grid_thw)
|
|
self.sample["vae_latent_shapes"].append((t, h, w))
|
|
# Use 3D-aware extrapolated position encoding.
|
|
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)
|
|
packed_timesteps, mse_loss_indexes = self.gen_timesteps(t, h, w, curr, num_vid_tokens)
|
|
self.sample["packed_timesteps"].extend(packed_timesteps)
|
|
vae_tensor = torch.randn([3, thw_video[0], thw_video[1], thw_video[2]], dtype=torch.float32) # Raw CTHW video, not latent.
|
|
self.sample["vae_video_tensors"].append(vae_tensor)
|
|
if self.data_config.text_template:
|
|
vit_num_tokens.append(num_vid_tokens)
|
|
text_template_assistant.append({"type": self.data_config.target_modality})
|
|
video_types.append("target_vae_video")
|
|
|
|
if text_template_user[0]['type']=='text':
|
|
text_template_user = text_template_user[1:] + text_template_user[:1]
|
|
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 ff2v_sample(self, idx: int) -> Dict[str, Any]:
|
|
"""Get a single sample."""
|
|
sample_modality, text_template_user, text_template_assistant, vit_num_tokens, video_types, search_text = [], [], [], [], [], ''
|
|
self.sample = self.set_sequence_status()
|
|
sample_lens = 0
|
|
sample = self.data[idx]
|
|
|
|
index = sample["index"]
|
|
data_sample = sample["data"]
|
|
additional_info = sample["data"]["additional_info"] if "additional_info" in sample["data"] else []
|
|
|
|
interleave_array, element_dtype_array, istarget_in_interleave = data_sample["interleave_array"], data_sample["element_dtype_array"], data_sample["istarget_in_interleave"]
|
|
interleave_array = list(interleave_array)
|
|
|
|
text_idx = next((i for i, dtype in enumerate(element_dtype_array) if dtype == "text"), None)
|
|
image_idx = next(
|
|
(
|
|
i
|
|
for i, (dtype, is_target) in enumerate(zip(element_dtype_array, istarget_in_interleave))
|
|
if dtype == "image" and is_target == 0
|
|
),
|
|
None,
|
|
)
|
|
if text_idx is not None and image_idx is not None:
|
|
interleave_array[text_idx] = self._maybe_enhance_i2v_prompt(
|
|
interleave_array[text_idx],
|
|
image_path=interleave_array[image_idx],
|
|
)
|
|
|
|
curr, curr_rope_id, curr_split_len, curr_video_grid_thw, video_sizes, caption_all, vae_image_tensor = 0, 0, 0, [], [], '', None
|
|
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:
|
|
vae_image_tensor = self.get_video_tensor_online(element, vision_stream="vae_video", element_dtype=element_dtype)
|
|
self.data_config.H, self.data_config.W = vae_image_tensor.shape[2], vae_image_tensor.shape[3]
|
|
self.frame_condition_idx = [0]
|
|
|
|
# Add the target VAE latent.
|
|
thw_video, thw_downsample = self.get_thw()
|
|
video_sizes.append(thw_video)
|
|
t, h, w = thw_downsample
|
|
num_vid_tokens = t * h * w
|
|
self.sample["vae_data_mode"].append("online")
|
|
spatial_merge_size = 2
|
|
vae_video_grid_thw = [
|
|
t,
|
|
h * spatial_merge_size,
|
|
w * spatial_merge_size,
|
|
]
|
|
curr_video_grid_thw.append(vae_video_grid_thw)
|
|
self.sample["vae_video_grid_thw"].append(vae_video_grid_thw)
|
|
self.sample["vae_latent_shapes"].append((t, h, w))
|
|
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)
|
|
packed_timesteps, mse_loss_indexes = self.gen_timesteps(t, h, w, curr, num_vid_tokens)
|
|
self.sample["packed_timesteps"].extend(packed_timesteps)
|
|
vae_tensor = torch.randn([3, thw_video[0], thw_video[1], thw_video[2]], dtype=torch.float32)
|
|
if vae_image_tensor is not None: # Fill in the first frame.
|
|
vae_tensor[:, :4, :, :] = vae_image_tensor[:, 0:1, :, :].repeat(1, 4, 1, 1)
|
|
else:
|
|
raise ValueError("vae_image_tensor of first frame is None")
|
|
self.sample["vae_video_tensors"].append(vae_tensor)
|
|
if self.data_config.text_template:
|
|
vit_num_tokens.append(num_vid_tokens)
|
|
text_template_assistant.append({"type": self.data_config.target_modality})
|
|
video_types.append("target_vae_video")
|
|
|
|
if len(text_template_user) > 0 and text_template_user[0]['type'] == 'text':
|
|
text_template_user = text_template_user[1:] + text_template_user[:1]
|
|
caption_instruction = generate_system_prompt(system_prompt_type=self.data_config.task, vision_type=self.data_config.target_modality)
|
|
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"]
|
|
|
|
self.sample["sample_task"] = torch.ones(sample_lens) * sample_task_map[self.sample_task]
|
|
self.sample["sample_modality"] = sample_modality
|
|
has_vit_video_grid = self.sample["vit_video_grid_thw"] != []
|
|
has_packed_vit_token_indexes = self.sample["packed_vit_token_indexes"] != []
|
|
if self.frame_condition_idx != []:
|
|
mse_loss_indexes_first = self.sample["mse_loss_indexes"][0]
|
|
self.sample["mse_loss_indexes"] = [idx + mse_loss_indexes_first for idx in mse_loss_indexes]
|
|
|
|
finalized_sample = self._finalize_sample(
|
|
sample_lens,
|
|
curr_video_grid_thw,
|
|
sample_type="gen",
|
|
sample=sample,
|
|
additional_fields={
|
|
"caption": caption_all,
|
|
"caption_cn": caption_all,
|
|
"additional_info": additional_info,
|
|
},
|
|
video_sizes=video_sizes,
|
|
)
|
|
if not has_vit_video_grid:
|
|
finalized_sample["vit_video_grid_thw"] = None
|
|
if not has_packed_vit_token_indexes:
|
|
finalized_sample["packed_vit_token_indexes"] = None
|
|
return finalized_sample
|
|
|
|
def render_template(self, instruction, text_template_assistant, text_template_user, vit_num_tokens, search_text=""):
|
|
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)
|
|
|
|
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]:
|
|
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 = 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:
|
|
if self.data_config.text_template:
|
|
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]
|
|
|
|
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':
|
|
caption_i = "You are a helpful assistant. " + caption_i
|
|
element = [caption_i, caption_q, caption_a]
|
|
|
|
caption_i, caption_q, caption_a = element[0], element[1], element[2]
|
|
|
|
text_template_assistant.append({"type": "text", "text": caption_a})
|
|
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]
|
|
else:
|
|
if isinstance(element, list):
|
|
element = element[-1]
|
|
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
|
|
|
|
elif element_dtype in ["image", "video"]:
|
|
|
|
vit_image_tensor = self.get_video_tensor_online(element, vision_stream="vit_video", element_dtype=element_dtype)
|
|
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]:
|
|
task = self.data_config.task
|
|
# Get target modality
|
|
if '_t' in task:
|
|
self.data_config.target_modality = 'text'
|
|
elif '2i' in task or 'image' in task:
|
|
self.data_config.target_modality = 'image'
|
|
else:
|
|
self.data_config.target_modality = 'video'
|
|
|
|
# Get sample
|
|
if task in ["t2i", "t2v"]: # Text-to-image or text-to-video
|
|
return self.t2v_sample(idx)
|
|
elif 'edit' in task: # Video Editing or Image Editing
|
|
self.sample_task = 'edit'
|
|
return self.tiv2v_sample(idx)
|
|
elif 'idip' in task: # Video IDIP, Image IDIP
|
|
self.sample_task = 'idip'
|
|
return self.tiv2v_sample(idx)
|
|
elif "i2v" in task: # Text-Image-to-Video
|
|
self.sample_task = 't2v'
|
|
return self.ff2v_sample(idx)
|
|
elif task in ["x2t", "x2t_image", "x2t_video"]: # Multi-modal Understanding
|
|
return self.x2t_sample(idx)
|
|
else:
|
|
raise ValueError(f"Unknown task: {task}")
|