# coding: utf-8 from abc import ABC, abstractmethod import os import random from typing import Any, Dict, Iterator, List, Optional import pyarrow.parquet as pq from PIL import Image import numpy as np from io import BytesIO import torch from data.datasets_factory.distributed_iterable_dataset import DistributedIterableDataset from data.parquet_utils import init_arrow_pf_fs, read_parquet_rows from data.transforms import VideoTransform from common.utils.logging import get_logger import decord import cv2 from decord import VideoReader from data.video.sampler.frames import FrameSamplerOutput from config.config_factory import TemplateArguments from data.common import parse_videochat2it_doubao_caption class TextCleaner: def __call__(self, text: Any) -> str: if text is None: return "" return str(text).strip() def _collect_target_positions(element_dtype_array: List[str], target_modality: str) -> List[int]: # Targets can only be selected from index >= 1, excluding index 0 return [i for i, t in enumerate(element_dtype_array) if i >= 1 and t == target_modality] def sample_task(task_type, task_type_rate): """ Sample one item from task_type according to task_type_rate. If all weights are zero or negative, fall back to uniform sampling. """ if len(task_type) != len(task_type_rate): raise ValueError("task_type and task_type_rate must have the same length") # Treat negative weights as 0 weights = [max(0.0, float(w)) for w in task_type_rate] if sum(weights) == 0.0: weights = [1.0] * len(task_type) # Fall back to uniform weights return random.choices(task_type, weights=weights, k=1)[0] def data_invert_text_image_pair(interleave_array, element_dtype_array, target_modality): # Handle a single image-text pair by swapping without considering position if len(element_dtype_array) == 2: if element_dtype_array[-1] != target_modality: interleave_array = interleave_array[::-1] element_dtype_array = element_dtype_array[::-1] return interleave_array, element_dtype_array class BaseMMParquetDataset(DistributedIterableDataset, ABC): def __init__( self, dataset_name: str, tokenizer: Any, data_dir_list: List[str], local_rank: int = 0, world_size: int = 1, num_workers: int = 8, data_status: Optional[Any] = None, **kwargs: Any, ): """ data_dir_list: list of data directories contains parquet files """ super().__init__(dataset_name, local_rank, world_size, num_workers) # Store config only and delay real initialization self.tokenizer = tokenizer self.data_dir_list = data_dir_list self.data_status = data_status self.seed = kwargs.get('seed', 42) self.caption_key = kwargs.get( 'caption_key', 'v3_0_long_internlm_caption_en_text' ) self.transform: VideoTransform = kwargs.get('transform') self.frame_sampler = kwargs.get("video_frame_sampler") # Video sampling self.vae_downsample = kwargs.get( 'vae_downsample', ( getattr(self.transform, 'stride_temporal', 4), getattr(self.transform, 'stride_spatial', 16), getattr(self.transform, 'stride_spatial', 16), ) ) self.max_bytes = kwargs.get('max_bytes', -1) self.logger = get_logger() # Mark as not initialized yet self.data_paths = kwargs.get('all_data_paths') self.cpu_count = os.cpu_count() or 1 self.apply_chat_template = kwargs.get('apply_chat_template', False) if self.apply_chat_template: self.chat_template = TemplateArguments().chat_template_T2I self.vision_stream = kwargs.get("vision_stream", "vae_video") # 'vae_video' | 'vit_video' self.vit_downsample = kwargs.get("vit_downsample", (2, 28, 28)) if kwargs.get('vit_transform') is not None: self.vit_transform: VideoTransform = kwargs.get('vit_transform') else: self.vit_transform: VideoTransform = kwargs.get('transform') self.text_cleaner = TextCleaner() self.dataset_type = kwargs.get("dataset_type", "interleave") self.force_last_as_gt_prob = kwargs.get("force_last_as_gt_prob", 0.0) self.N_target = kwargs.get("N_target", 1) self.N_target_random_prob = kwargs.get("N_target_random_prob", 0.0) self.max_num_split_vit, self.max_num_split_vae, self.max_num_split_text = kwargs.get("max_num_split", [1000, 1000, 1000]) self.is_image = kwargs.get("is_image", True) self.res_dump = kwargs.get("res_dump", "12fps_192p") self.data_mode = kwargs.get("data_mode", "online") self.text_template = kwargs.get("text_template", False) self.vision_cond_type = kwargs.get("vision_cond_type", ["vit"]) self.fbyf_group_interval = kwargs.get("fbyf_group_interval", -1) self.fbyf_type = kwargs.get("fbyf_type", "group") self.sample_task = kwargs.get("sample_task", "t2v") # Task identifier for joint multi-task training if "ocr" in self.dataset_type: self.data_filter = kwargs.get("data_filter", {}) self.save_video_image = kwargs.get("save_video_image", False) self.data_config = kwargs # ==== Hooks required or optionally overridden by subclasses ==== def select_columns(self) -> Optional[List[str]]: """Optional: return column names to read from parquet to reduce IO. None reads all columns.""" if "interleave" in self.dataset_type: return None # ["element_dtype_array", "interleave_array"] elif "ffhq" in self.dataset_type or "imagenet" in self.dataset_type: return ["tos_url", self.caption_key] # Select a subset of columns elif "hav" in self.dataset_type: return ["media_url", "properties"] elif "vertical" in self.dataset_type: return ["meta_url"] elif "audio_human" in self.dataset_type: return ["video_meta_url"] return None @staticmethod def _read_decord(video: VideoReader, frame_idx: List[int]) -> List[Image.Image]: # Use get_batch() instead of reading frames one by one for better performance frames_np = video.get_batch(frame_idx).asnumpy() return [Image.fromarray(frame) for frame in frames_np] def vision_token_count(self, video_tensor: torch.Tensor) -> int: _, T, H, W = video_tensor.shape if self.vision_stream == "vit_video": _T, _H, _W = self.vit_downsample return (T // _T) * (H // _H) * (W // _W) elif self.vision_stream == "vae_video": _T, _H, _W = self.vae_downsample return ((T // _T) + 1) * (H // _H) * (W // _W) else: raise ValueError(f"Unknown vision_stream: {self.vision_stream}") def get_thwc_url_new(self, media_url, worker_id): raise NotImplementedError("Remote media URLs are not supported. Use local parquet rows with embedded bytes.") 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": fps =24 try: fps = int(round(float(video_reader.get_avg_fps()))) except Exception: pass frames_info = { "clip_indices": [(0, total_frames)], # Left-closed, right-open interval; default is a single clip "fps": fps, # Default is 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) # Default DIT path video_tensor = self.vit_transform(video_frames) # fix: use List input if self.is_image: video_tensor = video_tensor.repeat(1, 2, 1, 1) # NOTE: Duplicate single images because the encoder temporal patch size is 2 # NOTE: Video length must be even if video_tensor.shape[1] % 2 == 1: last_frame = video_tensor[:, -1:, :, :] video_tensor = torch.cat([video_tensor, last_frame], dim=1) _, T, H, W = video_tensor.shape return (T, H, W) def get_video_tensor_online(self, media_url, vision_stream, worker_id=0, element_dtype="image", raw_bytes_input=False) -> torch.Tensor: self.vision_stream = vision_stream if raw_bytes_input: # raise NotImplementedError(f"raw_bytes_input must be True for {vision_stream}") video_stream = BytesIO(media_url) # # Method A: write directly to file; simplest debug code # from datetime import datetime # Import datetime module # timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") # with open(f"saved_image_{timestamp}.png", "wb") as f: # f.write(video_stream.getvalue()) # getvalue() returns all bytes in BytesIO # # Method A: write directly to file; simplest debug code # from datetime import datetime # Import datetime module # # Keep a high-precision timestamp: year-month-day_hour-minute-second-microsecond # timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") # # Use a video filename with .mp4 suffix and write video bytes # with open(f"saved_video_{timestamp}.mp4", "wb") as f: # f.write(video_stream.getvalue()) # video_stream is a BytesIO object containing video bytes else: raise NotImplementedError("Remote media URLs are not supported. Use raw_bytes_input=True with local parquet bytes.") if self.is_image and element_dtype == "image": image = Image.open(video_stream) if image.mode == "P": image = image.convert("RGBA") if image.mode == "RGBA": # Composite on a white background to remove transparency bg = Image.new("RGB", image.size, (255, 255, 255)) bg.paste(image, mask=image.split()[3]) # Use the alpha channel as the mask image = bg else: image = image.convert("RGB") video_frames = [image] # Save image if self.save_video_image: self.path = f"{self.path}.jpg" image.save(self.path, quality=95) print(f"Saved image to {self.path}") else: # for video video_reader = VideoReader(video_stream, ctx=decord.cpu(worker_id % self.cpu_count)) total_frames = len(video_reader) # Save video if self.save_video_image: fps = video_reader.get_avg_fps() width, height = video_reader[0].shape[1], video_reader[0].shape[0] self.path =f"{self.path}.mp4" # Saved video path # Save video with OpenCV fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(self.path, fourcc, fps, (width, height)) for frame in video_reader: frame_rgb = cv2.cvtColor(frame.asnumpy(), cv2.COLOR_BGR2RGB) # Convert BGR to RGB out.write(frame_rgb) out.release() print(f"Saved image to {self.path} with fps {fps}") sampler_name = self.frame_sampler.__class__.__name__ if sampler_name == "MultiClipsFrameSampler": fps =24 try: fps = int(round(float(video_reader.get_avg_fps()))) except Exception: pass frames_info = { "clip_indices": [(0, total_frames)], # Left-closed, right-open interval; default is a single clip "fps": fps, # Default is 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 self.is_image: video_tensor = video_tensor.repeat(1, 2, 1, 1) # NOTE: Duplicate single images because the encoder temporal patch size is 2 # NOTE: Video length must be even 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}") if not (self.is_image and element_dtype == "image"): del video_frames, video_reader, video_stream return video_tensor, self.vision_token_count(video_tensor) def get_video_tensor_offline(self, media_url, vision_stream, worker_id=0) -> torch.Tensor: self.vision_stream = vision_stream if vision_stream == "vae_video": video_tensor = media_url[0] # [t, h, w, c] num_token = video_tensor.shape[0] * video_tensor.shape[1] * video_tensor.shape[2] elif vision_stream == "vit_video": video_tensor = media_url[1] # [L, D] num_token = video_tensor.shape[0] if len(media_url) == 3 and vision_stream == "vit_video": if isinstance( media_url[2], str): # Get THW information thw = self.get_thwc_url_new(media_url[2], worker_id = worker_id) else: thw = media_url[2][1:] num_token_ = thw[0] * thw[1] * thw[2] // self.vit_downsample[0] // self.vit_downsample[1] // self.vit_downsample[2] if num_token_ != num_token: raise ValueError(f"Video tensor shape {video_tensor.shape} not match thw {thw}: {num_token_} != {num_token}") else: thw = None return video_tensor, num_token, thw def get_video_tensor(self, media_url, vision_stream, worker_id=0, element_dtype="image", raw_bytes_input=False) -> torch.Tensor: if isinstance(media_url, tuple): # offline video_tensor, num_tokens_, thw = self.get_video_tensor_offline(media_url, vision_stream=vision_stream) self.data_mode = "offline" video_tensor = [video_tensor] # Return as a list to distinguish this from online format else: # online video_tensor, num_tokens_ = self.get_video_tensor_online(media_url, vision_stream=vision_stream, worker_id=worker_id, element_dtype=element_dtype, raw_bytes_input=raw_bytes_input) self.data_mode = "online" thw = None return video_tensor, num_tokens_, thw # Get the sample count for each file during initialization def get_file_sample_counts(self, data_paths): sample_counts = [] for path in data_paths: fs = init_arrow_pf_fs(path) with fs.open_input_file(path) as f: fr = pq.ParquetFile(f) # Estimate or exactly compute the sample count in each file count = sum(fr.metadata.row_group(i).num_rows for i in range(fr.num_row_groups)) sample_counts.append(count) return sample_counts def get_condition_target_idx( self, element_dtype_array, ): if len(element_dtype_array) == 1 and self.target_modality in element_dtype_array: # A single element means there is no condition return [], [0], 1 target_pos_all = _collect_target_positions(element_dtype_array, self.target_modality) # Get target element positions excluding position 0 pos_all = list(range(len(element_dtype_array))) N_all = len(target_pos_all) if N_all == 0: # If no target type is available, fall back to all condition and empty target; upstream should drop this sample return None if random.random() < self.N_target_random_prob: # Randomly choose target count by probability; if N_target_random_prob is 0, default to self.N_target N_target = random.randint(1, N_all) else: N_target = self.N_target if self.target_modality in ["image", "video"] : N_target = min(N_target, N_all, self.max_num_split_vae) # Ensure target count does not exceed total count elif self.target_modality == "text": N_target = min(N_target, N_all, self.max_num_split_text) # --- Select target set --- choose_last = random.random() < self.force_last_as_gt_prob # Randomly decide whether to force the last item as target if choose_last: target_idx = target_pos_all[-N_target:] # Indexes of target elements else: target_last = random.randint(N_target - 1, N_all - 1) # Ensure the selected last target element index can cover N_target targets target_idx = target_pos_all[target_last - N_target + 1 : target_last + 1] # Add 1 because target_last is the target element index condition_idx = pos_all[: target_idx[0]] # Indexes of condition elements return condition_idx, target_idx, N_target @abstractmethod def _process_row(self, row, parquet_idx, row_group_id, row_idx, worker_id, parquet_file_path): pass def __iter__(self) -> Iterator[Dict[str, Any]]: data_paths_per_worker, worker_id = self.get_data_paths_per_worker() if self.data_status is not None: parquet_start_id = self.data_status[worker_id][0] row_group_start_id = self.data_status[worker_id][1] row_start_id = self.data_status[worker_id][2] + 1 else: parquet_start_id = 0 row_group_start_id = 0 row_start_id = 0 # log if data_paths_per_worker: self.logger.info( f"Rank-{self.local_rank} worker-{worker_id} dataset-{self.dataset_name}: " f"{len(data_paths_per_worker)} parquet files (first: {data_paths_per_worker[0]}, " f"last: {data_paths_per_worker[-1]}), " f"resuming at parquet#{parquet_start_id}, rg#{row_group_start_id}, row#{row_start_id}" ) else: self.logger.warning(f"Rank-{self.local_rank} worker-{worker_id} dataset-{self.dataset_name}: " "has 0 parquet files!") while True: data_paths_per_worker_ = data_paths_per_worker[parquet_start_id:] for parquet_idx, parquet_file_path in enumerate(data_paths_per_worker_, start=parquet_start_id): fs = init_arrow_pf_fs(parquet_file_path) with fs.open_input_file(parquet_file_path) as f: fr = pq.ParquetFile(f) row_group_ids = list(range(fr.num_row_groups)) row_group_ids_ = row_group_ids[row_group_start_id:] # Column pruning: subclasses can provide select_columns() cols = self.select_columns() # Default is None, which reads all columns for row_group_id in row_group_ids_: rows = read_parquet_rows(fr, row_group_id, columns=cols) rows = rows[row_start_id:] for row_idx, row in enumerate(rows, start=row_start_id): sample = self._process_row(row, parquet_idx, row_group_id, row_idx, worker_id, parquet_file_path) if sample: yield sample # self.logger.info(f"parquet_file_path: {parquet_file_path}, row_idx: {row_idx}, row_group_id: {row_group_id}, worker_id:{worker_id}, self.local_rank:{self.local_rank}") # Useful for locating bad data row_start_id = 0 row_group_start_id = 0 parquet_start_id = 0 if self.local_rank == 0: self.logger.info(f"{self.dataset_name} repeat in rank-{self.local_rank} worker-{worker_id}") pass def transform_row(self, row): if self.dataset_type == "text2video_general": video_bytes = row["video_bytes"] caption = row["caption"] interleave_array = [caption, video_bytes] if self.text_first else [video_bytes, caption] element_dtype_array = ["text", "video"] if self.text_first else ["video", "text"] elif self.dataset_type == "text2image_general": image_bytes = row["image_bytes"] caption = row["caption"] interleave_array = [caption, image_bytes] if self.text_first else [image_bytes, caption] element_dtype_array = ["text", "image"] if self.text_first else ["image", "text"] elif self.dataset_type == "x2t_general": if all(key in row for key in ["caption_i", "caption_q", "caption_a"]): caption = [row["caption_i"], row["caption_q"], row["caption_a"]] else: caption = parse_videochat2it_doubao_caption(row) if "image_bytes" in row: interleave_array = [row["image_bytes"], caption] if not self.text_first else [caption, row["image_bytes"]] element_dtype_array = ["image", "text"] if not self.text_first else ["text", "image"] elif "video_bytes" in row: interleave_array = [row["video_bytes"], caption] if not self.text_first else [caption, row["video_bytes"]] element_dtype_array = ["video", "text"] if not self.text_first else ["text", "video"] else: interleave_array = [caption] element_dtype_array = ["text"] elif self.dataset_type == "image2image": interleave_array = [row["caption"], row["input_image_bytes"], row["output_image_bytes"]] element_dtype_array = ["text", "image", "image"] self.force_last_as_gt_prob = 1 self.N_target = 1 self.N_target_random_prob = 0 self.sample_task = "edit" elif self.dataset_type == "image2image_online": interleave_array = [row["instruction"], row["input_image_url"], row["output_image_url"]] element_dtype_array = ["text", "image", "image"] self.force_last_as_gt_prob = 1 self.N_target = 1 self.N_target_random_prob = 0 self.sample_task = "edit" elif self.dataset_type == "video2video": interleave_array = [row["caption"], row["input_video_bytes"], row["output_video_bytes"]] element_dtype_array = ["text", "video", "video"] self.force_last_as_gt_prob = 1 self.N_target = 1 self.N_target_random_prob = 0 self.sample_task = "edit" else: raise ValueError(f"dataset_type {self.dataset_type} not supported") return interleave_array, np.array(element_dtype_array).astype(dtype=object)