"""Dataset core: image bookkeeping, bucketing and the BaseDataset/DatasetGroup hierarchy. This module owns the data-loading side of training: * ``IMAGE_EXTENSIONS`` / ``TEXT_ENCODER_OUTPUTS_CACHE_SUFFIX*`` constants * ``ImageInfo`` (per-image metadata), ``BucketManager`` (aspect-ratio buckets), ``BucketBatchIndex``, ``AugHelper`` (color augmentation) * ``BaseDataset`` (abstract dataset shared by DreamBooth / FineTuning / ControlNet) * ``DatasetGroup`` (concat of multiple datasets) * ``MinimalDataset`` and ``load_arbitrary_dataset`` (user-supplied dataset hook) * ``debug_dataset`` viewer plus ``glob_images`` / ``glob_images_pathlib`` helpers * ``split_train_val`` shared between training/validation dataset halves Leaf image helpers (``IMAGE_TRANSFORMS`` / ``load_image`` / ``get_crop_ltrb`` / ``trim_and_resize_if_required``) live in ``library.utils`` and are re-exported here for backward compatibility. The DreamBooth / FineTuning / ControlNet specializations of ``BaseDataset`` live in their dedicated modules: ``library.dreambooth_dataset``, ``library.finetuning_dataset`` and ``library.controlnet_dataset``. ``HIGH_VRAM`` (a mutable module-level flag toggled by ``enable_high_vram``) lives in ``library.accelerator_setup``; that module has no cycle with this one so it is imported directly. """ import glob import importlib import logging import math import os import pathlib import random import re from concurrent.futures import Future, ThreadPoolExecutor from typing import Any, Callable, Dict, List, NamedTuple, Optional, Sequence, Tuple, Union import cv2 import imagesize import numpy as np import torch from PIL import Image from accelerate import Accelerator from diffusers import AutoencoderKL from torchvision import transforms from tqdm import tqdm from transformers import CLIPTokenizer import library.model_util as model_util from library import accelerator_setup from library.device_utils import clean_memory_on_device from library.strategy_base import ( LatentsCachingStrategy, TextEncoderOutputsCachingStrategy, TextEncodingStrategy, TokenizeStrategy, ) from library.subset import ( BaseSubset, ControlNetSubset, DreamBoothSubset, FineTuningSubset, ) from library.utils import ( IMAGE_TRANSFORMS, get_crop_ltrb, load_image, resize_image, setup_logging, trim_and_resize_if_required, validate_interpolation_fn, ) setup_logging() logger = logging.getLogger(__name__) IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".webp", ".bmp", ".PNG", ".JPG", ".JPEG", ".WEBP", ".BMP"] try: import pillow_avif IMAGE_EXTENSIONS.extend([".avif", ".AVIF"]) except: pass # JPEG-XL on Linux try: from jxlpy import JXLImagePlugin from library.jpeg_xl_util import get_jxl_size IMAGE_EXTENSIONS.extend([".jxl", ".JXL"]) except: pass # JPEG-XL on Linux and Windows try: import pillow_jxl from library.jpeg_xl_util import get_jxl_size IMAGE_EXTENSIONS.extend([".jxl", ".JXL"]) except: pass TEXT_ENCODER_OUTPUTS_CACHE_SUFFIX = "_te_outputs.npz" TEXT_ENCODER_OUTPUTS_CACHE_SUFFIX_SD3 = "_sd3_te.npz" def split_train_val( paths: List[str], sizes: List[Optional[Tuple[int, int]]], is_training_dataset: bool, validation_split: float, validation_seed: int | None, ) -> Tuple[List[str], List[Optional[Tuple[int, int]]]]: """ Split the dataset into train and validation Shuffle the dataset based on the validation_seed or the current random seed. For example if the split of 0.2 of 100 images. [0:80] = 80 training images [80:] = 20 validation images """ dataset = list(zip(paths, sizes)) if validation_seed is not None: logging.info(f"Using validation seed: {validation_seed}") prevstate = random.getstate() random.seed(validation_seed) random.shuffle(dataset) random.setstate(prevstate) else: random.shuffle(dataset) paths, sizes = zip(*dataset) paths = list(paths) sizes = list(sizes) # Split the dataset between training and validation if is_training_dataset: # Training dataset we split to the first part split = math.ceil(len(paths) * (1 - validation_split)) return paths[0:split], sizes[0:split] else: # Validation dataset we split to the second part split = len(paths) - round(len(paths) * validation_split) return paths[split:], sizes[split:] class ImageInfo: def __init__( self, image_key: str, num_repeats: int, caption: str, is_reg: bool, absolute_path: str, caption_dropout_rate: float = 0.0 ) -> None: self.image_key: str = image_key self.num_repeats: int = num_repeats self.caption: str = caption self.is_reg: bool = is_reg self.absolute_path: str = absolute_path self.caption_dropout_rate: float = caption_dropout_rate self.image_size: Tuple[int, int] = None self.resized_size: Tuple[int, int] = None self.bucket_reso: Tuple[int, int] = None self.latents: Optional[torch.Tensor] = None self.latents_flipped: Optional[torch.Tensor] = None self.latents_npz: Optional[str] = None # set in cache_latents self.latents_original_size: Optional[Tuple[int, int]] = None # original image size, not latents size self.latents_crop_ltrb: Optional[Tuple[int, int]] = ( None # crop left top right bottom in original pixel size, not latents size ) self.cond_img_path: Optional[str] = None self.image: Optional[Image.Image] = None # optional, original PIL Image self.text_encoder_outputs_npz: Optional[str] = None # filename. set in cache_text_encoder_outputs # new self.text_encoder_outputs: Optional[List[torch.Tensor]] = None # old self.text_encoder_outputs1: Optional[torch.Tensor] = None self.text_encoder_outputs2: Optional[torch.Tensor] = None self.text_encoder_pool2: Optional[torch.Tensor] = None self.alpha_mask: Optional[torch.Tensor] = None # alpha mask can be flipped in runtime self.resize_interpolation: Optional[str] = None class BucketManager: def __init__(self, no_upscale, max_reso, min_size, max_size, reso_steps) -> None: if max_size is not None: if max_reso is not None: assert max_size >= max_reso[0], "the max_size should be larger than the width of max_reso" assert max_size >= max_reso[1], "the max_size should be larger than the height of max_reso" if min_size is not None: assert max_size >= min_size, "the max_size should be larger than the min_size" self.no_upscale = no_upscale if max_reso is None: self.max_reso = None self.max_area = None else: self.max_reso = max_reso self.max_area = max_reso[0] * max_reso[1] self.min_size = min_size self.max_size = max_size self.reso_steps = reso_steps self.resos = [] self.reso_to_id = {} self.buckets = [] # 前処理時は (image_key, image, original size, crop left/top)、学習時は image_key def add_image(self, reso, image_or_info): bucket_id = self.reso_to_id[reso] self.buckets[bucket_id].append(image_or_info) def shuffle(self): for bucket in self.buckets: random.shuffle(bucket) def sort(self): # 解像度順にソートする（表示時、メタデータ格納時の見栄えをよくするためだけ）。bucketsも入れ替えてreso_to_idも振り直す sorted_resos = self.resos.copy() sorted_resos.sort() sorted_buckets = [] sorted_reso_to_id = {} for i, reso in enumerate(sorted_resos): bucket_id = self.reso_to_id[reso] sorted_buckets.append(self.buckets[bucket_id]) sorted_reso_to_id[reso] = i self.resos = sorted_resos self.buckets = sorted_buckets self.reso_to_id = sorted_reso_to_id def make_buckets(self): resos = model_util.make_bucket_resolutions(self.max_reso, self.min_size, self.max_size, self.reso_steps) self.set_predefined_resos(resos) def set_predefined_resos(self, resos): # 規定サイズから選ぶ場合の解像度、aspect ratioの情報を格納しておく self.predefined_resos = resos.copy() self.predefined_resos_set = set(resos) self.predefined_aspect_ratios = np.array([w / h for w, h in resos]) def add_if_new_reso(self, reso): if reso not in self.reso_to_id: bucket_id = len(self.resos) self.reso_to_id[reso] = bucket_id self.resos.append(reso) self.buckets.append([]) # logger.info(reso, bucket_id, len(self.buckets)) def round_to_steps(self, x): x = int(x + 0.5) return x - x % self.reso_steps def select_bucket(self, image_width, image_height): aspect_ratio = image_width / image_height if not self.no_upscale: # 拡大および縮小を行う # 同じaspect ratioがあるかもしれないので（fine tuningで、no_upscale=Trueで前処理した場合）、解像度が同じものを優先する reso = (image_width, image_height) if reso in self.predefined_resos_set: pass else: ar_errors = self.predefined_aspect_ratios - aspect_ratio predefined_bucket_id = np.abs(ar_errors).argmin() # 当該解像度以外でaspect ratio errorが最も少ないもの reso = self.predefined_resos[predefined_bucket_id] ar_reso = reso[0] / reso[1] if aspect_ratio > ar_reso: # 横が長い→縦を合わせる scale = reso[1] / image_height else: scale = reso[0] / image_width resized_size = (int(image_width * scale + 0.5), int(image_height * scale + 0.5)) # logger.info(f"use predef, {image_width}, {image_height}, {reso}, {resized_size}") else: # 縮小のみを行う if image_width * image_height > self.max_area: # 画像が大きすぎるのでアスペクト比を保ったまま縮小することを前提にbucketを決める resized_width = math.sqrt(self.max_area * aspect_ratio) resized_height = self.max_area / resized_width assert abs(resized_width / resized_height - aspect_ratio) < 1e-2, "aspect is illegal" # リサイズ後の短辺または長辺をreso_steps単位にする：aspect ratioの差が少ないほうを選ぶ # 元のbucketingと同じロジック b_width_rounded = self.round_to_steps(resized_width) b_height_in_wr = self.round_to_steps(b_width_rounded / aspect_ratio) ar_width_rounded = b_width_rounded / b_height_in_wr b_height_rounded = self.round_to_steps(resized_height) b_width_in_hr = self.round_to_steps(b_height_rounded * aspect_ratio) ar_height_rounded = b_width_in_hr / b_height_rounded # logger.info(b_width_rounded, b_height_in_wr, ar_width_rounded) # logger.info(b_width_in_hr, b_height_rounded, ar_height_rounded) if abs(ar_width_rounded - aspect_ratio) < abs(ar_height_rounded - aspect_ratio): resized_size = (b_width_rounded, int(b_width_rounded / aspect_ratio + 0.5)) else: resized_size = (int(b_height_rounded * aspect_ratio + 0.5), b_height_rounded) # logger.info(resized_size) else: resized_size = (image_width, image_height) # リサイズは不要 # 画像のサイズ未満をbucketのサイズとする（paddingせずにcroppingする） bucket_width = resized_size[0] - resized_size[0] % self.reso_steps bucket_height = resized_size[1] - resized_size[1] % self.reso_steps # logger.info(f"use arbitrary {image_width}, {image_height}, {resized_size}, {bucket_width}, {bucket_height}") reso = (bucket_width, bucket_height) self.add_if_new_reso(reso) ar_error = (reso[0] / reso[1]) - aspect_ratio return reso, resized_size, ar_error @staticmethod def get_crop_ltrb(bucket_reso: Tuple[int, int], image_size: Tuple[int, int]): # implementation moved to library.utils.get_crop_ltrb; kept as a staticmethod for backward compatibility return get_crop_ltrb(bucket_reso, image_size) class BucketBatchIndex(NamedTuple): bucket_index: int bucket_batch_size: int batch_index: int class AugHelper: # albumentationsへの依存をなくしたがとりあえず同じinterfaceを持たせる def __init__(self): pass def color_aug(self, image: np.ndarray): # self.color_aug_method = albu.OneOf( # [ # albu.HueSaturationValue(8, 0, 0, p=0.5), # albu.RandomGamma((95, 105), p=0.5), # ], # p=0.33, # ) hue_shift_limit = 8 # remove dependency to albumentations if random.random() <= 0.33: if random.random() > 0.5: # hue shift hsv_img = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) hue_shift = random.uniform(-hue_shift_limit, hue_shift_limit) if hue_shift < 0: hue_shift = 180 + hue_shift hsv_img[:, :, 0] = (hsv_img[:, :, 0] + hue_shift) % 180 image = cv2.cvtColor(hsv_img, cv2.COLOR_HSV2BGR) else: # random gamma gamma = random.uniform(0.95, 1.05) image = np.clip(image**gamma, 0, 255).astype(np.uint8) return {"image": image} def get_augmentor(self, use_color_aug: bool): # -> Optional[Callable[[np.ndarray], Dict[str, np.ndarray]]]: return self.color_aug if use_color_aug else None class BaseDataset(torch.utils.data.Dataset): def __init__( self, resolution: Optional[Tuple[int, int]], network_multiplier: float, train_inpainting: bool, debug_dataset: bool, resize_interpolation: Optional[str] = None, skip_image_resolution: Optional[Tuple[int, int]] = None, ) -> None: super().__init__() # width/height is used when enable_bucket==False self.width, self.height = (None, None) if resolution is None else resolution self.network_multiplier = network_multiplier self.debug_dataset = debug_dataset self.subsets: List[Union[DreamBoothSubset, FineTuningSubset]] = [] self.token_padding_disabled = False self.tag_frequency = {} self.XTI_layers = None self.token_strings = None self.enable_bucket = False self.bucket_manager: BucketManager = None # not initialized self.min_bucket_reso = None self.max_bucket_reso = None self.bucket_reso_steps = None self.bucket_no_upscale = None self.bucket_info = None # for metadata self.current_epoch: int = 0 # インスタンスがepochごとに新しく作られるようなので外側から渡さないとダメ self.current_step: int = 0 self.max_train_steps: int = 0 self.seed: int = 0 # inpainting self.train_inpainting = train_inpainting # augmentation self.aug_helper = AugHelper() self.image_transforms = IMAGE_TRANSFORMS if resize_interpolation is not None: assert validate_interpolation_fn( resize_interpolation ), f'Resize interpolation "{resize_interpolation}" is not a valid interpolation' self.resize_interpolation = resize_interpolation self.skip_image_resolution = skip_image_resolution self.image_data: Dict[str, ImageInfo] = {} self.image_to_subset: Dict[str, Union[DreamBoothSubset, FineTuningSubset]] = {} self.replacements = {} self.tokenize_strategy = None self.text_encoder_output_caching_strategy = None self.latents_caching_strategy = None def set_current_strategies(self): self.tokenize_strategy = TokenizeStrategy.get_strategy() self.text_encoder_output_caching_strategy = TextEncoderOutputsCachingStrategy.get_strategy() self.latents_caching_strategy = LatentsCachingStrategy.get_strategy() def adjust_min_max_bucket_reso_by_steps( self, resolution: Tuple[int, int], min_bucket_reso: int, max_bucket_reso: int, bucket_reso_steps: int ) -> Tuple[int, int]: # make min/max bucket reso to be multiple of bucket_reso_steps if min_bucket_reso % bucket_reso_steps != 0: adjusted_min_bucket_reso = min_bucket_reso - min_bucket_reso % bucket_reso_steps logger.warning( f"min_bucket_reso is adjusted to be multiple of bucket_reso_steps" f" / min_bucket_resoがbucket_reso_stepsの倍数になるように調整されました: {min_bucket_reso} -> {adjusted_min_bucket_reso}" ) min_bucket_reso = adjusted_min_bucket_reso if max_bucket_reso % bucket_reso_steps != 0: adjusted_max_bucket_reso = max_bucket_reso + bucket_reso_steps - max_bucket_reso % bucket_reso_steps logger.warning( f"max_bucket_reso is adjusted to be multiple of bucket_reso_steps" f" / max_bucket_resoがbucket_reso_stepsの倍数になるように調整されました: {max_bucket_reso} -> {adjusted_max_bucket_reso}" ) max_bucket_reso = adjusted_max_bucket_reso assert ( min(resolution) >= min_bucket_reso ), f"min_bucket_reso must be equal or less than resolution / min_bucket_resoは最小解像度より大きくできません。解像度を大きくするかmin_bucket_resoを小さくしてください" assert ( max(resolution) <= max_bucket_reso ), f"max_bucket_reso must be equal or greater than resolution / max_bucket_resoは最大解像度より小さくできません。解像度を小さくするかmin_bucket_resoを大きくしてください" return min_bucket_reso, max_bucket_reso def set_seed(self, seed): self.seed = seed def set_current_epoch(self, epoch): if not self.current_epoch == epoch: # epochが切り替わったらバケツをシャッフルする if epoch > self.current_epoch: logger.info("epoch is incremented. current_epoch: {}, epoch: {}".format(self.current_epoch, epoch)) num_epochs = epoch - self.current_epoch for _ in range(num_epochs): self.current_epoch += 1 self.shuffle_buckets() # self.current_epoch seem to be set to 0 again in the next epoch. it may be caused by skipped_dataloader? else: logger.warning("epoch is not incremented. current_epoch: {}, epoch: {}".format(self.current_epoch, epoch)) self.current_epoch = epoch def set_current_step(self, step): self.current_step = step def set_max_train_steps(self, max_train_steps): self.max_train_steps = max_train_steps def set_tag_frequency(self, dir_name, captions): frequency_for_dir = self.tag_frequency.get(dir_name, {}) self.tag_frequency[dir_name] = frequency_for_dir for caption in captions: for tag in caption.split(","): tag = tag.strip() if tag: tag = tag.lower() frequency = frequency_for_dir.get(tag, 0) frequency_for_dir[tag] = frequency + 1 def disable_token_padding(self): self.token_padding_disabled = True def enable_XTI(self, layers=None, token_strings=None): self.XTI_layers = layers self.token_strings = token_strings def add_replacement(self, str_from, str_to): self.replacements[str_from] = str_to def process_caption(self, subset: BaseSubset, caption): # caption に prefix/suffix を付ける if subset.caption_prefix: caption = subset.caption_prefix + " " + caption if subset.caption_suffix: caption = caption + " " + subset.caption_suffix # dropoutの決定：tag dropがこのメソッド内にあるのでここで行うのが良い is_drop_out = subset.caption_dropout_rate > 0 and random.random() < subset.caption_dropout_rate is_drop_out = ( is_drop_out or subset.caption_dropout_every_n_epochs > 0 and self.current_epoch % subset.caption_dropout_every_n_epochs == 0 ) if is_drop_out: caption = "" else: # process wildcards if subset.enable_wildcard: # if caption is multiline, random choice one line if "\n" in caption: caption = random.choice(caption.split("\n")) # wildcard is like '{aaa|bbb|ccc...}' # escape the curly braces like {{ or }} replacer1 = "⦅" replacer2 = "⦆" while replacer1 in caption or replacer2 in caption: replacer1 += "⦅" replacer2 += "⦆" caption = caption.replace("{{", replacer1).replace("}}", replacer2) # replace the wildcard def replace_wildcard(match): return random.choice(match.group(1).split("|")) caption = re.sub(r"\{([^}]+)\}", replace_wildcard, caption) # unescape the curly braces caption = caption.replace(replacer1, "{").replace(replacer2, "}") else: # if caption is multiline, use the first line caption = caption.split("\n")[0] if subset.shuffle_caption or subset.token_warmup_step > 0 or subset.caption_tag_dropout_rate > 0: fixed_tokens = [] flex_tokens = [] fixed_suffix_tokens = [] if ( hasattr(subset, "keep_tokens_separator") and subset.keep_tokens_separator and subset.keep_tokens_separator in caption ): fixed_part, flex_part = caption.split(subset.keep_tokens_separator, 1) if subset.keep_tokens_separator in flex_part: flex_part, fixed_suffix_part = flex_part.split(subset.keep_tokens_separator, 1) fixed_suffix_tokens = [t.strip() for t in fixed_suffix_part.split(subset.caption_separator) if t.strip()] fixed_tokens = [t.strip() for t in fixed_part.split(subset.caption_separator) if t.strip()] flex_tokens = [t.strip() for t in flex_part.split(subset.caption_separator) if t.strip()] else: tokens = [t.strip() for t in caption.strip().split(subset.caption_separator)] flex_tokens = tokens[:] if subset.keep_tokens > 0: fixed_tokens = flex_tokens[: subset.keep_tokens] flex_tokens = tokens[subset.keep_tokens :] if subset.token_warmup_step < 1: # 初回に上書きする subset.token_warmup_step = math.floor(subset.token_warmup_step * self.max_train_steps) if subset.token_warmup_step and self.current_step < subset.token_warmup_step: tokens_len = ( math.floor( (self.current_step) * ((len(flex_tokens) - subset.token_warmup_min) / (subset.token_warmup_step)) ) + subset.token_warmup_min ) flex_tokens = flex_tokens[:tokens_len] def dropout_tags(tokens): if subset.caption_tag_dropout_rate <= 0: return tokens l = [] for token in tokens: if random.random() >= subset.caption_tag_dropout_rate: l.append(token) return l if subset.shuffle_caption: random.shuffle(flex_tokens) flex_tokens = dropout_tags(flex_tokens) caption = f"{subset.caption_separator} ".join(fixed_tokens + flex_tokens + fixed_suffix_tokens) # process secondary separator if subset.secondary_separator: caption = caption.replace(subset.secondary_separator, subset.caption_separator) # textual inversion対応 for str_from, str_to in self.replacements.items(): if str_from == "": # replace all if type(str_to) == list: caption = random.choice(str_to) else: caption = str_to else: caption = caption.replace(str_from, str_to) return caption def register_image(self, info: ImageInfo, subset: BaseSubset): self.image_data[info.image_key] = info self.image_to_subset[info.image_key] = subset def make_buckets(self): """ bucketingを行わない場合も呼び出し必須（ひとつだけbucketを作る） min_size and max_size are ignored when enable_bucket is False """ logger.info("loading image sizes.") for info in tqdm(self.image_data.values()): if info.image_size is None: info.image_size = self.get_image_size(info.absolute_path) # # run in parallel # max_workers = min(os.cpu_count(), len(self.image_data)) # TODO consider multi-gpu (processes) # with ThreadPoolExecutor(max_workers) as executor: # futures = [] # for info in tqdm(self.image_data.values(), desc="loading image sizes"): # if info.image_size is None: # def get_and_set_image_size(info): # info.image_size = self.get_image_size(info.absolute_path) # futures.append(executor.submit(get_and_set_image_size, info)) # # consume futures to reduce memory usage and prevent Ctrl-C hang # if len(futures) >= max_workers: # for future in futures: # future.result() # futures = [] # for future in futures: # future.result() if self.enable_bucket: logger.info("make buckets") else: logger.info("prepare dataset") # bucketを作成し、画像をbucketに振り分ける if self.enable_bucket: if self.bucket_manager is None: # fine tuningの場合でmetadataに定義がある場合は、すでに初期化済み self.bucket_manager = BucketManager( self.bucket_no_upscale, (self.width, self.height), self.min_bucket_reso, self.max_bucket_reso, self.bucket_reso_steps, ) if not self.bucket_no_upscale: self.bucket_manager.make_buckets() else: logger.warning( "min_bucket_reso and max_bucket_reso are ignored if bucket_no_upscale is set, because bucket reso is defined by image size automatically / bucket_no_upscaleが指定された場合は、bucketの解像度は画像サイズから自動計算されるため、min_bucket_resoとmax_bucket_resoは無視されます" ) img_ar_errors = [] for image_info in self.image_data.values(): image_width, image_height = image_info.image_size image_info.bucket_reso, image_info.resized_size, ar_error = self.bucket_manager.select_bucket( image_width, image_height ) # logger.info(image_info.image_key, image_info.bucket_reso) img_ar_errors.append(abs(ar_error)) self.bucket_manager.sort() else: self.bucket_manager = BucketManager(False, (self.width, self.height), None, None, None) self.bucket_manager.set_predefined_resos([(self.width, self.height)]) # ひとつの固定サイズbucketのみ for image_info in self.image_data.values(): image_width, image_height = image_info.image_size image_info.bucket_reso, image_info.resized_size, _ = self.bucket_manager.select_bucket(image_width, image_height) for image_info in self.image_data.values(): for _ in range(image_info.num_repeats): self.bucket_manager.add_image(image_info.bucket_reso, image_info.image_key) # bucket情報を表示、格納する if self.enable_bucket: self.bucket_info = {"buckets": {}} logger.info("number of images (including repeats) / 各bucketの画像枚数（繰り返し回数を含む）") for i, (reso, bucket) in enumerate(zip(self.bucket_manager.resos, self.bucket_manager.buckets)): count = len(bucket) if count > 0: self.bucket_info["buckets"][i] = {"resolution": reso, "count": len(bucket)} logger.info(f"bucket {i}: resolution {reso}, count: {len(bucket)}") if len(img_ar_errors) == 0: mean_img_ar_error = 0 # avoid NaN else: img_ar_errors = np.array(img_ar_errors) mean_img_ar_error = np.mean(np.abs(img_ar_errors)) self.bucket_info["mean_img_ar_error"] = mean_img_ar_error logger.info(f"mean ar error (without repeats): {mean_img_ar_error}") # データ参照用indexを作る。このindexはdatasetのshuffleに用いられる self.buckets_indices: List[BucketBatchIndex] = [] for bucket_index, bucket in enumerate(self.bucket_manager.buckets): batch_count = int(math.ceil(len(bucket) / self.batch_size)) for batch_index in range(batch_count): self.buckets_indices.append(BucketBatchIndex(bucket_index, self.batch_size, batch_index)) self.shuffle_buckets() self._length = len(self.buckets_indices) def shuffle_buckets(self): # set random seed for this epoch random.seed(self.seed + self.current_epoch) random.shuffle(self.buckets_indices) self.bucket_manager.shuffle() def verify_bucket_reso_steps(self, min_steps: int): assert self.bucket_reso_steps is None or self.bucket_reso_steps % min_steps == 0, ( f"bucket_reso_steps is {self.bucket_reso_steps}. it must be divisible by {min_steps}.\n" + f"bucket_reso_stepsが{self.bucket_reso_steps}です。{min_steps}で割り切れる必要があります" ) def is_latent_cacheable(self): return all([not subset.color_aug and not subset.random_crop for subset in self.subsets]) def is_text_encoder_output_cacheable(self, cache_supports_dropout: bool = False): return all( [ not ( subset.caption_dropout_rate > 0 and not cache_supports_dropout or subset.shuffle_caption or subset.token_warmup_step > 0 or subset.caption_tag_dropout_rate > 0 ) for subset in self.subsets ] ) def new_cache_latents(self, model: Any, accelerator: Accelerator): r""" a brand new method to cache latents. This method caches latents with caching strategy. normal cache_latents method is used by default, but this method is used when caching strategy is specified. """ logger.info("caching latents with caching strategy.") caching_strategy = LatentsCachingStrategy.get_strategy() image_infos = list(self.image_data.values()) # sort by resolution image_infos.sort(key=lambda info: info.bucket_reso[0] * info.bucket_reso[1]) # split by resolution and some conditions class Condition: def __init__(self, reso, flip_aug, alpha_mask, random_crop): self.reso = reso self.flip_aug = flip_aug self.alpha_mask = alpha_mask self.random_crop = random_crop def __eq__(self, other): return ( other is not None and self.reso == other.reso and self.flip_aug == other.flip_aug and self.alpha_mask == other.alpha_mask and self.random_crop == other.random_crop ) batch: List[ImageInfo] = [] current_condition = None # support multiple-gpus num_processes = accelerator.num_processes process_index = accelerator.process_index # define a function to submit a batch to cache def submit_batch(batch, cond): for info in batch: if info.image is not None and isinstance(info.image, Future): info.image = info.image.result() # future to image caching_strategy.cache_batch_latents(model, batch, cond.flip_aug, cond.alpha_mask, cond.random_crop) # remove image from memory for info in batch: info.image = None # define ThreadPoolExecutor to load images in parallel max_workers = min(os.cpu_count(), len(image_infos)) max_workers = max(1, max_workers // num_processes) # consider multi-gpu max_workers = min(max_workers, caching_strategy.batch_size) # max_workers should be less than batch_size executor = ThreadPoolExecutor(max_workers) try: # iterate images logger.info("caching latents...") for i, info in enumerate(tqdm(image_infos)): subset = self.image_to_subset[info.image_key] if info.latents_npz is not None: # fine tuning dataset continue # check disk cache exists and size of latents if caching_strategy.cache_to_disk: # info.latents_npz = os.path.splitext(info.absolute_path)[0] + file_suffix info.latents_npz = caching_strategy.get_latents_npz_path(info.absolute_path, info.image_size) # if the modulo of num_processes is not equal to process_index, skip caching # this makes each process cache different latents if i % num_processes != process_index: continue # print(f"{process_index}/{num_processes} {i}/{len(image_infos)} {info.latents_npz}") cache_available = caching_strategy.is_disk_cached_latents_expected( info.bucket_reso, info.latents_npz, subset.flip_aug, subset.alpha_mask ) if cache_available: # do not add to batch continue # if batch is not empty and condition is changed, flush the batch. Note that current_condition is not None if batch is not empty condition = Condition(info.bucket_reso, subset.flip_aug, subset.alpha_mask, subset.random_crop) if len(batch) > 0 and current_condition != condition: submit_batch(batch, current_condition) batch = [] if condition != current_condition and accelerator_setup.HIGH_VRAM: # even with high VRAM, if shape is changed clean_memory_on_device(accelerator.device) if info.image is None: # load image in parallel info.image = executor.submit(load_image, info.absolute_path, condition.alpha_mask) batch.append(info) current_condition = condition # if number of data in batch is enough, flush the batch if len(batch) >= caching_strategy.batch_size: submit_batch(batch, current_condition) batch = [] # current_condition = None # keep current_condition to avoid next `clean_memory_on_device` call if len(batch) > 0: submit_batch(batch, current_condition) finally: executor.shutdown() def new_cache_text_encoder_outputs(self, models: List[Any], accelerator: Accelerator): r""" a brand new method to cache text encoder outputs. This method caches text encoder outputs with caching strategy. """ tokenize_strategy = TokenizeStrategy.get_strategy() text_encoding_strategy = TextEncodingStrategy.get_strategy() caching_strategy = TextEncoderOutputsCachingStrategy.get_strategy() batch_size = caching_strategy.batch_size or self.batch_size logger.info("caching Text Encoder outputs with caching strategy.") image_infos = list(self.image_data.values()) # split by resolution batches = [] batch = [] # support multiple-gpus num_processes = accelerator.num_processes process_index = accelerator.process_index logger.info("checking cache validity...") for i, info in enumerate(tqdm(image_infos)): # check disk cache exists and size of text encoder outputs if caching_strategy.cache_to_disk: te_out_npz = caching_strategy.get_outputs_npz_path(info.absolute_path) info.text_encoder_outputs_npz = te_out_npz # set npz filename regardless of cache availability # if the modulo of num_processes is not equal to process_index, skip caching # this makes each process cache different text encoder outputs if i % num_processes != process_index: continue cache_available = caching_strategy.is_disk_cached_outputs_expected(te_out_npz) if cache_available: # do not add to batch continue batch.append(info) # if number of data in batch is enough, flush the batch if len(batch) >= batch_size: batches.append(batch) batch = [] if len(batch) > 0: batches.append(batch) if len(batches) == 0: logger.info("no Text Encoder outputs to cache") return # iterate batches logger.info("caching Text Encoder outputs...") for batch in tqdm(batches, smoothing=1, total=len(batches)): caching_strategy.cache_batch_outputs(tokenize_strategy, models, text_encoding_strategy, batch) def get_image_size(self, image_path): if image_path.endswith(".jxl") or image_path.endswith(".JXL"): return get_jxl_size(image_path) # return imagesize.get(image_path) image_size = imagesize.get(image_path) if image_size[0] <= 0: # imagesize doesn't work for some images, so use PIL as a fallback try: with Image.open(image_path) as img: image_size = img.size except Exception as e: logger.warning(f"failed to get image size: {image_path}, error: {e}") image_size = (0, 0) return image_size def load_image_with_face_info(self, subset: BaseSubset, image_path: str, alpha_mask=False): img = load_image(image_path, alpha_mask) face_cx = face_cy = face_w = face_h = 0 if subset.face_crop_aug_range is not None: tokens = os.path.splitext(os.path.basename(image_path))[0].split("_") if len(tokens) >= 5: face_cx = int(tokens[-4]) face_cy = int(tokens[-3]) face_w = int(tokens[-2]) face_h = int(tokens[-1]) return img, face_cx, face_cy, face_w, face_h # いい感じに切り出す def crop_target(self, subset: BaseSubset, image, face_cx, face_cy, face_w, face_h): height, width = image.shape[0:2] if height == self.height and width == self.width: return image # 画像サイズはsizeより大きいのでリサイズする face_size = max(face_w, face_h) size = min(self.height, self.width) # 短いほう min_scale = max(self.height / height, self.width / width) # 画像がモデル入力サイズぴったりになる倍率（最小の倍率） min_scale = min(1.0, max(min_scale, size / (face_size * subset.face_crop_aug_range[1]))) # 指定した顔最小サイズ max_scale = min(1.0, max(min_scale, size / (face_size * subset.face_crop_aug_range[0]))) # 指定した顔最大サイズ if min_scale >= max_scale: # range指定がmin==max scale = min_scale else: scale = random.uniform(min_scale, max_scale) nh = int(height * scale + 0.5) nw = int(width * scale + 0.5) assert nh >= self.height and nw >= self.width, f"internal error. small scale {scale}, {width}*{height}" image = resize_image(image, width, height, nw, nh, subset.resize_interpolation) face_cx = int(face_cx * scale + 0.5) face_cy = int(face_cy * scale + 0.5) height, width = nh, nw # 顔を中心として448*640とかへ切り出す for axis, (target_size, length, face_p) in enumerate(zip((self.height, self.width), (height, width), (face_cy, face_cx))): p1 = face_p - target_size // 2 # 顔を中心に持ってくるための切り出し位置 if subset.random_crop: # 背景も含めるために顔を中心に置く確率を高めつつずらす range = max(length - face_p, face_p) # 画像の端から顔中心までの距離の長いほう p1 = p1 + (random.randint(0, range) + random.randint(0, range)) - range # -range ~ +range までのいい感じの乱数 else: # range指定があるときのみ、すこしだけランダムに（わりと適当） if subset.face_crop_aug_range[0] != subset.face_crop_aug_range[1]: if face_size > size // 10 and face_size >= 40: p1 = p1 + random.randint(-face_size // 20, +face_size // 20) p1 = max(0, min(p1, length - target_size)) if axis == 0: image = image[p1 : p1 + target_size, :] else: image = image[:, p1 : p1 + target_size] return image def __len__(self): return self._length def __getitem__(self, index): bucket = self.bucket_manager.buckets[self.buckets_indices[index].bucket_index] bucket_batch_size = self.buckets_indices[index].bucket_batch_size image_index = self.buckets_indices[index].batch_index * bucket_batch_size loss_weights = [] captions = [] input_ids_list = [] latents_list = [] alpha_mask_list = [] images = [] original_sizes_hw = [] crop_top_lefts = [] target_sizes_hw = [] flippeds = [] # 変数名が微妙 text_encoder_outputs_list = [] custom_attributes = [] masks = [] masked_images = [] for image_key in bucket[image_index : image_index + bucket_batch_size]: image_info = self.image_data[image_key] subset = self.image_to_subset[image_key] custom_attributes.append(subset.custom_attributes) # in case of fine tuning, is_reg is always False loss_weights.append(self.prior_loss_weight if image_info.is_reg else 1.0) flipped = subset.flip_aug and random.random() < 0.5 # not flipped or flipped with 50% chance # image/latentsを処理する if image_info.latents is not None: # cache_latents=Trueの場合 original_size = image_info.latents_original_size crop_ltrb = image_info.latents_crop_ltrb # calc values later if flipped if not flipped: latents = image_info.latents alpha_mask = image_info.alpha_mask else: latents = image_info.latents_flipped alpha_mask = None if image_info.alpha_mask is None else torch.flip(image_info.alpha_mask, [1]) image = None elif image_info.latents_npz is not None: # FineTuningDatasetまたはcache_latents_to_disk=Trueの場合 latents, original_size, crop_ltrb, flipped_latents, alpha_mask = ( self.latents_caching_strategy.load_latents_from_disk(image_info.latents_npz, image_info.bucket_reso) ) if flipped: latents = flipped_latents alpha_mask = None if alpha_mask is None else alpha_mask[:, ::-1].copy() # copy to avoid negative stride problem del flipped_latents latents = torch.FloatTensor(latents) if alpha_mask is not None: alpha_mask = torch.FloatTensor(alpha_mask) image = None else: # 画像を読み込み、必要ならcropする img, face_cx, face_cy, face_w, face_h = self.load_image_with_face_info( subset, image_info.absolute_path, subset.alpha_mask ) im_h, im_w = img.shape[0:2] if self.enable_bucket: img, original_size, crop_ltrb = trim_and_resize_if_required( subset.random_crop, img, image_info.bucket_reso, image_info.resized_size, resize_interpolation=image_info.resize_interpolation, ) else: if face_cx > 0: # 顔位置情報あり img = self.crop_target(subset, img, face_cx, face_cy, face_w, face_h) elif im_h > self.height or im_w > self.width: assert ( subset.random_crop ), f"image too large, but cropping and bucketing are disabled / 画像サイズが大きいのでface_crop_aug_rangeかrandom_crop、またはbucketを有効にしてください: {image_info.absolute_path}" if im_h > self.height: p = random.randint(0, im_h - self.height) img = img[p : p + self.height] if im_w > self.width: p = random.randint(0, im_w - self.width) img = img[:, p : p + self.width] im_h, im_w = img.shape[0:2] assert ( im_h == self.height and im_w == self.width ), f"image size is small / 画像サイズが小さいようです: {image_info.absolute_path}" original_size = [im_w, im_h] crop_ltrb = (0, 0, 0, 0) # augmentation aug = self.aug_helper.get_augmentor(subset.color_aug) if aug is not None: # augment RGB channels only img_rgb = img[:, :, :3] img_rgb = aug(image=img_rgb)["image"] img[:, :, :3] = img_rgb if flipped: img = img[:, ::-1, :].copy() # copy to avoid negative stride problem if subset.alpha_mask: if img.shape[2] == 4: alpha_mask = img[:, :, 3] # [H,W] alpha_mask = alpha_mask.astype(np.float32) / 255.0 # 0.0~1.0 alpha_mask = torch.FloatTensor(alpha_mask) else: alpha_mask = torch.ones((img.shape[0], img.shape[1]), dtype=torch.float32) else: alpha_mask = None img = img[:, :, :3] # remove alpha channel if self.train_inpainting: pil_image = transforms.functional.to_pil_image(img) mask = self.random_mask(pil_image.size) mask, masked_image = self.prepare_mask_and_masked_image(pil_image, mask) masks.append(mask) masked_images.append(masked_image) latents = None image = self.image_transforms(img) # -1.0~1.0のtorch.Tensorになる del img images.append(image) latents_list.append(latents) alpha_mask_list.append(alpha_mask) target_size = (image.shape[2], image.shape[1]) if image is not None else (latents.shape[2] * 8, latents.shape[1] * 8) if not flipped: crop_left_top = (crop_ltrb[0], crop_ltrb[1]) else: # crop_ltrb[2] is right, so target_size[0] - crop_ltrb[2] is left in flipped image crop_left_top = (target_size[0] - crop_ltrb[2], crop_ltrb[1]) original_sizes_hw.append((int(original_size[1]), int(original_size[0]))) crop_top_lefts.append((int(crop_left_top[1]), int(crop_left_top[0]))) target_sizes_hw.append((int(target_size[1]), int(target_size[0]))) flippeds.append(flipped) # captionとtext encoder outputを処理する caption = image_info.caption # default tokenization_required = ( self.text_encoder_output_caching_strategy is None or self.text_encoder_output_caching_strategy.is_partial ) text_encoder_outputs = None input_ids = None if image_info.text_encoder_outputs is not None: # cached text_encoder_outputs = image_info.text_encoder_outputs elif image_info.text_encoder_outputs_npz is not None: # on disk text_encoder_outputs = self.text_encoder_output_caching_strategy.load_outputs_npz( image_info.text_encoder_outputs_npz ) else: tokenization_required = True text_encoder_outputs_list.append(text_encoder_outputs) if tokenization_required: caption = self.process_caption(subset, image_info.caption) input_ids = [ids[0] for ids in self.tokenize_strategy.tokenize(caption)] # remove batch dimension input_ids_list.append(input_ids) captions.append(caption) def none_or_stack_elements(tensors_list, converter): # [[clip_l, clip_g, t5xxl], [clip_l, clip_g, t5xxl], ...] -> [torch.stack(clip_l), torch.stack(clip_g), torch.stack(t5xxl)] if len(tensors_list) == 0 or tensors_list[0] == None or len(tensors_list[0]) == 0 or tensors_list[0][0] is None: return None # old implementation without padding: all elements must have same length # return [torch.stack([converter(x[i]) for x in tensors_list]) for i in range(len(tensors_list[0]))] # new implementation with padding support result = [] for i in range(len(tensors_list[0])): tensors = [x[i] for x in tensors_list] if tensors[0].ndim == 0: # scalar value: e.g. ocr mask result.append(torch.stack([converter(x[i]) for x in tensors_list])) continue min_len = min([len(x) for x in tensors]) max_len = max([len(x) for x in tensors]) if min_len == max_len: # no padding result.append(torch.stack([converter(x) for x in tensors])) else: # padding tensors = [converter(x) for x in tensors] if tensors[0].ndim == 1: # input_ids or mask result.append(torch.stack([(torch.nn.functional.pad(x, (0, max_len - x.shape[0]))) for x in tensors])) else: # text encoder outputs result.append(torch.stack([(torch.nn.functional.pad(x, (0, 0, 0, max_len - x.shape[0]))) for x in tensors])) return result # set example example = {} example["custom_attributes"] = custom_attributes # may be list of empty dict example["loss_weights"] = torch.FloatTensor(loss_weights) example["text_encoder_outputs_list"] = none_or_stack_elements(text_encoder_outputs_list, torch.FloatTensor) example["input_ids_list"] = none_or_stack_elements(input_ids_list, lambda x: x) # if one of alpha_masks is not None, we need to replace None with ones none_or_not = [x is None for x in alpha_mask_list] if all(none_or_not): example["alpha_masks"] = None elif any(none_or_not): for i in range(len(alpha_mask_list)): if alpha_mask_list[i] is None: if images[i] is not None: alpha_mask_list[i] = torch.ones((images[i].shape[1], images[i].shape[2]), dtype=torch.float32) else: alpha_mask_list[i] = torch.ones( (latents_list[i].shape[1] * 8, latents_list[i].shape[2] * 8), dtype=torch.float32 ) example["alpha_masks"] = torch.stack(alpha_mask_list) else: example["alpha_masks"] = torch.stack(alpha_mask_list) if images[0] is not None: images = torch.stack(images) images = images.to(memory_format=torch.contiguous_format).float() else: images = None example["images"] = images example["masks"] = torch.stack(masks) if masks else None example["masked_images"] = torch.stack(masked_images) if masked_images else None example["latents"] = torch.stack(latents_list) if latents_list[0] is not None else None example["captions"] = captions example["original_sizes_hw"] = torch.stack([torch.LongTensor(x) for x in original_sizes_hw]) example["crop_top_lefts"] = torch.stack([torch.LongTensor(x) for x in crop_top_lefts]) example["target_sizes_hw"] = torch.stack([torch.LongTensor(x) for x in target_sizes_hw]) example["flippeds"] = flippeds example["network_multipliers"] = torch.FloatTensor([self.network_multiplier] * len(captions)) if self.debug_dataset: example["image_keys"] = bucket[image_index : image_index + self.batch_size] return example @staticmethod def prepare_mask_and_masked_image(image, mask): image = np.array(image.convert("RGB")) image = image.transpose(2, 0, 1) # HWC -> CHW image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0 mask = np.array(mask.convert("L")) mask = mask.astype(np.float32) / 255.0 mask = mask[None] # 1,H,W mask[mask < 0.5] = 0 mask[mask >= 0.5] = 1 mask = torch.from_numpy(mask) masked_image = image * (mask < 0.5) return mask, masked_image # generate random masks @staticmethod def random_mask(im_shape): from library.mask_generator import random_mask as _random_mask w, h = im_shape return _random_mask(w, h) class DatasetGroup(torch.utils.data.ConcatDataset): # DreamBoothDataset / FineTuningDataset still live in library.train_util (PR-1d will move them). def __init__(self, datasets: 'Sequence[Union[DreamBoothDataset, FineTuningDataset]]'): self.datasets: 'List[Union[DreamBoothDataset, FineTuningDataset]]' super().__init__(datasets) self.image_data = {} self.num_train_images = 0 self.num_reg_images = 0 # simply concat together # TODO: handling image_data key duplication among dataset # In practical, this is not the big issue because image_data is accessed from outside of dataset only for debug_dataset. for dataset in datasets: self.image_data.update(dataset.image_data) self.num_train_images += dataset.num_train_images self.num_reg_images += dataset.num_reg_images def add_replacement(self, str_from, str_to): for dataset in self.datasets: dataset.add_replacement(str_from, str_to) # def make_buckets(self): # for dataset in self.datasets: # dataset.make_buckets() def set_text_encoder_output_caching_strategy(self, strategy: TextEncoderOutputsCachingStrategy): """ DataLoader is run in multiple processes, so we need to set the strategy manually. """ for dataset in self.datasets: dataset.set_text_encoder_output_caching_strategy(strategy) def enable_XTI(self, *args, **kwargs): for dataset in self.datasets: dataset.enable_XTI(*args, **kwargs) def new_cache_latents(self, model: Any, accelerator: Accelerator): for i, dataset in enumerate(self.datasets): logger.info(f"[Dataset {i}]") dataset.new_cache_latents(model, accelerator) accelerator.wait_for_everyone() def new_cache_text_encoder_outputs(self, models: List[Any], accelerator: Accelerator): for i, dataset in enumerate(self.datasets): logger.info(f"[Dataset {i}]") dataset.new_cache_text_encoder_outputs(models, accelerator) accelerator.wait_for_everyone() def verify_bucket_reso_steps(self, min_steps: int): for dataset in self.datasets: dataset.verify_bucket_reso_steps(min_steps) def get_resolutions(self) -> List[Tuple[int, int]]: return [(dataset.width, dataset.height) for dataset in self.datasets] def is_latent_cacheable(self) -> bool: return all([dataset.is_latent_cacheable() for dataset in self.datasets]) def is_text_encoder_output_cacheable(self, cache_supports_dropout: bool = False) -> bool: return all([dataset.is_text_encoder_output_cacheable(cache_supports_dropout) for dataset in self.datasets]) def set_current_strategies(self): for dataset in self.datasets: dataset.set_current_strategies() def set_current_epoch(self, epoch): for dataset in self.datasets: dataset.set_current_epoch(epoch) def set_current_step(self, step): for dataset in self.datasets: dataset.set_current_step(step) def set_max_train_steps(self, max_train_steps): for dataset in self.datasets: dataset.set_max_train_steps(max_train_steps) def disable_token_padding(self): for dataset in self.datasets: dataset.disable_token_padding() def debug_dataset(train_dataset, show_input_ids=False): logger.info(f"Total dataset length (steps) / データセットの長さ（ステップ数）: {len(train_dataset)}") logger.info( "`S` for next step, `E` for next epoch no. , Escape for exit. / Sキーで次のステップ、Eキーで次のエポック、Escキーで中断、終了します" ) epoch = 1 while True: logger.info(f"") logger.info(f"epoch: {epoch}") steps = (epoch - 1) * len(train_dataset) + 1 indices = list(range(len(train_dataset))) random.shuffle(indices) k = 0 for i, idx in enumerate(indices): train_dataset.set_current_epoch(epoch) train_dataset.set_current_step(steps) logger.info(f"steps: {steps} ({i + 1}/{len(train_dataset)})") example = train_dataset[idx] if example["latents"] is not None: logger.info(f"sample has latents from npz file: {example['latents'].size()}") for j, (ik, cap, lw, orgsz, crptl, trgsz, flpdz) in enumerate( zip( example["image_keys"], example["captions"], example["loss_weights"], # example["input_ids"], example["original_sizes_hw"], example["crop_top_lefts"], example["target_sizes_hw"], example["flippeds"], ) ): logger.info( f'{ik}, size: {train_dataset.image_data[ik].image_size}, loss weight: {lw}, caption: "{cap}", original size: {orgsz}, crop top left: {crptl}, target size: {trgsz}, flipped: {flpdz}' ) if "network_multipliers" in example: logger.info(f"network multiplier: {example['network_multipliers'][j]}") if "custom_attributes" in example: logger.info(f"custom attributes: {example['custom_attributes'][j]}") # if show_input_ids: # logger.info(f"input ids: {iid}") # if "input_ids2" in example: # logger.info(f"input ids2: {example['input_ids2'][j]}") if example["images"] is not None: im = example["images"][j] logger.info(f"image size: {im.size()}") im = ((im.numpy() + 1.0) * 127.5).astype(np.uint8) im = np.transpose(im, (1, 2, 0)) # c,H,W -> H,W,c im = im[:, :, ::-1] # RGB -> BGR (OpenCV) if "conditioning_images" in example or "masked_images" in example: cond_img = example["conditioning_images"][j] if "conditioning_images" in example else example["masked_images"][j] logger.info(f"conditioning image size: {cond_img.size()}") cond_img = ((cond_img.numpy() + 1.0) * 127.5).astype(np.uint8) cond_img = np.transpose(cond_img, (1, 2, 0)) cond_img = cond_img[:, :, ::-1] if os.name == "nt": cv2.imshow("cond_img", cond_img) if "alpha_masks" in example and example["alpha_masks"] is not None: alpha_mask = example["alpha_masks"][j] logger.info(f"alpha mask size: {alpha_mask.size()}") alpha_mask = (alpha_mask.numpy() * 255.0).astype(np.uint8) if os.name == "nt": cv2.imshow("alpha_mask", alpha_mask) if os.name == "nt": # only windows cv2.imshow("img", im) k = cv2.waitKey() cv2.destroyAllWindows() if k == 27 or k == ord("s") or k == ord("e"): break steps += 1 if k == ord("e"): break if k == 27 or (example["images"] is None and i >= 8): k = 27 break if k == 27: break epoch += 1 def glob_images(directory, base="*"): img_paths = [] for ext in IMAGE_EXTENSIONS: if base == "*": img_paths.extend(glob.glob(os.path.join(glob.escape(directory), base + ext))) else: img_paths.extend(glob.glob(glob.escape(os.path.join(directory, base + ext)))) img_paths = list(set(img_paths)) # 重複を排除 img_paths.sort() return img_paths def glob_images_pathlib(dir_path, recursive): image_paths = [] if recursive: for ext in IMAGE_EXTENSIONS: image_paths += list(dir_path.rglob("*" + ext)) else: for ext in IMAGE_EXTENSIONS: image_paths += list(dir_path.glob("*" + ext)) image_paths = list(set(image_paths)) # 重複を排除 image_paths.sort() return image_paths class MinimalDataset(BaseDataset): def __init__(self, resolution, network_multiplier, train_inpainting=False, debug_dataset=False): super().__init__(resolution, network_multiplier, train_inpainting, debug_dataset) self.num_train_images = 0 # update in subclass self.num_reg_images = 0 # update in subclass self.datasets = [self] self.batch_size = 1 # update in subclass self.subsets = [self] self.num_repeats = 1 # update in subclass if needed self.img_count = 1 # update in subclass if needed self.bucket_info = {} self.is_reg = False self.image_dir = "dummy" # for metadata def verify_bucket_reso_steps(self, min_steps: int): pass def is_latent_cacheable(self) -> bool: return False def __len__(self): raise NotImplementedError # override to avoid shuffling buckets def set_current_epoch(self, epoch): self.current_epoch = epoch def __getitem__(self, idx): r""" The subclass may have image_data for debug_dataset, which is a dict of ImageInfo objects. Returns: example like this: for i in range(batch_size): image_key = ... # whatever hashable image_keys.append(image_key) image = ... # PIL Image img_tensor = self.image_transforms(img) images.append(img_tensor) caption = ... # str input_ids = [ids[0] for ids in self.tokenize_strategy.tokenize(caption)] input_ids_list.append(input_ids) captions.append(caption) images = torch.stack(images, dim=0) input_ids_list = torch.stack(input_ids_list, dim=0) example = { "images": images, "input_ids": input_ids_list, "captions": captions, # for debug_dataset "latents": None, "image_keys": image_keys, # for debug_dataset "loss_weights": torch.ones(batch_size, dtype=torch.float32), } return example """ raise NotImplementedError def get_resolutions(self) -> List[Tuple[int, int]]: return [] def load_arbitrary_dataset(args, tokenizer=None) -> MinimalDataset: module = ".".join(args.dataset_class.split(".")[:-1]) dataset_class = args.dataset_class.split(".")[-1] module = importlib.import_module(module) dataset_class = getattr(module, dataset_class) train_dataset_group: MinimalDataset = dataset_class(tokenizer, args.max_token_length, args.resolution, args.debug_dataset) return train_dataset_group # collate_fn 用 epoch, step は multiprocessing.Value class collator_class: def __init__(self, epoch, step, dataset): self.current_epoch = epoch self.current_step = step self.dataset = dataset # not used if worker_info is not None, in case of multiprocessing def __call__(self, examples): worker_info = torch.utils.data.get_worker_info() # worker_info is None in the main process if worker_info is not None: dataset = worker_info.dataset else: dataset = self.dataset # set epoch and step dataset.set_current_epoch(self.current_epoch.value) dataset.set_current_step(self.current_step.value) return examples[0]