import os import numpy as np import torch import random from torchvision.datasets.folder import default_loader from diffusion.data.datasets.InternalData import InternalData, InternalDataSigma from diffusion.data.builder import get_data_path, DATASETS from diffusion.utils.logger import get_root_logger import torchvision.transforms as T from torchvision.transforms.functional import InterpolationMode from diffusion.data.datasets.utils import * def get_closest_ratio(height: float, width: float, ratios: dict): aspect_ratio = height / width closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - aspect_ratio)) return ratios[closest_ratio], float(closest_ratio) @DATASETS.register_module() class InternalDataMS(InternalData): def __init__(self, root, image_list_json='data_info.json', transform=None, resolution=256, sample_subset=None, load_vae_feat=False, input_size=32, patch_size=2, mask_ratio=0.0, mask_type='null', load_mask_index=False, real_prompt_ratio=1.0, max_length=120, config=None, **kwargs): self.root = get_data_path(root) self.transform = transform self.load_vae_feat = load_vae_feat self.ori_imgs_nums = 0 self.resolution = resolution self.N = int(resolution // (input_size // patch_size)) self.mask_ratio = mask_ratio self.load_mask_index = load_mask_index self.mask_type = mask_type self.real_prompt_ratio = real_prompt_ratio self.max_lenth = max_length self.base_size = int(kwargs['aspect_ratio_type'].split('_')[-1]) self.aspect_ratio = eval(kwargs.pop('aspect_ratio_type')) # base aspect ratio self.meta_data_clean = [] self.img_samples = [] self.txt_feat_samples = [] self.vae_feat_samples = [] self.mask_index_samples = [] self.ratio_index = {} self.ratio_nums = {} # self.weight_dtype = torch.float16 if self.real_prompt_ratio > 0 else torch.float32 for k, v in self.aspect_ratio.items(): self.ratio_index[float(k)] = [] # used for self.getitem self.ratio_nums[float(k)] = 0 # used for batch-sampler image_list_json = image_list_json if isinstance(image_list_json, list) else [image_list_json] for json_file in image_list_json: meta_data = self.load_json(os.path.join(self.root, json_file)) self.ori_imgs_nums += len(meta_data) meta_data_clean = [item for item in meta_data if item['ratio'] <= 4] self.meta_data_clean.extend(meta_data_clean) self.img_samples.extend([os.path.join(self.root.replace('InternData', "InternImgs"), item['path']) for item in meta_data_clean]) self.txt_feat_samples.extend([os.path.join(self.root, 'caption_features', '_'.join(item['path'].rsplit('/', 1)).replace('.png', '.npz')) for item in meta_data_clean]) self.vae_feat_samples.extend([os.path.join(self.root, f'img_vae_fatures_{resolution}_multiscale/ms', '_'.join(item['path'].rsplit('/', 1)).replace('.png', '.npy')) for item in meta_data_clean]) # Set loader and extensions if load_vae_feat: self.transform = None self.loader = self.vae_feat_loader else: self.loader = default_loader if sample_subset is not None: self.sample_subset(sample_subset) # sample dataset for local debug # scan the dataset for ratio static for i, info in enumerate(self.meta_data_clean[:len(self.meta_data_clean)//3]): ori_h, ori_w = info['height'], info['width'] closest_size, closest_ratio = get_closest_ratio(ori_h, ori_w, self.aspect_ratio) self.ratio_nums[closest_ratio] += 1 if len(self.ratio_index[closest_ratio]) == 0: self.ratio_index[closest_ratio].append(i) # print(self.ratio_nums) logger = get_root_logger() if config is None else get_root_logger(os.path.join(config.work_dir, 'train_log.log')) logger.info(f"T5 max token length: {self.max_lenth}") def getdata(self, index): img_path = self.img_samples[index] npz_path = self.txt_feat_samples[index] npy_path = self.vae_feat_samples[index] ori_h, ori_w = self.meta_data_clean[index]['height'], self.meta_data_clean[index]['width'] # Calculate the closest aspect ratio and resize & crop image[w, h] closest_size, closest_ratio = get_closest_ratio(ori_h, ori_w, self.aspect_ratio) closest_size = list(map(lambda x: int(x), closest_size)) self.closest_ratio = closest_ratio if self.load_vae_feat: try: img = self.loader(npy_path) if index not in self.ratio_index[closest_ratio]: self.ratio_index[closest_ratio].append(index) except Exception: index = random.choice(self.ratio_index[closest_ratio]) return self.getdata(index) h, w = (img.shape[1], img.shape[2]) assert h, w == (ori_h//8, ori_w//8) else: img = self.loader(img_path) h, w = (img.size[1], img.size[0]) assert h, w == (ori_h, ori_w) data_info = {'img_hw': torch.tensor([ori_h, ori_w], dtype=torch.float32)} data_info['aspect_ratio'] = closest_ratio data_info["mask_type"] = self.mask_type txt_info = np.load(npz_path) txt_fea = torch.from_numpy(txt_info['caption_feature']) attention_mask = torch.ones(1, 1, txt_fea.shape[1]) if 'attention_mask' in txt_info.keys(): attention_mask = torch.from_numpy(txt_info['attention_mask'])[None] if not self.load_vae_feat: if closest_size[0] / ori_h > closest_size[1] / ori_w: resize_size = closest_size[0], int(ori_w * closest_size[0] / ori_h) else: resize_size = int(ori_h * closest_size[1] / ori_w), closest_size[1] self.transform = T.Compose([ T.Lambda(lambda img: img.convert('RGB')), T.Resize(resize_size, interpolation=InterpolationMode.BICUBIC), # Image.BICUBIC T.CenterCrop(closest_size), T.ToTensor(), T.Normalize([.5], [.5]), ]) if self.transform: img = self.transform(img) return img, txt_fea, attention_mask, data_info def __getitem__(self, idx): for _ in range(20): try: return self.getdata(idx) except Exception as e: print(f"Error details: {str(e)}") idx = random.choice(self.ratio_index[self.closest_ratio]) raise RuntimeError('Too many bad data.') @DATASETS.register_module() class InternalDataMSSigma(InternalDataSigma): def __init__(self, root, image_list_json='data_info.json', transform=None, resolution=256, sample_subset=None, load_vae_feat=False, load_t5_feat=False, input_size=32, patch_size=2, mask_ratio=0.0, mask_type='null', load_mask_index=False, real_prompt_ratio=1.0, max_length=300, config=None, **kwargs): self.root = get_data_path(root) self.transform = transform self.load_vae_feat = load_vae_feat self.load_t5_feat = load_t5_feat self.ori_imgs_nums = 0 self.resolution = resolution self.N = int(resolution // (input_size // patch_size)) self.mask_ratio = mask_ratio self.load_mask_index = load_mask_index self.mask_type = mask_type self.real_prompt_ratio = real_prompt_ratio self.max_lenth = max_length self.base_size = int(kwargs['aspect_ratio_type'].split('_')[-1]) self.aspect_ratio = eval(kwargs.pop('aspect_ratio_type')) # base aspect ratio self.meta_data_clean = [] self.img_samples = [] self.txt_samples = [] self.sharegpt4v_txt_samples = [] self.txt_feat_samples = [] self.vae_feat_samples = [] self.mask_index_samples = [] self.ratio_index = {} self.ratio_nums = {} self.gpt4v_txt_feat_samples = [] self.weight_dtype = torch.float16 if self.real_prompt_ratio > 0 else torch.float32 self.interpolate_model = InterpolationMode.BICUBIC if self.aspect_ratio in [ASPECT_RATIO_2048, ASPECT_RATIO_2880]: self.interpolate_model = InterpolationMode.LANCZOS suffix = '' for k, v in self.aspect_ratio.items(): self.ratio_index[float(k)] = [] # used for self.getitem self.ratio_nums[float(k)] = 0 # used for batch-sampler logger = get_root_logger() if config is None else get_root_logger(os.path.join(config.work_dir, 'train_log.log')) logger.info(f"T5 max token length: {self.max_lenth}") logger.info(f"ratio of real user prompt: {self.real_prompt_ratio}") image_list_json = image_list_json if isinstance(image_list_json, list) else [image_list_json] for json_file in image_list_json: meta_data = self.load_json(os.path.join(self.root, json_file)) logger.info(f"{json_file} data volume: {len(meta_data)}") self.ori_imgs_nums += len(meta_data) meta_data_clean = [item for item in meta_data if item['ratio'] <= 4.5] self.meta_data_clean.extend(meta_data_clean) self.img_samples.extend([ os.path.join(self.root.replace('InternData'+suffix, 'InternImgs'), item['path']) for item in meta_data_clean ]) self.txt_samples.extend([item['prompt'] for item in meta_data_clean]) self.sharegpt4v_txt_samples.extend([item['sharegpt4v'] if 'sharegpt4v' in item else '' for item in meta_data_clean]) self.txt_feat_samples.extend([ os.path.join( self.root, 'caption_features_new', '_'.join(item['path'].rsplit('/', 1)).replace('.png', '.npz') ) for item in meta_data_clean ]) self.gpt4v_txt_feat_samples.extend([ os.path.join( self.root, 'sharegpt4v_caption_features_new', '_'.join(item['path'].rsplit('/', 1)).replace('.png', '.npz') ) for item in meta_data_clean ]) self.vae_feat_samples.extend( [ os.path.join( self.root + suffix, f'img_sdxl_vae_features_{resolution}resolution_ms_new', '_'.join(item['path'].rsplit('/', 1)).replace('.png', '.npy') ) for item in meta_data_clean ]) if self.real_prompt_ratio < 1: assert len(self.sharegpt4v_txt_samples[0]) != 0 # Set loader and extensions if load_vae_feat: self.transform = None self.loader = self.vae_feat_loader else: self.loader = default_loader if sample_subset is not None: self.sample_subset(sample_subset) # sample dataset for local debug # scan the dataset for ratio static for i, info in enumerate(self.meta_data_clean[:len(self.meta_data_clean)//3]): ori_h, ori_w = info['height'], info['width'] closest_size, closest_ratio = get_closest_ratio(ori_h, ori_w, self.aspect_ratio) self.ratio_nums[closest_ratio] += 1 if len(self.ratio_index[closest_ratio]) == 0: self.ratio_index[closest_ratio].append(i) def getdata(self, index): img_path = self.img_samples[index] real_prompt = random.random() < self.real_prompt_ratio npz_path = self.txt_feat_samples[index] if real_prompt else self.gpt4v_txt_feat_samples[index] txt = self.txt_samples[index] if real_prompt else self.sharegpt4v_txt_samples[index] npy_path = self.vae_feat_samples[index] data_info = {} ori_h, ori_w = self.meta_data_clean[index]['height'], self.meta_data_clean[index]['width'] # Calculate the closest aspect ratio and resize & crop image[w, h] closest_size, closest_ratio = get_closest_ratio(ori_h, ori_w, self.aspect_ratio) closest_size = list(map(lambda x: int(x), closest_size)) self.closest_ratio = closest_ratio if self.load_vae_feat: img = self.loader(npy_path) if index not in self.ratio_index[closest_ratio]: self.ratio_index[closest_ratio].append(index) h, w = (img.shape[1], img.shape[2]) assert h, w == (ori_h//8, ori_w//8) else: img = self.loader(img_path) h, w = (img.size[1], img.size[0]) assert h, w == (ori_h, ori_w) data_info['img_hw'] = torch.tensor([ori_h, ori_w], dtype=torch.float32) data_info['aspect_ratio'] = closest_ratio data_info["mask_type"] = self.mask_type attention_mask = torch.ones(1, 1, self.max_lenth) if self.load_t5_feat: txt_info = np.load(npz_path) txt_fea = torch.from_numpy(txt_info['caption_feature']) if 'attention_mask' in txt_info.keys(): attention_mask = torch.from_numpy(txt_info['attention_mask'])[None] if txt_fea.shape[1] != self.max_lenth: txt_fea = torch.cat([txt_fea, txt_fea[:, -1:].repeat(1, self.max_lenth-txt_fea.shape[1], 1)], dim=1).to(self.weight_dtype) attention_mask = torch.cat([attention_mask, torch.zeros(1, 1, self.max_lenth-attention_mask.shape[-1])], dim=-1) else: txt_fea = txt if not self.load_vae_feat: if closest_size[0] / ori_h > closest_size[1] / ori_w: resize_size = closest_size[0], int(ori_w * closest_size[0] / ori_h) else: resize_size = int(ori_h * closest_size[1] / ori_w), closest_size[1] self.transform = T.Compose([ T.Lambda(lambda img: img.convert('RGB')), T.Resize(resize_size, interpolation=self.interpolate_model), # Image.BICUBIC T.CenterCrop(closest_size), T.ToTensor(), T.Normalize([.5], [.5]), ]) if self.transform: img = self.transform(img) return img, txt_fea, attention_mask.to(torch.int16), data_info def __getitem__(self, idx): for _ in range(20): try: data = self.getdata(idx) return data except Exception as e: print(f"Error details: {str(e)}") idx = random.choice(self.ratio_index[self.closest_ratio]) raise RuntimeError('Too many bad data.')