wehub-resource-sync
246 lines
13 KiB
Python
246 lines
13 KiB
Python
# coding: utf-8
|
|
|
|
|
|
from typing import Any, Dict, Optional
|
|
import random
|
|
from data.base_mm_parquet_dataset import BaseMMParquetDataset, sample_task, data_invert_text_image_pair
|
|
from data.common import generate_system_prompt
|
|
from data.system_prompt_render import render_qwenvl_prompt, expand_and_index_by_token_ids_new
|
|
|
|
|
|
class X2VInterleaveIterableDataset(BaseMMParquetDataset):
|
|
def __init__(self, *args, **kwargs):
|
|
kwargs.setdefault("vision_stream", "vae_video") # NOTE: Must be vae_video, vit_video, or another supported stream
|
|
super().__init__(*args, **kwargs)
|
|
|
|
self.target_modality = kwargs.get("target_modality", "image") # Default is "image"; target element type is "text" for I2T
|
|
self.task_type = kwargs.get("task_type", "t2v")
|
|
self.task_type_rate = kwargs.get("task_type_rate", 1)
|
|
self.text_first = True
|
|
self.L_video_group = kwargs.get("L_video_group", 1000) # Number of groups used to split a video; values larger than T disable grouping
|
|
self.random_choose_group = kwargs.get("random_choose_group", False) # Whether to randomly choose video groups
|
|
self.raw_bytes_input = kwargs.get("raw_bytes_input", False) # Whether image and video inputs are raw bytes
|
|
if self.local_rank == 0:
|
|
self.logger.info(f"X2VInterleaveIterableDataset raw_bytes_input: {self.raw_bytes_input}")
|
|
|
|
def _process_row(self, row: Any, parquet_idx: int, row_group_id: int, row_idx: int, worker_id: int, parquet_file_path: str) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Process one row.
|
|
Return None if processing fails or the data is invalid.
|
|
"""
|
|
try:
|
|
if self.dataset_type == "interleave":
|
|
interleave_array, element_dtype_array = row["interleave_array"], row["element_dtype_array"]
|
|
else:
|
|
# Local text2video_general / text2image_general data also goes through the shared transform_row path.
|
|
try:
|
|
interleave_array, element_dtype_array = self.transform_row(row)
|
|
except Exception as e:
|
|
if "ocr" not in self.dataset_type: # OCR has filtering; skip this log
|
|
self.logger.warning(f"Warning transform row: {e} in self.dataset_type: {self.dataset_type}")
|
|
return None
|
|
|
|
interleave_array, element_dtype_array = data_invert_text_image_pair(interleave_array, element_dtype_array, self.target_modality)
|
|
|
|
try:
|
|
condition_idx, target_idx, N_target = self.get_condition_target_idx(element_dtype_array)
|
|
except Exception as e:
|
|
self.logger.warning(f"Warning processing row: num of target element {self.target_modality} is 0, element_dtype_array is {element_dtype_array}")
|
|
return None
|
|
|
|
# Select task type
|
|
if isinstance(self.task_type, str):
|
|
task_type_sample = self.task_type
|
|
elif "text" not in element_dtype_array[condition_idx]:
|
|
task_type_sample = "v2v"
|
|
elif isinstance(self.task_type, list):
|
|
task_type_sample = sample_task(self.task_type, self.task_type_rate)
|
|
|
|
if task_type_sample == 'flf2v': # First and Last Frame to Video Generation
|
|
self.frame_condition_idx = [0, -1]
|
|
elif task_type_sample == 'ff2v': # First Frame to Video Generation
|
|
self.frame_condition_idx = [0]
|
|
else:
|
|
self.frame_condition_idx = []
|
|
|
|
num_tokens, sequence_plan, text_ids_list, video_tensor_list = [], [], [], [] # sequence_plan stores loss, cfg, and other marker metadata
|
|
text_template_user, text_template_assistant, vit_num_tokens, search_text = [], [], [], ""
|
|
num_split_vit, num_split_vae, num_split_text = 0, 0, 0
|
|
is_offline = False
|
|
for idx in range(target_idx[-1] + 1):
|
|
if idx in condition_idx: # Process condition element
|
|
if task_type_sample in ["t2v", "tv2v", "vt2v", "flf2v", "ff2v"] and element_dtype_array[idx] == "text": # Text condition element (required)
|
|
if num_split_text >= self.max_num_split_text:
|
|
continue
|
|
caption = interleave_array[idx]
|
|
caption = self.text_cleaner(caption)
|
|
|
|
if not caption:
|
|
continue
|
|
if self.text_template:
|
|
if random.random() > self.data_config['text_cond_dropout_prob']: # HACK: cfg handling; drop text condition
|
|
text_template_user.append({"type": "text", "text": caption})
|
|
search_text = caption
|
|
else:
|
|
text_ids = self.tokenizer.encode(caption)
|
|
num_tokens.append(len(text_ids))
|
|
text_ids_list.append(text_ids)
|
|
sequence_plan.append(
|
|
{
|
|
"type": "text",
|
|
"enable_cfg": 1,
|
|
"loss": 0,
|
|
"special_token_loss": 0, # NOTE: Special tokens are provided manually and are not predicted
|
|
"special_token_label": None,
|
|
}
|
|
)
|
|
num_split_text += 1
|
|
elif task_type_sample in ["v2v", "tv2v", "vt2v", "flf2v", "ff2v"] and element_dtype_array[idx] in ["image", "video"]: # Process visual condition element (optional)
|
|
if num_split_vit >= self.max_num_split_vit:
|
|
continue
|
|
if self.text_template and random.random() < self.data_config['vit_cond_dropout_prob']: # HACK: cfg handling; drop vit condition here, which also drops vae and vit conditions
|
|
continue
|
|
|
|
media_url = interleave_array[idx]
|
|
if "vit" in self.vision_cond_type:
|
|
# if self.text_template and random.random() < self.data_config['vit_cond_dropout_prob']: # HACK: cfg handling; drop vit condition
|
|
# pass
|
|
# else:
|
|
video_tensor, num_tokens_, thw = self.get_video_tensor(media_url, vision_stream="vit_video", element_dtype=element_dtype_array[idx], raw_bytes_input=self.raw_bytes_input)
|
|
video_tensor_list.append(video_tensor)
|
|
sequence_plan.append(
|
|
{
|
|
"type": "vit_video",
|
|
"enable_cfg": 1,
|
|
"loss": 0,
|
|
"special_token_loss": 0, # NOTE: Special tokens are provided manually and are not predicted
|
|
"special_token_label": None, # eos
|
|
"thw": thw,
|
|
"apply_text_template": self.text_template, # Default is false
|
|
}
|
|
)
|
|
num_tokens.append(num_tokens_) # NOTE: Count vision tokens
|
|
num_split_vit += 1
|
|
|
|
if self.text_template:
|
|
text_template_user.append({"type": element_dtype_array[idx]})
|
|
vit_num_tokens.append(num_tokens_)
|
|
|
|
del video_tensor
|
|
if "vae_Nsplit" in self.vision_cond_type:
|
|
video_tensor, num_tokens_, thw = self.get_video_tensor(media_url, vision_stream="vae_video", element_dtype=element_dtype_array[idx], raw_bytes_input=self.raw_bytes_input)
|
|
sequence_plan.append(
|
|
{
|
|
"type": "vae_video",
|
|
"enable_cfg": 1,
|
|
"loss": 0,
|
|
"special_token_loss": 0, # NOTE: Special tokens are provided manually and are not predicted
|
|
"special_token_label": None, # eos
|
|
"thw": thw,
|
|
"apply_text_template": self.text_template, # Default is false
|
|
}
|
|
)
|
|
video_tensor_list.append(video_tensor)
|
|
num_tokens.append(num_tokens_) # NOTE: Count vision tokens
|
|
num_split_vae += 1
|
|
|
|
if self.text_template:
|
|
text_template_user.append({"type": element_dtype_array[idx]})
|
|
vit_num_tokens.append(num_tokens_)
|
|
|
|
del video_tensor
|
|
|
|
elif idx in target_idx: # Process target element
|
|
# if num_split_vae >= self.max_num_split_vae:
|
|
# continue
|
|
media_url = interleave_array[idx]
|
|
video_tensor, num_tokens_, thw = self.get_video_tensor(media_url, vision_stream="vae_video", element_dtype=element_dtype_array[idx], raw_bytes_input=self.raw_bytes_input)
|
|
|
|
if isinstance(video_tensor, list):
|
|
is_offline = True
|
|
|
|
# Support group-by-group generation
|
|
if is_offline:
|
|
_ = video_tensor[0].shape[0]
|
|
else:
|
|
_ = (video_tensor.shape[1] - 1) // self.vae_downsample[0] + 1
|
|
|
|
video_tensor_list.append(video_tensor)
|
|
del video_tensor
|
|
num_tokens.append(num_tokens_) # NOTE: Count vision tokens
|
|
sequence_plan.append(
|
|
{
|
|
"type": "vae_video",
|
|
"enable_cfg": 0,
|
|
"loss": 1,
|
|
"special_token_loss": 0, # NOTE: Special tokens are provided manually and are not predicted
|
|
"special_token_label": None,
|
|
"frame_condition_idx": self.frame_condition_idx,
|
|
"L_video_group": self.L_video_group,
|
|
"random_choose_group": self.random_choose_group,
|
|
"thw": thw,
|
|
"apply_text_template": self.text_template, # Default is false
|
|
}
|
|
#"N_key_frame": self.N_key_frame,
|
|
)
|
|
num_split_vae += 1
|
|
|
|
if self.text_template:
|
|
text_template_assistant.append({"type": element_dtype_array[idx]})
|
|
vit_num_tokens.append(num_tokens_)
|
|
|
|
# Handle video before text:
|
|
text_template_user = text_template_user[1:] + text_template_user[:1] # HACK
|
|
|
|
|
|
messages = [
|
|
{
|
|
"role": "user",
|
|
"content": text_template_user, # Original usage
|
|
},
|
|
{
|
|
"role": "assistant",
|
|
"content": text_template_assistant,
|
|
},
|
|
]
|
|
|
|
caption_all = render_qwenvl_prompt(messages, default_system=generate_system_prompt(self.sample_task, vision_type=element_dtype_array[idx]), 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
|
|
)
|
|
|
|
text_ids_list.append(all_token_id)
|
|
sequence_plan.append(
|
|
{
|
|
"type": "text_template",
|
|
"enable_cfg": 0,
|
|
"loss": 0,
|
|
"special_token_loss": 0,
|
|
"special_token_label": None,
|
|
"packed_label_index": tgt_index,
|
|
"spans_index": spans_index,
|
|
'cap_index': search_index,
|
|
}
|
|
)
|
|
|
|
sample = dict(
|
|
video_tensor_list=video_tensor_list,
|
|
text_ids_list=text_ids_list,
|
|
num_tokens=sum(num_tokens) if not self.text_template else len(all_token_id),
|
|
sequence_plan=sequence_plan,
|
|
sample_task=self.sample_task,
|
|
N_target=N_target,
|
|
data_indexes={
|
|
"data_indexes": [parquet_idx, row_group_id, row_idx],
|
|
"worker_id": worker_id,
|
|
"dataset_name": self.dataset_name,
|
|
},
|
|
)
|
|
return sample
|
|
|
|
except Exception as e:
|
|
# Keep a top-level catch-all for unexpected errors
|
|
self.logger.warning(f"Error processing row: {e} in rg#{row_group_id}, {parquet_file_path}")
|
|
return None
|