Spaces:
Running
on
Zero
Running
on
Zero
File size: 5,913 Bytes
e6af450 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
# SPDX-License-Identifier: Apache-2.0
import io
import json
import pyarrow.parquet as pq
import random
from PIL import Image
from .data_utils import pil_img2rgb
from .distributed_iterable_dataset import DistributedIterableDataset
from .parquet_utils import get_parquet_data_paths, init_arrow_pf_fs
Image.MAX_IMAGE_PIXELS = 20_000_000
class T2IIterableDataset(DistributedIterableDataset):
def __init__(
self, dataset_name, transform, tokenizer, data_dir_list, num_used_data,
local_rank=0, world_size=1, num_workers=8, data_status=None,
):
"""
data_dir_list: list of data directories contains parquet files
num_used_data: list of number of sampled data paths for each data directory
"""
super().__init__(dataset_name, local_rank, world_size, num_workers)
self.transform = transform
self.tokenizer = tokenizer
self.data_status = data_status
self.data_paths = self.get_data_paths(data_dir_list, num_used_data)
self.set_epoch()
def get_data_paths(self, data_dir_list, num_used_data):
return get_parquet_data_paths(data_dir_list, num_used_data)
def __iter__(self):
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
transform_stride = self.transform.stride
print(
f"rank-{self.local_rank} worker-{worker_id} dataset-{self.dataset_name}: "
f"resuming data at parquet#{parquet_start_id}, rg#{row_group_start_id}, row#{row_start_id}"
)
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:]
for row_group_id in row_group_ids_:
df = fr.read_row_group(row_group_id).to_pandas()
df = df.iloc[row_start_id:]
for row_idx, row in df.iterrows():
num_tokens = 0
try:
image_byte = row['image']
image = pil_img2rgb(Image.open(io.BytesIO(image_byte)))
except Exception as e:
print(f'Error: {e} in rg#{row_group_id}, {parquet_file_path}')
continue
image_tensor = self.transform(image)
height, width = image_tensor.shape[1:]
num_tokens += width * height // transform_stride ** 2
try:
caption_dict = row['captions']
caption_dict = json.loads(caption_dict)
except Exception as e:
print(f'Error: {e} in rg#{row_group_id}, {parquet_file_path}')
continue
caps_token = [self.tokenizer.encode(v) for _, v in caption_dict.items()]
if len(caps_token) == 0:
print(f'no caption in rg#{row_group_id}, {parquet_file_path}')
caption_token = self.tokenizer.encode(' ')
else:
caption_token = random.choice(caps_token)
sequence_plan, text_ids_list = [], []
text_ids = caption_token
num_tokens += len(caption_token)
text_ids_list.append(text_ids)
sequence_plan.append({
'type': 'text',
'enable_cfg': 1,
'loss': 0,
'special_token_loss': 0,
'special_token_label': None,
})
sequence_plan.append({
'type': 'vae_image',
'enable_cfg': 0,
'loss': 1,
'special_token_loss': 0,
'special_token_label': None,
})
sample = dict(
image_tensor_list=[image_tensor],
text_ids_list=text_ids_list,
num_tokens=num_tokens,
sequence_plan=sequence_plan,
data_indexes={
"data_indexes": [parquet_idx, row_group_id, row_idx],
"worker_id": worker_id,
"dataset_name": self.dataset_name,
}
)
yield sample
row_start_id = 0
row_group_start_id = 0
parquet_start_id = 0
print(f"{self.dataset_name} repeat in rank-{self.local_rank} worker-{worker_id}")
|