from abc import ABC, abstractmethod import copy import json import logging import os from pathlib import Path import random from time import sleep import traceback import warnings import pandas as pd from tqdm import tqdm import h5py import torch.distributed as dist from torch.utils.data import Dataset import yaml logger = logging.getLogger(__name__) class DataBriefReportException(Exception): def __init__(self, message=None): self.message = message def __str__(self): return f"{self.__class__}: {self.message}" class DataNoReportException(Exception): def __init__(self, message=None): self.message = message def __str__(self): return f"{self.__class__}: {self.message}" class ItemProcessor(ABC): @abstractmethod def process_item(self, data_item, training_mode=False): raise NotImplementedError def is_huggingface_path(path: str) -> bool: # Heuristic: Hugging Face dataset paths are in format "user/dataset" # and not an existing local file or directory. return ("/" in path and not os.path.exists(path) and not "booru" in path) or (os.path.exists(path) and os.path.isdir(path)) global_log_count = 0 def log_every_n(n, msg): global global_log_count if global_log_count % n == 0: logger.warning(msg) global_log_count += 1 class MyDataset(Dataset): def __init__(self, config_path, item_processor: ItemProcessor, cache_on_disk=False): logger.info(f"read dataset config from {config_path}") with open(config_path, "r") as f: self.config = yaml.load(f, Loader=yaml.FullLoader) logger.info("DATASET CONFIG:") logger.info(self.config) self.cache_on_disk = cache_on_disk if self.cache_on_disk: cache_dir = self._get_cache_dir(config_path) if int(os.environ["LOCAL_RANK"]) == 0: local_rank = dist.get_rank() print(f"Building cache on rank {local_rank}") self._collect_annotations_and_save_to_cache(cache_dir) dist.barrier() ann, group_indice_range = self._load_annotations_from_cache(cache_dir) else: cache_dir = None ann, group_indice_range = self._collect_annotations() self.ann = ann self.group_indices = {key: list(range(val[0], val[1])) for key, val in group_indice_range.items()} logger.info(f"total length: {len(self)}") self.item_processor = item_processor def __len__(self): return len(self.ann) def _collect_annotations(self): meta_type_to_caption_type = { "image_text" : "prompt", "image_nl_caption" : "sentence", "image_alttext" : "alttext", "default" : "prompt", "super_high_quality_caption" : "super_high_quality_caption", "image_tags" : "tags", } switchable_keys = ["prompt", "sentence", "alttext", "super_high_quality_caption", "tags"] group_ann = {} for meta in self.config["META"]: meta_path, meta_type = meta["path"], meta.get("type", "default") meta_key = meta_type_to_caption_type.get(meta_type, "prompt") logger.info(f"Reading {meta_path} with type {meta_type} and key {meta_key}") if is_huggingface_path(meta_path): raise NotImplementedError("Hugging Face datasets are not supported in this minimal example.") else: meta_ext = os.path.splitext(meta_path)[-1] if meta_ext == ".json": # with open(meta_path) as f: # meta_l = json.load(f) with open(meta_path, 'r') as json_file: f = json_file.read() meta_l = json.loads(f) elif meta_ext == ".jsonl": meta_l = [] with open(meta_path) as f: for i, line in tqdm(enumerate(f), desc=f"Reading {meta_path}"): try: read_result = json.loads(line) if isinstance(read_result, dict): for key in switchable_keys: if key in read_result and meta_key != key: read_result[meta_key] = read_result[key] read_result.pop(key) break if read_result[meta_key].strip(): meta_l.append(read_result) else: logger.error(f"Empty prompt in {meta_path} line {i}, file: {meta_path}") log_every_n(10000, f"line {i}: {read_result}") else: raise ValueError(f"Expected a dictionary, got {type(read_result)} for {meta_path} line {i}") except json.decoder.JSONDecodeError as e: logger.error(f"Error decoding the following jsonl line ({i}):\n{line.rstrip()}") raise e elif meta_ext == ".parquet": meta_l = [] df = pd.read_parquet(meta_path) # Read the Parquet file into a DataFrame pq_cols = meta.get("pq_cols", None) if pq_cols is not None: cols = pq_cols.split(",") else: cols = None if cols: if "index" not in cols: raise ValueError(f"The 'index' column must be included in the 'pq_cols' list., in {meta_path}") if not all([col in df.columns for col in cols]): raise ValueError(f"Columns in 'pq_cols' must be present in the Parquet file., in {meta_path}") for _, row in tqdm(df.iterrows(), total=len(df), desc=f"Reading {meta_path}"): # Pull the 'index' column (whatever column indicates image index/id) index_val = row["index"] # For each *other* column in the row, if not None/NaN, use it as "prompt" for col in df.columns: if col == "index": continue if cols: if col not in cols: continue # Skip if the value is None or NaN if pd.notna(row[col]) and str(row[col]): log_every_n(10000, f"{meta_key}: {row[col]}") meta_l.append({ "image_path": f"danbooru://{index_val}" if not os.path.exists(index_val) and "://" not in str(index_val) else str(index_val), meta_key: str(row[col]) # Cast to str in case it's not a string }) else: raise NotImplementedError( f'Unknown meta file extension: "{meta_ext}". ' f"Currently, .json, .jsonl, .parquet (with index column + caption columns) are supported. " "If you are using a supported format, please set the file extension so that the proper parsing " "routine can be called." ) logger.info(f"{meta_path}, type{meta_type}: len {len(meta_l)}") if "ratio" in meta: random.seed(0) meta_l = random.sample(meta_l, int(len(meta_l) * meta["ratio"])) logger.info(f"sample (ratio = {meta['ratio']}) {len(meta_l)} items") if "root" in meta: for item in meta_l: for path_key in ["path", "image_url", "image", "image_path"]: if path_key in item: item[path_key] = os.path.join(meta["root"], item[path_key]) if meta_type not in group_ann: group_ann[meta_type] = [] group_ann[meta_type] += meta_l ann = sum(list(group_ann.values()), start=[]) group_indice_range = {} start_pos = 0 for meta_type, meta_l in group_ann.items(): group_indice_range[meta_type] = [start_pos, start_pos + len(meta_l)] start_pos = start_pos + len(meta_l) return ann, group_indice_range def _collect_annotations_and_save_to_cache(self, cache_dir): if (Path(cache_dir) / "data.h5").exists() and (Path(cache_dir) / "ready").exists(): # off-the-shelf annotation cache exists warnings.warn( f"Use existing h5 data cache: {Path(cache_dir)}\n" f"Note: if the actual data defined by the data config has changed since your last run, " f"please delete the cache manually and re-run this experiment, or the data actually used " f"will not be updated" ) return Path(cache_dir).mkdir(parents=True, exist_ok=True) ann, group_indice_range = self._collect_annotations() # when cache on disk, rank0 saves items to an h5 file serialized_ann = [json.dumps(_) for _ in ann] logger.info(f"start to build data cache to: {Path(cache_dir)}") with h5py.File(Path(cache_dir) / "data.h5", "w") as file: dt = h5py.vlen_dtype(str) h5_ann = file.create_dataset("ann", (len(serialized_ann),), dtype=dt) h5_ann[:] = serialized_ann file.create_dataset("group_indice_range", data=json.dumps(group_indice_range)) with open(Path(cache_dir) / "ready", "w") as f: f.write("ready") logger.info(f"data cache built") @staticmethod def _get_cache_dir(config_path): config_identifier = config_path disallowed_chars = ["/", "\\", ".", "?", "!"] for _ in disallowed_chars: config_identifier = config_identifier.replace(_, "-") cache_dir = f"./accessory_data_cache/{config_identifier}" return cache_dir @staticmethod def _load_annotations_from_cache(cache_dir): while not (Path(cache_dir) / "ready").exists(): # cache has not yet been completed by rank 0 assert int(os.environ["LOCAL_RANK"]) != 0 sleep(1) cache_file = h5py.File(Path(cache_dir) / "data.h5", "r") annotations = cache_file["ann"] group_indice_range = json.loads(cache_file["group_indice_range"].asstr()[()]) return annotations, group_indice_range def get_item_func(self, index): data_item = self.ann[index] if self.cache_on_disk: data_item = json.loads(data_item) else: data_item = copy.deepcopy(data_item) return self.item_processor.process_item(data_item, training_mode=True) def __getitem__(self, index): try: return self.get_item_func(index) except Exception as e: if isinstance(e, DataNoReportException): pass elif isinstance(e, DataBriefReportException): logger.info(e) else: logger.info( f"Item {index} errored, annotation:\n" f"{self.ann[index]}\n" f"Error:\n" f"{traceback.format_exc()}" ) for group_name, indices_this_group in self.group_indices.items(): if indices_this_group[0] <= index <= indices_this_group[-1]: if index == indices_this_group[0]: new_index = indices_this_group[-1] else: new_index = index - 1 return self[new_index] raise RuntimeError def groups(self): return list(self.group_indices.values())