chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,14 @@
|
||||
# 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
|
||||
@@ -0,0 +1,497 @@
|
||||
# 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)
|
||||
+225
@@ -0,0 +1,225 @@
|
||||
# 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 torch
|
||||
import random
|
||||
import io
|
||||
from einops import rearrange
|
||||
from typing import List
|
||||
import torch.nn.functional as F
|
||||
import re
|
||||
import numpy as np
|
||||
|
||||
RE_ZH = re.compile(r"[\u4e00-\u9fff]")
|
||||
RE_EN = re.compile(r"[A-Za-z]")
|
||||
|
||||
def generate_system_prompt(system_prompt_type="caption", vision_type="video"):
|
||||
if system_prompt_type == "caption":
|
||||
str_list = [
|
||||
f"Generate a detailed and accurate description of the {vision_type}, including all the key moments and visual details.",
|
||||
f"Write an in-depth depiction of the {vision_type}, covering all its aspects.",
|
||||
f"Write an exhaustive depiction of the given {vision_type}, capturing its essence and key moments.",
|
||||
f"Describe the key features of the input {vision_type}, including color, shape, size, texture, objects, background.",
|
||||
]
|
||||
elif system_prompt_type == "t2v" or system_prompt_type == "i2v":
|
||||
str_list = [f"Describe the {vision_type} by detailing the color, quantity, visible text, shape, size, texture, spatial relationships and motion/camera movements of the objects and background:"]
|
||||
elif system_prompt_type == "t2i":
|
||||
str_list = [f"Describe the {vision_type} by detailing the color, quantity, text, shape, size, texture, spatial relationships of the objects and background:"]
|
||||
elif "edit" in system_prompt_type:
|
||||
str_list = [f"Describe the key features of the input {vision_type} (color, shape, size, texture, objects, background), then explain how the user’s text instruction should alter or modify the {vision_type}. Generate a new {vision_type} that meets the user’s requirements while maintaining consistency with the original input where appropriate."]
|
||||
elif "idip" in system_prompt_type:
|
||||
str_list = [f"Describe the key features of the input image (color, shape, size, texture, objects, background, style), then incorporate the user’s text description to generate a new {vision_type} that satisfies the user’s requirements while preserving the essential identity and object or style information from the reference input."]
|
||||
elif 'maze' in system_prompt_type:
|
||||
str_list = [
|
||||
"Describe the key elements of the input maze image (layout, white path, black walls, blue star, red flag, and overall background), then generate a 2D animation. The blue star should slide smoothly along the white path, stop exactly on the red flag, and then acquire a trophy. Ensure the blue star never crosses or enters the black maze walls. Keep the camera as a static top-down view showing the entire maze."
|
||||
]
|
||||
|
||||
return random.choice(str_list)
|
||||
|
||||
|
||||
def shift_position_ids(
|
||||
position_ids: torch.Tensor,
|
||||
pos_shift: any,
|
||||
attn_modes: List[str],
|
||||
split_lens: int,
|
||||
shift_attn_mode=["full_noise", "full"],
|
||||
pro_type=None,
|
||||
i_sample_task=None,
|
||||
i_sample_modality=None,
|
||||
) -> torch.Tensor:
|
||||
curr_split = 0
|
||||
for i, attn_mode in enumerate(attn_modes):
|
||||
if attn_mode in shift_attn_mode:
|
||||
if pro_type == 10: # Related to sample_modality.
|
||||
if position_ids[:, :, i_sample_modality == 4].sum() != 0:
|
||||
pos_shift_type4 = 1000 - position_ids[:, :, i_sample_modality == 4][0, 0, 0]
|
||||
position_ids[0, :, i_sample_modality == 4] += pos_shift_type4
|
||||
if position_ids[:, :, i_sample_modality == 3].sum() != 0:
|
||||
pos_shift_type3 = 2000 - position_ids[:, :, i_sample_modality == 3][0, 0, 0]
|
||||
position_ids[0, :, i_sample_modality == 3] += pos_shift_type3
|
||||
if position_ids[:, :, i_sample_modality == 2].sum() != 0 and sum(i_sample_modality == 2) == sum(i_sample_modality == 1):
|
||||
position_ids[:, :, i_sample_modality == 1] = position_ids[:, :, i_sample_modality == 2]
|
||||
|
||||
curr_split += split_lens[i]
|
||||
|
||||
return position_ids
|
||||
|
||||
def detect_lang_simple(s: str) -> str:
|
||||
"""
|
||||
Fast heuristic: return 'zh' if Chinese is present, 'en' if English letters are present, otherwise 'other'.
|
||||
Useful for quick routing. If both are present, this returns 'zh'; adjust if needed.
|
||||
"""
|
||||
# Remove digits before detection
|
||||
s_without_digits = re.sub(r'\d+', '', s)
|
||||
|
||||
if RE_ZH.search(s_without_digits):
|
||||
return "zh"
|
||||
if RE_EN.search(s_without_digits):
|
||||
return "en"
|
||||
return "other"
|
||||
|
||||
def map_to_nearest_aspect_ratio(h, w, target_resolution=256):
|
||||
"""
|
||||
Map h and w to the closest preset aspect ratio and return adjusted h and w near the target resolution.
|
||||
|
||||
Preset ratios: ["21:9", "16:9", "4:3", "1:1", "3:4", "9:16"].
|
||||
target_resolution: Base target resolution, default 256.
|
||||
"""
|
||||
# Precompute all preset aspect ratios as width / height
|
||||
PRESET_RATIOS = [21 / 9, 16 / 9, 4 / 3, 1 / 1, 3 / 4, 9 / 16]
|
||||
|
||||
# Compute the original aspect ratio
|
||||
original_ratio = w / h
|
||||
|
||||
# Find the closest preset ratio
|
||||
min_index = min(range(len(PRESET_RATIOS)), key=lambda i: abs(original_ratio - PRESET_RATIOS[i]))
|
||||
best_ratio = PRESET_RATIOS[min_index]
|
||||
|
||||
# Compute scale so the longer side is close to the target resolution
|
||||
if best_ratio >= 1: # Landscape: width >= height
|
||||
scale = target_resolution / best_ratio
|
||||
adjusted_w = round(target_resolution)
|
||||
adjusted_h = round(scale)
|
||||
else: # Portrait: height > width
|
||||
scale = target_resolution
|
||||
adjusted_h = round(target_resolution)
|
||||
adjusted_w = round(scale * best_ratio)
|
||||
|
||||
return adjusted_h, adjusted_w
|
||||
|
||||
|
||||
def concat_resize_tensor_list(video_latents: List[torch.Tensor], dim: int = 0, is_offline: bool = False, max_num_frames: int = 121) -> torch.Tensor:
|
||||
"""
|
||||
Concatenate tensors along dim; resize H/W of tensors with different sizes to match target.
|
||||
- tensors: Non-empty list; all tensors must have the same ndim.
|
||||
- dim: Concatenation axis; negative values are supported.
|
||||
- pad_value: Padding value, default 0.0.
|
||||
Returns: Concatenated tensor.
|
||||
"""
|
||||
|
||||
if is_offline:
|
||||
H, W = video_latents[-1].shape[-3], video_latents[-1].shape[-2]
|
||||
else:
|
||||
H, W = video_latents[-1].shape[-2], video_latents[-1].shape[-1]
|
||||
|
||||
padded_video_latents = []
|
||||
num_frames_target = video_latents[-1].shape[dim]
|
||||
|
||||
num_frames_all = num_frames_target
|
||||
for index, video_latent in enumerate(video_latents):
|
||||
if index != len(video_latents) - 1 and num_frames_all + video_latent.shape[dim] > max_num_frames: # Avoid producing videos longer than MAX_NUM_FRAMES
|
||||
continue
|
||||
num_frames_all += video_latent.shape[dim]
|
||||
if is_offline:
|
||||
# video_latent:[t,h,w,c] -> [t,c,h,w]
|
||||
video_latent = rearrange(video_latent, "t h w c -> t c h w")
|
||||
|
||||
h, w = video_latent.shape[-2], video_latent.shape[-1]
|
||||
if h != H or w != W:
|
||||
video_latent = F.interpolate(video_latent, size=(H, W), mode="bilinear", align_corners=False)
|
||||
padded_video_latents.append(video_latent)
|
||||
|
||||
padded_video_latents = torch.cat(padded_video_latents, dim=dim)
|
||||
|
||||
if is_offline:
|
||||
# padded_video_latents: [t,c,h,w] -> [t,h,w,c]
|
||||
padded_video_latents = rearrange(padded_video_latents, "t c h w -> t h w c")
|
||||
|
||||
return padded_video_latents
|
||||
|
||||
|
||||
def concat_pad_tensor_list(video_latents: List[torch.Tensor], dim: int = 0, pad_value: float = 0.0, max_num_frames: int = 121) -> torch.Tensor:
|
||||
"""
|
||||
Concatenate tensors along dim; pad other axes to the maximum length on each axis with pad_value.
|
||||
- tensors: Non-empty list; all tensors must have the same ndim.
|
||||
- dim: Concatenation axis; negative values are supported.
|
||||
- pad_value: Padding value, default 0.0.
|
||||
Returns: Concatenated tensor.
|
||||
"""
|
||||
video_sizes = [item.shape for item in video_latents]
|
||||
max_video_size = [max(item) for item in list(zip(*video_sizes))]
|
||||
padded_video_latents = []
|
||||
num_frames_target = video_latents[-1].shape[dim]
|
||||
|
||||
num_frames_all = num_frames_target
|
||||
for index, video_latent in enumerate(video_latents):
|
||||
if index != len(video_latents) - 1 and num_frames_all + video_latent.shape[dim] > max_num_frames: # Avoid producing videos longer than MAX_NUM_FRAMES
|
||||
continue
|
||||
num_frames_all += video_latent.shape[dim]
|
||||
max_video_size[dim] = video_latent.shape[dim]
|
||||
padded_video_latent = torch.zeros(max_video_size)
|
||||
n1, n2, n3, n4 = video_latent.shape
|
||||
padded_video_latent[:n1, :n2, :n3, :n4] = video_latent
|
||||
padded_video_latents.append(padded_video_latent)
|
||||
|
||||
padded_video_latents = torch.cat(padded_video_latents, dim=dim)
|
||||
|
||||
return padded_video_latents
|
||||
|
||||
|
||||
def parse_videochat2it_doubao_caption(row):
|
||||
try:
|
||||
IQA_i = "View the video attentively and provide a suitable answer to the posed question."
|
||||
rewrite_VQA = json.loads(row['rewrite_VQA'])
|
||||
|
||||
IQA_q = rewrite_VQA['question'] if 'question' in rewrite_VQA.keys() else rewrite_VQA['Question']
|
||||
IQA_a = rewrite_VQA['final_answer']
|
||||
IQA_resoning = rewrite_VQA['reasoning']
|
||||
# Combine reasoning and final_answer as the final answer
|
||||
|
||||
# if random.random() < 0.5: # 50% chance to include the reasoning process
|
||||
# IQA_a = IQA_a + '\n' + IQA_resoning
|
||||
# IQA_i = IQA_i + ' Please provide the reasoning process for selecting the correct answer.'
|
||||
|
||||
if 'options' not in IQA_q and 'Options' not in IQA_q: # When the question does not contain options
|
||||
try:
|
||||
options = rewrite_VQA['options']
|
||||
except:
|
||||
options = rewrite_VQA['Options']
|
||||
if options == []:
|
||||
return [IQA_i, IQA_q, IQA_a]
|
||||
elif isinstance(options,list):
|
||||
options = '\n'.join(options)
|
||||
elif isinstance(options,dict):
|
||||
options_str = [key + ' ' + value if value not in key else key for key,value in options.items()]
|
||||
options = '\n'.join(options_str)
|
||||
|
||||
IQA_q = IQA_q + '\nOptions:\n' + options # Add options to the question
|
||||
return [IQA_i, IQA_q, IQA_a]
|
||||
except:
|
||||
if 'rewrite_VQA' in row.keys():
|
||||
raise ValueError(f"wrong rewrite_VQA in {row['rewrite_VQA']}")
|
||||
else:
|
||||
raise ValueError(f"wrong rewrite_VQA in {row}")
|
||||
@@ -0,0 +1,230 @@
|
||||
# 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
|
||||
|
||||
"""
|
||||
Data helpers used by inference (`inference_lance.py`, `ValidationDataset`) and the
|
||||
Lance model core (`modeling/lance/lance.py`).
|
||||
|
||||
Exported utilities:
|
||||
- Position id helpers (image / video, interpolate / extrapolate)
|
||||
- Patchify helpers (image + video-with-merge)
|
||||
- create_sparse_mask : flex-attention sparse mask builder
|
||||
- add_special_tokens : register chat / vision tokens on a tokenizer
|
||||
- len2weight : CE loss reweighting factor
|
||||
"""
|
||||
|
||||
from einops import rearrange
|
||||
|
||||
import torch
|
||||
from torch.nn.attention.flex_attention import or_masks, and_masks
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Position id helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_flattened_position_ids_interpolate_video(num_frames, img_h, img_w, patch_size, max_num_frames, max_num_patches_per_side):
|
||||
num_patches_h, num_patches_w = img_h // patch_size, img_w // patch_size
|
||||
# temporal
|
||||
boundaries_t = torch.arange(1 / max_num_frames, 1.0, 1 / max_num_frames)
|
||||
fractional_coords_t = torch.arange(0, 1 - 1e-6, 1 / num_frames)
|
||||
bucket_coords_t = torch.bucketize(fractional_coords_t, boundaries_t, right=True)
|
||||
# spatial
|
||||
boundaries_s = torch.arange(1 / max_num_patches_per_side, 1.0, 1 / max_num_patches_per_side)
|
||||
fractional_coords_h = torch.arange(0, 1 - 1e-6, 1 / num_patches_h)
|
||||
fractional_coords_w = torch.arange(0, 1 - 1e-6, 1 / num_patches_w)
|
||||
bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries_s, right=True)
|
||||
bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries_s, right=True)
|
||||
pos_ids = (
|
||||
bucket_coords_t[:, None, None] * max_num_patches_per_side * max_num_patches_per_side
|
||||
+ bucket_coords_h[None, :, None] * max_num_patches_per_side
|
||||
+ bucket_coords_w[None, None, :]
|
||||
).flatten()
|
||||
return pos_ids
|
||||
|
||||
|
||||
def get_flattened_position_ids_extrapolate_video(t, h, w, max_latent_size):
|
||||
"""
|
||||
Defaults:
|
||||
num_frames = 7 (corresponding to 25 frames)
|
||||
max_num_patches_per_side = 64
|
||||
"""
|
||||
coords_t = torch.arange(0, t)
|
||||
coords_h = torch.arange(0, h)
|
||||
coords_w = torch.arange(0, w)
|
||||
pos_ids = (
|
||||
coords_t[:, None, None] * max_latent_size * max_latent_size
|
||||
+ coords_h[None, :, None] * max_latent_size
|
||||
+ coords_w[None, None, :]
|
||||
).flatten()
|
||||
return pos_ids
|
||||
|
||||
|
||||
def get_flattened_position_ids_extrapolate(img_h, img_w, patch_size, max_num_patches_per_side):
|
||||
num_patches_h, num_patches_w = img_h // patch_size, img_w // patch_size
|
||||
coords_h = torch.arange(0, num_patches_h)
|
||||
coords_w = torch.arange(0, num_patches_w)
|
||||
pos_ids = (coords_h[:, None] * max_num_patches_per_side + coords_w).flatten()
|
||||
return pos_ids
|
||||
|
||||
|
||||
def get_flattened_position_ids_interpolate(img_h, img_w, patch_size, max_num_patches_per_side):
|
||||
num_patches_h, num_patches_w = img_h // patch_size, img_w // patch_size
|
||||
boundaries = torch.arange(1 / max_num_patches_per_side, 1.0, 1 / max_num_patches_per_side)
|
||||
fractional_coords_h = torch.arange(0, 1 - 1e-6, 1 / num_patches_h)
|
||||
fractional_coords_w = torch.arange(0, 1 - 1e-6, 1 / num_patches_w)
|
||||
bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries, right=True)
|
||||
bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries, right=True)
|
||||
pos_ids = (bucket_coords_h[:, None] * max_num_patches_per_side + bucket_coords_w).flatten()
|
||||
return pos_ids
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Patchify helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def patchify(image, patch_size):
|
||||
p = patch_size
|
||||
c, h, w = image.shape
|
||||
assert h % p == 0 and w % p == 0
|
||||
image = image.reshape(c, h // p, p, w // p, p)
|
||||
image = torch.einsum("chpwq->hwpqc", image)
|
||||
image = image.reshape(-1, p**2 * c)
|
||||
return image
|
||||
|
||||
|
||||
def patchify_video_with_merge(video, spatial_patch_size, temporal_patch_size, merge_size=2):
|
||||
"""
|
||||
Args:
|
||||
video: Tensor of shape [C, T, H, W]
|
||||
spatial_patch_size: patch size for H/W
|
||||
temporal_patch_size: patch size for T
|
||||
merge_size: merging factor for spatial grid (fixed at 2)
|
||||
|
||||
Returns:
|
||||
patches: Tensor of shape [num_patches, patch_dim]
|
||||
"""
|
||||
video = rearrange(video, "C T H W -> T C H W")
|
||||
T, C, H, W = video.shape
|
||||
p, tp, ms = spatial_patch_size, temporal_patch_size, merge_size
|
||||
|
||||
gt, gh, gw = T // tp, H // p, W // p
|
||||
video = video.reshape(gt, tp, C, gh // ms, ms, p, gw // ms, ms, p)
|
||||
video = video.permute(0, 3, 6, 4, 7, 2, 1, 5, 8)
|
||||
patches = video.reshape(gt * gh * gw, C * tp * p * p)
|
||||
return patches
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Sparse attention mask (flex-attention)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def create_sparse_mask(document_lens, split_lens, attn_modes, device):
|
||||
def causal_mask(b, h, q_idx, kv_idx):
|
||||
return q_idx >= kv_idx
|
||||
|
||||
def full_and_noise_mask(b, h, q_idx, kv_idx):
|
||||
return (full_and_noise_seq_id[q_idx] == full_and_noise_seq_id[kv_idx]) & (full_and_noise_seq_id[q_idx] >= 0)
|
||||
|
||||
def remove_noise_mask(b, h, q_idx, kv_idx):
|
||||
return ~((noise_seq_id[kv_idx] >= 0) & (noise_seq_id[q_idx] != noise_seq_id[kv_idx]))
|
||||
|
||||
def sample_mask(b, h, q_idx, kv_idx):
|
||||
return document_id[q_idx] == document_id[kv_idx]
|
||||
|
||||
full_and_noise_tmp = []
|
||||
noise_tmp = []
|
||||
|
||||
for i, (length, mode) in enumerate(zip(split_lens, attn_modes)):
|
||||
value = i if mode in ["full", "noise"] else -1
|
||||
full_and_noise_tmp.extend([value] * length)
|
||||
value_noise = i if mode == "noise" else -1
|
||||
noise_tmp.extend([value_noise] * length)
|
||||
|
||||
full_and_noise_seq_id = torch.Tensor(full_and_noise_tmp).to(device)
|
||||
noise_seq_id = torch.Tensor(noise_tmp).to(device)
|
||||
|
||||
document_id = torch.cat([torch.full((l,), i) for i, l in enumerate(document_lens, start=1)]).to(device)
|
||||
|
||||
return and_masks(or_masks(causal_mask, full_and_noise_mask), remove_noise_mask, sample_mask)
|
||||
|
||||
def prepare_attention_mask_per_sample(split_lens, attn_modes, device="cpu"):
|
||||
"""
|
||||
nested_split_lens: A list of N lists of ints. Each int indicates the length of a split within
|
||||
a sample, where each sample contains multiple splits with different attn modes.
|
||||
nested_attn_modes: whether to use full attn in each split.
|
||||
"""
|
||||
sample_len = sum(split_lens)
|
||||
attention_mask = torch.zeros((sample_len, sample_len), dtype=torch.bool, device=device)
|
||||
|
||||
csum = 0
|
||||
for s, attn_mode in zip(split_lens, attn_modes):
|
||||
assert attn_mode in ["causal", "full", "noise"]
|
||||
if attn_mode == "causal":
|
||||
attention_mask[csum : csum + s, csum : csum + s] = torch.ones((s, s), device=device).tril()
|
||||
attention_mask[csum : csum + s, :csum] = 1
|
||||
else:
|
||||
attention_mask[csum : csum + s, csum : csum + s] = torch.ones((s, s))
|
||||
attention_mask[csum : csum + s, :csum] = 1
|
||||
csum += s
|
||||
|
||||
csum = 0
|
||||
for s, attn_mode in zip(split_lens, attn_modes):
|
||||
if attn_mode == "noise":
|
||||
attention_mask[:, csum : csum + s] = torch.zeros((sample_len, s))
|
||||
attention_mask[csum : csum + s, csum : csum + s] = torch.ones((s, s))
|
||||
csum += s
|
||||
|
||||
attention_mask = torch.zeros_like(attention_mask, dtype=torch.float).masked_fill_(~attention_mask, float("-inf"))
|
||||
|
||||
return attention_mask
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Tokenizer / loss helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def add_special_tokens(tokenizer):
|
||||
all_special_tokens = []
|
||||
for k, v in tokenizer.special_tokens_map.items():
|
||||
if isinstance(v, str):
|
||||
all_special_tokens.append(v)
|
||||
elif isinstance(v, list):
|
||||
all_special_tokens += v
|
||||
|
||||
new_tokens = []
|
||||
for tok in ("<|im_start|>", "<|im_end|>", "<|vision_start|>", "<|vision_end|>"):
|
||||
if tok not in all_special_tokens:
|
||||
new_tokens.append(tok)
|
||||
|
||||
num_new_tokens = tokenizer.add_tokens(new_tokens)
|
||||
new_token_ids = dict(
|
||||
bos_token_id=tokenizer.convert_tokens_to_ids("<|im_start|>"),
|
||||
eos_token_id=tokenizer.convert_tokens_to_ids("<|im_end|>"),
|
||||
start_of_image=tokenizer.convert_tokens_to_ids("<|vision_start|>"),
|
||||
end_of_image=tokenizer.convert_tokens_to_ids("<|vision_end|>"),
|
||||
)
|
||||
return tokenizer, new_token_ids, num_new_tokens
|
||||
|
||||
|
||||
def len2weight(x, loss_reduction="square"):
|
||||
if x == 0:
|
||||
return x
|
||||
if loss_reduction == "token":
|
||||
return 1
|
||||
if loss_reduction == "sample":
|
||||
return 1 / x
|
||||
if loss_reduction == "square":
|
||||
return 1 / (x**0.5)
|
||||
raise NotImplementedError(loss_reduction)
|
||||
@@ -0,0 +1,80 @@
|
||||
# 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
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, Tuple
|
||||
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataConfig:
|
||||
"""
|
||||
DataConfig 版本，其中 vae_downsample 是一个三元组。
|
||||
"""
|
||||
grouped_datasets: Dict[str, Any] = field(default_factory=dict)
|
||||
text_cond_dropout_prob: float = 0.1
|
||||
vit_cond_dropout_prob: float = 0.4
|
||||
vae_cond_dropout_prob: float = 0.1
|
||||
|
||||
# vae_downsample stores temporal, height, and width downsampling factors.
|
||||
vae_downsample: Tuple[int, int, int] = (4, 16, 16)
|
||||
|
||||
max_latent_size: int = 64 # by ModelArguments
|
||||
vit_patch_size: int = 14 # by ModelArguments
|
||||
vit_patch_size_temporal: int = 2 # by ModelArguments
|
||||
vit_max_num_patch_per_side: int = 70 # by ModelArguments
|
||||
max_num_frames: int = 25 # by ModelArguments
|
||||
|
||||
latent_patch_size: int = None # by ModelArguments
|
||||
|
||||
@classmethod
|
||||
def from_yaml(cls, file_path: str) -> 'DataConfig':
|
||||
"""从 YAML/JSON 文件创建 DataConfig 实例"""
|
||||
with open(file_path, "r") as stream:
|
||||
data = yaml.safe_load(stream)
|
||||
return cls(grouped_datasets=data)
|
||||
|
||||
|
||||
class SimpleCustomBatch:
|
||||
def __init__(self, batch):
|
||||
data = batch[0]
|
||||
for key, value in data.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
def pin_memory(self):
|
||||
for key, value in self.__dict__.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
setattr(self, key, value.pin_memory())
|
||||
elif isinstance(value, list) and value and all(isinstance(i, torch.Tensor) for i in value):
|
||||
setattr(self, key, [i.pin_memory() for i in value])
|
||||
return self
|
||||
|
||||
def cuda(self, device):
|
||||
for key, value in self.__dict__.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
setattr(self, key, value.to(device))
|
||||
elif isinstance(value, list) and value and all(isinstance(i, torch.Tensor) for i in value):
|
||||
setattr(self, key, [i.to(device) for i in value])
|
||||
return self
|
||||
|
||||
def to_dict(self):
|
||||
return self.__dict__.copy()
|
||||
|
||||
|
||||
# Top-level function so it can be pickled.
|
||||
def simple_custom_collate(batch):
|
||||
return SimpleCustomBatch(batch)
|
||||
@@ -0,0 +1,998 @@
|
||||
# coding: utf-8
|
||||
|
||||
|
||||
import os
|
||||
EXP_HW_20250819 = os.environ.get("EXP_HW_20250819", "False").lower() == "true"
|
||||
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Tuple
|
||||
import numpy as np
|
||||
import torch
|
||||
import yaml
|
||||
from itertools import chain
|
||||
import sys
|
||||
|
||||
from .data_utils import (
|
||||
get_flattened_position_ids_interpolate_video,
|
||||
get_flattened_position_ids_extrapolate_video,
|
||||
len2weight,
|
||||
prepare_attention_mask_per_sample,
|
||||
patchify_video_with_merge,
|
||||
)
|
||||
from .dataset_info import DATASET_REGISTRY
|
||||
from data.video.sampler.utils import FRAME_SAMPLER_TYPES
|
||||
from data.transforms import ImageTransform, VideoTransform
|
||||
from data.video.video_utils import FrameSampler
|
||||
from data.parquet_utils import get_parquet_data_paths_balanced
|
||||
from common.utils.logging import get_logger
|
||||
from common.utils.basic import get_global_rank
|
||||
import bisect
|
||||
|
||||
sample_task_map = {
|
||||
't2v': 0,
|
||||
'idip': 1,
|
||||
'edit': 2,
|
||||
'refedit': 3,
|
||||
'maze':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
|
||||
}
|
||||
|
||||
@dataclass
|
||||
class DataConfig:
|
||||
"""
|
||||
DataConfig variant where vae_downsample is a 3-tuple.
|
||||
"""
|
||||
grouped_datasets: Dict[str, Any]
|
||||
text_cond_dropout_prob: float = 0.1
|
||||
vit_cond_dropout_prob: float = 0.4
|
||||
vae_cond_dropout_prob: float = 0.1
|
||||
|
||||
# Use a 3-tuple for vae_downsample: temporal, height, and width downsampling rates
|
||||
vae_downsample: Tuple[int, int, int] = (4, 16, 16)
|
||||
|
||||
max_latent_size: int = 64 # by ModelArguments
|
||||
vit_patch_size: int = 14 # by ModelArguments
|
||||
vit_patch_size_temporal: int = 2 # by ModelArguments
|
||||
vit_max_num_patch_per_side: int = 70 # by ModelArguments
|
||||
max_num_frames: int = 25 # by ModelArguments
|
||||
|
||||
latent_patch_size: int = None # by ModelArguments
|
||||
|
||||
@classmethod
|
||||
def from_yaml(cls, file_path: str) -> 'DataConfig':
|
||||
"""
|
||||
Create a DataConfig instance from a YAML file.
|
||||
"""
|
||||
with open(file_path, "r") as stream:
|
||||
data = yaml.safe_load(stream)
|
||||
return cls(grouped_datasets=data)
|
||||
|
||||
class PackedDataset(torch.utils.data.IterableDataset):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data_config: DataConfig,
|
||||
tokenizer,
|
||||
special_tokens,
|
||||
local_rank,
|
||||
world_size,
|
||||
num_workers,
|
||||
expected_num_tokens=32768, # Expected packed sequence length
|
||||
max_num_tokens_per_sample=16384, # Maximum length of a single sample
|
||||
max_num_tokens=36864, # Hard limit for packed sequence length
|
||||
prefer_buffer_before=16384, # Sample-length threshold for preferring the buffer
|
||||
max_buffer_size=50, # Maximum buffer capacity
|
||||
interpolate_pos=False,
|
||||
use_flex=False,
|
||||
data_status=None,
|
||||
apply_chat_template=False,
|
||||
image_token_id=151655,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
self.expected_num_tokens = expected_num_tokens
|
||||
self.max_num_tokens_per_sample = max_num_tokens_per_sample
|
||||
self.prefer_buffer_before = prefer_buffer_before
|
||||
self.max_num_tokens = max_num_tokens
|
||||
self.max_buffer_size = max_buffer_size
|
||||
self.tokenizer = tokenizer
|
||||
self.local_rank = local_rank
|
||||
self.world_size = world_size
|
||||
self.num_workers = num_workers
|
||||
self.use_flex = use_flex
|
||||
self.data_config: DataConfig = data_config
|
||||
self.max_num_latent_frames = self.data_config.max_num_frames // self.data_config.vae_downsample[0] + 1
|
||||
self.apply_chat_template = apply_chat_template
|
||||
self.cfg_type = kwargs.get("cfg_type", 0)
|
||||
self.cfg_uncond_token_id = kwargs.get("cfg_uncond_token_id", 151643)
|
||||
self.image_token_id = image_token_id
|
||||
self.data_args = kwargs
|
||||
self.expected_num_ce_loss_tokens = kwargs.get("expected_num_ce_loss_tokens", 1000000) # Default 1000000 effectively means unlimited
|
||||
self.N_key_frame = kwargs.get("N_key_frame", -1)
|
||||
self.fbyf_type = kwargs.get("fbyf_type", "group")
|
||||
self.fbyf_group_interval = kwargs.get("fbyf_group_interval", -1)
|
||||
self.incre_time_pro = kwargs.get("incre_time_pro", 0)
|
||||
self.require_und_gen = kwargs.get("require_und_gen", False)
|
||||
|
||||
# NOTE: Add special tokens such as <|im_start|> and <|im_end|> here
|
||||
for k, v in special_tokens.items():
|
||||
setattr(self, k, v)
|
||||
|
||||
# add a logger
|
||||
self.logger = get_logger()
|
||||
# self.log_rank0 = lambda msg: self.logger.info(msg) if get_global_rank() == 0 else None # Log only on rank 0
|
||||
|
||||
grouped_datasets, is_mandatory, grouped_weights, data_type = self.build_datasets(data_config.grouped_datasets, data_status)
|
||||
self.grouped_datasets = grouped_datasets
|
||||
# self.dataset_iters = [iter(dataset) for dataset in grouped_datasets] # this is dataset iter
|
||||
self.dataset_iters = None # Delay creation until __iter__
|
||||
self.is_mandatory = is_mandatory
|
||||
self.grouped_weights = grouped_weights
|
||||
self.interpolate_pos = interpolate_pos
|
||||
self.data_type = data_type
|
||||
|
||||
# Log only on rank 0; instance methods can be pickled
|
||||
def log_rank0(self, msg: str):
|
||||
if get_global_rank() == 0:
|
||||
self.logger.info(msg)
|
||||
|
||||
# Avoid sending non-serializable objects such as loggers or iterators to child processes
|
||||
def __getstate__(self):
|
||||
state = self.__dict__.copy()
|
||||
state["logger"] = None
|
||||
state["dataset_iters"] = None
|
||||
return state
|
||||
|
||||
def __setstate__(self, state):
|
||||
self.__dict__.update(state)
|
||||
self.logger = get_logger()
|
||||
self.dataset_iters = None
|
||||
|
||||
def set_epoch(self, seed):
|
||||
for dataset in self.grouped_datasets:
|
||||
dataset.set_epoch(seed)
|
||||
|
||||
# for dataset in self.grouped_datasets:
|
||||
# if hasattr(dataset, 'seed'):
|
||||
# dataset.set_epoch(dataset.seed)
|
||||
# else:
|
||||
# dataset.set_epoch(seed)
|
||||
|
||||
def set_sequence_status(self):
|
||||
sequence_status = dict(
|
||||
curr = 0, # Pointer
|
||||
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 arguments
|
||||
vae_video_latent = [], # for vae video latent offline
|
||||
vae_data_mode = [], # offline or online
|
||||
vit_data_mode = [], # offline or online
|
||||
# sample_task for joint training
|
||||
sample_task = [],
|
||||
sample_modality = [],
|
||||
)
|
||||
return sequence_status
|
||||
|
||||
def to_tensor(self, sequence_status: Dict[str, Any]):
|
||||
data = dict(
|
||||
sequence_length=sum(sequence_status['sample_lens']),
|
||||
sample_lens=sequence_status['sample_lens'],
|
||||
sample_type=sequence_status['sample_type'],
|
||||
sample_N_target=sequence_status['sample_N_target'],
|
||||
packed_text_ids=torch.tensor(sequence_status['packed_text_ids']),
|
||||
packed_text_indexes=torch.tensor(sequence_status['packed_text_indexes']),
|
||||
packed_position_ids=torch.tensor(sequence_status['packed_position_ids']),
|
||||
vit_data_mode=sequence_status['vit_data_mode'],
|
||||
video_grid_thw=sequence_status['video_grid_thw'],
|
||||
sample_task=torch.tensor(sequence_status['sample_task']),
|
||||
sample_modality=torch.tensor(sequence_status['sample_modality']),
|
||||
)
|
||||
|
||||
data['vae_data_mode'] = sequence_status['vae_data_mode']
|
||||
if not self.use_flex:
|
||||
data['nested_attention_masks'] = sequence_status['nested_attention_masks']
|
||||
else:
|
||||
sequence_len = data['sequence_length']
|
||||
pad_len = self.max_num_tokens - sequence_len
|
||||
assert pad_len >= 0, f"pad_len must be greater than 0, but got {pad_len}" # !!!
|
||||
data['split_lens'] = sequence_status['split_lens'] + [pad_len]
|
||||
data['attn_modes'] = sequence_status['attn_modes'] + ['causal']
|
||||
data['sample_lens'] += [pad_len]
|
||||
data['sample_type'] += ['pad']
|
||||
data['sample_N_target'] += [0]
|
||||
|
||||
# if the model has a convnet vae (e.g., as visual tokenizer)
|
||||
|
||||
data['padded_videos'] = sequence_status.pop('vae_video_tensors')
|
||||
if len(data['padded_videos']) > 0:
|
||||
# Pack as dynamic resolution
|
||||
# NOTE: The following keys are shared between image and video for now
|
||||
if 'patchified_vae_latent_shapes' not in data:
|
||||
data['patchified_vae_latent_shapes'] = sequence_status['vae_latent_shapes']
|
||||
data['packed_latent_position_ids'] = torch.cat(sequence_status['packed_latent_position_ids'], dim=0)
|
||||
data['packed_vae_token_indexes'] = torch.tensor(sequence_status['packed_vae_token_indexes'])
|
||||
|
||||
# process for offline data: padding
|
||||
if len(sequence_status["vae_video_latent"]) > 0:
|
||||
video_latents = sequence_status.pop("vae_video_latent")
|
||||
video_sizes = [item.shape for item in video_latents]
|
||||
max_video_size = [max(item) for item in list(zip(*video_sizes))]
|
||||
padded_videos_latent = torch.zeros(size=(len(video_latents), *max_video_size))
|
||||
for i, video_latent in enumerate(video_latents): # [T, H, W, C]
|
||||
t, h, w, c = video_latent.shape
|
||||
padded_videos_latent[i, :t, :h, :w, :c] = video_latent
|
||||
data["padded_latent"] = padded_videos_latent
|
||||
# NOTE: The following keys are shared between image and video for now
|
||||
if "patchified_vae_latent_shapes" not in data:
|
||||
data["patchified_vae_latent_shapes"] = sequence_status["vae_latent_shapes"]
|
||||
data["packed_latent_position_ids"] = torch.cat(sequence_status["packed_latent_position_ids"], dim=0)
|
||||
data["packed_vae_token_indexes"] = torch.tensor(sequence_status["packed_vae_token_indexes"])
|
||||
|
||||
# if the model has a vit (e.g., as visual tokenizer)
|
||||
if len(sequence_status['packed_vit_tokens']) > 0:
|
||||
data['packed_vit_tokens'] = sequence_status.pop('packed_vit_tokens')
|
||||
# data['packed_vit_tokens'] = torch.cat(sequence_status['packed_vit_tokens'], dim=0)
|
||||
data['packed_vit_position_ids'] = torch.cat(sequence_status['packed_vit_position_ids'], dim=0)
|
||||
data['packed_vit_token_indexes'] = torch.tensor(sequence_status['packed_vit_token_indexes'])
|
||||
data['vit_token_seqlens'] = torch.tensor(sequence_status['vit_token_seqlens'])
|
||||
|
||||
# Pack as dynamic resolution
|
||||
data['padded_videos_vit'] = sequence_status.pop('vit_video_tensors')
|
||||
|
||||
# if the model is required to perform visual generation
|
||||
if len(sequence_status['packed_timesteps']) > 0:
|
||||
data['packed_timesteps'] = torch.tensor(sequence_status['packed_timesteps'])
|
||||
data['mse_loss_indexes'] = torch.tensor(sequence_status['mse_loss_indexes'])
|
||||
|
||||
# if the model is required to perform text generation
|
||||
if len(sequence_status['packed_label_ids']) > 0:
|
||||
data['packed_label_ids'] = torch.tensor(sequence_status['packed_label_ids'])
|
||||
data['ce_loss_indexes'] = torch.tensor(sequence_status['ce_loss_indexes'])
|
||||
data['ce_loss_weights'] = torch.tensor(sequence_status['ce_loss_weights'])
|
||||
|
||||
if len(sequence_status['vae_video_grid_thw']) > 0:
|
||||
data['vae_video_grid_thw'] = torch.tensor(sequence_status['vae_video_grid_thw'])
|
||||
if len(sequence_status['vit_video_grid_thw']) > 0:
|
||||
data['vit_video_grid_thw'] = torch.tensor(sequence_status['vit_video_grid_thw'])
|
||||
|
||||
# Memory optimization: release sequence_status contents that are no longer needed
|
||||
sequence_status.clear()
|
||||
return data
|
||||
|
||||
def build_datasets(self, datasets_metainfo, data_status):
|
||||
datasets = []
|
||||
is_mandatory = []
|
||||
grouped_weights = []
|
||||
data_type = []
|
||||
for grouped_dataset_name, dataset_args in datasets_metainfo.items():
|
||||
if grouped_dataset_name.startswith('D'): # Handle the new multi-level nested logic
|
||||
grouped_dataset_name, dataset_args = list(dataset_args.items())[0]
|
||||
if '2t' in grouped_dataset_name:
|
||||
data_type.append('x2t')
|
||||
else:
|
||||
data_type.append('x2v')
|
||||
is_mandatory.append(dataset_args.pop('is_mandatory', False))
|
||||
grouped_weights.append(dataset_args.pop('weight', 0.0))
|
||||
|
||||
if 'frame_sampler_args' in dataset_args.keys():
|
||||
# NOTE: NOT for video
|
||||
frame_sampler = FrameSampler(**dataset_args.pop('frame_sampler_args'))
|
||||
dataset_args['frame_sampler'] = frame_sampler
|
||||
if 'image_transform_args' in dataset_args.keys(): # TODO: deprecate this
|
||||
transform = ImageTransform(**dataset_args.pop('image_transform_args'))
|
||||
dataset_args['transform'] = transform
|
||||
if 'video_transform_args' in dataset_args.keys():
|
||||
# video
|
||||
transform = VideoTransform(**dataset_args.pop('video_transform_args'))
|
||||
dataset_args['transform'] = transform
|
||||
dataset_args['vae_downsample'] = self.data_config.vae_downsample # fix: pass this in; TODO: consider vae_downsample and vit_downsample, low priority
|
||||
|
||||
# Add the video frame sampler here
|
||||
if 'video_frame_sampler_args' in dataset_args:
|
||||
dataset_args['res_dump'] = dataset_args['video_frame_sampler_args']['res_dump'] if 'res_dump' in dataset_args['video_frame_sampler_args'].keys() else ''
|
||||
|
||||
video_frame_sampler_args = dataset_args.pop('video_frame_sampler_args')
|
||||
|
||||
video_frame_sampler = FRAME_SAMPLER_TYPES[video_frame_sampler_args.get("type", "fixed")](**video_frame_sampler_args.get("params", {}))
|
||||
|
||||
dataset_args['video_frame_sampler'] = video_frame_sampler
|
||||
|
||||
if 'vit_video_transform_args' in dataset_args.keys():
|
||||
# video
|
||||
vit_transform = VideoTransform(**dataset_args.pop('vit_video_transform_args'))
|
||||
dataset_args['vit_transform'] = vit_transform
|
||||
|
||||
elif 'vit_image_transform_args' in dataset_args.keys(): # TODO: deprecate this
|
||||
vit_transform = ImageTransform(**dataset_args.pop('vit_image_transform_args'))
|
||||
dataset_args['vit_transform'] = vit_transform
|
||||
|
||||
assert 'dataset_names' in dataset_args, dataset_args.keys() or "missing 'dataset_names'"
|
||||
dataset_names = dataset_args.pop('dataset_names') # NOTE: Pay attention to this pop pattern
|
||||
dataset_args['data_dir_list'] = []
|
||||
|
||||
# Iterate and build datasets
|
||||
for item, meta_info in dataset_names.items():
|
||||
if self.local_rank == 0:
|
||||
self.logger.info(f'Preparing Dataset {grouped_dataset_name}/{item}')
|
||||
|
||||
data_dir = meta_info['data_dir']
|
||||
if isinstance(data_dir, str):
|
||||
# If it is a path
|
||||
dataset_args['data_dir_list'].append(meta_info['data_dir'])
|
||||
elif isinstance(data_dir, list):
|
||||
# If it is a list
|
||||
dataset_args['data_dir_list'].extend(data_dir)
|
||||
else:
|
||||
raise Exception(f'Unknown data_dir type {type(data_dir)}')
|
||||
|
||||
# NOTE: Collect all paths at the outer level, then pass them in
|
||||
all_data_paths = get_parquet_data_paths_balanced(
|
||||
data_dir_list=dataset_args.get('data_dir_list'),
|
||||
rank=self.local_rank,
|
||||
world_size=self.world_size,
|
||||
num_repeat=dataset_args.get('num_repeat', 1),
|
||||
)
|
||||
|
||||
if 'all_data_paths' in dataset_args.keys():
|
||||
dataset_args['all_data_paths'].extend(all_data_paths)
|
||||
else:
|
||||
dataset_args['all_data_paths'] = all_data_paths
|
||||
|
||||
resume_data_status = dataset_args.pop('resume_data_status', True)
|
||||
if data_status is not None and grouped_dataset_name in data_status.keys() and resume_data_status:
|
||||
data_status_per_group = data_status[grouped_dataset_name]
|
||||
else:
|
||||
data_status_per_group = None
|
||||
|
||||
dataset_args.update(self.data_args)
|
||||
|
||||
# Update
|
||||
dataset_args['vit_cond_dropout_prob'] = self.data_config.vit_cond_dropout_prob
|
||||
dataset_args['text_cond_dropout_prob'] = self.data_config.text_cond_dropout_prob
|
||||
dataset_args['vae_cond_dropout_prob'] = self.data_config.vae_cond_dropout_prob
|
||||
|
||||
dataset = DATASET_REGISTRY[grouped_dataset_name](
|
||||
dataset_name=grouped_dataset_name,
|
||||
tokenizer=self.tokenizer,
|
||||
local_rank=self.local_rank,
|
||||
world_size=self.world_size,
|
||||
num_workers=self.num_workers,
|
||||
data_status=data_status_per_group,
|
||||
apply_chat_template=self.apply_chat_template,
|
||||
**dataset_args,
|
||||
)
|
||||
datasets.append(dataset)
|
||||
|
||||
return datasets, is_mandatory, grouped_weights, data_type
|
||||
|
||||
# Add the video processing branch in pack_sequence
|
||||
def pack_sequence(self, sample: Dict[str, Any], sequence_status: Dict[str, Any]):
|
||||
image_tensor_list = sample.get('image_tensor_list', []) # just for debug
|
||||
video_tensor_list = sample.get('video_tensor_list', [])
|
||||
video_latent_list = sample.get('video_latent_list', [])
|
||||
sample_N_target = sample.get('N_target', 1)
|
||||
text_ids_list = sample['text_ids_list']
|
||||
sequence_plan = sample['sequence_plan']
|
||||
sample_task = sample.get('sample_task', 't2v')
|
||||
|
||||
split_lens, attn_modes = list(), list()
|
||||
curr = sequence_status['curr']
|
||||
curr_rope_id = 0
|
||||
sample_lens = 0
|
||||
sample_type = ''
|
||||
curr_split_idx = sequence_status['curr']
|
||||
apply_text_template = False
|
||||
curr_video_grid_thw = []
|
||||
|
||||
for item in sequence_plan:
|
||||
split_start = item.get('split_start', True)
|
||||
if split_start:
|
||||
curr_split_len = 0
|
||||
|
||||
# TODO: add more item types to help classification
|
||||
if item['type'] == 'text':
|
||||
sample_type = 'und' # This is overwritten, so only the last item takes effect
|
||||
text_ids = text_ids_list.pop(0)
|
||||
if item['enable_cfg'] == 1 and random.random() < self.data_config.text_cond_dropout_prob:
|
||||
|
||||
if self.cfg_type == 0: # 0 fully removes the text condition
|
||||
continue
|
||||
elif self.cfg_type == 1: # 1 keeps only special tokens
|
||||
text_ids = []
|
||||
elif self.cfg_type == 2: # 2 keeps special tokens and replaces middle text tokens with uncond_token
|
||||
text_ids = [self.cfg_uncond_token_id] * len(text_ids)
|
||||
|
||||
if not item.get('special_token_start_nouse'): # When special_token_start_nouse is None or False
|
||||
shifted_text_ids = [self.bos_token_id] + text_ids # NOTE: self.bos_token_id=151644 <|im_start|>
|
||||
else:
|
||||
shifted_text_ids = text_ids
|
||||
|
||||
sequence_status['packed_text_ids'].extend(shifted_text_ids)
|
||||
sequence_status['packed_text_indexes'].extend(range(curr, curr + len(shifted_text_ids)))
|
||||
|
||||
# NOTE: item['loss'] == 1 identifies understanding vs generation
|
||||
if item['loss'] == 1:
|
||||
loss_token_shift = item.get('loss_token_shift') or 0
|
||||
sequence_status['ce_loss_indexes'].extend(range(curr - loss_token_shift, curr + len(shifted_text_ids)))
|
||||
sequence_status['ce_loss_weights'].extend([len2weight(len(shifted_text_ids) + loss_token_shift)] * (len(shifted_text_ids) + loss_token_shift))
|
||||
sequence_status['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
|
||||
if not item.get('special_token_end_nouse'):
|
||||
sequence_status['packed_text_ids'].append(self.eos_token_id)
|
||||
sequence_status['packed_text_indexes'].append(curr)
|
||||
if item['special_token_loss'] == 1: # <|im_end|> may have loss
|
||||
sequence_status['ce_loss_indexes'].append(curr)
|
||||
sequence_status['ce_loss_weights'].append(1.0)
|
||||
sequence_status['packed_label_ids'].append(item['special_token_label'])
|
||||
curr += 1
|
||||
curr_split_len += 1
|
||||
|
||||
# update sequence status
|
||||
attn_modes.append("causal")
|
||||
#if self.apply_chat_template:
|
||||
sequence_status['packed_position_ids'].extend(range(curr_rope_id, curr_rope_id + curr_split_len))
|
||||
curr_rope_id += curr_split_len
|
||||
|
||||
sequence_status['sample_modality'].extend([modality_map[item['type']]] * curr_split_len)
|
||||
|
||||
elif item["type"] == "text_template":
|
||||
apply_text_template = True
|
||||
text_ids = text_ids_list.pop(0)
|
||||
# The current template only applies to UND, so ignore cfg for now
|
||||
|
||||
sequence_status["packed_text_ids"].extend(text_ids)
|
||||
sample_lens = len(text_ids)
|
||||
spans_index = item.get("spans_index", None) # Vision padding tokens as a list; each item is an index list for the corresponding video/image pad tokens
|
||||
curr_sample_modality = []
|
||||
caption_index = item.get("cap_index", []) # Indexes for text excluding the system prompt
|
||||
|
||||
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] # Indexes of the first and last '<|video_pad|>'
|
||||
sequence_status["packed_text_indexes"].extend(range(curr, vision_start))
|
||||
if (vision_start - 1) - curr != 0: # Confirm there is a text split before vision; HACK: changed from the LLaVA version
|
||||
curr_split_len = (vision_start - 1) - curr
|
||||
sequence_status["packed_position_ids"].extend(
|
||||
range(curr_rope_id, curr_rope_id + curr_split_len)
|
||||
) # Note: use vision_start - 1, not vision_start, because vision_start is the starting token of the video split
|
||||
curr_rope_id += curr_split_len
|
||||
curr_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: interleaved text is not supported; text must be contiguous
|
||||
split_len_1 = caption_index[0] - curr # Length of the system prompt before text
|
||||
split_len_2 = len(caption_index) # Text length
|
||||
split_len_3 = curr_split_len - split_len_1 - split_len_2 # Length of the system prompt after text
|
||||
|
||||
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]
|
||||
attn_modes.extend(["causal"] * len(split_len_text))
|
||||
split_lens.extend(split_len_text)
|
||||
else:
|
||||
attn_modes.append("causal")
|
||||
split_lens.append(curr_split_len)
|
||||
|
||||
|
||||
curr_split_len = len(span_index) + 2
|
||||
if sequence_plan[video_id]["type"] == 'vit_video':
|
||||
sequence_status["packed_vit_token_indexes"].extend(range(vision_start, vision_end + 1))
|
||||
attn_modes.append("full") # TODO: add this check if the GEN branch also uses templates
|
||||
curr_sample_modality.extend([modality_map['ref_vit']] * curr_split_len)
|
||||
elif sequence_plan[video_id]["type"] == 'vae_video':
|
||||
sequence_status["packed_vae_token_indexes"].extend(range(vision_start, vision_end + 1))
|
||||
if sequence_plan[video_id]["loss"] == 0:
|
||||
attn_modes.append("full_noise") # TODO: add this check if the GEN branch also uses templates
|
||||
if sample_task == 'edit':
|
||||
curr_sample_modality.extend([modality_map['ref_source']] * curr_split_len)
|
||||
elif sample_task == 'idip' or sample_task == 'maze':
|
||||
curr_sample_modality.extend([modality_map['ref_image']] * curr_split_len)
|
||||
else:
|
||||
if frame_condition_idx == []:
|
||||
sequence_status["mse_loss_indexes"].extend(range(vision_start, vision_end + 1))
|
||||
else: # Support f2v; multiple target videos are not currently supported because different videos need separate THW values
|
||||
frame_condition_indexes = []
|
||||
mse_loss_indexes = list(range(vision_start, vision_end + 1))
|
||||
for idx in frame_condition_idx:
|
||||
if idx == -1:
|
||||
idx = t - 1
|
||||
if idx == 1:
|
||||
continue # Skip when there are only two frames to avoid using identical condition frames for all frames
|
||||
frame_condition_indexes.extend(mse_loss_indexes[idx * h * w : (idx + 1) * h * w])
|
||||
mse_loss_indexes = sorted(list(set(mse_loss_indexes) - set(frame_condition_indexes)))
|
||||
sequence_status["mse_loss_indexes"].extend(mse_loss_indexes)
|
||||
|
||||
|
||||
attn_modes.append("noise") # TODO: add this check if the GEN branch also uses templates
|
||||
sample_type = "gen" # This is overwritten, so only the last item takes effect
|
||||
curr_sample_modality.extend([modality_map['noise']] * curr_split_len)
|
||||
|
||||
sequence_status["packed_position_ids"].extend([curr_rope_id] * curr_split_len)
|
||||
#attn_modes.append("full") # TODO: add this check if the GEN branch also uses templates
|
||||
split_lens.append(len(span_index) + 2)
|
||||
curr = vision_end + 1 # Index of the '<|vision_end|>' token
|
||||
curr_rope_id += 1
|
||||
sequence_status["packed_text_indexes"].append(curr)
|
||||
curr += 1 # Starting token of the next sequence
|
||||
|
||||
len_split_last = sample_lens - (curr - curr_split_idx) if spans_index != [] else len(text_ids)
|
||||
if len_split_last != 0: # A trailing text segment remains
|
||||
split_lens.append(len_split_last)
|
||||
sequence_status["packed_text_indexes"].extend(range(curr, curr + len_split_last))
|
||||
sequence_status["packed_position_ids"].extend(range(curr_rope_id, curr_rope_id + len_split_last))
|
||||
attn_modes.append("causal")
|
||||
curr_sample_modality.extend([modality_map['system_prompt']] * len_split_last)
|
||||
|
||||
if item["loss"] == 1: # This marks an understanding task and requires CE loss
|
||||
packed_label_index = item.get("packed_label_index", [])[:]
|
||||
sequence_status["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()
|
||||
sequence_status["ce_loss_indexes"].extend(ce_loss_indexes)
|
||||
sequence_status["ce_loss_weights"].extend([len2weight(len(packed_label_index))] * (len(packed_label_index)))
|
||||
sample_type = "und" # This is overwritten, so only the last item takes effect
|
||||
# else: # This marks a generation task
|
||||
# sample_type = "gen" # This is overwritten, so only the last item takes effect
|
||||
#packed_label_index = item.get("packed_label_index", [])[1:-2]
|
||||
|
||||
|
||||
# Get caption indexes in text and update their sample_modality
|
||||
if caption_index != []:
|
||||
curr_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
|
||||
sequence_status['sample_modality'].extend(curr_sample_modality)
|
||||
|
||||
elif item['type'] == 'vae_video':
|
||||
sample_type = 'gen'
|
||||
video_tensor = video_tensor_list.pop(0) # CTHW
|
||||
if item['enable_cfg'] == 1 and random.random() < self.data_config.vae_cond_dropout_prob:
|
||||
# FIXME: fix VAE dropout in video2video. TODO: confirm whether enable_cfg is needed in the vae_video branch
|
||||
# curr_rope_id += 1
|
||||
continue
|
||||
|
||||
apply_text_template = item.get('apply_text_template', False)
|
||||
|
||||
|
||||
num_special_tokens = 0
|
||||
# Add <|startofimage|> token shared by video and image. TODO: split image and video special tokens?
|
||||
if not apply_text_template:
|
||||
sequence_status['packed_text_ids'].append(self.start_of_image) # self.start_of_image=151652, <|vision_start|>
|
||||
sequence_status['packed_text_indexes'].append(curr)
|
||||
curr += 1
|
||||
curr_split_len += 1
|
||||
num_special_tokens += 1
|
||||
|
||||
# In online mode, video_tensor is a tensor; in offline mode, it is a list [latent]
|
||||
if isinstance(video_tensor, torch.Tensor): # online
|
||||
# Preprocess video
|
||||
|
||||
sequence_status["vae_video_tensors"].append(video_tensor) # Raw CTHW video, not latent
|
||||
|
||||
# Assume video_tensor has shape (C, T, H, W)
|
||||
_, T, H, W = video_tensor.shape
|
||||
_T, _H, _W = self.data_config.vae_downsample # NOTE: Absolute-scale downsample including patchification
|
||||
t = (T - 1) // _T + 1 # k*N+1; normally the T dimension is not patchified. Update this if T patchification is needed
|
||||
h = H // _H
|
||||
w = W // _W
|
||||
sequence_status["vae_data_mode"].append("online")
|
||||
else: # offline
|
||||
# video_latent = video_tensor[0] # [[T, H, W, C]]
|
||||
sequence_status["vae_video_tensors"].append(video_tensor[0]) # video_tensor[0] is a tensor with shape [T, H, W, C]
|
||||
# Assume video_latent has shape [T, H, W, C]
|
||||
T, H, W, _ = video_tensor[0].shape
|
||||
_T, _H, _W = self.data_config.latent_patch_size # NOTE: Only includes patchification
|
||||
|
||||
t = T // _T # k*N+1; normally the T dimension is not patchified. Update this if T patchification is needed
|
||||
h = H // _H
|
||||
w = W // _W
|
||||
|
||||
sequence_status["vae_data_mode"].append("offline")
|
||||
|
||||
spatial_merge_size = 2 # TODO: must spatial_merge_size always be 2?
|
||||
vae_video_grid_thw = [
|
||||
t,
|
||||
h * spatial_merge_size,
|
||||
w * spatial_merge_size,
|
||||
] # Qwen-VL RoPE divides by spatial_merge_size by default to match VIT processing, so VAE needs an extra multiply by spatial_merge_size
|
||||
|
||||
if EXP_HW_20250819: # HACK: temporary experiment
|
||||
vae_video_grid_thw = [1, 2, 2]
|
||||
|
||||
sequence_status["vae_video_grid_thw"].append(vae_video_grid_thw)
|
||||
curr_video_grid_thw.append(vae_video_grid_thw)
|
||||
|
||||
# Use the native (t, h, w) latent shape
|
||||
sequence_status['vae_latent_shapes'].append((t, h, w))
|
||||
|
||||
# Use the 3D-aware position encoding function
|
||||
if self.interpolate_pos:
|
||||
# Interpolation
|
||||
packed_latent_position_ids = get_flattened_position_ids_interpolate_video(
|
||||
t, h, w, 1, # Latent-space patch size is 1
|
||||
max_num_frames=self.max_num_latent_frames,
|
||||
max_num_patches_per_side=self.data_config.max_latent_size
|
||||
)
|
||||
else:
|
||||
# Extrapolation
|
||||
packed_latent_position_ids = get_flattened_position_ids_extrapolate_video(
|
||||
t, h, w,
|
||||
max_latent_size=self.data_config.max_latent_size
|
||||
)
|
||||
|
||||
sequence_status['packed_latent_position_ids'].append(packed_latent_position_ids)
|
||||
|
||||
num_vid_tokens = t * h * w
|
||||
|
||||
if not apply_text_template:
|
||||
sequence_status['packed_vae_token_indexes'].extend(range(curr, curr + num_vid_tokens))
|
||||
|
||||
if item["loss"] == 1:
|
||||
pro_fbyf = False
|
||||
if split_start:
|
||||
timestep = np.random.randn() # NOTE: A sigmoid is applied outside
|
||||
|
||||
frame_condition_idx = item.get("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 # Skip when there are only two frames to avoid using identical condition frames for all frames
|
||||
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 apply_text_template:
|
||||
sequence_status["mse_loss_indexes"].extend(mse_loss_indexes) # range(curr, curr + num_vid_tokens))
|
||||
else:
|
||||
if self.incre_time_pro <= 0:
|
||||
timestep = float("-inf")
|
||||
else:
|
||||
timestep = 0.
|
||||
packed_timesteps = [timestep] * num_vid_tokens
|
||||
|
||||
sequence_status['packed_timesteps'].extend(packed_timesteps)
|
||||
|
||||
|
||||
if not apply_text_template:
|
||||
curr += num_vid_tokens
|
||||
curr_split_len += num_vid_tokens
|
||||
sequence_status["packed_text_ids"].extend([self.image_token_id] * num_vid_tokens)
|
||||
|
||||
# Add <|endofimage|> token
|
||||
sequence_status['packed_text_ids'].append(self.end_of_image) # self.end_of_image=151653, <|vision_end|>
|
||||
sequence_status['packed_text_indexes'].append(curr)
|
||||
# <|endofimage|> may have loss
|
||||
if item['special_token_loss'] == 1:
|
||||
sequence_status['ce_loss_indexes'].append(curr)
|
||||
sequence_status['ce_loss_weights'].append(1.0)
|
||||
sequence_status['packed_label_ids'].append(item['special_token_label'])
|
||||
curr += 1
|
||||
curr_split_len += 1
|
||||
num_special_tokens += 1
|
||||
|
||||
# Update sequence status
|
||||
if split_start:
|
||||
if item['loss'] == 1 and 'frame_delta' not in item.keys(): # Is frame_delta for multi-frame generation?
|
||||
attn_modes.append("noise")
|
||||
else:
|
||||
# attn_modes.append("full")
|
||||
attn_modes.append("full_noise")
|
||||
|
||||
sequence_status['packed_position_ids'].extend([curr_rope_id] * (num_vid_tokens + num_special_tokens)) # NOTE: why is RoPE fixed?
|
||||
if 'frame_delta' in item.keys():
|
||||
curr_rope_id += item['frame_delta']
|
||||
elif item['loss'] == 0:
|
||||
curr_rope_id += 1
|
||||
|
||||
# update sample sequence modality
|
||||
if item['loss'] == 1:
|
||||
sequence_status['sample_modality'].extend([modality_map['noise']] * curr_split_len)
|
||||
elif item['loss'] == 0 and sample_task == 'edit':
|
||||
sequence_status['sample_modality'].extend([modality_map['ref_source']] * curr_split_len)
|
||||
elif item['loss'] == 0 and (sample_task == 'idip' or sample_task == 'maze'):
|
||||
sequence_status['sample_modality'].extend([modality_map['ref_image']] * curr_split_len)
|
||||
|
||||
del video_tensor, packed_timesteps, packed_latent_position_ids
|
||||
|
||||
elif item['type'] == 'vit_video':
|
||||
apply_text_template = item.get('apply_text_template', False)
|
||||
video_tensor = video_tensor_list.pop(0) # CTHW
|
||||
if item['enable_cfg'] == 1 and random.random() < self.data_config.vit_cond_dropout_prob and not apply_text_template:
|
||||
# FIXME: fix VAE dropout in video2video. TODO: confirm whether enable_cfg is needed in the vit_video branch
|
||||
curr_rope_id += 1 # HACK ？？？？
|
||||
continue
|
||||
|
||||
# Add <|startofimage|> token shared by video and image. TODO: split image and video special tokens?
|
||||
if not apply_text_template:
|
||||
sequence_status['packed_text_ids'].append(self.start_of_image) # 151652, <|vision_start|>
|
||||
sequence_status['packed_text_indexes'].append(curr)
|
||||
curr += 1
|
||||
curr_split_len += 1
|
||||
|
||||
# In online mode, video_tensor is a tensor; in offline mode, it is a list [latent]
|
||||
if isinstance(video_tensor, torch.Tensor): # online
|
||||
sequence_status['vit_video_tensors'].append(video_tensor) # Raw CTHW video, not latent; used only for validation visualization
|
||||
|
||||
# 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 also needs merging: 2x2 merges into 1, hardcoded temporarily
|
||||
t, h, w = video_tensor.size(1), video_tensor.size(2), video_tensor.size(3)
|
||||
|
||||
sequence_status['packed_vit_tokens'].append(vit_tokens)
|
||||
sequence_status['vit_data_mode'].append('online')
|
||||
|
||||
del vit_tokens
|
||||
|
||||
else: # offline
|
||||
# video_latent, video_tensor = video_tensor # [L,D]
|
||||
# sequence_status['vit_video_tensors'].append(video_tensor[0])
|
||||
num_video_tokens = video_tensor[0].shape[0]
|
||||
|
||||
sequence_status['packed_vit_tokens'].append(video_tensor[0])
|
||||
sequence_status['vit_data_mode'].append('offline')
|
||||
thw = item.get('thw',None)
|
||||
if thw is not None:
|
||||
t, h, w = thw
|
||||
else:
|
||||
t = None
|
||||
|
||||
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]
|
||||
sequence_status["vit_video_grid_thw"].append(vit_video_grid_thw)
|
||||
curr_video_grid_thw.append(vit_video_grid_thw)
|
||||
|
||||
sequence_status['vit_token_seqlens'].append(num_video_tokens)
|
||||
sequence_status['packed_vit_position_ids'].append(torch.zeros(num_video_tokens)) # TODO: this may not always be 0; multiple VIT sequences may be problematic, but this is currently only used in vit model:siglip
|
||||
|
||||
if not apply_text_template:
|
||||
sequence_status['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
|
||||
sequence_status['packed_text_ids'].extend([self.image_token_id]*num_video_tokens)
|
||||
|
||||
# add a <|endofimage|> token
|
||||
sequence_status['packed_text_ids'].append(self.end_of_image) # 151653, <|vision_end|>
|
||||
sequence_status['packed_text_indexes'].append(curr)
|
||||
if item['special_token_loss'] == 1: # <|endofimage|> may have loss
|
||||
sequence_status['ce_loss_indexes'].append(curr)
|
||||
sequence_status['ce_loss_weights'].append(1.0)
|
||||
sequence_status['packed_label_ids'].append(item['special_token_label'])
|
||||
curr += 1
|
||||
curr_split_len += 1
|
||||
sequence_status['packed_position_ids'].extend([curr_rope_id] * curr_split_len)
|
||||
curr_rope_id += 1
|
||||
|
||||
# update sequence status
|
||||
attn_modes.append("full")
|
||||
|
||||
sequence_status['sample_modality'].extend([modality_map['ref_vit']] * curr_split_len)
|
||||
|
||||
del video_tensor
|
||||
|
||||
if item.get('split_end', True) and not apply_text_template:
|
||||
if isinstance(curr_split_len, list):
|
||||
split_lens.extend(curr_split_len)
|
||||
sample_lens += sum(curr_split_len)
|
||||
else:
|
||||
split_lens.append(curr_split_len)
|
||||
sample_lens += curr_split_len
|
||||
|
||||
sequence_status['curr'] = curr
|
||||
sequence_status['sample_lens'].append(sample_lens) # sample_lens is the length of a pair, e.g. <text, video>
|
||||
sequence_status['sample_type'].append(sample_type)
|
||||
sequence_status['sample_N_target'].append(sample_N_target)
|
||||
sequence_status['video_grid_thw'].append(torch.tensor(curr_video_grid_thw)) # video_grid_thw is a sequence split by sample
|
||||
# prepare attention mask
|
||||
if not self.use_flex:
|
||||
sequence_status['nested_attention_masks'].append(
|
||||
prepare_attention_mask_per_sample(split_lens, attn_modes)
|
||||
)
|
||||
else:
|
||||
sequence_status['split_lens'].extend(split_lens)
|
||||
sequence_status['attn_modes'].extend(attn_modes)
|
||||
# add sample_task
|
||||
sequence_status['sample_task'].extend([sample_task_map[sample_task]] * sample_lens)
|
||||
return sequence_status
|
||||
|
||||
def sample_pro(self, sample, sequence_status, batch_data_indexes, buffer, sample_from_buffer=False, skip_count=0):
|
||||
|
||||
num_tokens = sample['num_tokens'] + 2 * len(sample['sequence_plan']) # NOTE: 2*2 special tokens
|
||||
is_pro = False
|
||||
if num_tokens < self.max_num_tokens_per_sample and sequence_status['curr'] + num_tokens < self.max_num_tokens :
|
||||
sequence_status = self.pack_sequence(sample, sequence_status)
|
||||
batch_data_indexes.append(sample['data_indexes'])
|
||||
is_pro = True
|
||||
elif sequence_status['curr'] + num_tokens > self.max_num_tokens:
|
||||
if len(buffer) < self.max_buffer_size and not sample_from_buffer:
|
||||
buffer.append(sample)
|
||||
else:
|
||||
self.logger.info(f"skip a sample with length {num_tokens}")
|
||||
skip_count += 1
|
||||
else:
|
||||
self.logger.info(f"skip a sample with length {num_tokens}")
|
||||
skip_count += 1
|
||||
return sequence_status, batch_data_indexes, buffer, is_pro, skip_count
|
||||
|
||||
|
||||
def __iter__(self):
|
||||
# NOTE: Core logic repeatedly calls pack_sequence to pack multimodal data
|
||||
# Create underlying dataset iterators at each iteration; spawn/fork safe
|
||||
self.dataset_iters = [iter(dataset) for dataset in self.grouped_datasets]
|
||||
|
||||
total_weights = sum(self.grouped_weights)
|
||||
assert total_weights > 0.0
|
||||
group_cumprobs = [sum(self.grouped_weights[: i + 1]) / total_weights for i in range(len(self.grouped_weights))]
|
||||
sequence_status = self.set_sequence_status()
|
||||
batch_data_indexes = []
|
||||
|
||||
buffer = []
|
||||
data_type_pro = ['x2t', 'x2v']
|
||||
skip_count = 0
|
||||
max_skips_before_reset = 30 # Threshold for resetting sequence_status
|
||||
while True:
|
||||
# Ensure at least one sample from each group
|
||||
if sequence_status['curr'] == 0:
|
||||
for group_index, group_iter in enumerate(self.dataset_iters):
|
||||
if self.is_mandatory[group_index]:
|
||||
while True:
|
||||
sample = next(group_iter)
|
||||
# if a sample is too long, skip it
|
||||
sequence_status, batch_data_indexes, buffer, is_pro, skip_count = self.sample_pro(sample, sequence_status, batch_data_indexes, buffer, skip_count=skip_count)
|
||||
if self.data_type[group_index] in data_type_pro:
|
||||
data_type_pro.remove(self.data_type[group_index])
|
||||
if is_pro:
|
||||
break
|
||||
if self.require_und_gen and 'x2t' in self.data_type and 'x2v' in self.data_type: # NOTE: In joint UND + GEN training, sequences with only one sample may cause communication failures
|
||||
while True:
|
||||
if data_type_pro == []:
|
||||
break
|
||||
n = random.random()
|
||||
group_index = bisect.bisect_left(group_cumprobs, n)
|
||||
|
||||
if self.data_type[group_index] in data_type_pro:
|
||||
sample = next(self.dataset_iters[group_index])
|
||||
sequence_status, batch_data_indexes, buffer, is_pro, skip_count = self.sample_pro(sample, sequence_status, batch_data_indexes, buffer, skip_count=skip_count)
|
||||
if is_pro:
|
||||
data_type_pro.remove(self.data_type[group_index])
|
||||
|
||||
if skip_count >= max_skips_before_reset: # NOTE: If 30 samples are skipped consecutively, break and reset sequence_status
|
||||
break
|
||||
|
||||
if skip_count >= max_skips_before_reset: # Reset sequence_status
|
||||
self.logger.info(f"Too many skips ({skip_count}), resetting sequence_status")
|
||||
sequence_status = self.set_sequence_status()
|
||||
batch_data_indexes = []
|
||||
data_type_pro = ['x2t', 'x2v']
|
||||
skip_count = 0
|
||||
continue
|
||||
|
||||
|
||||
if sequence_status['curr'] < self.prefer_buffer_before and len(buffer) > 0:
|
||||
sample = buffer.pop(0)
|
||||
sample_from_buffer = True
|
||||
else:
|
||||
# sample normally across all groups
|
||||
n = random.random()
|
||||
group_index = bisect.bisect_left(group_cumprobs, n)
|
||||
sample = next(self.dataset_iters[group_index])
|
||||
sample_from_buffer = False
|
||||
|
||||
# if a sample is too long, skip it
|
||||
num_tokens = sample['num_tokens'] + 2 * len(sample['sequence_plan'])
|
||||
if num_tokens > self.max_num_tokens_per_sample:
|
||||
self.logger.info(f"skip a sample with length {num_tokens}")
|
||||
continue
|
||||
|
||||
|
||||
if sequence_status['curr'] + num_tokens > self.max_num_tokens or len(sequence_status['ce_loss_indexes']) > self.expected_num_ce_loss_tokens:
|
||||
if len(buffer) < self.max_buffer_size and not sample_from_buffer:
|
||||
buffer.append(sample)
|
||||
# elif self.require_und_gen:
|
||||
# sample_type = 1
|
||||
|
||||
else:
|
||||
# self.logger.info(f"Yielding data with length {sum(sequence_status['sample_lens'])}")
|
||||
data = self.to_tensor(sequence_status)
|
||||
data['batch_data_indexes'] = batch_data_indexes
|
||||
data_type_pro = ['x2t', 'x2v']
|
||||
yield data
|
||||
# Reset sequence_status after yield
|
||||
sequence_status = self.set_sequence_status()
|
||||
batch_data_indexes = []
|
||||
continue
|
||||
|
||||
sequence_status = self.pack_sequence(sample, sequence_status)
|
||||
batch_data_indexes.append(sample['data_indexes'])
|
||||
|
||||
if sequence_status['curr'] >= self.expected_num_tokens or len(sequence_status['ce_loss_indexes']) >= self.expected_num_ce_loss_tokens:
|
||||
data = self.to_tensor(sequence_status)
|
||||
data['batch_data_indexes'] = batch_data_indexes
|
||||
data_type_pro = ['x2t', 'x2v']
|
||||
yield data
|
||||
# Reset sequence_status after yield
|
||||
sequence_status = self.set_sequence_status()
|
||||
batch_data_indexes = []
|
||||
|
||||
|
||||
class SimpleCustomBatch:
|
||||
def __init__(self, batch):
|
||||
data = batch[0]
|
||||
for key, value in data.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
def pin_memory(self):
|
||||
for key, value in self.__dict__.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
setattr(self, key, value.pin_memory())
|
||||
elif isinstance(value, list) and value and all(isinstance(i, torch.Tensor) for i in value):
|
||||
setattr(self, key, [i.pin_memory() for i in value])
|
||||
|
||||
return self
|
||||
|
||||
def cuda(self, device):
|
||||
for key, value in self.__dict__.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
setattr(self, key, value.to(device))
|
||||
elif isinstance(value, list) and value and all(isinstance(i, torch.Tensor) for i in value):
|
||||
setattr(self, key, [i.to(device) for i in value])
|
||||
|
||||
return self
|
||||
|
||||
def to_dict(self):
|
||||
return self.__dict__.copy()
|
||||
|
||||
|
||||
def collate_wrapper():
|
||||
def collate_fn(batch):
|
||||
return SimpleCustomBatch(batch)
|
||||
return collate_fn
|
||||
|
||||
# Top-level function; pickleable
|
||||
def simple_custom_collate(batch):
|
||||
return SimpleCustomBatch(batch)
|
||||
@@ -0,0 +1,14 @@
|
||||
# coding: utf-8
|
||||
|
||||
from data.datasets_factory.dataset_x2v_interleave import X2VInterleaveIterableDataset
|
||||
from data.datasets_factory.dataset_x2v_interleave_local import X2VInterleaveLocalIterableDataset
|
||||
from data.datasets_factory.dataset_x2t_interleave import X2TInterleaveIterableDataset
|
||||
from data.datasets_factory.dataset_x2t_interleave_local import X2TInterleaveLocalIterableDataset
|
||||
|
||||
DATASET_REGISTRY = {
|
||||
# GEN
|
||||
'x2v_interleave': X2VInterleaveIterableDataset,
|
||||
'x2v_interleave_local': X2VInterleaveLocalIterableDataset,
|
||||
'x2t_interleave': X2TInterleaveIterableDataset,
|
||||
'x2t_interleave_local': X2TInterleaveLocalIterableDataset,
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
# 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
|
||||
|
||||
"""自定义数据集入口。"""
|
||||
|
||||
from .validation_dataset import ValidationDataset
|
||||
|
||||
__all__ = ["ValidationDataset"]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,228 @@
|
||||
# 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, detect_lang_simple
|
||||
from data.system_prompt_render import render_qwenvl_prompt, expand_and_index_by_token_ids_new
|
||||
|
||||
class X2TInterleaveIterableDataset(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 = "text" # Target element type
|
||||
self.task_type = kwargs.get("task_type", "v2t")
|
||||
self.task_type_rate = kwargs.get("task_type_rate", 1)
|
||||
self.text_first = False
|
||||
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"X2TInterleaveIterableDataset 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:
|
||||
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
|
||||
|
||||
condition_modalities = [element_dtype_array[i] for i in condition_idx]
|
||||
|
||||
# Select task type
|
||||
if "text" not in condition_modalities:
|
||||
task_type_sample = "v2t"
|
||||
elif "image" not in condition_modalities and "video" not in condition_modalities:
|
||||
task_type_sample = "t2t"
|
||||
elif isinstance(self.task_type, list):
|
||||
task_type_sample = sample_task(self.task_type, self.task_type_rate)
|
||||
else:
|
||||
task_type_sample = self.task_type
|
||||
|
||||
num_tokens, sequence_plan, text_ids_list, video_tensor_list, target_captions = [], [], [], [], [] # sequence_plan stores loss, cfg, and other marker metadata
|
||||
num_split_vit, num_split_vae, num_split_text = 0, 0, 0
|
||||
text_template_user, text_template_assistant, vit_num_tokens = [], [], []
|
||||
for idx in range(target_idx[-1] + 1):
|
||||
if idx in condition_idx: # Process condition element
|
||||
if task_type_sample in ["v2t", "tv2t", "vt2t"] and element_dtype_array[idx] in ["image","video"]: # Visual condition element
|
||||
if num_split_vit >= self.max_num_split_vit:
|
||||
continue
|
||||
media_url = interleave_array[idx]
|
||||
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)
|
||||
|
||||
if self.text_template:
|
||||
text_template_user.append({"type": element_dtype_array[idx]})
|
||||
vit_num_tokens.append(num_tokens_)
|
||||
else:
|
||||
num_tokens.append(num_tokens_) # NOTE: Count vision tokens
|
||||
|
||||
sequence_plan.append(
|
||||
{
|
||||
"type": "vit_video",
|
||||
"enable_cfg": 0,
|
||||
"loss": 0,
|
||||
"special_token_loss": 0, # NOTE: Special tokens are provided manually and are not predicted
|
||||
"special_token_label": None, # eos
|
||||
"apply_text_template": self.text_template, # Default is false
|
||||
"thw": thw,
|
||||
}
|
||||
)
|
||||
num_split_vit += 1
|
||||
elif task_type_sample in ["t2t", "tv2t", "vt2t"] and element_dtype_array[idx] == "text": # Process text condition element
|
||||
if num_split_text >= (self.max_num_split_text - N_target):
|
||||
continue
|
||||
caption = self.text_cleaner(interleave_array[idx])
|
||||
if not caption:
|
||||
continue
|
||||
|
||||
if self.text_template:
|
||||
text_template_user.append({"type": "text", "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": 0,
|
||||
"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 idx in target_idx: # Process target element
|
||||
if num_split_text >= self.max_num_split_text:
|
||||
continue
|
||||
caption = interleave_array[idx]
|
||||
target_captions.append(caption)
|
||||
num_split_text += 1
|
||||
|
||||
if len(target_captions) == 0:
|
||||
return None
|
||||
|
||||
target_caption = target_captions[-1]
|
||||
if (
|
||||
isinstance(target_caption, str)
|
||||
and detect_lang_simple(target_caption) != "en"
|
||||
) or (
|
||||
not isinstance(target_caption, str)
|
||||
and detect_lang_simple(target_caption[1]) != "en"
|
||||
and target_caption[1] != ""
|
||||
):
|
||||
self.logger.warning(f"Wrong caption: {target_caption}")
|
||||
return None
|
||||
|
||||
if self.text_template:
|
||||
if isinstance(target_caption, str): # Only one text item
|
||||
if "video" in condition_modalities:
|
||||
vision_type = 'video'
|
||||
else:
|
||||
vision_type = 'image'
|
||||
caption_a = target_caption
|
||||
caption_i = generate_system_prompt(system_prompt_type='caption', vision_type=vision_type)
|
||||
caption_q = ""
|
||||
else:
|
||||
caption_i, caption_q, caption_a = target_caption[0], target_caption[1], target_caption[2]
|
||||
|
||||
if self.data_config.get('system_prompt_type') == 'SP2':
|
||||
caption_q = caption_i + " " + caption_q
|
||||
caption_i = "You are a helpful assistant. "
|
||||
elif self.data_config.get('system_prompt_type') == 'SP1':
|
||||
# SP1: assistant
|
||||
caption_i = "You are a helpful assistant. " + caption_i
|
||||
|
||||
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, # Original usage
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": text_template_assistant,
|
||||
},
|
||||
]
|
||||
caption_all = render_qwenvl_prompt(messages, default_system=caption_i, include_assistant_content=True) # NOTE: Whether to add '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=""
|
||||
)
|
||||
|
||||
assert len(all_token_id[tgt_index[0] :]) == len(tgt_index)
|
||||
|
||||
num_tokens.append(len(all_token_id))
|
||||
|
||||
text_ids_list.append(all_token_id)
|
||||
sequence_plan.append(
|
||||
{
|
||||
"type": "text_template",
|
||||
"enable_cfg": 0,
|
||||
"loss": 1,
|
||||
"special_token_loss": 0,
|
||||
"special_token_label": None,
|
||||
"packed_label_index": tgt_index,
|
||||
"spans_index": spans_index,
|
||||
}
|
||||
)
|
||||
else:
|
||||
caption = " ".join(self.text_cleaner(item) for item in target_captions if self.text_cleaner(item))
|
||||
caption = self.text_cleaner(caption)
|
||||
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": 0,
|
||||
"loss": 1,
|
||||
"special_token_loss": 0,
|
||||
"special_token_label": None,
|
||||
}
|
||||
)
|
||||
|
||||
sample = dict(
|
||||
video_tensor_list=video_tensor_list,
|
||||
text_ids_list=text_ids_list,
|
||||
num_tokens=sum(num_tokens),
|
||||
sequence_plan=sequence_plan,
|
||||
N_target=1, # In theory, N_target is always 1 for the UND branch
|
||||
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
|
||||
@@ -0,0 +1,90 @@
|
||||
# coding: utf-8
|
||||
|
||||
import os
|
||||
import pyarrow.parquet as pq
|
||||
import torch
|
||||
|
||||
from data.parquet_utils import init_arrow_fs, read_parquet_rows
|
||||
from data.datasets_factory.dataset_x2t_interleave import X2TInterleaveIterableDataset
|
||||
|
||||
|
||||
class X2TInterleaveLocalIterableDataset(X2TInterleaveIterableDataset):
|
||||
def __init__(self, *args, **kwargs):
|
||||
# Debug/local smoke-test dataset: read local parquet rows with raw image/video bytes.
|
||||
kwargs.setdefault("raw_bytes_input", True)
|
||||
super().__init__(*args, **kwargs)
|
||||
# Debug-only option for quick smoke tests on tiny local datasets:
|
||||
# repeat the same parquet file list to avoid very short epochs / empty shards.
|
||||
self.debug_parquet_repeat = max(1, int(kwargs.get("debug_parquet_repeat", 1)))
|
||||
self.data_paths = (kwargs.get("all_data_paths") or self.data_dir_list) * self.debug_parquet_repeat
|
||||
self._validate_local_data_paths()
|
||||
|
||||
def _validate_local_data_paths(self):
|
||||
for path in self.data_paths:
|
||||
if str(path).startswith(("hdfs://", "http://", "https://")):
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__} only supports local parquet files with embedded bytes, got: {path}"
|
||||
)
|
||||
if not os.path.isfile(path):
|
||||
raise FileNotFoundError(f"Local parquet file not found: {path}")
|
||||
|
||||
def _validate_local_row(self, row):
|
||||
if self.dataset_type == "x2t_general":
|
||||
if "image_bytes" not in row and "video_bytes" not in row:
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__} expects x2t_general rows to contain image_bytes or video_bytes"
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__} only supports dataset_type=x2t_general for local byte reading, got {self.dataset_type}"
|
||||
)
|
||||
|
||||
def set_epoch(self, seed=42):
|
||||
if not self.data_paths:
|
||||
return
|
||||
|
||||
data_paths = sorted(self.data_paths)
|
||||
self.rng.seed(seed)
|
||||
self.rng.shuffle(data_paths)
|
||||
self.data_paths_per_rank = data_paths
|
||||
self.num_files_per_rank = len(data_paths)
|
||||
|
||||
def __iter__(self):
|
||||
if not hasattr(self, "data_paths_per_rank"):
|
||||
self.set_epoch(self.seed)
|
||||
|
||||
info = torch.utils.data.get_worker_info()
|
||||
worker_id = 0 if info is None else info.id
|
||||
num_workers = 1 if info is None else info.num_workers
|
||||
global_worker_id = self.local_rank * num_workers + worker_id
|
||||
global_num_workers = max(1, self.world_size * num_workers)
|
||||
|
||||
while True:
|
||||
for parquet_idx, parquet_file_path in enumerate(self.data_paths_per_rank):
|
||||
fs = init_arrow_fs(parquet_file_path)
|
||||
with fs.open_input_file(parquet_file_path) as f:
|
||||
parquet_file = pq.ParquetFile(f)
|
||||
for row_group_id in range(parquet_file.num_row_groups):
|
||||
rows = read_parquet_rows(parquet_file, row_group_id)
|
||||
for row_idx, row in enumerate(rows):
|
||||
if row_idx % global_num_workers != global_worker_id:
|
||||
continue
|
||||
try:
|
||||
self._validate_local_row(row)
|
||||
except Exception as e:
|
||||
self.logger.warning(
|
||||
f"Invalid local x2t row in rg#{row_group_id} row#{row_idx} {parquet_file_path}: {e}"
|
||||
)
|
||||
continue
|
||||
sample = self._process_row(
|
||||
row=row,
|
||||
parquet_idx=parquet_idx,
|
||||
row_group_id=row_group_id,
|
||||
row_idx=row_idx,
|
||||
worker_id=worker_id,
|
||||
parquet_file_path=parquet_file_path,
|
||||
)
|
||||
if sample:
|
||||
yield sample
|
||||
|
||||
self.logger.info(f"{self.dataset_name} repeat in rank-{self.local_rank} worker-{worker_id}")
|
||||
@@ -0,0 +1,245 @@
|
||||
# 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
|
||||
@@ -0,0 +1,100 @@
|
||||
# coding: utf-8
|
||||
|
||||
import os
|
||||
import pyarrow.parquet as pq
|
||||
import torch
|
||||
|
||||
from data.parquet_utils import init_arrow_fs, read_parquet_rows
|
||||
from data.datasets_factory.dataset_x2v_interleave import X2VInterleaveIterableDataset
|
||||
|
||||
|
||||
class X2VInterleaveLocalIterableDataset(X2VInterleaveIterableDataset):
|
||||
def __init__(self, *args, **kwargs):
|
||||
# Debug/local smoke-test dataset: read local parquet rows with raw image/video bytes.
|
||||
kwargs.setdefault("raw_bytes_input", True)
|
||||
super().__init__(*args, **kwargs)
|
||||
# Debug-only option for quick smoke tests on tiny local datasets:
|
||||
# repeat the same parquet file list to avoid very short epochs / empty shards.
|
||||
self.debug_parquet_repeat = max(1, int(kwargs.get("debug_parquet_repeat", 1)))
|
||||
self.data_paths = (kwargs.get("all_data_paths") or self.data_dir_list) * self.debug_parquet_repeat
|
||||
self._validate_local_data_paths()
|
||||
|
||||
def _validate_local_data_paths(self):
|
||||
for path in self.data_paths:
|
||||
if str(path).startswith(("hdfs://", "http://", "https://")):
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__} only supports local parquet files with embedded bytes, got: {path}"
|
||||
)
|
||||
if not os.path.isfile(path):
|
||||
raise FileNotFoundError(f"Local parquet file not found: {path}")
|
||||
|
||||
def _validate_local_row(self, row):
|
||||
dataset_type_to_required_keys = {
|
||||
"text2image_general": {"caption", "image_bytes"},
|
||||
"text2video_general": {"caption", "video_bytes"},
|
||||
"image2image": {"caption", "input_image_bytes", "output_image_bytes"},
|
||||
"video2video": {"caption", "input_video_bytes", "output_video_bytes"},
|
||||
}
|
||||
required_keys = dataset_type_to_required_keys.get(self.dataset_type)
|
||||
if required_keys is None:
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__} only supports local byte dataset types "
|
||||
f"{sorted(dataset_type_to_required_keys)}, got {self.dataset_type}"
|
||||
)
|
||||
|
||||
missing = sorted(key for key in required_keys if key not in row)
|
||||
if missing:
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__} expects dataset_type={self.dataset_type} rows to contain keys {sorted(required_keys)}, "
|
||||
f"missing {missing}"
|
||||
)
|
||||
|
||||
def set_epoch(self, seed=42):
|
||||
if not self.data_paths:
|
||||
return
|
||||
|
||||
data_paths = sorted(self.data_paths)
|
||||
self.rng.seed(seed)
|
||||
self.rng.shuffle(data_paths)
|
||||
self.data_paths_per_rank = data_paths
|
||||
self.num_files_per_rank = len(data_paths)
|
||||
|
||||
def __iter__(self):
|
||||
if not hasattr(self, "data_paths_per_rank"):
|
||||
self.set_epoch(self.seed)
|
||||
|
||||
info = torch.utils.data.get_worker_info()
|
||||
worker_id = 0 if info is None else info.id
|
||||
num_workers = 1 if info is None else info.num_workers
|
||||
global_worker_id = self.local_rank * num_workers + worker_id
|
||||
global_num_workers = max(1, self.world_size * num_workers)
|
||||
|
||||
while True:
|
||||
for parquet_idx, parquet_file_path in enumerate(self.data_paths_per_rank):
|
||||
fs = init_arrow_fs(parquet_file_path)
|
||||
with fs.open_input_file(parquet_file_path) as f:
|
||||
parquet_file = pq.ParquetFile(f)
|
||||
for row_group_id in range(parquet_file.num_row_groups):
|
||||
rows = read_parquet_rows(parquet_file, row_group_id)
|
||||
for row_idx, row in enumerate(rows):
|
||||
if row_idx % global_num_workers != global_worker_id:
|
||||
continue
|
||||
try:
|
||||
self._validate_local_row(row)
|
||||
except Exception as e:
|
||||
self.logger.warning(
|
||||
f"Invalid local x2v row in rg#{row_group_id} row#{row_idx} {parquet_file_path}: {e}"
|
||||
)
|
||||
continue
|
||||
sample = self._process_row(
|
||||
row=row,
|
||||
parquet_idx=parquet_idx,
|
||||
row_group_id=row_group_id,
|
||||
row_idx=row_idx,
|
||||
worker_id=worker_id,
|
||||
parquet_file_path=parquet_file_path,
|
||||
)
|
||||
if sample:
|
||||
yield sample
|
||||
|
||||
self.logger.info(f"{self.dataset_name} repeat in rank-{self.local_rank} worker-{worker_id}")
|
||||
@@ -0,0 +1,69 @@
|
||||
# coding: utf-8
|
||||
|
||||
import os
|
||||
import random
|
||||
import torch
|
||||
from common.utils.logging import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
class DistributedIterableDataset(torch.utils.data.IterableDataset):
|
||||
def __init__(self, dataset_name, local_rank=0, world_size=1, num_workers=8):
|
||||
self.dataset_name = dataset_name
|
||||
self.local_rank = local_rank
|
||||
self.world_size = world_size
|
||||
self.num_workers = num_workers
|
||||
self.rng = random.Random()
|
||||
self.data_paths = None
|
||||
|
||||
def get_data_paths(self, *args, **kwargs):
|
||||
raise NotImplementedError
|
||||
|
||||
def set_epoch(self, seed=42):
|
||||
# Process granularity: one shard per rank
|
||||
if self.data_paths is None:
|
||||
return
|
||||
|
||||
if isinstance(self.data_paths[0], tuple):
|
||||
data_paths = sorted(self.data_paths, key=lambda x: (x[0], x[1]))
|
||||
elif isinstance(self.data_paths[0], str):
|
||||
data_paths = sorted(self.data_paths)
|
||||
else:
|
||||
raise ValueError(f"Unknown data_paths type: {type(self.data_paths[0])}")
|
||||
|
||||
self.rng.seed(seed)
|
||||
self.rng.shuffle(data_paths)
|
||||
|
||||
num_files_per_rank = len(data_paths) // self.world_size
|
||||
local_start = self.local_rank * num_files_per_rank
|
||||
local_end = (self.local_rank + 1) * num_files_per_rank
|
||||
self.num_files_per_rank = num_files_per_rank
|
||||
self.data_paths_per_rank = data_paths[local_start:local_end]
|
||||
|
||||
# ================== Add this log line ==================
|
||||
if self.data_paths_per_rank and self.local_rank == 0: # Ensure the list is non-empty and log only on rank 0
|
||||
logger.info(f"[Rank-Split-Check] Rank {self.local_rank} got {len(self.data_paths_per_rank)} files. "
|
||||
f"First file: {os.path.basename(self.data_paths_per_rank[0])}")
|
||||
# =======================================================
|
||||
|
||||
def get_data_paths_per_worker(self):
|
||||
# Worker granularity: one shard per worker process
|
||||
if self.data_paths is None:
|
||||
return None
|
||||
|
||||
info = torch.utils.data.get_worker_info()
|
||||
if info is None:
|
||||
# Single worker: Use all files assigned to the rank
|
||||
return self.data_paths_per_rank, 0
|
||||
|
||||
worker_id = info.id
|
||||
num_files_per_worker = self.num_files_per_rank // info.num_workers
|
||||
start = num_files_per_worker * worker_id
|
||||
end = num_files_per_worker * (worker_id + 1)
|
||||
data_paths_per_worker = self.data_paths_per_rank[start:end]
|
||||
|
||||
# NOTE: The reverse order ::-1 is probably unnecessary
|
||||
return data_paths_per_worker, worker_id
|
||||
|
||||
def __iter__(self):
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,91 @@
|
||||
# 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
|
||||
|
||||
"""
|
||||
Generic parquet/data file helpers that do not depend on cluster-specific settings.
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
from typing import Iterable, List, Optional
|
||||
|
||||
import pyarrow.fs as pf
|
||||
import torch.distributed as dist
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
SUPPORTED_DATA_EXTENSIONS = (".parquet", ".json", ".jsonl")
|
||||
|
||||
|
||||
def select_supported_files(file_paths: Iterable[str]) -> List[str]:
|
||||
file_paths = list(file_paths)
|
||||
parquet_files = [path for path in file_paths if path.endswith(".parquet")]
|
||||
if parquet_files:
|
||||
return parquet_files
|
||||
return [path for path in file_paths if path.endswith((".json", ".jsonl"))]
|
||||
|
||||
|
||||
def list_local_supported_data_files(data_dir: str) -> List[str]:
|
||||
if os.path.isfile(data_dir):
|
||||
if not data_dir.endswith(SUPPORTED_DATA_EXTENSIONS):
|
||||
raise ValueError(f"Unsupported data file: {data_dir}")
|
||||
return [data_dir]
|
||||
|
||||
if not os.path.isdir(data_dir):
|
||||
raise FileNotFoundError(f"Data path does not exist: {data_dir}")
|
||||
|
||||
files = [
|
||||
os.path.join(data_dir, name)
|
||||
for name in os.listdir(data_dir)
|
||||
if name.endswith(SUPPORTED_DATA_EXTENSIONS)
|
||||
]
|
||||
return select_supported_files(files)
|
||||
|
||||
|
||||
def init_arrow_fs(parquet_file_path: str) -> pf.FileSystem:
|
||||
if parquet_file_path.startswith("hdfs://"):
|
||||
raise ValueError(f"Only local parquet files are supported, got remote path: {parquet_file_path}")
|
||||
return pf.LocalFileSystem()
|
||||
|
||||
|
||||
def init_arrow_pf_fs(parquet_file_path: str) -> pf.FileSystem:
|
||||
return init_arrow_fs(parquet_file_path)
|
||||
|
||||
|
||||
def get_parquet_data_paths_balanced(data_dir_list, rank=0, world_size=1, num_repeat=1):
|
||||
all_data_paths = []
|
||||
if rank == 0:
|
||||
logger.info("Rank 0 is gathering local data file paths...")
|
||||
for data_dir in data_dir_list:
|
||||
all_data_paths.extend(sorted(list_local_supported_data_files(data_dir)))
|
||||
logger.info(f"Total local data files found: {len(all_data_paths)}")
|
||||
|
||||
if num_repeat > 1:
|
||||
original_len = len(all_data_paths)
|
||||
all_data_paths = all_data_paths * num_repeat
|
||||
logger.info(f"Repeating dataset {num_repeat} times, from {original_len} to {len(all_data_paths)} files.")
|
||||
|
||||
if world_size > 1 and dist.is_available() and dist.is_initialized():
|
||||
object_list = [all_data_paths]
|
||||
dist.broadcast_object_list(object_list, src=0)
|
||||
return object_list[0]
|
||||
|
||||
return all_data_paths
|
||||
|
||||
|
||||
def read_parquet_rows(parquet_file, row_group_id: int, columns: Optional[List[str]] = None):
|
||||
return parquet_file.read_row_group(row_group_id, columns=columns).to_pylist()
|
||||
@@ -0,0 +1,193 @@
|
||||
# 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.
|
||||
|
||||
from __future__ import annotations
|
||||
from typing import Any, Dict, List, Tuple
|
||||
from jinja2 import Environment, BaseLoader
|
||||
|
||||
|
||||
JINJA_PROMPT_TMPL = (
|
||||
"<|im_start|>system\n"
|
||||
"{{ system_prompt }}<|im_end|>\n"
|
||||
"{% for m in msgs -%}"
|
||||
"<|im_start|>{{ m.role }}\n"
|
||||
"{% if not (m.role == 'assistant' and not include_assistant_content) -%}"
|
||||
"{{ m.content | render_mm_list }}"
|
||||
"{% endif -%}"
|
||||
"{% if (not (loop.last and m.role == 'assistant')) or include_assistant_content -%}"
|
||||
"<|im_end|>\n"
|
||||
"{% endif -%}"
|
||||
"{% endfor -%}"
|
||||
)
|
||||
|
||||
VS, VE = "<|vision_start|>", "<|vision_end|>"
|
||||
VP, IP = "<|video_pad|>", "<|image_pad|>"
|
||||
|
||||
def expand_and_index_by_token_ids_new(
|
||||
rendered_text: str,
|
||||
tokens: List[int],
|
||||
tokenizer,
|
||||
target_text: str = "",
|
||||
search_text: str = "",
|
||||
) -> Tuple[str, List[int], List[List[int]], List[int]]:
|
||||
"""
|
||||
Returns:
|
||||
new_rendered_text: expanded text
|
||||
all_token_id : token ids of new_rendered_text
|
||||
spans_index : indexes of each pad block in all_token_id, in occurrence order, e.g. [[100..199], [350..549], ...]
|
||||
tgt_index : indexes of target_text in all_token_id, or [] if not found
|
||||
"""
|
||||
vs_ids = tokenizer(VS, add_special_tokens=False)["input_ids"]
|
||||
ve_ids = tokenizer(VE, add_special_tokens=False)["input_ids"]
|
||||
vp_ids = tokenizer(VP, add_special_tokens=False)["input_ids"]
|
||||
ip_ids = tokenizer(IP, add_special_tokens=False)["input_ids"]
|
||||
|
||||
enc = tokenizer(rendered_text, add_special_tokens=False)
|
||||
base_ids = enc["input_ids"]
|
||||
|
||||
# ---------- 1) Scan VP/IP in occurrence order and expand each to K copies, recording pad block metadata ----------
|
||||
# find all VS positions and pair them with nearest VE after each VS
|
||||
|
||||
all_ids: List[int] = []
|
||||
spans_index: List[List[int]] = []
|
||||
|
||||
i = 0 # Scan pointer for base_ids.
|
||||
tk_ptr = 0 # Pointer for tokens(K).
|
||||
|
||||
while True:
|
||||
try:
|
||||
vs_positions_ = base_ids[i:].index(vs_ids[0]) + i
|
||||
except:
|
||||
all_ids.extend(base_ids[i:])
|
||||
break
|
||||
all_ids.extend(base_ids[i: vs_positions_])
|
||||
i = vs_positions_ + 3
|
||||
|
||||
# Expand the sequence and insert placeholder ids into pad_ids.
|
||||
pad_ids = base_ids[vs_positions_ + 1:vs_positions_ + 2]
|
||||
K = int(tokens[tk_ptr])
|
||||
start, end = len(all_ids) + 1, len(all_ids) + 1 + K
|
||||
all_ids.extend(vs_ids + pad_ids * K + ve_ids)
|
||||
tk_ptr += 1
|
||||
|
||||
# Collect indexes of each pad token block in all_token_id, in occurrence order, e.g. [[100..199], [350..549], ...].
|
||||
#start, end = vs_positions_ + 1, vs_positions_ + 1 + K
|
||||
spans_index.append(list(range(start, end)))
|
||||
|
||||
tgt_index: List[int] = []
|
||||
if target_text:
|
||||
tgt_ids_identify = tokenizer(target_text, add_special_tokens=False)["input_ids"]
|
||||
i = 0 # Scan pointer for base_ids.
|
||||
|
||||
while i < len(all_ids):
|
||||
tgt_positions_ = all_ids[i:].index(tgt_ids_identify[0]) + i
|
||||
if all_ids[tgt_positions_+len(tgt_ids_identify)-1] == tgt_ids_identify[-1]:
|
||||
tgt_index = list(range(tgt_positions_+len(tgt_ids_identify), len(all_ids)))
|
||||
break
|
||||
else:
|
||||
i = tgt_positions_ + 1
|
||||
|
||||
search_index: List[int] = []
|
||||
if search_text:
|
||||
search_ids_identify = tokenizer(search_text, add_special_tokens=False)["input_ids"]
|
||||
i = 0 # Scan pointer for base_ids.
|
||||
|
||||
while i < len(all_ids):
|
||||
search_positions_ = all_ids[i:].index(search_ids_identify[0]) + i
|
||||
if all_ids[search_positions_:search_positions_+len(search_ids_identify)] == search_ids_identify:
|
||||
search_index = list(range(search_positions_, search_positions_+len(search_ids_identify)))
|
||||
break
|
||||
else:
|
||||
i = search_positions_ + 1
|
||||
|
||||
|
||||
return all_ids, spans_index, tgt_index, search_index
|
||||
|
||||
def _extract_system_prompt(messages: List[Dict[str, Any]], default_system: str) -> str:
|
||||
for m in messages:
|
||||
if m.get("role") == "system":
|
||||
c = m.get("content", "")
|
||||
if isinstance(c, str):
|
||||
return c
|
||||
if isinstance(c, list):
|
||||
texts = [it.get("text", "") for it in c if isinstance(it, dict) and it.get("type") == "text"]
|
||||
if texts:
|
||||
return "".join(texts)
|
||||
return default_system
|
||||
|
||||
|
||||
def _normalize_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
norm: List[Dict[str, Any]] = []
|
||||
for m in messages:
|
||||
role = m.get("role")
|
||||
if role == "system":
|
||||
continue
|
||||
c = m.get("content", "")
|
||||
if isinstance(c, str):
|
||||
items = [{"type": "text", "text": c}]
|
||||
elif isinstance(c, list):
|
||||
items = c
|
||||
else:
|
||||
items = []
|
||||
norm.append({"role": role, "content": items})
|
||||
return norm
|
||||
|
||||
|
||||
def render_qwenvl_prompt(
|
||||
messages: List[Dict[str, Any]],
|
||||
default_system: str = "You are a helpful assistant.",
|
||||
include_assistant_content: bool = False, # Key option: whether to render assistant text.
|
||||
force_video_pad: bool = False,
|
||||
) -> str:
|
||||
system_prompt = _extract_system_prompt(messages, default_system)
|
||||
msgs = _normalize_messages(messages)
|
||||
|
||||
def _render_mm_list(items: Any) -> str:
|
||||
if isinstance(items, str):
|
||||
return items
|
||||
if not isinstance(items, list):
|
||||
return ""
|
||||
parts: List[str] = []
|
||||
for it in items:
|
||||
if not isinstance(it, dict):
|
||||
continue
|
||||
t = it.get("type")
|
||||
if t == "text":
|
||||
parts.append(it.get("text", ""))
|
||||
elif t == "image":
|
||||
if force_video_pad:
|
||||
parts.append("<|vision_start|><|image_pad|><|vision_end|>")
|
||||
else:
|
||||
parts.append("<|vision_start|><|video_pad|><|vision_end|>")
|
||||
elif t == "video":
|
||||
parts.append("<|vision_start|><|video_pad|><|vision_end|>")
|
||||
# Other modalities can be added here.
|
||||
return "".join(parts)
|
||||
|
||||
env = Environment(
|
||||
loader=BaseLoader(),
|
||||
autoescape=False,
|
||||
trim_blocks=True, # Remove newlines after block endings.
|
||||
lstrip_blocks=True, # Remove whitespace before block starts.
|
||||
newline_sequence="\n",
|
||||
keep_trailing_newline=False,
|
||||
)
|
||||
env.filters["render_mm_list"] = _render_mm_list
|
||||
template = env.from_string(JINJA_PROMPT_TMPL)
|
||||
|
||||
return template.render(
|
||||
system_prompt=system_prompt,
|
||||
msgs=msgs,
|
||||
include_assistant_content=include_assistant_content,
|
||||
)
|
||||
+32
@@ -0,0 +1,32 @@
|
||||
# coding: utf-8
|
||||
|
||||
class TosWrapper:
|
||||
"""Placeholder for removed remote object storage support.
|
||||
|
||||
The public training framework only supports local parquet files with media
|
||||
bytes embedded in each row. Remote TOS access has been intentionally
|
||||
removed from the default code path.
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _raise_removed():
|
||||
raise NotImplementedError(
|
||||
"Remote TOS access is not supported. Use local parquet files with embedded image/video bytes."
|
||||
)
|
||||
|
||||
def get_obj(self, *args, **kwargs):
|
||||
self._raise_removed()
|
||||
|
||||
def get_obj_by_url(self, *args, **kwargs):
|
||||
self._raise_removed()
|
||||
|
||||
def get_obj_meta_by_url(self, *args, **kwargs):
|
||||
self._raise_removed()
|
||||
|
||||
def put_obj(self, *args, **kwargs):
|
||||
self._raise_removed()
|
||||
|
||||
__all__ = ["TosWrapper"]
|
||||
@@ -0,0 +1,380 @@
|
||||
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import random
|
||||
from PIL import Image
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from torchvision import transforms
|
||||
from torchvision.transforms import functional as F
|
||||
from torchvision.transforms import InterpolationMode, Compose, Normalize
|
||||
|
||||
from .video.transforms.na_resize import NaResize
|
||||
from .video.transforms.divisible_crop import DivisibleCrop
|
||||
from .video.transforms.rearrange import Rearrange
|
||||
|
||||
|
||||
class MaxLongEdgeMinShortEdgeResize(torch.nn.Module):
|
||||
"""Resize the input image so that its longest side and shortest side are within a specified range,
|
||||
ensuring that both sides are divisible by a specified stride.
|
||||
|
||||
Args:
|
||||
max_size (int): Maximum size for the longest edge of the image.
|
||||
min_size (int): Minimum size for the shortest edge of the image.
|
||||
stride (int): Value by which the height and width of the image must be divisible.
|
||||
max_pixels (int): Maximum pixels for the full image.
|
||||
interpolation (InterpolationMode): Desired interpolation enum defined by
|
||||
:class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``.
|
||||
If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.NEAREST_EXACT``,
|
||||
``InterpolationMode.BILINEAR``, and ``InterpolationMode.BICUBIC`` are supported.
|
||||
The corresponding Pillow integer constants, e.g., ``PIL.Image.BILINEAR`` are also accepted.
|
||||
antialias (bool, optional): Whether to apply antialiasing (default is True).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_size: int,
|
||||
min_size: int,
|
||||
stride: int,
|
||||
max_pixels: int,
|
||||
interpolation=InterpolationMode.BICUBIC,
|
||||
antialias=True
|
||||
):
|
||||
super().__init__()
|
||||
self.max_size = max_size
|
||||
self.min_size = min_size
|
||||
self.stride = stride
|
||||
self.max_pixels = max_pixels
|
||||
self.interpolation = interpolation
|
||||
self.antialias = antialias
|
||||
|
||||
def _make_divisible(self, value, stride):
|
||||
"""Ensure the value is divisible by the stride."""
|
||||
return max(stride, int(round(value / stride) * stride))
|
||||
|
||||
def _apply_scale(self, width, height, scale):
|
||||
new_width = round(width * scale)
|
||||
new_height = round(height * scale)
|
||||
new_width = self._make_divisible(new_width, self.stride)
|
||||
new_height = self._make_divisible(new_height, self.stride)
|
||||
return new_width, new_height
|
||||
|
||||
def forward(self, img, img_num=1):
|
||||
"""
|
||||
Args:
|
||||
img (PIL Image): Image to be resized.
|
||||
img_num (int): Number of images, used to change max_tokens.
|
||||
Returns:
|
||||
PIL Image or Tensor: Rescaled image with divisible dimensions.
|
||||
"""
|
||||
if isinstance(img, torch.Tensor):
|
||||
height, width = img.shape[-2:]
|
||||
else:
|
||||
width, height = img.size
|
||||
|
||||
scale = min(self.max_size / max(width, height), 1.0)
|
||||
scale = max(scale, self.min_size / min(width, height))
|
||||
new_width, new_height = self._apply_scale(width, height, scale)
|
||||
|
||||
# Ensure the number of pixels does not exceed max_pixels
|
||||
if new_width * new_height > self.max_pixels / img_num:
|
||||
scale = self.max_pixels / img_num / (new_width * new_height)
|
||||
new_width, new_height = self._apply_scale(new_width, new_height, scale)
|
||||
|
||||
# Ensure longest edge does not exceed max_size
|
||||
if max(new_width, new_height) > self.max_size:
|
||||
scale = self.max_size / max(new_width, new_height)
|
||||
new_width, new_height = self._apply_scale(new_width, new_height, scale)
|
||||
|
||||
return F.resize(img, (new_height, new_width), self.interpolation, antialias=self.antialias)
|
||||
|
||||
|
||||
class ImageTransform:
|
||||
def __init__(
|
||||
self,
|
||||
max_image_size,
|
||||
min_image_size,
|
||||
image_stride,
|
||||
max_pixels=14*14*9*1024,
|
||||
image_mean=[0.5, 0.5, 0.5],
|
||||
image_std=[0.5, 0.5, 0.5]
|
||||
):
|
||||
self.stride = image_stride
|
||||
|
||||
self.resize_transform = MaxLongEdgeMinShortEdgeResize(
|
||||
max_size=max_image_size,
|
||||
min_size=min_image_size,
|
||||
stride=image_stride,
|
||||
max_pixels=max_pixels,
|
||||
)
|
||||
self.to_tensor_transform = transforms.ToTensor()
|
||||
self.normalize_transform = transforms.Normalize(mean=image_mean, std=image_std, inplace=True)
|
||||
|
||||
def __call__(self, img, img_num=1):
|
||||
img = self.resize_transform(img, img_num=img_num)
|
||||
img = self.to_tensor_transform(img)
|
||||
img = self.normalize_transform(img)
|
||||
return img
|
||||
|
||||
|
||||
class VideoTransform:
|
||||
def __init__(
|
||||
self,
|
||||
resolution=640,
|
||||
mode="area",
|
||||
divisible_crop_size=16,
|
||||
aspect_ratios=("21:9", "16:9", "4:3", "1:1", "3:4", "9:16"),
|
||||
stride_spatial=16,
|
||||
stride_temporal=4,
|
||||
mean=0.5,
|
||||
std=0.5,
|
||||
**kwargs
|
||||
):
|
||||
self.transform = Compose(
|
||||
[
|
||||
NaResize(
|
||||
resolution=resolution,
|
||||
mode=mode,
|
||||
downsample_only=True,
|
||||
stride=stride_spatial,
|
||||
# NOTE: aspect_ratios are only for `bucket` resize.
|
||||
aspect_ratios=aspect_ratios,
|
||||
),
|
||||
DivisibleCrop(divisible_crop_size),
|
||||
Normalize(mean, std),
|
||||
Rearrange("t c h w -> c t h w"),
|
||||
]
|
||||
)
|
||||
# self.stride = divisible_crop_size if isinstance(divisible_crop_size, int) else divisible_crop_size[0]
|
||||
self.stride_spatial = stride_spatial
|
||||
self.stride_temporal = stride_temporal
|
||||
|
||||
def __call__(self, video):
|
||||
return self.transform(video)
|
||||
|
||||
|
||||
class VisualTransform:
|
||||
def __init__(
|
||||
self,
|
||||
max_frame_size,
|
||||
min_frame_size,
|
||||
image_stride,
|
||||
max_pixels=14*14*9*1024,
|
||||
image_mean=[0.5, 0.5, 0.5],
|
||||
image_std=[0.5, 0.5, 0.5]
|
||||
):
|
||||
self.stride = image_stride
|
||||
self.resize_transform = MaxLongEdgeMinShortEdgeResize(
|
||||
max_size=max_frame_size,
|
||||
min_size=min_frame_size,
|
||||
stride=image_stride,
|
||||
max_pixels=max_pixels,
|
||||
)
|
||||
self.to_tensor_transform = transforms.ToTensor()
|
||||
self.normalize_transform = transforms.Normalize(mean=image_mean, std=image_std, inplace=True)
|
||||
|
||||
def _process_single(self, img, img_num=1):
|
||||
img = self.resize_transform(img, img_num=img_num)
|
||||
img = self.to_tensor_transform(img)
|
||||
img = self.normalize_transform(img)
|
||||
return img
|
||||
|
||||
def __call__(self, img, img_num=1):
|
||||
# --- Video sequence processing ---
|
||||
if isinstance(img, (list, tuple)):
|
||||
# List of PIL.Image or tensors
|
||||
out = torch.stack([self._process_single(frame, img_num=img_num) for frame in img]) # [T, C, H, W]
|
||||
out = out.permute(1, 0, 2, 3) # [C, T, H, W]
|
||||
return out
|
||||
elif isinstance(img, np.ndarray) and img.ndim == 4:
|
||||
# numpy array: [T, H, W, C] or [T, C, H, W]
|
||||
frames = [img[i] for i in range(img.shape[0])]
|
||||
processed_frames = [self._process_single(Image.fromarray(frame) if frame.shape[-1] in [3, 4] else frame, img_num=img_num)
|
||||
for frame in frames]
|
||||
out = torch.stack(processed_frames) # [T, C, H, W]
|
||||
out = out.permute(1, 0, 2, 3) # [C, T, H, W]
|
||||
return out
|
||||
elif isinstance(img, torch.Tensor) and img.ndim == 4:
|
||||
# torch tensor: [T, C, H, W] or [T, H, W, C]
|
||||
frames = [img[i] for i in range(img.shape[0])]
|
||||
processed_frames = [self._process_single(frame, img_num=img_num) for frame in frames]
|
||||
out = torch.stack(processed_frames) # [T, C, H, W]
|
||||
out = out.permute(1, 0, 2, 3) # [C, T, H, W]
|
||||
return out
|
||||
else:
|
||||
# Single frame
|
||||
return self._process_single(img, img_num=img_num)
|
||||
|
||||
|
||||
def decolorization(image):
|
||||
gray_image = image.convert('L')
|
||||
return Image.merge(image.mode, [gray_image] * 3) if image.mode in ('RGB', 'L') else gray_image
|
||||
|
||||
|
||||
def downscale(image, scale_factor):
|
||||
new_width = int(round(image.width * scale_factor))
|
||||
new_height = int(round(image.height * scale_factor))
|
||||
new_width = max(1, new_width)
|
||||
new_height = max(1, new_height)
|
||||
return image.resize((new_width, new_height), resample=Image.BICUBIC)
|
||||
|
||||
|
||||
def crop(image, crop_factors):
|
||||
target_h, target_w = crop_factors
|
||||
img_w, img_h = image.size
|
||||
|
||||
if target_h > img_h or target_w > img_w:
|
||||
raise ValueError("Crop size exceeds image dimensions")
|
||||
|
||||
x = random.randint(0, img_w - target_w)
|
||||
y = random.randint(0, img_h - target_h)
|
||||
|
||||
return image.crop((x, y, x + target_w, y + target_h)), [[x, y], [x + target_w, y + target_h]]
|
||||
|
||||
|
||||
def motion_blur_opencv(image, kernel_size=15, angle=0):
|
||||
# Linear kernel
|
||||
kernel = np.zeros((kernel_size, kernel_size), dtype=np.float32)
|
||||
kernel[kernel_size // 2, :] = np.ones(kernel_size, dtype=np.float32)
|
||||
|
||||
# Rotation kernel
|
||||
center = (kernel_size / 2 - 0.5, kernel_size / 2 - 0.5)
|
||||
M = cv2.getRotationMatrix2D(center, angle, 1)
|
||||
rotated_kernel = cv2.warpAffine(kernel, M, (kernel_size, kernel_size))
|
||||
|
||||
# Normalize kernel
|
||||
rotated_kernel /= rotated_kernel.sum() if rotated_kernel.sum() != 0 else 1
|
||||
|
||||
img = np.array(image)
|
||||
if img.ndim == 2:
|
||||
blurred = cv2.filter2D(img, -1, rotated_kernel, borderType=cv2.BORDER_REFLECT)
|
||||
else:
|
||||
# For color images, convolve each channel independently
|
||||
blurred = np.zeros_like(img)
|
||||
for c in range(img.shape[2]):
|
||||
blurred[..., c] = cv2.filter2D(img[..., c], -1, rotated_kernel, borderType=cv2.BORDER_REFLECT)
|
||||
|
||||
return Image.fromarray(blurred.astype(np.uint8))
|
||||
|
||||
|
||||
def shuffle_patch(image, num_splits, gap_size=2):
|
||||
"""Split an image into patches, allowing non-divisible sizes, shuffle them, and stitch them with gaps."""
|
||||
h_splits, w_splits = num_splits
|
||||
img_w, img_h = image.size
|
||||
|
||||
base_patch_h = img_h // h_splits
|
||||
patch_heights = [base_patch_h] * (h_splits - 1)
|
||||
patch_heights.append(img_h - sum(patch_heights))
|
||||
|
||||
base_patch_w = img_w // w_splits
|
||||
patch_widths = [base_patch_w] * (w_splits - 1)
|
||||
patch_widths.append(img_w - sum(patch_widths))
|
||||
|
||||
patches = []
|
||||
current_y = 0
|
||||
for i in range(h_splits):
|
||||
current_x = 0
|
||||
patch_h = patch_heights[i]
|
||||
for j in range(w_splits):
|
||||
patch_w = patch_widths[j]
|
||||
patch = image.crop((current_x, current_y, current_x + patch_w, current_y + patch_h))
|
||||
patches.append(patch)
|
||||
current_x += patch_w
|
||||
current_y += patch_h
|
||||
|
||||
random.shuffle(patches)
|
||||
|
||||
total_width = sum(patch_widths) + (w_splits - 1) * gap_size
|
||||
total_height = sum(patch_heights) + (h_splits - 1) * gap_size
|
||||
new_image = Image.new(image.mode, (total_width, total_height), color=(255, 255, 255))
|
||||
|
||||
current_y = 0 # Starting Y coordinate of the current row
|
||||
patch_idx = 0 # Current patch index
|
||||
for i in range(h_splits):
|
||||
current_x = 0 # Starting X coordinate of the current column
|
||||
patch_h = patch_heights[i] # Patch height for the current row
|
||||
for j in range(w_splits):
|
||||
# Fetch the shuffled patch
|
||||
patch = patches[patch_idx]
|
||||
patch_w = patch_widths[j] # Patch width for the current column
|
||||
# Paste the patch with top-left corner at (current_x, current_y)
|
||||
new_image.paste(patch, (current_x, current_y))
|
||||
# Update X coordinate: next patch starts after current patch width plus gap
|
||||
current_x += patch_w + gap_size
|
||||
patch_idx += 1
|
||||
# Update Y coordinate: next row starts after current row height plus gap
|
||||
current_y += patch_h + gap_size
|
||||
|
||||
return new_image
|
||||
|
||||
|
||||
def inpainting(image, num_splits, blank_ratio=0.3, blank_color=(255, 255, 255)):
|
||||
"""
|
||||
Split an image and randomly blank out patches for inpainting tasks.
|
||||
|
||||
Args:
|
||||
image: Input PIL.Image in RGB mode.
|
||||
h_splits: Number of row splits.
|
||||
w_splits: Number of column splits.
|
||||
blank_ratio: Ratio of blank patches, from 0 to 1.
|
||||
blank_color: RGB color for blank regions, e.g. white (255, 255, 255).
|
||||
|
||||
Returns:
|
||||
Processed and stitched PIL.Image.
|
||||
"""
|
||||
h_splits, w_splits = num_splits
|
||||
img_w, img_h = image.size
|
||||
|
||||
base_patch_h = img_h // h_splits
|
||||
patch_heights = [base_patch_h] * (h_splits - 1)
|
||||
patch_heights.append(img_h - sum(patch_heights))
|
||||
|
||||
base_patch_w = img_w // w_splits
|
||||
patch_widths = [base_patch_w] * (w_splits - 1)
|
||||
patch_widths.append(img_w - sum(patch_widths))
|
||||
|
||||
patches = []
|
||||
current_y = 0
|
||||
for i in range(h_splits):
|
||||
current_x = 0
|
||||
patch_h = patch_heights[i]
|
||||
for j in range(w_splits):
|
||||
patch_w = patch_widths[j]
|
||||
patch = image.crop((current_x, current_y, current_x + patch_w, current_y + patch_h))
|
||||
patches.append(patch)
|
||||
current_x += patch_w
|
||||
current_y += patch_h
|
||||
|
||||
total_patches = h_splits * w_splits
|
||||
num_blank = int(total_patches * blank_ratio)
|
||||
num_blank = max(0, min(num_blank, total_patches))
|
||||
blank_indices = random.sample(range(total_patches), num_blank)
|
||||
|
||||
processed_patches = []
|
||||
for idx, patch in enumerate(patches):
|
||||
if idx in blank_indices:
|
||||
blank_patch = Image.new("RGB", patch.size, color=blank_color)
|
||||
processed_patches.append(blank_patch)
|
||||
else:
|
||||
processed_patches.append(patch)
|
||||
|
||||
# Create the result image with the same size as the original
|
||||
result_image = Image.new("RGB", (img_w, img_h))
|
||||
current_y = 0
|
||||
patch_idx = 0
|
||||
for i in range(h_splits):
|
||||
current_x = 0
|
||||
patch_h = patch_heights[i]
|
||||
for j in range(w_splits):
|
||||
# Fetch the processed patch
|
||||
patch = processed_patches[patch_idx]
|
||||
patch_w = patch_widths[j]
|
||||
# Paste it back to the original position
|
||||
result_image.paste(patch, (current_x, current_y))
|
||||
current_x += patch_w
|
||||
patch_idx += 1
|
||||
current_y += patch_h
|
||||
|
||||
return result_image
|
||||
@@ -0,0 +1,141 @@
|
||||
# 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
|
||||
|
||||
from typing import Any, Dict, List, Literal, NamedTuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
class FrameSamplerOutput(NamedTuple):
|
||||
indices: List[int]
|
||||
additional_info: Dict[str, Any]
|
||||
|
||||
|
||||
class MultiClipsFrameSampler:
|
||||
"""
|
||||
Deterministic sampler used by Lance inference for image/video inputs.
|
||||
|
||||
The inference dataset always builds a single clip covering the full video.
|
||||
This sampler keeps the public behavior that matters for inference: sample
|
||||
at a target FPS, optionally clamp to max_duration, and return a frame count
|
||||
compatible with the VAE temporal downsample factor.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
temporal: int = 4,
|
||||
sample_fps: int = 12,
|
||||
truncate: bool = False,
|
||||
max_duration: int = 12,
|
||||
length_type: Literal["kn", "kn+1"] = "kn+1",
|
||||
assert_seconds: bool = True,
|
||||
):
|
||||
self.temporal = temporal
|
||||
self.sample_fps = sample_fps
|
||||
self.truncate = truncate
|
||||
self.max_duration = max_duration
|
||||
self.length_type = length_type
|
||||
self.assert_seconds = assert_seconds
|
||||
|
||||
def __call__(self, frames_info: Dict[str, Any]) -> FrameSamplerOutput:
|
||||
clip_indices = frames_info["clip_indices"]
|
||||
origin_fps = frames_info["fps"]
|
||||
|
||||
if self.truncate:
|
||||
clip_indices = self.truncate_to_bucket(clip_indices, origin_fps)
|
||||
|
||||
if self.assert_seconds:
|
||||
duration_sec = int(round(sum((end - start) / origin_fps for start, end in clip_indices)))
|
||||
if not self.truncate:
|
||||
duration_sec = min(duration_sec, self.max_duration)
|
||||
n_frames = duration_sec * self.sample_fps
|
||||
if self.length_type == "kn+1":
|
||||
n_frames += 1
|
||||
else:
|
||||
duration = sum((end - start) / origin_fps for start, end in clip_indices)
|
||||
if not self.truncate:
|
||||
duration = min(duration, self.max_duration)
|
||||
n_frames = int(round(duration * self.sample_fps))
|
||||
if self.length_type == "kn+1":
|
||||
if n_frames % self.temporal != 0:
|
||||
n_frames = n_frames // self.temporal * self.temporal + 1
|
||||
else:
|
||||
n_frames = n_frames // self.temporal * self.temporal + 1 - self.temporal
|
||||
|
||||
clip_n_frames = self.split_n_frames_by_clip(n_frames, clip_indices)
|
||||
sample_indices = self.sample_frame_indices(clip_indices, clip_n_frames)
|
||||
clip_n_latent_frames = [(n + self.temporal - 1) // self.temporal for n in clip_n_frames]
|
||||
|
||||
return FrameSamplerOutput(
|
||||
indices=sample_indices,
|
||||
additional_info={
|
||||
"clip_n_frames": clip_n_frames,
|
||||
"clip_n_latent_frames": clip_n_latent_frames,
|
||||
},
|
||||
)
|
||||
|
||||
def truncate_to_bucket(self, clip_indices, fps):
|
||||
clip_indices = [tuple(index) for index in clip_indices]
|
||||
durations = [(end - start) / fps for start, end in clip_indices]
|
||||
duration = sum(durations)
|
||||
max_duration = min(int(duration), self.max_duration)
|
||||
cutoff = duration - max_duration
|
||||
if cutoff <= 0:
|
||||
return clip_indices
|
||||
|
||||
if durations[-1] - cutoff > durations[0] - cutoff:
|
||||
start, end = clip_indices[-1]
|
||||
end = min(round((durations[-1] - cutoff) * fps), end) + start
|
||||
clip_indices[-1] = (start, end)
|
||||
else:
|
||||
start, end = clip_indices[0]
|
||||
start = max(end - round((durations[0] - cutoff) * fps), start)
|
||||
clip_indices[0] = (start, end)
|
||||
return clip_indices
|
||||
|
||||
def split_n_frames_by_clip(self, n_frames, clip_indices):
|
||||
n_latent_frames = n_frames // self.temporal
|
||||
clip_lengths = [end - start for start, end in clip_indices]
|
||||
total_length = sum(clip_lengths)
|
||||
clip_n_latent_frames = [int(length / total_length * n_latent_frames) for length in clip_lengths]
|
||||
n_remains = n_latent_frames - sum(clip_n_latent_frames)
|
||||
for i in range(n_remains):
|
||||
clip_n_latent_frames[i] += 1
|
||||
clip_n_frames = [n * self.temporal for n in clip_n_latent_frames]
|
||||
if self.length_type == "kn+1":
|
||||
clip_n_frames[0] += 1
|
||||
return clip_n_frames
|
||||
|
||||
@staticmethod
|
||||
def sample_frame_indices(clip_indices, clip_n_frames):
|
||||
shift_clip_indices = []
|
||||
accum_n_frames = 0
|
||||
for start, end in clip_indices:
|
||||
shift_start, shift_end = accum_n_frames, accum_n_frames + (end - start)
|
||||
shift_clip_indices.append((shift_start, shift_end))
|
||||
accum_n_frames += end - start
|
||||
|
||||
all_sample_indices = []
|
||||
for i, ((start, end), (shift_start, shift_end), n_frames) in enumerate(
|
||||
zip(clip_indices, shift_clip_indices, clip_n_frames)
|
||||
):
|
||||
indices = np.arange(start, end)
|
||||
next_shift_start = shift_clip_indices[i + 1][0] if i < len(clip_indices) - 1 else shift_end
|
||||
shift_sample_indices = (
|
||||
np.linspace(shift_start, next_shift_start - 1, n_frames, dtype=int) - shift_start
|
||||
)
|
||||
all_sample_indices.extend(indices[shift_sample_indices].tolist())
|
||||
|
||||
return all_sample_indices
|
||||
@@ -0,0 +1,20 @@
|
||||
# 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
|
||||
|
||||
from .frames import MultiClipsFrameSampler
|
||||
|
||||
FRAME_SAMPLER_TYPES = {
|
||||
"multi_clips": MultiClipsFrameSampler,
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
# 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.
|
||||
import math
|
||||
from typing import List, Union
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torchvision.transforms import functional as TVF
|
||||
from torchvision.transforms.functional import InterpolationMode, to_tensor
|
||||
|
||||
|
||||
class AreaResize:
|
||||
def __init__(
|
||||
self,
|
||||
max_area: float,
|
||||
downsample_only: bool = False,
|
||||
interpolation: InterpolationMode = InterpolationMode.BICUBIC,
|
||||
):
|
||||
self.max_area = max_area
|
||||
self.downsample_only = downsample_only
|
||||
self.interpolation = interpolation
|
||||
|
||||
def __call__(self, image: Union[torch.Tensor, Image.Image, List[Image.Image]]):
|
||||
|
||||
if isinstance(image, torch.Tensor):
|
||||
height, width = image.shape[-2:]
|
||||
elif isinstance(image, Image.Image):
|
||||
width, height = image.size
|
||||
elif isinstance(image, list) and isinstance(image[0], Image.Image):
|
||||
width, height = image[0].size
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
scale = math.sqrt(self.max_area / (height * width))
|
||||
|
||||
# keep original height and width for small pictures.
|
||||
scale = 1 if scale >= 1 and self.downsample_only else scale
|
||||
|
||||
resized_height, resized_width = round(height * scale), round(width * scale)
|
||||
|
||||
if isinstance(image, list) and isinstance(image[0], Image.Image):
|
||||
image = torch.stack(
|
||||
[
|
||||
to_tensor(
|
||||
TVF.resize(
|
||||
_image,
|
||||
size=(resized_height, resized_width),
|
||||
interpolation=self.interpolation,
|
||||
)
|
||||
)
|
||||
for _image in image
|
||||
]
|
||||
)
|
||||
else:
|
||||
image = TVF.resize(
|
||||
image,
|
||||
size=(resized_height, resized_width),
|
||||
interpolation=self.interpolation,
|
||||
)
|
||||
if isinstance(image, Image.Image):
|
||||
image = to_tensor(image)
|
||||
return image
|
||||
@@ -0,0 +1,126 @@
|
||||
# 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 math
|
||||
from typing import List, Tuple, Union
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torchvision.transforms import RandomResizedCrop
|
||||
from torchvision.transforms.functional import InterpolationMode, to_tensor
|
||||
|
||||
|
||||
class BucketResize:
|
||||
def __init__(
|
||||
self,
|
||||
max_area: float,
|
||||
interpolation: InterpolationMode = InterpolationMode.LANCZOS,
|
||||
aspect_ratios: List[str] = None,
|
||||
stride: Union[int, Tuple[int]] = None,
|
||||
):
|
||||
self.max_area = max_area
|
||||
self.interpolation = interpolation
|
||||
|
||||
assert aspect_ratios and stride, "`aspect_ratios` or `stride` not given!"
|
||||
self.buckets, self.bucket_ratios = self.init_buckets(aspect_ratios, max_area, stride)
|
||||
self.bucket_resize = {
|
||||
# NOTICE: despite the name, this setting performs a deterministic center crop.
|
||||
# bucket: (h,w)
|
||||
bucket: RandomResizedCrop(
|
||||
size=(bucket[0], bucket[1]),
|
||||
scale=(1, 1),
|
||||
ratio=(bucket_ratio, bucket_ratio),
|
||||
interpolation=self.interpolation,
|
||||
)
|
||||
for bucket, bucket_ratio in zip(self.buckets, self.bucket_ratios)
|
||||
}
|
||||
|
||||
def __call__(self, image: Union[torch.Tensor, Image.Image, List[Image.Image]]):
|
||||
|
||||
if isinstance(image, torch.Tensor):
|
||||
height, width = image.shape[-2:]
|
||||
elif isinstance(image, Image.Image):
|
||||
width, height = image.size
|
||||
elif isinstance(image, list) and isinstance(image[0], Image.Image):
|
||||
width, height = image[0].size
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
bucket = self.find_nearest_bucket(width, height)
|
||||
resizer = self.bucket_resize[bucket]
|
||||
|
||||
if isinstance(image, list) and isinstance(image[0], Image.Image):
|
||||
return torch.stack([to_tensor(resizer(_image)) for _image in image])
|
||||
else:
|
||||
image = resizer(image)
|
||||
if isinstance(image, Image.Image):
|
||||
image = to_tensor(image)
|
||||
return image
|
||||
|
||||
def find_nearest_bucket(self, width, height):
|
||||
"""
|
||||
找到与给定图片最近的bucket尺寸
|
||||
"""
|
||||
image_ratio = width / height
|
||||
diff = np.abs(image_ratio - self.bucket_ratios)
|
||||
index = diff.argmin()
|
||||
return self.buckets[index]
|
||||
|
||||
@staticmethod
|
||||
def init_buckets(aspect_ratio_names, max_area, stride):
|
||||
"""
|
||||
指定一些列最接近给定宽高比和面积的,同时整除vae降采样和patch_size倍数的宽高
|
||||
"""
|
||||
if not isinstance(stride, (tuple, list)):
|
||||
stride = (stride, stride)
|
||||
height_factor, width_factor = stride
|
||||
|
||||
buckets, bucket_ratios = [], []
|
||||
for name in aspect_ratio_names:
|
||||
w, h = (int(v) for v in name.split(":"))
|
||||
aspect_ratio = w / h
|
||||
|
||||
resize_width1 = math.sqrt(max_area * aspect_ratio)
|
||||
bucket_width1 = round(resize_width1 / width_factor) * width_factor
|
||||
resize_height1 = bucket_width1 / aspect_ratio
|
||||
bucket_height1 = round(resize_height1 / height_factor) * height_factor
|
||||
bucket_ratio1 = bucket_width1 / bucket_height1
|
||||
bucket_area1 = bucket_width1 * bucket_height1
|
||||
|
||||
resize_height2 = math.sqrt(max_area / aspect_ratio)
|
||||
bucket_height2 = round(resize_height2 / height_factor) * height_factor
|
||||
resize_width2 = bucket_height2 * aspect_ratio
|
||||
bucket_width2 = round(resize_width2 / width_factor) * width_factor
|
||||
bucket_ratio2 = bucket_width2 / bucket_height2
|
||||
bucket_area2 = bucket_width2 * bucket_height2
|
||||
|
||||
if abs(bucket_ratio1 - aspect_ratio) < abs(bucket_ratio2 - aspect_ratio):
|
||||
bucket_width, bucket_height = bucket_width1, bucket_height1
|
||||
elif abs(bucket_ratio1 - aspect_ratio) > abs(bucket_ratio2 - aspect_ratio):
|
||||
bucket_width, bucket_height = bucket_width2, bucket_height2
|
||||
else:
|
||||
if abs(bucket_area1 - max_area) <= abs(bucket_area2 - max_area):
|
||||
bucket_width, bucket_height = bucket_width1, bucket_height1
|
||||
else:
|
||||
bucket_width, bucket_height = bucket_width2, bucket_height2
|
||||
|
||||
bucket_ratio = bucket_width / bucket_height
|
||||
|
||||
buckets.append((bucket_height, bucket_width))
|
||||
bucket_ratios.append(bucket_ratio)
|
||||
|
||||
bucket_ratios = np.array(bucket_ratios)
|
||||
|
||||
return buckets, bucket_ratios
|
||||
@@ -0,0 +1,39 @@
|
||||
# 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.
|
||||
from typing import Union
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torchvision.transforms import functional as TVF
|
||||
|
||||
|
||||
class DivisibleCrop:
|
||||
def __init__(self, factor):
|
||||
if not isinstance(factor, tuple):
|
||||
factor = (factor, factor)
|
||||
|
||||
self.height_factor, self.width_factor = factor[0], factor[1]
|
||||
|
||||
def __call__(self, image: Union[torch.Tensor, Image.Image]):
|
||||
if isinstance(image, torch.Tensor):
|
||||
height, width = image.shape[-2:]
|
||||
elif isinstance(image, Image.Image):
|
||||
width, height = image.size
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
cropped_height = height - (height % self.height_factor)
|
||||
cropped_width = width - (width % self.width_factor)
|
||||
|
||||
image = TVF.center_crop(img=image, output_size=(cropped_height, cropped_width))
|
||||
return image
|
||||
@@ -0,0 +1,56 @@
|
||||
# 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.
|
||||
from typing import Literal
|
||||
from torchvision.transforms import CenterCrop, Compose, InterpolationMode, Resize
|
||||
|
||||
from .area_resize import AreaResize
|
||||
from .bucket_resize import BucketResize
|
||||
|
||||
def NaResize(
|
||||
resolution: int,
|
||||
mode: Literal["area", "square", "bucket"],
|
||||
downsample_only: bool,
|
||||
interpolation: InterpolationMode = InterpolationMode.BICUBIC,
|
||||
**kwargs,
|
||||
):
|
||||
if mode == "area":
|
||||
return AreaResize(
|
||||
max_area=resolution**2,
|
||||
downsample_only=downsample_only,
|
||||
interpolation=interpolation,
|
||||
)
|
||||
elif mode == "square":
|
||||
return Compose(
|
||||
[
|
||||
Resize(
|
||||
size=resolution,
|
||||
interpolation=interpolation,
|
||||
),
|
||||
CenterCrop(resolution),
|
||||
]
|
||||
)
|
||||
elif mode == "bucket":
|
||||
aspect_ratios = kwargs.get("aspect_ratios", ["21:9", "16:9", "4:3", "1:1", "3:4", "9:16"])
|
||||
stride = kwargs.get("stride", 16)
|
||||
return Compose(
|
||||
[
|
||||
BucketResize(
|
||||
max_area=resolution**2,
|
||||
interpolation=interpolation,
|
||||
aspect_ratios=aspect_ratios,
|
||||
stride=stride,
|
||||
)
|
||||
]
|
||||
)
|
||||
raise ValueError(f"Unknown resize mode: {mode}")
|
||||
@@ -0,0 +1,23 @@
|
||||
# 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.
|
||||
from einops import rearrange
|
||||
|
||||
|
||||
class Rearrange:
|
||||
def __init__(self, pattern: str, **kwargs):
|
||||
self.pattern = pattern
|
||||
self.kwargs = kwargs
|
||||
|
||||
def __call__(self, x):
|
||||
return rearrange(x, self.pattern, **self.kwargs)
|
||||
@@ -0,0 +1,165 @@
|
||||
# Copyright (c) 2023 OpenGVLab
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: MIT
|
||||
#
|
||||
# This file has been modified by ByteDance Ltd. and/or its affiliates. on 2025-05-20.
|
||||
#
|
||||
# Original file was released under MIT, with the full license text
|
||||
# available at https://github.com/OpenGVLab/InternVL/blob/main/LICENSE.
|
||||
#
|
||||
# This modified file is released under the same license.
|
||||
|
||||
|
||||
import io
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
|
||||
import numpy as np
|
||||
import decord
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def get_frame_indices(num_frames, vlen, sample='rand', fix_start=None, input_fps=1, max_num_frames=-1):
|
||||
if sample in ['rand', 'middle']: # uniform sampling
|
||||
acc_samples = min(num_frames, vlen)
|
||||
# split the video into `acc_samples` intervals, and sample from each interval.
|
||||
intervals = np.linspace(start=0, stop=vlen, num=acc_samples + 1).astype(int)
|
||||
ranges = []
|
||||
for idx, interv in enumerate(intervals[:-1]):
|
||||
ranges.append((interv, intervals[idx + 1] - 1))
|
||||
if sample == 'rand':
|
||||
try:
|
||||
frame_indices = [random.choice(range(x[0], x[1])) for x in ranges]
|
||||
except:
|
||||
frame_indices = np.random.permutation(vlen)[:acc_samples]
|
||||
frame_indices.sort()
|
||||
frame_indices = list(frame_indices)
|
||||
elif fix_start is not None:
|
||||
frame_indices = [x[0] + fix_start for x in ranges]
|
||||
elif sample == 'middle':
|
||||
frame_indices = [(x[0] + x[1]) // 2 for x in ranges]
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
if len(frame_indices) < num_frames: # padded with last frame
|
||||
padded_frame_indices = [frame_indices[-1]] * num_frames
|
||||
padded_frame_indices[:len(frame_indices)] = frame_indices
|
||||
frame_indices = padded_frame_indices
|
||||
elif 'fps' in sample: # fps0.5, sequentially sample frames at 0.5 fps
|
||||
output_fps = float(sample[3:])
|
||||
duration = float(vlen) / input_fps
|
||||
delta = 1 / output_fps # gap between frames, this is also the clip length each frame represents
|
||||
frame_seconds = np.arange(0 + delta / 2, duration + delta / 2, delta)
|
||||
frame_indices = np.around(frame_seconds * input_fps).astype(int)
|
||||
frame_indices = [e for e in frame_indices if e < vlen]
|
||||
if max_num_frames > 0 and len(frame_indices) > max_num_frames:
|
||||
frame_indices = frame_indices[:max_num_frames]
|
||||
else:
|
||||
raise ValueError
|
||||
return frame_indices
|
||||
|
||||
|
||||
def read_frames_decord(video_path, num_frames, sample='rand', fix_start=None, clip=None, min_num_frames=4):
|
||||
video_reader = decord.VideoReader(video_path, num_threads=1)
|
||||
vlen = len(video_reader)
|
||||
fps = video_reader.get_avg_fps()
|
||||
duration = vlen / float(fps)
|
||||
if clip:
|
||||
start, end = clip
|
||||
duration = end - start
|
||||
vlen = int(duration * fps)
|
||||
start_index = int(start * fps)
|
||||
|
||||
t_num_frames = np.random.randint(min_num_frames, num_frames + 1)
|
||||
|
||||
frame_indices = get_frame_indices(
|
||||
t_num_frames, vlen, sample=sample, fix_start=fix_start,
|
||||
input_fps=fps
|
||||
)
|
||||
if clip:
|
||||
frame_indices = [f + start_index for f in frame_indices]
|
||||
frames = video_reader.get_batch(frame_indices).asnumpy() # (T, H, W, C), np.uint8
|
||||
frames = [Image.fromarray(frames[i]) for i in range(frames.shape[0])]
|
||||
return frames
|
||||
|
||||
|
||||
def extract_frame_number(filename):
|
||||
# Extract the numeric part from the filename using regular expressions
|
||||
match = re.search(r'_(\d+).jpg$', filename)
|
||||
return int(match.group(1)) if match else -1
|
||||
|
||||
|
||||
def sort_frames(frame_paths):
|
||||
# Extract filenames from each path and sort by their numeric part
|
||||
return sorted(frame_paths, key=lambda x: extract_frame_number(os.path.basename(x)))
|
||||
|
||||
|
||||
def read_frames_folder(video_path, num_frames, sample='rand', fix_start=None, min_num_frames=4):
|
||||
image_list = sort_frames(list(os.listdir(video_path)))
|
||||
frames = []
|
||||
for image in image_list:
|
||||
fp = os.path.join(video_path, image)
|
||||
frame = Image.open(fp).convert('RGB')
|
||||
frames.append(frame)
|
||||
vlen = len(frames)
|
||||
|
||||
t_num_frames = np.random.randint(min_num_frames, num_frames + 1)
|
||||
|
||||
if vlen > t_num_frames:
|
||||
frame_indices = get_frame_indices(
|
||||
t_num_frames, vlen, sample=sample, fix_start=fix_start
|
||||
)
|
||||
frames = [frames[i] for i in frame_indices]
|
||||
return frames
|
||||
|
||||
|
||||
class FrameSampler:
|
||||
def __init__(self, max_num_frames=-1, min_num_frames=8, sample='rand'):
|
||||
self.max_num_frames = max_num_frames
|
||||
self.min_num_frames = min_num_frames
|
||||
self.sample = sample
|
||||
|
||||
def __call__(self, file_name):
|
||||
fn = read_frames_folder if file_name.endswith('/') else read_frames_decord
|
||||
frames = fn(file_name, num_frames=self.max_num_frames, min_num_frames=self.min_num_frames, sample=self.sample)
|
||||
return frames
|
||||
|
||||
|
||||
def decode_video_byte(video_bytes):
|
||||
video_stream = io.BytesIO(video_bytes)
|
||||
vr = decord.VideoReader(video_stream)
|
||||
return vr
|
||||
|
||||
|
||||
def sample_mp4_frames(mp4_p, n_frames=None, fps=None, return_frame_indices=False, random_sample=False):
|
||||
if isinstance(mp4_p, str):
|
||||
vr = decord.VideoReader(mp4_p, num_threads=1)
|
||||
elif isinstance(mp4_p, decord.video_reader.VideoReader):
|
||||
vr = mp4_p
|
||||
video_fps = vr.get_avg_fps() # Get the video frame rate
|
||||
video_duration = len(vr) / video_fps
|
||||
if n_frames is not None:
|
||||
if random_sample:
|
||||
frame_indices = sorted(random.sample(range(len(vr)), n_frames))
|
||||
else:
|
||||
frame_indices = np.linspace(0, len(vr)-1, n_frames, dtype=int).tolist()
|
||||
else:
|
||||
frame_indices = [int(i) for i in np.arange(0, len(vr)-1, video_fps/fps)]
|
||||
frames = vr.get_batch(frame_indices).asnumpy() # Convert to a numpy array
|
||||
frames = [Image.fromarray(frame).convert("RGB") for frame in frames]
|
||||
if not return_frame_indices:
|
||||
return frames, video_duration
|
||||
else:
|
||||
return frames, video_duration, frame_indices
|
||||
|
||||
|
||||
def sample_mp4_frames_by_indices(mp4_p, frame_indices: list):
|
||||
if isinstance(mp4_p, str):
|
||||
vr = decord.VideoReader(mp4_p, num_threads=1)
|
||||
elif isinstance(mp4_p, decord.video_reader.VideoReader):
|
||||
vr = mp4_p
|
||||
# sample the frames in frame_indices
|
||||
frames = vr.get_batch(frame_indices).asnumpy() # Convert to a numpy array
|
||||
frames = [Image.fromarray(frame).convert("RGB") for frame in frames]
|
||||
return frames
|
||||
Reference in New Issue
Block a user