from email.mime import image import os import random import typing from pathlib import Path from typing import Optional import subprocess import datasets import torch from numpy import pad from PIL import Image, ImageFile from tensordict import TensorDict from torchvision import transforms from decoupled_utils import get_world_size import time import re import shutil from constants import UNIDISC_DIR from decoupled_utils import barrier, get_rank, gprint, is_local_main_process, is_main_process, is_torch_cuda_available, is_torch_xla_available, rprint from models.datasets.webdataset_utils import get_data import hashlib from decoupled_utils import sanitize_filename from omegaconf import OmegaConf, read_write from models.datasets.misc_image_datasets import * from copy import deepcopy from datasets import Dataset, DatasetDict import numpy as np from PIL import Image import json ImageFile.LOAD_TRUNCATED_IMAGES = True import torch from torch.utils.data import Subset def split_dataset(dataset, n: int, m: int): # Ensure m is valid if m < 0 or m >= n: raise ValueError(f"m must be between 0 and {n-1}, but got {m}.") # Calculate the size of each subset total_len = len(dataset) subset_size = total_len // n remainder = total_len % n # Calculate the start and end index of the m-th subset start_idx = m * subset_size + min(m, remainder) end_idx = start_idx + subset_size + (1 if m < remainder else 0) # Return the m-th subset indices = list(range(start_idx, end_idx)) if isinstance(dataset, torch.utils.data.Dataset): return Subset(dataset, indices) else: return dataset[slice(start_idx, end_idx)] def get_webdataset_indexed(config, tokenizer, transform, cond_transform, n_samples, name, should_tokenize=False): should_tokenize = ("tokenize" in name) or should_tokenize import wids # You need to use the custom sampler!! custom_ignore_func_dict = { "pixelprose": lambda x: len(x[".txt"]) > 400, } valid_func = None for k in custom_ignore_func_dict.keys(): if k in name: valid_func = custom_ignore_func_dict[k] break from dataloader import tokenize_text def process(x, idx): data = {} if "mmc4" in name: print(x['.json']['image_info'][0]) breakpoint() img = x[".jpg"].convert("RGB") data["is_valid"] = True if valid_func is not None and valid_func(x) is False: print(f"Invalid") data["is_valid"] = False data["img"] = transform(img) if cond_transform is not None: data["cond_img"] = cond_transform(x[".jpg"].convert("RGB")) if data["img"].shape[0] != 3: raise Exception(f"Image shape: {data['img'].shape}, {x['.jpg'].size}, {x['.jpg'].mode}") if "pixelprose" in name: before = x[".txt"] x[".txt"] = re.sub(r"This image displays.*?(?=[a-zA-Z0-9])", "", x[".txt"]) if abs(len(before) - len(x[".txt"])) > 100: data["is_valid"] = False if not "imagenet" in name: if should_tokenize: tokens = tokenize_text(tokenizer, config.data.block_size, x[".txt"]) data["input_ids"] = tokens["input_ids"] data["attention_mask"] = tokens["attention_mask"].float() else: data[".txt"] = x[".txt"] data["idx"] = idx return data disable_split = False if isinstance(config.data.raw_data_dir, str) and '*' in config.data.raw_data_dir: import glob index_path = sorted(glob.glob(config.data.raw_data_dir)) if not index_path: raise ValueError(f"No files found matching the pattern: {config.data.raw_data_dir:}") print(f"Expanded glob pattern to {len(index_path)} files") if os.getenv("SLURM_ARRAY_TASK_COUNT", None) is not None: index_path = split_dataset(index_path, int(os.getenv("SLURM_ARRAY_TASK_COUNT")), int(os.getenv("SLURM_ARRAY_TASK_ID"))) print(f"After splitting, dataset is length {len(index_path)}") shards = [] for shard in index_path: shards.append({"url": shard, "nsamples": wids.wids.compute_num_samples(shard)}) print(f"Shard: {shard}") index_path = shards disable_split = True elif Path(config.data.raw_data_dir).is_file(): index_path = config.data.raw_data_dir else: default_path = Path(config.data.raw_data_dir) / "index.json" shard_path = Path(config.data.raw_data_dir) / "shardindex.json" index_path = str(default_path if default_path.exists() else shard_path) assert getattr(config.data, "shard_list_path", None) is None, "shard_list_path is deprecated, use raw_data_dir instead" dataset = wids.ShardListDataset(index_path) # lru_size=20 dataset = CustomTransformDataset(dataset, process) if n_samples is not None: from torch.utils.data import Subset indices = torch.randperm(len(dataset))[:n_samples] dataset = Subset(dataset, indices) if config.data.split_dataset and not disable_split: gprint(f"Original dataset was length {len(dataset)}") dataset = split_dataset(dataset, int(os.getenv("SLURM_ARRAY_TASK_COUNT")), int(os.getenv("SLURM_ARRAY_TASK_ID"))) gprint(f"After splitting, dataset is length {len(dataset)}") return dataset def _copy_data(src_path, dst_path, use_rsync=True): dst_path.mkdir(parents=True, exist_ok=True) if use_rsync: rprint(f"Rsyncing data from {src_path} to {dst_path}") rsync_command = ["rsync", "-av", str(src_path) + "/", str(dst_path) + "/"] try: result = subprocess.run(rsync_command, check=True, capture_output=True, text=True) rprint(f"Rsync output: {result.stdout}") rprint(f"Successfully rsynced data from {src_path} to {dst_path}") except subprocess.CalledProcessError as e: rprint(f"Rsync failed: {e}") rprint(f"Rsync stderr: {e.stderr}") raise else: rprint(f"Copying tensordict from {src_path} to {dst_path}") shutil.copytree(src_path, dst_path) rprint(f"Copied tensordict from {src_path} to {dst_path}") def copy_data(shm_path, src_path, dst_path): shm_path.mkdir(parents=True, exist_ok=True) use_rsync = True if not dst_path.exists() or use_rsync: _copy_data(src_path, dst_path, use_rsync=use_rsync) else: src_files = sum(1 for _ in src_path.rglob('*')) dst_files = sum(1 for _ in dst_path.rglob('*')) src_size = sum(f.stat().st_size for f in src_path.rglob('*') if f.is_file()) dst_size = sum(f.stat().st_size for f in dst_path.rglob('*') if f.is_file()) size_diff_percent = abs(src_size - dst_size) / max(src_size, dst_size) * 100 if src_files != dst_files or size_diff_percent > 10: rprint(f"Src files: {src_files}, Dst files: {dst_files} contain different number of files, {src_size} {dst_size}, size difference {size_diff_percent}, Deleting {dst_path}") shutil.rmtree(dst_path) rprint(f"Deleted {dst_path}, copying from {src_path}") _copy_data(src_path, dst_path, use_rsync=False) rprint(f"Deleted and re-copied tensordict from {src_path} to {dst_path}") else: rprint(f"Tensordict already exists at {dst_path}, loading from there") def get_tensordict(config, path, dataset_idx, dataset_name=None): parquet_files = list(Path(path).glob('*.arrow')) if parquet_files: # Does not load into memory by default from datasets import load_from_disk dataset = load_from_disk(path) rprint(f"Loaded {len(dataset)} samples from {path} as parquet") return dataset if getattr(config.data, "force_dummy_tensordict", False): return get_dummy_tensordict(config, 1000000, dataset_idx=dataset_idx) if config.data.move_tensordict_to_shm: assert config.data.keep_tensordict_on_disk is True shm_path = Path(getattr(config.data, "shm_path", Path("/dev/shm") / Path.home().name)) src_path = Path(path) dst_path = shm_path / (dataset_name if dataset_name is not None else src_path.name) if getattr(config.data, "skip_copy_tensordict_to_shm", False): gprint(f"Skipping copy of tensordict to SHM") elif is_torch_xla_available(): if is_main_process(): copy_data(shm_path, src_path, dst_path) barrier() if not is_main_process(): import time from torch_xla._internal import tpu host_ip = tpu.get_worker_ips()[0] file_dst_path = Path(shm_path) src_path_on_host = f"aswerdlow@{host_ip}:{file_dst_path}" gprint(f"Copying data from {src_path_on_host} to {file_dst_path}") file_dst_path.mkdir(parents=True, exist_ok=True) max_retries = 5 retry_delay = 15 for attempt in range(max_retries): try: gprint(f"After main copy, rsyncing data from {src_path_on_host} to {file_dst_path}") command = f'bash {(UNIDISC_DIR / "scripts/rsync_data.sh").resolve()} {src_path_on_host}/ {file_dst_path}/' os.environ.pop('SSH_AUTH_SOCK', None) # Breaks without this gprint(command) subprocess.run(command, shell=True, check=True) gprint(f"Successfully rsynced data from {src_path_on_host} to {file_dst_path}") break except subprocess.CalledProcessError as e: if attempt < max_retries - 1: gprint(f"Rsync attempt {attempt + 1} failed. Retrying in {retry_delay} seconds..., {e}") time.sleep(retry_delay) retry_delay *= 2 else: raise RuntimeError(f"Failed to rsync data after {max_retries} attempts: {e}") gprint(f"Finished rsyncing data from {src_path_on_host} to {file_dst_path}") barrier() else: if is_local_main_process(): copy_data(shm_path, src_path, dst_path) # For now we assume we are on SPMD and there is only one process per worker [node] if not is_torch_xla_available(): barrier() else: dst_path = Path(path) path = dst_path data = TensorDict.load_memmap(path) if config.data.keep_tensordict_on_disk: rprint(f"Keeping tensordict on disk at {path}") else: data = data.clone() # Move to CPU memory rprint(f"Loaded {len(data)} samples from {path}") return data def get_dummy_tensordict(config, dataset_size, txt_length=None, img_length=None, dataset_idx=0): if img_length is None: img_length = config.model.img_length if txt_length is None: txt_length = config.model.txt_length return TensorDict( { "input_ids": torch.ones(dataset_size, config.model.length, dtype=torch.int32), "attention_mask": torch.ones(dataset_size, config.model.length, dtype=torch.bool), "img_input_ids": torch.ones(dataset_size, img_length, dtype=torch.int16), "txt_input_ids": torch.ones(dataset_size, txt_length, dtype=torch.int32), "txt_attention_mask": torch.ones(dataset_size, txt_length, dtype=torch.bool), "idx": torch.arange(dataset_size, dtype=torch.int32).view(-1, 1), "dataset_idx": torch.full((dataset_size,), fill_value=dataset_idx, dtype=torch.int32), "write_flag": torch.zeros(dataset_size, 1, dtype=torch.bool), }, batch_size=[dataset_size], ) def get_token_dataset(config, name, is_train, n_samples, n_duplicate, tokenizer): assert getattr(config.data, "token_data_dir", None) is None, "token_data_dir is deprecated, use data_dir_train and data_dir_val instead" if "dummy" in name: return get_dummy_tensordict(config, n_samples if n_samples is not None else 100000) data_key = ( config.data.data_dir_train if is_train else (config.data.data_dir_val if config.data.data_dir_val is not None else config.data.data_dir_train) ) image_datasets_key = getattr(config.data, "image_data_train", None) if is_train else getattr(config.data, "image_data_val", None) if config.data.use_weighted_tensordict_sampler: _dataset_cls = MultipleTensorDictDataset _datasets = [get_tensordict(config, x['dir'], dataset_idx=i, dataset_name=x['name']) for i, x in enumerate(data_key)] _weights = [x['weight'] for x in data_key] _dataset_names = [x['name'] for x in data_key] _kwargs = dict() _kwargs["config"] = config _kwargs["tokenizer"] = tokenizer if any(not isinstance(x, TensorDict) for x in _datasets): _kwargs["returns_parquet"] = True elif getattr(config.data, "add_text_to_weighted_sampler", False): from datasets import load_dataset, interleave_datasets rprint("Loading smollm datasets") ds1 = load_dataset("HuggingFaceTB/smollm-corpus", "cosmopedia-v2", split="train", cache_dir=config.data.cache_dir, streaming=True) ds2 = load_dataset("HuggingFaceTB/smollm-corpus", "fineweb-edu-dedup", split="train", cache_dir=config.data.cache_dir, streaming=True) # DKYoon/SlimPajama-6B, "cerebras/SlimPajama-627B" ds3 = load_dataset("DKYoon/SlimPajama-6B", split="train", cache_dir=config.data.cache_dir, streaming=True) ds4 = load_dataset("bigcode/starcoderdata", split="train", cache_dir=config.data.cache_dir, streaming=True) rprint(f"Finished loading datasets") if getattr(config.data, "code_only", False): _dataset = ds4 else: _dataset = interleave_datasets([ds1, ds2, ds3, ds4], probabilities=[0.3, 0.3, 0.2, 0.2], seed=config.seed) rprint(f"Finished interleaving datasets") _datasets.append(_dataset) _weights.append(1) _dataset_names.append("SlimPajama-627B") _kwargs["returns_tokenized_text"] = True rprint(f"Finished creating dataset") elif image_datasets_key is not None: returns_raw_images = False tokenize_vqvae_in_dataloader = False allow_label = False for key in image_datasets_key: _key = OmegaConf.to_object(key) if _key.get("raw_images", False) or config.data.force_raw_images_in_multiple_tensordict: rprint(f"WARNING!!! Using raw images") returns_raw_images = True if _key.get("tokenize_vqvae_in_dataloader", False): tokenize_vqvae_in_dataloader = True if _key.get("allow_label", False): rprint(f"WARNING!!! Using allow_label") allow_label = True if config.data.force_raw_images_in_multiple_tensordict: tokenize_vqvae_in_dataloader = False _key["tokenize_vqvae_in_dataloader"] = False _key["disable_text_modality"] = True image_config = OmegaConf.merge(deepcopy(config), { "data": { **{k:v for k,v in _key.items() if k not in {"dir", "weight", "name", "raw_images"}} }, } ) image_dataset = get_image_dataset( mode="train" if is_train else "val", config=image_config, tokenizer=tokenizer, allow_aug=is_train, force_aug=False, name=key["train"] if is_train else key["val"], ) _datasets.append(image_dataset) _weights.append(key["weight"]) _dataset_names.append(key["name"]) _kwargs["returns_raw_images"] = returns_raw_images _kwargs["returns_tokenize_vqvae_in_dataloader"] = tokenize_vqvae_in_dataloader _kwargs["allow_label"] = allow_label if n_samples is not None: if getattr(config.data, "force_no_shuffle_tensordict", False): _datasets = [data[:n_samples] for data in _datasets] else: _datasets = [data[torch.randperm(len(data), generator=torch.Generator().manual_seed(config.seed))[:n_samples]] for data in _datasets] rprint(f"Sampled {n_samples} samples from {len(_datasets)}, is_train: {is_train}.") data = _dataset_cls(datasets=_datasets, weights=_weights, dataset_names=_dataset_names, **_kwargs) else: data = get_tensordict(config, data_key, 0) if n_samples is not None: if getattr(config.data, "force_no_shuffle_tensordict", False): indices = list(range(n_samples)) else: indices = torch.randperm(len(data), generator=torch.Generator().manual_seed(config.seed))[:n_samples] data = data[indices] rprint(f"Sampled {n_samples} samples from {len(data)}, is_train: {is_train}. First 2 indices: {indices[:2]}") if n_duplicate is not None: data = torch.cat([data for _ in range(n_duplicate)], dim=0) rprint(f"Duplicated {n_duplicate} times, now {len(data)} samples") return data class UnpairedDatasetWrapper(torch.utils.data.Dataset): def __init__(self, img_dataset, txt_dataset, shuffle=True): self.img_dataset = img_dataset self.txt_dataset = txt_dataset self.shuffle = shuffle def __len__(self): if self.shuffle: return min(len(self.img_dataset), len(self.txt_dataset)) else: return max(len(self.img_dataset), len(self.txt_dataset)) def __getitem__(self, idx): while True: try: if self.shuffle: img_idx = torch.randint(0, len(self.img_dataset), (1,)).item() txt_idx = torch.randint(0, len(self.txt_dataset), (1,)).item() else: txt_idx = idx img_idx = idx % len(self.img_dataset) return dict(**self.img_dataset[img_idx], **self.txt_dataset[txt_idx]) except Exception as e: gprint(e) import traceback traceback.print_exc() idx = (idx + 1) % len(self) def get_unpaired_dataset(config=None, tokenizer=None, mode="train", **kwargs): image_dataset = get_image_dataset(config=config, mode=mode, tokenizer=tokenizer, **kwargs) from models.datasets.text_datasets import get_text_dataset text_dataset = get_text_dataset( dataset_name=getattr(config.data, "txt_train", "text8") if mode == "train" else getattr(config.data, "txt_val", "text8"), tokenizer=tokenizer, mode="test" if (mode == "validation" and getattr(config.data, "txt_val", "text8") == "lm1b") else mode, wrap=config.data.wrap, block_size=config.model.txt_length, # Intentional cache_dir=config.data.cache_dir, num_proc=config.data.num_proc, streaming=config.data.streaming, ) return UnpairedDatasetWrapper(image_dataset, text_dataset, shuffle=getattr(config.data, "force_disable_shuffle", False) is False) def get_transform(resolution, orig_mode, allow_aug, force_aug, aggressive_aug=False): if orig_mode == "train" and (allow_aug or force_aug): if aggressive_aug: rprint("Using aggressive augmentations") transform = transforms.Compose( [ transforms.RandomResizedCrop((resolution, resolution), scale=(0.8, 1.0), ratio=(0.97, 1.03)), transforms.RandomHorizontalFlip(1.0 if force_aug else 0.5), transforms.ToTensor(), ] ) else: transform = transforms.Compose( [ transforms.Resize(resolution, interpolation=transforms.InterpolationMode.BILINEAR), transforms.RandomCrop((resolution, resolution)), transforms.RandomHorizontalFlip(1.0 if force_aug else 0.5), transforms.ToTensor(), ] ) else: transform = transforms.Compose( [ transforms.Resize(resolution, interpolation=transforms.InterpolationMode.BILINEAR), transforms.CenterCrop((resolution, resolution)), transforms.ToTensor(), ] ) return transform def load_vqvae_from_cache(config, full_cache_path): global_cache_parent = os.environ.get("DIFFUSION_DATA_DIR", None) if global_cache_parent is not None: global_full_cache_path = Path(global_cache_parent) / full_cache_path.relative_to(Path(config.data.cache_dir)) gprint(f"Checking global cache path: {global_full_cache_path}") if global_full_cache_path.exists() and len(list(global_full_cache_path.iterdir())) > 0: gprint(f"Loading data from global cache: {global_full_cache_path}") full_cache_path = global_full_cache_path if not (full_cache_path.exists() and len(list(full_cache_path.iterdir())) > 0): gprint(f"Cache path {full_cache_path} does not exist or is empty") return None gprint(f"Loading data from: {full_cache_path}, found {len(list(full_cache_path.iterdir()))} shards") ret = [] kwargs = dict() if config.data.keep_hf_dataset_in_memory: kwargs["keep_in_memory"] = True if config.loader.num_workers > 0: for _ in range(5): gprint(f"WARNING!!!! Keeping dataset in memory and num_workers > 0, this will cause excessive memory usage") else: gprint(f"Loading datasets into memory") for folder in full_cache_path.iterdir(): if folder.is_dir(): ret.append(datasets.load_from_disk(folder, **kwargs)) ret = datasets.concatenate_datasets(ret).with_format("torch") gprint(f"Loaded data from cache: {full_cache_path} with {len(ret)} samples") return ret def get_vqvae_dataloader(config, name, split): cache_key = f'vqvae_tokenized_{name}_{split}_{config.data.resolution}' vae_ckpt_hash = "" if hasattr(config.model, "use_custom_vae_ckpt") and config.model.use_custom_vae_ckpt: vae_ckpt_hash = hashlib.md5(str(Path(config.model.use_custom_vae_ckpt).name).encode()).hexdigest()[:8] cache_key += f"_{vae_ckpt_hash}" if hasattr(config.model, "vae_type") and config.model.vae_type != "VQ-16": cache_key += f"_{config.model.vae_type}" if getattr(config.data, "vqvae_cache_suffix", None) is not None: cache_key += f"_{config.data.vqvae_cache_suffix}" cache_dir = config.data.cache_dir full_cache_path = Path(cache_dir) / "tokens" / sanitize_filename(cache_key) return full_cache_path def get_image_dataset(mode, config, tokenizer, allow_aug=True, force_aug=False, name=None, **kwargs): rprint(f"Getting image dataset with mode {mode}") if getattr(config.data, "tokenizers_parallelism", None) is not None: rprint(f"Setting tokenizers parallelism to {config.data.tokenizers_parallelism}") os.environ["TOKENIZERS_PARALLELISM"] = "false" if config.data.tokenizers_parallelism is False else "true" resolution = config.data.resolution name = name or config.data.train streaming = config.data.streaming precache = config.data.precache dynamic = streaming or precache is False orig_mode = mode block_size = getattr(config.data, "block_size", 1024) is_train = orig_mode == "train" n_duplicate_train = getattr(config.data, "n_duplicate_train", None) n_duplicate_val = getattr(config.data, "n_duplicate_val", None) n_duplicate = n_duplicate_train if is_train else n_duplicate_val n_val_samples = getattr(config.data, "n_val_samples", None) n_train_samples = getattr(config.data, "n_train_samples", None) n_samples = n_train_samples if is_train else n_val_samples raw_data_dir = getattr(config.data, "raw_data_dir", getattr(config.data, "data_dir", None)) rprint(f"Data dir is {raw_data_dir}") unified_model = getattr(config.model, "unified_model", False) and getattr(config.data, "unpaired", False) is False cond_resolution = getattr(config.data, "cond_resolution", None) if "sora" in name: return get_sora_dataset(config=config, tokenizer=tokenizer, **kwargs) elif "tokens" in name: print(f"Loading token dataset {name}") assert config.data.use_token_dataset, "data.use_token_dataset must be true to load token datasets" return get_token_dataset(config, name, is_train, n_samples, n_duplicate, tokenizer) dataset_splits = { "cassiekang/cub200_dataset": ( "train" if ((orig_mode == "train" and n_train_samples is not None) or (orig_mode != "train" and n_val_samples is not None)) else "train+test" ), "nlphuji/flickr30k": "test", "richwardle/reduced-imagenet": "train", "tglcourse/lsun_church_train": "train" if is_train else "test", "pixparse/cc12m-wds": "train", "imagenet": "train" if is_train else "val", "imagefolder": "train" if is_train else "validation", "ILSVRC/imagenet-1k": "train" if is_train else "validation", "pouya-haghi/imagenet-subset": "validation", "laion/clevr-webdataset": "train" if is_train else "validation", "pcuenq/lsun-bedrooms": "train" if is_train else "test", "facebook/winoground": "test", "sayakpaul/coco-30-val-2014": "train" } split = dataset_splits[name] if name in dataset_splits else "train" if n_samples is not None: split = f"{split}[:{n_samples}]" extra_kwargs = dict() cache_dir = Path(config.data.cache_dir) cache_dir.mkdir(parents=True, exist_ok=True) if "HF_HUB_DATASETS_TOKEN" in os.environ: extra_kwargs["token"] = os.environ["HF_HUB_DATASETS_TOKEN"] if name == "mmc4" or name == "cambrian": from unidisc.tokenizers.tokenize_interleaved import JsonlDataset dataset = JsonlDataset(glob_pattern=config.data.raw_data_dir) if n_samples is not None: from torch.utils.data import Subset indices = list(range(len(dataset)))[:n_samples] dataset = Subset(dataset, indices) if config.data.split_dataset: if getattr(config.data, "split_dataset_total_count", None) is not None and \ getattr(config.data, "split_dataset_cur_idx", None) is not None: gprint(f"Splitting dataset into {config.data.split_dataset_total_count} shards, original length {len(dataset)}") dataset = split_dataset(dataset, config.data.split_dataset_total_count, config.data.split_dataset_cur_idx) gprint(f"Original dataset was length {len(dataset)}") total_count, cur_idx = int(os.getenv("SLURM_ARRAY_TASK_COUNT")), int(os.getenv("SLURM_ARRAY_TASK_ID")) dataset = split_dataset(dataset, total_count, cur_idx) gprint(f"After splitting, dataset is length {len(dataset)}") return dataset if name == "imagefolder": from datasets.data_files import DataFilesDict with open(config.data.train_data_dir, "r") as f: train_txt = [f"{config.data.data_dir}/{line.strip()}" for line in f.readlines()] with open(config.data.val_data_dir, "r") as f: val_txt = [f"{config.data.data_dir}/{line.strip()}" for line in f.readlines()] data_files = DataFilesDict({"train": train_txt, "validation": val_txt}) extra_kwargs["data_files"] = data_files if config.data.tokenize_vqvae_in_dataloader and not getattr(config.data, "allow_aug_vqvae_dataloader", False): rprint(f"WARNING!!!! Disabling augmentations for VQVAE dataloader") allow_aug = False force_aug = False transform = get_transform(resolution, orig_mode, allow_aug, force_aug, getattr(config.data, "aggressive_aug", False)) if cond_resolution is not None: cond_transform = get_transform(cond_resolution, orig_mode, allow_aug, force_aug) else: cond_transform = None if kwargs.get("transform", None) is not None: rprint(f"Using transform from kwargs: {kwargs['transform']}") transform = kwargs.pop("transform") if name == "torchvision_imagenet": from torchvision.datasets import ImageFolder raw_data_dir = Path(config.data.raw_data_dir) raw_data_dir = raw_data_dir / "train" if orig_mode == "train" else raw_data_dir / "val" dataset = ImageFolder(raw_data_dir, transform=transform) dataset = CustomTransformDataset(dataset, lambda x, idx: {"img": x[0], "label": x[1]}) return dataset if "pixparse/cc12m-wds-fast" in name or "pixparse/cc3m-wds-fast" in name or "indexed" in name: return get_webdataset_indexed(config, tokenizer, transform, cond_transform, n_samples, name, should_tokenize=True) if name == "vggface2": dataset = VGGFace( Path(raw_data_dir), is_train, transform=transform, filter_resolution=(resolution - 48), cond_transform=cond_transform, v2=getattr(config.data, "add_vggface_v2_attributes", False), ) rprint(f"VGGFace2 has size {len(dataset)}") return dataset if name == "cub2011_custom": from models.datasets.cub200 import TextDataset dataset = TextDataset(data_dir='/path/to/cub200/birds', split='train' if is_train else 'test') return dataset wds_config = OmegaConf.create( { "train_data": None, "val_data": None, "dataset_type": "webdataset", "train_data_upsampling_factors": None, "batch_size": config.loader.batch_size if mode == "train" else config.loader.eval_batch_size, "workers": config.loader.num_workers, "distributed": True, "seed": config.seed, "val_num_samples": None, "train_num_samples": config.data.webdataset_train_num_samples, "val_num_samples": config.data.webdataset_val_num_samples, "world_size": config.trainer.devices * config.trainer.num_nodes, "block_size": block_size, } ) if config.data.dataset_type == "webdataset": clean_brace_escape = lambda x: x.replace("[", "{").replace("]", "}") wds_config.train_data = clean_brace_escape(config.data.webdataset_train_data) wds_config.val_data = clean_brace_escape(config.data.webdataset_val_data) if getattr(config.data, "webdataset_prefix", None) is not None: wds_config.train_data = config.data.webdataset_prefix.replace("LITERALQUOTE", "'").replace("LITERALSPACE", " ") + wds_config.train_data wds_config.val_data = config.data.webdataset_prefix.replace("LITERALQUOTE", "'").replace("LITERALSPACE", " ") + wds_config.val_data if getattr(config.data, "webdataset_postfix", None) is not None: wds_config.train_data = wds_config.train_data + config.data.webdataset_postfix.replace("LITERALQUOTE", "'").replace("LITERALSPACE", " ") wds_config.val_data = wds_config.val_data + config.data.webdataset_postfix.replace("LITERALQUOTE", "'").replace("LITERALSPACE", " ") return get_data(wds_config, (transform, transform), epoch=0, tokenizer=tokenizer) if name == "laion400m": # TODO: Debug if these configs are correct!!!! Not fully sure how the webdataset sharded dataloader should work. wds_config.train_data = "/grogu/datasets/laion400m/dataset/{00000..00625}.tar" wds_config.val_data = "/grogu/datasets/laion400m/dataset/{00000..00625}.tar" return get_data(wds_config, (transform, transform), epoch=0, tokenizer=tokenizer) elif name == "cc12m_3m": # TODO: Debug if these configs are correct!!!! Not fully sure how the webdataset sharded dataloader should work. wds_config.train_data = config.data.raw_data_dir + "/cc3m-train-{0000..0575}.tar" wds_config.val_data = config.data.raw_data_dir + "/cc3m-validation-{0000..0015}.tar" return get_data(wds_config, (transform, transform), epoch=0, tokenizer=tokenizer) elif name == "facecaption": if getattr(config.data, "webdataset_iterable", False): wds_config.train_data = "/grogu/user/mprabhud/data/diffusion/facecaption/{00000..00001}.tar" wds_config.val_data = "/grogu/user/mprabhud/data/diffusion/facecaption/{00000..00001}.tar" return get_data(wds_config, (transform, transform), epoch=0, tokenizer=tokenizer) elif getattr(config.data, "webdataset_indexed", False) is False: return get_webdataset_indexed(config, tokenizer, transform, cond_transform, n_samples, name, should_tokenize=True) else: raise Exception("Unknown webdataset type") # hf webdataset if name == "pixparse/cc12m-wds": extra_kwargs["data_dir"] = config.data.raw_data_dir if name == "generated_images": extra_kwargs["data_files"] = {"train": getattr(config.data, "parquet_path", None)} if name != "imagefolder": rprint(f"Loading dataset {name}, split={split}, streaming={streaming}, cache_dir={cache_dir}, extra_kwargs={extra_kwargs}, dynamic={dynamic}") load_map = {"pixparse/cc12m-wds": "webdataset", "laion400m": "webdataset", "generated_images": "parquet"} load_name = load_map.get(name, name) if streaming is False: extra_kwargs["num_proc"] = 16 if config.data.tokenize_vqvae_in_dataloader: full_cache_path = get_vqvae_dataloader(config, name, split) _ret = load_vqvae_from_cache(config, full_cache_path) if _ret is not None: return _ret from model import get_image_batch, get_vae rank = get_rank() vae = get_vae(config, device="cpu").eval() vae.to(f"cuda:{rank}") def tokenize_vqvae(batch): device = f"cuda:{rank}" img_input_ids = get_image_batch(config, vae, batch, device) batch.pop("img") batch["img_input_ids"] = img_input_ids return batch if config.data.keep_hf_dataset_in_memory: extra_kwargs["keep_in_memory"] = True gprint(f"WARNING!!!! Keeping dataset in memory") if name == "geneval": def create_blank_image(): return Image.new("RGB", (resolution, resolution), color=(255, 255, 255)) # https://github.com/djghosh13/geneval/blob/main/prompts/generation_prompts.txt prompts_path = Path.home() / ".cache" / "unidisc" / "geneval_generation_prompts.txt" if not prompts_path.exists(): prompts_path.parent.mkdir(parents=True, exist_ok=True) import urllib.request urllib.request.urlretrieve( "https://raw.githubusercontent.com/djghosh13/geneval/main/prompts/generation_prompts.txt", prompts_path ) with open(prompts_path, "r") as f: captions = [line.strip() for line in f.readlines()] dataset = Dataset.from_dict({ "caption": captions, "image": [ create_blank_image() for i in range(len(captions)) ], }) elif name == "MJHQ": def create_blank_image(): return Image.new("RGB", (resolution, resolution), color=(255, 255, 255)) prompts_path = Path.home() / ".cache" / "unidisc" / "MJHQ_meta_data.json" if not prompts_path.exists(): prompts_path.parent.mkdir(parents=True, exist_ok=True) import urllib.request urllib.request.urlretrieve( "https://huggingface.co/datasets/playgroundai/MJHQ-30K/resolve/main/meta_data.json", prompts_path ) with open(prompts_path, "r") as f: data = json.load(f) captions = [item["prompt"] for item in data.values()] dataset = Dataset.from_dict({ "caption": captions, "image": [ create_blank_image() for i in range(len(captions)) ], }) else: dataset = datasets.load_dataset(load_name, split=split, streaming=streaming, cache_dir=cache_dir, **extra_kwargs) dataset_keys = { "cassiekang/cub200_dataset": ("image", "text"), "Andron00e/CUB200-custom": ("image",), "nlphuji/flickr30k": ("image", "caption"), "ILSVRC/imagenet-1k": ("image", "label"), "richwardle/reduced-imagenet": ("image",), "tglcourse/lsun_church_train": ("image",), "imagefolder": ("image",), "pixparse/cc12m-wds": ("jpg", "txt"), "pravsels/FFHQ_1024": ("image",), "pravsels/SFHQ_256": ("image",), "jxie/celeba-hq": ("image",), "tglcourse/lsun_church_train": ("image",), "pouya-haghi/imagenet-subset": ("image",), "DeepLearner101/ImageNetSubsetValidate": ("image",), "PixArt-alpha/SAM-LLaVA-Captions10M": ("__key__", "txt"), "generated_images": ("__key__", "caption"), "laion/clevr-webdataset": ("jpg","txt"), "pcuenq/lsun-bedrooms": ("image",), "facebook/winoground": ("image_0", "image_1", "caption_0", "caption_1"), "sayakpaul/coco-30-val-2014": ("image", "caption"), "geneval": ("image", "caption"), "MJHQ": ("image", "caption"), } from dataloader import tokenize_text def preprocess_images(example, index: typing.Optional[typing.Any] = None): data = {} if dataset_keys[name][0] == "__key__": images = [] is_valid = [] for key, _image_path in zip(example[dataset_keys[name][0]], example["image_path"]): img_path = ( (Path(config.data.raw_data_dir) / key).with_suffix(".jpg") if not key.endswith(".jpg") else (Path(config.data.raw_data_dir) / key) ) allow_relative = False if Path(_image_path).exists() and Path(_image_path).stat().st_size > 0: img = Image.open(_image_path) is_valid.append(True) elif allow_relative and img_path.exists() and img_path.stat().st_size > 0: img = Image.open(img_path) is_valid.append(True) else: img = Image.new("RGB", (resolution, resolution), color=(255, 255, 255)) is_valid.append(False) images.append(img) data["is_valid"] = is_valid if sum(data["is_valid"]) < len(data["is_valid"]): gprint(f"WARNING!!! Found {len(data['is_valid']) - sum(data['is_valid'])} invalid images") else: images = [image.convert("RGB") for image in example[dataset_keys[name][0]]] data["img"] = [transform(image) for image in images] if cond_resolution is not None: data["cond_img"] = [cond_transform(image) for image in images] if index is not None: data["idx"] = index if "idx" in example: data["idx"] = example["idx"] if dynamic and dataset_keys[name][0] is not None: data["img"] = torch.stack(data["img"]) if "label" in example: data["label"] = example["label"] if (unified_model or getattr(config.data, "txt_only", False)) and not getattr(config.data, "disable_text_modality", False): tokenizer.padding_side = "right" tokenizer.truncation_side = "right" if name == "facebook/winoground": caption_0 = example["caption_0"] caption_1 = example["caption_1"] img_0 = example["image_0"] img_1 = example["image_1"] # tokenize and store captions separately tokens_0 = tokenize_text(tokenizer, block_size, caption_0) tokens_1 = tokenize_text(tokenizer, block_size, caption_1) data["caption_0_input_ids"] = tokens_0["input_ids"] data["caption_0_attention_mask"] = tokens_0["attention_mask"].float() data["caption_1_input_ids"] = tokens_1["input_ids"] data["caption_1_attention_mask"] = tokens_1["attention_mask"].float() # convert img_0 and img_1 which are lists of PIL images to tensors # convert some rgba pil images to rgb data["img_0"] = torch.stack([transform(img.convert("RGB")) for img in img_0]) data["img_1"] = torch.stack([transform(img.convert("RGB")) for img in img_1]) else: text_data = example[dataset_keys[name][1]] if isinstance(text_data[0], list): # Flickr has a list of captions for each image text_data = [random.choice(_data) for _data in text_data] tokens = tokenize_text(tokenizer, block_size, text_data) data["input_ids"] = tokens["input_ids"] data["attention_mask"] = tokens["attention_mask"].float() return data if precache is False: tokenized_dataset = dataset.with_transform(preprocess_images) else: extra_kwargs = dict() if streaming is False: extra_kwargs["load_from_cache_file"] = True else: if name == "pixparse/cc12m-wds": extra_kwargs["remove_columns"] = ["__key__", "jpg", "__url__", "json", "txt"] elif name == "ILSVRC/imagenet-1k": extra_kwargs["remove_columns"] = ["image"] tokenized_dataset = dataset.map(preprocess_images, batched=True, with_indices=True, **extra_kwargs) allowed_column_names = ["img", "input_ids", "attention_mask", "tokens", "text", "idx"] current_column_names = tokenized_dataset.column_names if current_column_names is not None: for column_name in current_column_names: if column_name not in allowed_column_names: tokenized_dataset = tokenized_dataset.remove_columns(column_name) if n_duplicate is not None: tokenized_dataset = datasets.concatenate_datasets([tokenized_dataset] * n_duplicate) ret = tokenized_dataset if dynamic else tokenized_dataset.with_format("torch") if isinstance(dataset, torch.utils.data.IterableDataset) or "cc12m" in name: ret = ResilientIterableDatasetWrapper(ret) if config.data.tokenize_vqvae_in_dataloader: assert config.data.force_mp_spawn ret = ret.shard(num_shards=get_world_size(), index=get_rank(), contiguous=True, keep_in_memory=True) gprint(f"Rank {rank} has {len(ret)} samples. World size is {get_world_size()}") ret = ret.map(tokenize_vqvae, batch_size=getattr(config.data, "vqvae_batch_size", 128), batched=True, keep_in_memory=True) ret.reset_format() allowed_column_names = ["img_input_ids"] map_column_list = getattr(config.data, "map_columns", None) if map_column_list is not None: for old_column_name, new_column_name in map_column_list.items(): ret = ret.rename_column(old_column_name, new_column_name) if getattr(config.data, "allow_label", False): allowed_column_names.append("label") if getattr(config.data, "allowed_columns_vqvae_dataloader", None): allowed_column_names.extend(list(config.data.allowed_columns_vqvae_dataloader)) current_column_names = ret.column_names if current_column_names is not None: for column_name in current_column_names: if column_name not in allowed_column_names: ret = ret.remove_columns(column_name) rank_cache_path = full_cache_path / f"rank_{rank}" gprint(f"Rank {rank} has saved to {rank_cache_path} with {len(ret)} samples") ret.save_to_disk(rank_cache_path) barrier() gprint(f"Rank {rank} has finished saving to {rank_cache_path}. Sleeping for a bit. You may want to Ctrl+C now") time.sleep(60 * 30) ret = load_vqvae_from_cache(config, full_cache_path) gprint(f"Rank {rank} has finished loading from file: {rank_cache_path}") return ret