import copy
import glob
import io
import json
import math
import os
import random
import re
from dataclasses import dataclass
from typing import Dict, List, Optional, Sequence
import pyarrow.parquet as pq
import torch
import transformers
import yaml
from PIL import Image, ImageFile
from torch.utils.data import Dataset
from torchvision.transforms import v2
from torchvision import transforms
from datasets import load_dataset, concatenate_datasets
from blip3o.constants import (
DEFAULT_IM_END_TOKEN,
DEFAULT_IM_START_TOKEN,
DEFAULT_IMAGE_TOKEN,
IGNORE_INDEX,
IMAGE_TOKEN_INDEX,
)
from blip3o.utils import rank0_print
import random
ImageFile.LOAD_TRUNCATED_IMAGES = True
## target transform for sana
target_transform = v2.Compose(
[
v2.Resize(512),
v2.CenterCrop(512),
v2.ToImage(),
v2.ToDtype(torch.float32, scale=True),
v2.Normalize([0.5], [0.5]),
]
)
def expand2square(pil_img, background_color):
width, height = pil_img.size
if width == height:
return pil_img
elif width > height:
result = Image.new(pil_img.mode, (width, width), background_color)
result.paste(pil_img, (0, (width - height) // 2))
return result
else:
result = Image.new(pil_img.mode, (height, height), background_color)
result.paste(pil_img, ((height - width) // 2, 0))
return result
def preprocess_multimodal(sources: Sequence[str], data_args) -> Dict:
is_multimodal = data_args.is_multimodal
if not is_multimodal:
return sources
for source in sources:
for sentence in source:
replace_token = DEFAULT_IMAGE_TOKEN # "<image>"
# NOTE: only add im_start_end when image generation
if data_args.mm_use_im_start_end and sentence['from'] == 'gpt':
replace_token = DEFAULT_IM_START_TOKEN + replace_token + DEFAULT_IM_END_TOKEN
sentence["value"] = sentence["value"].replace(DEFAULT_IMAGE_TOKEN, replace_token)
# For videoInstruct-100k noisy_data. TODO: Ask Yuanhan to clean the data instead of leaving the noise code here.
sentence["value"] = sentence["value"].replace("QA_GT_caption_based_noisy", "")
return sources
def preprocess_qwen(sources, tokenizer: transformers.PreTrainedTokenizer, has_image: bool = False, max_len=2048, system_message: str = "You are a helpful assistant.") -> Dict:
# roles = {"human": "<|im_start|>user", "gpt": "<|im_start|>assistant"}
roles = {"human": "user", "gpt": "assistant"}
#tokenizer = copy.deepcopy(tokenizer)
# When there is actually an image, we add the image tokens as a special token
if 'image_token_index' not in globals():
tokenizer.add_tokens(["<image>"], special_tokens=True)
global image_token_index
image_token_index = tokenizer.convert_tokens_to_ids("<image>") # 217210
# if has_image:
# tokenizer.add_tokens(["<image>"], special_tokens=True)
# image_token_index = tokenizer.convert_tokens_to_ids("<image>")
im_start, im_end = tokenizer.additional_special_tokens_ids[:2]
# unmask_tokens = ["<|im_start|>", "<|im_start|>", "\n"]
unmask_tokens_idx = [198, im_start, im_end] # [198, 151644, 151645]
# nl_tokens = tokenizer("\n").input_ids
# Reset Qwen chat templates so that it won't include system message every time we apply
chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
tokenizer.chat_template = chat_template
# _system = tokenizer("system").input_ids + nl_tokens
# _user = tokenizer("user").input_ids + nl_tokens
# _assistant = tokenizer("assistant").input_ids + nl_tokens
# Apply prompt templates
input_ids, targets = [], []
for i, source in enumerate(sources):
if roles[source[0]["from"]] != roles["human"]:
source = source[1:]
input_id, target = [], []
# New version, use apply chat template
# Build system message for each sentence
input_id += tokenizer.apply_chat_template([{"role" : "system", "content" : system_message}])
# target += [IGNORE_INDEX] * len(input_id)
target += input_id
for conv in source:
# Make sure blip3o data can load
try:
role = conv["role"]
content = conv["content"]
except:
role = conv["from"]
content = conv["value"]
role = roles.get(role, role)
conv = [{"role" : role, "content" : content}] # Qwen's format
encode_id = tokenizer.apply_chat_template(conv)
# '<|im_start|>user\n<image>\nPlease reconstruct the given image based on the image content: a photography of a woman with blonde hair and a bow on her head<|im_end|>\n'
# '<|im_start|>assistant\n<im_start><image><im_end><|im_end|>\n'
# <im_start>151669 <im_end>151670 != <|im_start|>151644 <|im_end|>151645
# decoded_text = tokenizer.decode(encode_id, skip_special_tokens=False)
input_id += encode_id
if role in ["user", "system"]:
# target += [IGNORE_INDEX] * len(encode_id)
target += encode_id
else:
target += encode_id
assert len(input_id) == len(target), f"{len(input_id)} != {len(target)}"
for idx, encode_id in enumerate(input_id):
if encode_id in unmask_tokens_idx: # ["<|im_start|>", "<|im_start|>", "\n"] does not work
target[idx] = encode_id
if encode_id == image_token_index:
input_id[idx] = IMAGE_TOKEN_INDEX # Replaces all <image>217210 with -200
input_ids.append(input_id)
targets.append(target)
input_ids = torch.tensor(input_ids, dtype=torch.long)
targets = torch.tensor(targets, dtype=torch.long)
return dict(
input_ids=input_ids,
labels=targets,
)
class LazySupervisedMixDataset(Dataset):
"""Dataset for supervised fine-tuning."""
def __init__(
self,
tokenizer: transformers.PreTrainedTokenizer,
data_path: str,
data_args
):
super(LazySupervisedMixDataset, self).__init__()
self.data_args = data_args
list_data_dict = []
if os.path.isdir(data_path):
files = glob.glob(os.path.join(data_path, "**", "*.tar"), recursive=True)
else:
files = glob.glob(data_path, recursive=True)
train_dataset = load_dataset("webdataset", data_files=files, split="train", num_proc=1, cache_dir='')
train_dataset = train_dataset.rename_column("jpg", "image")
train_dataset = train_dataset.add_column('type', len(train_dataset) * ['T2I'])
train_dataset = train_dataset.remove_columns([col for col in train_dataset.column_names if not col in (
["image", "txt", "type"])])
print(f"Finished loading image {len(train_dataset)}")
list_data_dict.append(train_dataset)
if len(list_data_dict) > 1:
list_data_dict = concatenate_datasets(list_data_dict)
else:
list_data_dict = list_data_dict[0]
list_data_dict = list_data_dict.shuffle(seed=42)
rank0_print(f"Totoal number of training instance: {len(list_data_dict)}")
self.tokenizer = tokenizer
self.list_data_dict = list_data_dict
self.modality = torch.tensor(0) # 0 is for understanding task, 1 is for generation task
def __len__(self):
return len(self.list_data_dict)
def process_image(self, image):
processor = self.data_args.image_processor
image_size = image.size
image = processor.preprocess(image, return_tensors="pt")["pixel_values"][0]
return image, image_size, self.modality
def process_target_image(self, image):
image = target_transform(image)
return image
@property
def lengths(self):
length_list = []
for sample in self.list_data_dict:
img_tokens = 128 if "image" in sample else 0
length_list.append(sum(len(conv["value"].split()) for conv in sample["conversations"]) + img_tokens)
return length_list
@property
def modality_lengths(self):
length_list = []
for sample in self.list_data_dict:
cur_len = sum(len(conv["value"].split()) for conv in sample["conversations"])
cur_len = cur_len if "image" in sample else -cur_len
length_list.append(cur_len)
return length_list
def __getitem__(self, i) -> Dict[str, torch.Tensor]:
while True:
sources = self.list_data_dict[i]
if sources["type"] == "T2I":
sources["conversations"] = [
{"from": "human", "value": f"Please generate image based on the following caption: {sources['txt']}"},
{"from": "gpt", "value": "<image>"},
]
elif sources["type"] == "I2I":
sources["conversations"] = [
{
"from": "human",
"value": f"<image>\nPlease reconstruct the given image.",
},
{"from": "gpt", "value": ""},
]
else:
raise ValueError("Unknown source type. Please check the 'type' in 'sources'.")
if "image" in sources:
if sources["type"] == "T2I" or sources["type"] == "I2I":
image_files = self.list_data_dict[i]["image"]
if not isinstance(image_files, list):
image_files = [image_files]
images = []
for img in image_files:
try:
if sources["type"] == "T2I" or sources["type"] == "I2I":
img = img.convert("RGB")
else:
raise ValueError("Unknown source type. Please check the 'type' in 'sources'.")
images.append(img)
except Exception as e:
print(f"Error opening image {img}: {e}")
images = None
break # Skip to the next image if there's an error
## test if can apply img_process
if not images is None:
try:
process_images = [self.process_image(f) for f in images]
except Exception as e:
print(f"Error wrong number of channels: {e}")
images = None
# If no valid images were found, randomly pick another item
if images is None:
print(sources)
print(f"Warning: Invalid image!")
i = random.randint(0, len(self.list_data_dict) - 1)
continue
sources = preprocess_multimodal(copy.deepcopy([sources["conversations"]]), self.data_args)
else:
sources = copy.deepcopy([sources["conversations"]])
data_dict = preprocess_qwen(sources, self.tokenizer, has_image=("image" in self.list_data_dict[i]))
if isinstance(i, int):
data_dict = dict(input_ids=data_dict["input_ids"][0], labels=data_dict["labels"][0])
# image exist in the data
if "image" in self.list_data_dict[i]:
data_dict["image"] = process_images
data_dict["target_image"] = [self.process_target_image(f) for f in images]
data_dict["ids"] = self.list_data_dict[i]["id"] if "id" in self.list_data_dict[i] else "unk"
return data_dict
class LazySupervisedRestoreDataset(Dataset):
"""Dataset for supervised fine-tuning."""
def __init__(
self,
tokenizer: transformers.PreTrainedTokenizer,
data_path: str,
data_args
):
super(LazySupervisedRestoreDataset, self).__init__()
self.data_args = data_args
datasets_to_merge = []
with open(data_path, 'r', encoding='utf-8') as f:
for line in f:
path = line.strip()
if path and os.path.isdir(path):
parquet_files = glob.glob(os.path.join(path, "**", "*.parquet"), recursive=True)
tar_files = glob.glob(os.path.join(path, "**", "*.tar"), recursive=True)
if parquet_files:
dataset = load_dataset("parquet", data_files=parquet_files, split="train", num_proc=1, cache_dir='/data/zgq/yaozhengjian/Datasets/FFHQ/cache')
if "text" in dataset.column_names:
dataset = dataset.rename_column("text", "txt")
required_columns = ["image", "txt"]
columns_to_keep = [col for col in dataset.column_names if col in required_columns]
dataset = dataset.select_columns(columns_to_keep)
dataset = dataset.add_column('type', len(dataset) * ['I2I'])
datasets_to_merge.append(dataset)
if tar_files:
dataset = load_dataset("webdataset", data_files=tar_files, split="train", num_proc=1, cache_dir='/data/zgq/yaozhengjian/Datasets/UniWorld-V1/data/BLIP3o-60k/webdataset')
dataset = dataset.rename_column("jpg", "image")
dataset = dataset.add_column('type', len(dataset) * ['I2I'])
dataset = dataset.remove_columns([col for col in dataset.column_names if not col in (
["image", "txt", "type"])])
datasets_to_merge.append(dataset)
if datasets_to_merge:
train_dataset = concatenate_datasets(datasets_to_merge)
else:
raise ValueError(f"No valid parquet files found in paths from {data_path}")
train_dataset = train_dataset.shuffle(seed=42)
rank0_print(f"Total number of training instance: {len(train_dataset)}") # 110848
self.tokenizer = tokenizer
self.list_data_dict = train_dataset
self.modality = torch.tensor(1) # 0 is for understanding task, 1 is for generation task
self.degradation_params = {
'gt_size': 512,
'in_size': 512,
'use_motion_kernel': False,
'blur_kernel_size': 41,
'blur_sigma': [1, 15],
# 'downsample_range': [1, 30],
'downsample_range': [1, 15],
'noise_range': [0, 20],
'jpeg_range': [30, 90]
}
def __len__(self):
return len(self.list_data_dict)
def process_image(self, image):
# === Image Processor Configuration ===
# {'_processor_class': None, 'do_resize': True, 'size': (384, 384), 'resample': <Resampling.BICUBIC: 3>, 'do_rescale': True, 'rescale_factor': 0.00392156862745098, 'do_normalize': True, 'image_mean': [0.5, 0.5, 0.5], 'image_std': [0.5, 0.5, 0.5], 'do_convert_rgb': None, 'crop_size': {'height': 384, 'width': 384}, 'image_processor_type': 'SiglipImageProcessor'}
# key 384*384
# =====================================
processor = self.data_args.image_processor
image_size = image.size
image = processor.preprocess(image, return_tensors="pt")["pixel_values"][0]
return image, image_size, self.modality
def process_target_image(self, image):
image = target_transform(image)
return image
@property
def lengths(self):
length_list = []
for sample in self.list_data_dict:
img_tokens = 128 if "image" in sample else 0
length_list.append(sum(len(conv["value"].split()) for conv in sample["conversations"]) + img_tokens)
return length_list
@property
def modality_lengths(self):
length_list = []
for sample in self.list_data_dict:
cur_len = sum(len(conv["value"].split()) for conv in sample["conversations"])
cur_len = cur_len if "image" in sample else -cur_len
length_list.append(cur_len)
return length_list
def __getitem__(self, i) -> Dict[str, torch.Tensor]:
while True:
sources = self.list_data_dict[i]
if sources["type"] == "T2I":
sources["conversations"] = [
{"from": "human", "value": f"Please generate image based on the following caption: {sources['txt']}"},
{"from": "gpt", "value": "<image>"},
]
elif sources["type"] == "I2I":
sources["conversations"] = [
{
"from": "human",
"value": f"<image>\nPlease reconstruct the given image based on the image content: {sources['txt']}" if random.random() < 0.9 else "<image>\nPlease reconstruct the given image.",
},
{"from": "gpt", "value": "<image>"},
]
# }, # Without caption
# {"from": "gpt", "value": "<image>"},
# ]
else:
raise ValueError("Unknown source type. Please check the 'type' in 'sources'.")
if "image" in sources:
if sources["type"] == "T2I" or sources["type"] == "I2I":
image_files = self.list_data_dict[i]["image"]
if not isinstance(image_files, list):
image_files = [image_files]
images = []
for img in image_files:
try:
if sources["type"] == "T2I" or sources["type"] == "I2I":
img = img.convert("RGB") # PIL.Image.Image
else:
raise ValueError("Unknown source type. Please check the 'type' in 'sources'.")
images.append(img)
except Exception as e:
print(f"Error opening image {img}: {e}")
images = None
break # Skip to the next image if there's an error
## test if can apply img_process
if not images is None:
try:
from .image_degradation import degrade_image
if random.random() < 0.2:
degrade_images = images # Directly use the original image for reconstruction task
sources["conversations"] = [
{
"from": "human",
"value": f"<image>\nThis is a side mission. Please complete the auxiliary reconstruction task by duplicating the given high-definition image.",
},
{"from": "gpt", "value": "<image>"},
]
else:
degrade_images = [degrade_image(img, **self.degradation_params) for img in images]
images_for_llm = degrade_images + images
process_images = [self.process_image(f) for f in images_for_llm]
except Exception as e:
print(f"Error wrong number of channels: {e}")
images = None
# If no valid images were found, randomly pick another item
if images is None:
print(sources)
print(f"Warning: Invalid image!")
i = random.randint(0, len(self.list_data_dict) - 1)
continue
sources = preprocess_multimodal(copy.deepcopy([sources["conversations"]]), self.data_args)
else:
sources = copy.deepcopy([sources["conversations"]])
data_dict = preprocess_qwen(sources, self.tokenizer, has_image=("image" in self.list_data_dict[i]))
if isinstance(i, int):
data_dict = dict(input_ids=data_dict["input_ids"][0], labels=data_dict["labels"][0])
# image exist in the data
if "image" in self.list_data_dict[i]:
data_dict["image"] = process_images # Siglip images [(torch.Size([3, 384, 384]), (512, 512), tensor(1)), .....]
data_dict["target_image"] = [self.process_target_image(f) for f in images] # 512x512 [torch.Size([3, 512, 512])]
data_dict["detailed_condition"] = [self.process_target_image(f) for f in degrade_images] # [torch.Size([3, 512, 512])]
data_dict["ids"] = self.list_data_dict[i]["id"] if "id" in self.list_data_dict[i] else "unk"
return data_dict
@dataclass
class DataCollatorForSupervisedDataset(object):
"""Collate examples for supervised fine-tuning."""
tokenizer: transformers.PreTrainedTokenizer
def pad_sequence(self, input_ids, batch_first, padding_value):
if self.tokenizer.padding_side == "left":
input_ids = [torch.flip(_input_ids, [0]) for _input_ids in input_ids]
input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=batch_first, padding_value=padding_value)
if self.tokenizer.padding_side == "left":
input_ids = torch.flip(input_ids, [1])
return input_ids
def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
input_ids = [_input_ids[: self.tokenizer.model_max_length] for _input_ids in input_ids]
labels = [_labels[: self.tokenizer.model_max_length] for _labels in labels]
if self.tokenizer.pad_token_id is None: # "<|endoftext|>"151643
self.tokenizer.pad_token_id = 0 # This gets the best result. Don't know why.
input_ids = self.pad_sequence(input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id)
labels = self.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
batch = dict(input_ids=input_ids, labels=labels.long() if labels.dtype == torch.int32 else labels, attention_mask=input_ids.ne(self.tokenizer.pad_token_id))
if "image" in instances[0]:
images = [instance["image"] for instance in instances]
batch["image_sizes"] = [im[1] for im_list in images for im in im_list]
batch["modalities"] = [im[2] for im_list in images for im in im_list]
images = [im[0] for im_list in images for im in im_list]
batch["images"] = images
target_images = [instance["target_image"][0] for instance in instances]
target_images = torch.stack(target_images, dim=0) if target_images else None # [B, 3, 512, 512]
batch["target_images"] = target_images
if "detailed_condition" in instances[0]:
detailed_conditions = [instance["detailed_condition"][0] for instance in instances] # target_image
detailed_conditions = torch.stack(detailed_conditions, dim=0) if detailed_conditions else None # [B, 3, 512, 512]
batch["detailed_conditions"] = detailed_conditions
if "prompt" in instances[0]:
batch["prompts"] = [instance["prompt"] for instance in instances]
return batch
def get_dataset_cls(name):
if name == 'mix':
dataset_cls = LazySupervisedMixDataset
elif name == 'restore':
dataset_cls = LazySupervisedRestoreDataset
else:
raise ValueError(f'Unknown dataset class {name}')
return dataset_cls
def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizer, data_args) -> Dict:
"""Make dataset and collator for supervised fine-tuning."""
dataset_cls = get_dataset_cls(data_args.dataset_cls)
train_dataset = dataset_cls(tokenizer=tokenizer, data_path=data_args.data_path, data_args=data_args)
data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)