Allex21/LT
0
1# common functions for training2 3import argparse4import ast5import asyncio6import datetime7import importlib8import json9import logging10import pathlib11import re12import shutil13import time14from typing import (15 Dict,16 List,17 NamedTuple,18 Optional,19 Sequence,20 Tuple,21 Union,22)23from accelerate import Accelerator, InitProcessGroupKwargs, DistributedDataParallelKwargs, PartialState24import glob25import math26import os27import random28import hashlib29import subprocess30from io import BytesIO31import toml32 33from tqdm import tqdm34 35import torch36from library.device_utils import init_ipex, clean_memory_on_device37 38init_ipex()39 40from torch.nn.parallel import DistributedDataParallel as DDP41from torch.optim import Optimizer42from torchvision import transforms43from transformers import CLIPTokenizer, CLIPTextModel, CLIPTextModelWithProjection44import transformers45from diffusers.optimization import (46 SchedulerType as DiffusersSchedulerType,47 TYPE_TO_SCHEDULER_FUNCTION as DIFFUSERS_TYPE_TO_SCHEDULER_FUNCTION,48)49from transformers.optimization import SchedulerType, TYPE_TO_SCHEDULER_FUNCTION50from diffusers import (51 StableDiffusionPipeline,52 DDPMScheduler,53 EulerAncestralDiscreteScheduler,54 DPMSolverMultistepScheduler,55 DPMSolverSinglestepScheduler,56 LMSDiscreteScheduler,57 PNDMScheduler,58 DDIMScheduler,59 EulerDiscreteScheduler,60 HeunDiscreteScheduler,61 KDPM2DiscreteScheduler,62 KDPM2AncestralDiscreteScheduler,63 AutoencoderKL,64)65from library import custom_train_functions66from library.original_unet import UNet2DConditionModel67from huggingface_hub import hf_hub_download68import numpy as np69from PIL import Image70import imagesize71import cv272import safetensors.torch73from library.lpw_stable_diffusion import StableDiffusionLongPromptWeightingPipeline74import library.model_util as model_util75import library.huggingface_util as huggingface_util76import library.sai_model_spec as sai_model_spec77import library.deepspeed_utils as deepspeed_utils78from library.utils import setup_logging, pil_resize79 80setup_logging()81import logging82 83logger = logging.getLogger(__name__)84# from library.attention_processors import FlashAttnProcessor85# from library.hypernetwork import replace_attentions_for_hypernetwork86from library.original_unet import UNet2DConditionModel87 88# Tokenizer: checkpointから読み込むのではなくあらかじめ提供されているものを使う89TOKENIZER_PATH = "openai/clip-vit-large-patch14"90V2_STABLE_DIFFUSION_PATH = "stabilityai/stable-diffusion-2" # ここからtokenizerだけ使う v2とv2.1はtokenizer仕様は同じ91 92HIGH_VRAM = False93 94# checkpointファイル名95EPOCH_STATE_NAME = "{}-{:06d}-state"96EPOCH_FILE_NAME = "{}-{:06d}"97EPOCH_DIFFUSERS_DIR_NAME = "{}-{:06d}"98LAST_STATE_NAME = "{}-state"99DEFAULT_EPOCH_NAME = "epoch"100DEFAULT_LAST_OUTPUT_NAME = "last"101 102DEFAULT_STEP_NAME = "at"103STEP_STATE_NAME = "{}-step{:08d}-state"104STEP_FILE_NAME = "{}-step{:08d}"105STEP_DIFFUSERS_DIR_NAME = "{}-step{:08d}"106 107# region dataset108 109IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".webp", ".bmp", ".PNG", ".JPG", ".JPEG", ".WEBP", ".BMP"]110 111try:112 import pillow_avif113 114 IMAGE_EXTENSIONS.extend([".avif", ".AVIF"])115except:116 pass117 118# JPEG-XL on Linux119try:120 from jxlpy import JXLImagePlugin121 122 IMAGE_EXTENSIONS.extend([".jxl", ".JXL"])123except:124 pass125 126# JPEG-XL on Windows127try:128 import pillow_jxl129 130 IMAGE_EXTENSIONS.extend([".jxl", ".JXL"])131except:132 pass133 134IMAGE_TRANSFORMS = transforms.Compose(135 [136 transforms.ToTensor(),137 transforms.Normalize([0.5], [0.5]),138 ]139)140 141TEXT_ENCODER_OUTPUTS_CACHE_SUFFIX = "_te_outputs.npz"142 143 144class ImageInfo:145 def __init__(self, image_key: str, num_repeats: int, caption: str, is_reg: bool, absolute_path: str) -> None:146 self.image_key: str = image_key147 self.num_repeats: int = num_repeats148 self.caption: str = caption149 self.is_reg: bool = is_reg150 self.absolute_path: str = absolute_path151 self.image_size: Tuple[int, int] = None152 self.resized_size: Tuple[int, int] = None153 self.bucket_reso: Tuple[int, int] = None154 self.latents: torch.Tensor = None155 self.latents_flipped: torch.Tensor = None156 self.latents_npz: str = None157 self.latents_original_size: Tuple[int, int] = None # original image size, not latents size158 self.latents_crop_ltrb: Tuple[int, int] = None # crop left top right bottom in original pixel size, not latents size159 self.cond_img_path: str = None160 self.image: Optional[Image.Image] = None # optional, original PIL Image161 # SDXL, optional162 self.text_encoder_outputs_npz: Optional[str] = None163 self.text_encoder_outputs1: Optional[torch.Tensor] = None164 self.text_encoder_outputs2: Optional[torch.Tensor] = None165 self.text_encoder_pool2: Optional[torch.Tensor] = None166 self.alpha_mask: Optional[torch.Tensor] = None # alpha mask can be flipped in runtime167 168 169class BucketManager:170 def __init__(self, no_upscale, max_reso, min_size, max_size, reso_steps) -> None:171 if max_size is not None:172 if max_reso is not None:173 assert max_size >= max_reso[0], "the max_size should be larger than the width of max_reso"174 assert max_size >= max_reso[1], "the max_size should be larger than the height of max_reso"175 if min_size is not None:176 assert max_size >= min_size, "the max_size should be larger than the min_size"177 178 self.no_upscale = no_upscale179 if max_reso is None:180 self.max_reso = None181 self.max_area = None182 else:183 self.max_reso = max_reso184 self.max_area = max_reso[0] * max_reso[1]185 self.min_size = min_size186 self.max_size = max_size187 self.reso_steps = reso_steps188 189 self.resos = []190 self.reso_to_id = {}191 self.buckets = [] # 前処理時は (image_key, image, original size, crop left/top)、学習時は image_key192 193 def add_image(self, reso, image_or_info):194 bucket_id = self.reso_to_id[reso]195 self.buckets[bucket_id].append(image_or_info)196 197 def shuffle(self):198 for bucket in self.buckets:199 random.shuffle(bucket)200 201 def sort(self):202 # 解像度順にソートする（表示時、メタデータ格納時の見栄えをよくするためだけ）。bucketsも入れ替えてreso_to_idも振り直す203 sorted_resos = self.resos.copy()204 sorted_resos.sort()205 206 sorted_buckets = []207 sorted_reso_to_id = {}208 for i, reso in enumerate(sorted_resos):209 bucket_id = self.reso_to_id[reso]210 sorted_buckets.append(self.buckets[bucket_id])211 sorted_reso_to_id[reso] = i212 213 self.resos = sorted_resos214 self.buckets = sorted_buckets215 self.reso_to_id = sorted_reso_to_id216 217 def make_buckets(self):218 resos = model_util.make_bucket_resolutions(self.max_reso, self.min_size, self.max_size, self.reso_steps)219 self.set_predefined_resos(resos)220 221 def set_predefined_resos(self, resos):222 # 規定サイズから選ぶ場合の解像度、aspect ratioの情報を格納しておく223 self.predefined_resos = resos.copy()224 self.predefined_resos_set = set(resos)225 self.predefined_aspect_ratios = np.array([w / h for w, h in resos])226 227 def add_if_new_reso(self, reso):228 if reso not in self.reso_to_id:229 bucket_id = len(self.resos)230 self.reso_to_id[reso] = bucket_id231 self.resos.append(reso)232 self.buckets.append([])233 # logger.info(reso, bucket_id, len(self.buckets))234 235 def round_to_steps(self, x):236 x = int(x + 0.5)237 return x - x % self.reso_steps238 239 def select_bucket(self, image_width, image_height):240 aspect_ratio = image_width / image_height241 if not self.no_upscale:242 # 拡大および縮小を行う243 # 同じaspect ratioがあるかもしれないので（fine tuningで、no_upscale=Trueで前処理した場合）、解像度が同じものを優先する244 reso = (image_width, image_height)245 if reso in self.predefined_resos_set:246 pass247 else:248 ar_errors = self.predefined_aspect_ratios - aspect_ratio249 predefined_bucket_id = np.abs(ar_errors).argmin() # 当該解像度以外でaspect ratio errorが最も少ないもの250 reso = self.predefined_resos[predefined_bucket_id]251 252 ar_reso = reso[0] / reso[1]253 if aspect_ratio > ar_reso: # 横が長い→縦を合わせる254 scale = reso[1] / image_height255 else:256 scale = reso[0] / image_width257 258 resized_size = (int(image_width * scale + 0.5), int(image_height * scale + 0.5))259 # logger.info(f"use predef, {image_width}, {image_height}, {reso}, {resized_size}")260 else:261 # 縮小のみを行う262 if image_width * image_height > self.max_area:263 # 画像が大きすぎるのでアスペクト比を保ったまま縮小することを前提にbucketを決める264 resized_width = math.sqrt(self.max_area * aspect_ratio)265 resized_height = self.max_area / resized_width266 assert abs(resized_width / resized_height - aspect_ratio) < 1e-2, "aspect is illegal"267 268 # リサイズ後の短辺または長辺をreso_steps単位にする：aspect ratioの差が少ないほうを選ぶ269 # 元のbucketingと同じロジック270 b_width_rounded = self.round_to_steps(resized_width)271 b_height_in_wr = self.round_to_steps(b_width_rounded / aspect_ratio)272 ar_width_rounded = b_width_rounded / b_height_in_wr273 274 b_height_rounded = self.round_to_steps(resized_height)275 b_width_in_hr = self.round_to_steps(b_height_rounded * aspect_ratio)276 ar_height_rounded = b_width_in_hr / b_height_rounded277 278 # logger.info(b_width_rounded, b_height_in_wr, ar_width_rounded)279 # logger.info(b_width_in_hr, b_height_rounded, ar_height_rounded)280 281 if abs(ar_width_rounded - aspect_ratio) < abs(ar_height_rounded - aspect_ratio):282 resized_size = (b_width_rounded, int(b_width_rounded / aspect_ratio + 0.5))283 else:284 resized_size = (int(b_height_rounded * aspect_ratio + 0.5), b_height_rounded)285 # logger.info(resized_size)286 else:287 resized_size = (image_width, image_height) # リサイズは不要288 289 # 画像のサイズ未満をbucketのサイズとする（paddingせずにcroppingする）290 bucket_width = resized_size[0] - resized_size[0] % self.reso_steps291 bucket_height = resized_size[1] - resized_size[1] % self.reso_steps292 # logger.info(f"use arbitrary {image_width}, {image_height}, {resized_size}, {bucket_width}, {bucket_height}")293 294 reso = (bucket_width, bucket_height)295 296 self.add_if_new_reso(reso)297 298 ar_error = (reso[0] / reso[1]) - aspect_ratio299 return reso, resized_size, ar_error300 301 @staticmethod302 def get_crop_ltrb(bucket_reso: Tuple[int, int], image_size: Tuple[int, int]):303 # Stability AIの前処理に合わせてcrop left/topを計算する。crop rightはflipのaugmentationのために求める304 # Calculate crop left/top according to the preprocessing of Stability AI. Crop right is calculated for flip augmentation.305 306 bucket_ar = bucket_reso[0] / bucket_reso[1]307 image_ar = image_size[0] / image_size[1]308 if bucket_ar > image_ar:309 # bucketのほうが横長→縦を合わせる310 resized_width = bucket_reso[1] * image_ar311 resized_height = bucket_reso[1]312 else:313 resized_width = bucket_reso[0]314 resized_height = bucket_reso[0] / image_ar315 crop_left = (bucket_reso[0] - resized_width) // 2316 crop_top = (bucket_reso[1] - resized_height) // 2317 crop_right = crop_left + resized_width318 crop_bottom = crop_top + resized_height319 return crop_left, crop_top, crop_right, crop_bottom320 321 322class BucketBatchIndex(NamedTuple):323 bucket_index: int324 bucket_batch_size: int325 batch_index: int326 327 328class AugHelper:329 # albumentationsへの依存をなくしたがとりあえず同じinterfaceを持たせる330 331 def __init__(self):332 pass333 334 def color_aug(self, image: np.ndarray):335 # self.color_aug_method = albu.OneOf(336 # [337 # albu.HueSaturationValue(8, 0, 0, p=0.5),338 # albu.RandomGamma((95, 105), p=0.5),339 # ],340 # p=0.33,341 # )342 hue_shift_limit = 8343 344 # remove dependency to albumentations345 if random.random() <= 0.33:346 if random.random() > 0.5:347 # hue shift348 hsv_img = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)349 hue_shift = random.uniform(-hue_shift_limit, hue_shift_limit)350 if hue_shift < 0:351 hue_shift = 180 + hue_shift352 hsv_img[:, :, 0] = (hsv_img[:, :, 0] + hue_shift) % 180353 image = cv2.cvtColor(hsv_img, cv2.COLOR_HSV2BGR)354 else:355 # random gamma356 gamma = random.uniform(0.95, 1.05)357 image = np.clip(image**gamma, 0, 255).astype(np.uint8)358 359 return {"image": image}360 361 def get_augmentor(self, use_color_aug: bool): # -> Optional[Callable[[np.ndarray], Dict[str, np.ndarray]]]:362 return self.color_aug if use_color_aug else None363 364 365class BaseSubset:366 def __init__(367 self,368 image_dir: Optional[str],369 alpha_mask: Optional[bool],370 num_repeats: int,371 shuffle_caption: bool,372 caption_separator: str,373 keep_tokens: int,374 keep_tokens_separator: str,375 secondary_separator: Optional[str],376 enable_wildcard: bool,377 color_aug: bool,378 flip_aug: bool,379 face_crop_aug_range: Optional[Tuple[float, float]],380 random_crop: bool,381 caption_dropout_rate: float,382 caption_dropout_every_n_epochs: int,383 caption_tag_dropout_rate: float,384 caption_prefix: Optional[str],385 caption_suffix: Optional[str],386 token_warmup_min: int,387 token_warmup_step: Union[float, int],388 ) -> None:389 self.image_dir = image_dir390 self.alpha_mask = alpha_mask if alpha_mask is not None else False391 self.num_repeats = num_repeats392 self.shuffle_caption = shuffle_caption393 self.caption_separator = caption_separator394 self.keep_tokens = keep_tokens395 self.keep_tokens_separator = keep_tokens_separator396 self.secondary_separator = secondary_separator397 self.enable_wildcard = enable_wildcard398 self.color_aug = color_aug399 self.flip_aug = flip_aug400 self.face_crop_aug_range = face_crop_aug_range401 self.random_crop = random_crop402 self.caption_dropout_rate = caption_dropout_rate403 self.caption_dropout_every_n_epochs = caption_dropout_every_n_epochs404 self.caption_tag_dropout_rate = caption_tag_dropout_rate405 self.caption_prefix = caption_prefix406 self.caption_suffix = caption_suffix407 408 self.token_warmup_min = token_warmup_min # step=0におけるタグの数409 self.token_warmup_step = token_warmup_step # N（N<1ならN*max_train_steps）ステップ目でタグの数が最大になる410 411 self.img_count = 0412 413 414class DreamBoothSubset(BaseSubset):415 def __init__(416 self,417 image_dir: str,418 is_reg: bool,419 class_tokens: Optional[str],420 caption_extension: str,421 cache_info: bool,422 alpha_mask: bool,423 num_repeats,424 shuffle_caption,425 caption_separator: str,426 keep_tokens,427 keep_tokens_separator,428 secondary_separator,429 enable_wildcard,430 color_aug,431 flip_aug,432 face_crop_aug_range,433 random_crop,434 caption_dropout_rate,435 caption_dropout_every_n_epochs,436 caption_tag_dropout_rate,437 caption_prefix,438 caption_suffix,439 token_warmup_min,440 token_warmup_step,441 ) -> None:442 assert image_dir is not None, "image_dir must be specified / image_dirは指定が必須です"443 444 super().__init__(445 image_dir,446 alpha_mask,447 num_repeats,448 shuffle_caption,449 caption_separator,450 keep_tokens,451 keep_tokens_separator,452 secondary_separator,453 enable_wildcard,454 color_aug,455 flip_aug,456 face_crop_aug_range,457 random_crop,458 caption_dropout_rate,459 caption_dropout_every_n_epochs,460 caption_tag_dropout_rate,461 caption_prefix,462 caption_suffix,463 token_warmup_min,464 token_warmup_step,465 )466 467 self.is_reg = is_reg468 self.class_tokens = class_tokens469 self.caption_extension = caption_extension470 if self.caption_extension and not self.caption_extension.startswith("."):471 self.caption_extension = "." + self.caption_extension472 self.cache_info = cache_info473 474 def __eq__(self, other) -> bool:475 if not isinstance(other, DreamBoothSubset):476 return NotImplemented477 return self.image_dir == other.image_dir478 479 480class FineTuningSubset(BaseSubset):481 def __init__(482 self,483 image_dir,484 metadata_file: str,485 alpha_mask: bool,486 num_repeats,487 shuffle_caption,488 caption_separator,489 keep_tokens,490 keep_tokens_separator,491 secondary_separator,492 enable_wildcard,493 color_aug,494 flip_aug,495 face_crop_aug_range,496 random_crop,497 caption_dropout_rate,498 caption_dropout_every_n_epochs,499 caption_tag_dropout_rate,500 caption_prefix,501 caption_suffix,502 token_warmup_min,503 token_warmup_step,504 ) -> None:505 assert metadata_file is not None, "metadata_file must be specified / metadata_fileは指定が必須です"506 507 super().__init__(508 image_dir,509 alpha_mask,510 num_repeats,511 shuffle_caption,512 caption_separator,513 keep_tokens,514 keep_tokens_separator,515 secondary_separator,516 enable_wildcard,517 color_aug,518 flip_aug,519 face_crop_aug_range,520 random_crop,521 caption_dropout_rate,522 caption_dropout_every_n_epochs,523 caption_tag_dropout_rate,524 caption_prefix,525 caption_suffix,526 token_warmup_min,527 token_warmup_step,528 )529 530 self.metadata_file = metadata_file531 532 def __eq__(self, other) -> bool:533 if not isinstance(other, FineTuningSubset):534 return NotImplemented535 return self.metadata_file == other.metadata_file536 537 538class ControlNetSubset(BaseSubset):539 def __init__(540 self,541 image_dir: str,542 conditioning_data_dir: str,543 caption_extension: str,544 cache_info: bool,545 num_repeats,546 shuffle_caption,547 caption_separator,548 keep_tokens,549 keep_tokens_separator,550 secondary_separator,551 enable_wildcard,552 color_aug,553 flip_aug,554 face_crop_aug_range,555 random_crop,556 caption_dropout_rate,557 caption_dropout_every_n_epochs,558 caption_tag_dropout_rate,559 caption_prefix,560 caption_suffix,561 token_warmup_min,562 token_warmup_step,563 ) -> None:564 assert image_dir is not None, "image_dir must be specified / image_dirは指定が必須です"565 566 super().__init__(567 image_dir,568 False, # alpha_mask569 num_repeats,570 shuffle_caption,571 caption_separator,572 keep_tokens,573 keep_tokens_separator,574 secondary_separator,575 enable_wildcard,576 color_aug,577 flip_aug,578 face_crop_aug_range,579 random_crop,580 caption_dropout_rate,581 caption_dropout_every_n_epochs,582 caption_tag_dropout_rate,583 caption_prefix,584 caption_suffix,585 token_warmup_min,586 token_warmup_step,587 )588 589 self.conditioning_data_dir = conditioning_data_dir590 self.caption_extension = caption_extension591 if self.caption_extension and not self.caption_extension.startswith("."):592 self.caption_extension = "." + self.caption_extension593 self.cache_info = cache_info594 595 def __eq__(self, other) -> bool:596 if not isinstance(other, ControlNetSubset):597 return NotImplemented598 return self.image_dir == other.image_dir and self.conditioning_data_dir == other.conditioning_data_dir599 600 601class BaseDataset(torch.utils.data.Dataset):602 def __init__(603 self,604 tokenizer: Union[CLIPTokenizer, List[CLIPTokenizer]],605 max_token_length: int,606 resolution: Optional[Tuple[int, int]],607 network_multiplier: float,608 debug_dataset: bool,609 ) -> None:610 super().__init__()611 612 self.tokenizers = tokenizer if isinstance(tokenizer, list) else [tokenizer]613 614 self.max_token_length = max_token_length615 # width/height is used when enable_bucket==False616 self.width, self.height = (None, None) if resolution is None else resolution617 self.network_multiplier = network_multiplier618 self.debug_dataset = debug_dataset619 620 self.subsets: List[Union[DreamBoothSubset, FineTuningSubset]] = []621 622 self.token_padding_disabled = False623 self.tag_frequency = {}624 self.XTI_layers = None625 self.token_strings = None626 627 self.enable_bucket = False628 self.bucket_manager: BucketManager = None # not initialized629 self.min_bucket_reso = None630 self.max_bucket_reso = None631 self.bucket_reso_steps = None632 self.bucket_no_upscale = None633 self.bucket_info = None # for metadata634 635 self.tokenizer_max_length = self.tokenizers[0].model_max_length if max_token_length is None else max_token_length + 2636 637 self.current_epoch: int = 0 # インスタンスがepochごとに新しく作られるようなので外側から渡さないとダメ638 639 self.current_step: int = 0640 self.max_train_steps: int = 0641 self.seed: int = 0642 643 # augmentation644 self.aug_helper = AugHelper()645 646 self.image_transforms = IMAGE_TRANSFORMS647 648 self.image_data: Dict[str, ImageInfo] = {}649 self.image_to_subset: Dict[str, Union[DreamBoothSubset, FineTuningSubset]] = {}650 651 self.replacements = {}652 653 # caching654 self.caching_mode = None # None, 'latents', 'text'655 656 def adjust_min_max_bucket_reso_by_steps(657 self, resolution: Tuple[int, int], min_bucket_reso: int, max_bucket_reso: int, bucket_reso_steps: int658 ) -> Tuple[int, int]:659 # make min/max bucket reso to be multiple of bucket_reso_steps660 if min_bucket_reso % bucket_reso_steps != 0:661 adjusted_min_bucket_reso = min_bucket_reso - min_bucket_reso % bucket_reso_steps662 logger.warning(663 f"min_bucket_reso is adjusted to be multiple of bucket_reso_steps"664 f" / min_bucket_resoがbucket_reso_stepsの倍数になるように調整されました: {min_bucket_reso} -> {adjusted_min_bucket_reso}"665 )666 min_bucket_reso = adjusted_min_bucket_reso667 if max_bucket_reso % bucket_reso_steps != 0:668 adjusted_max_bucket_reso = max_bucket_reso + bucket_reso_steps - max_bucket_reso % bucket_reso_steps669 logger.warning(670 f"max_bucket_reso is adjusted to be multiple of bucket_reso_steps"671 f" / max_bucket_resoがbucket_reso_stepsの倍数になるように調整されました: {max_bucket_reso} -> {adjusted_max_bucket_reso}"672 )673 max_bucket_reso = adjusted_max_bucket_reso674 675 assert (676 min(resolution) >= min_bucket_reso677 ), f"min_bucket_reso must be equal or less than resolution / min_bucket_resoは最小解像度より大きくできません。解像度を大きくするかmin_bucket_resoを小さくしてください"678 assert (679 max(resolution) <= max_bucket_reso680 ), f"max_bucket_reso must be equal or greater than resolution / max_bucket_resoは最大解像度より小さくできません。解像度を小さくするかmin_bucket_resoを大きくしてください"681 682 return min_bucket_reso, max_bucket_reso683 684 def set_seed(self, seed):685 self.seed = seed686 687 def set_caching_mode(self, mode):688 self.caching_mode = mode689 690 def set_current_epoch(self, epoch):691 if not self.current_epoch == epoch: # epochが切り替わったらバケツをシャッフルする692 if epoch > self.current_epoch:693 logger.info("epoch is incremented. current_epoch: {}, epoch: {}".format(self.current_epoch, epoch))694 num_epochs = epoch - self.current_epoch695 for _ in range(num_epochs):696 self.current_epoch += 1697 self.shuffle_buckets()698 # self.current_epoch seem to be set to 0 again in the next epoch. it may be caused by skipped_dataloader?699 else:700 logger.warning("epoch is not incremented. current_epoch: {}, epoch: {}".format(self.current_epoch, epoch))701 self.current_epoch = epoch702 703 def set_current_step(self, step):704 self.current_step = step705 706 def set_max_train_steps(self, max_train_steps):707 self.max_train_steps = max_train_steps708 709 def set_tag_frequency(self, dir_name, captions):710 frequency_for_dir = self.tag_frequency.get(dir_name, {})711 self.tag_frequency[dir_name] = frequency_for_dir712 for caption in captions:713 for tag in caption.split(","):714 tag = tag.strip()715 if tag:716 tag = tag.lower()717 frequency = frequency_for_dir.get(tag, 0)718 frequency_for_dir[tag] = frequency + 1719 720 def disable_token_padding(self):721 self.token_padding_disabled = True722 723 def enable_XTI(self, layers=None, token_strings=None):724 self.XTI_layers = layers725 self.token_strings = token_strings726 727 def add_replacement(self, str_from, str_to):728 self.replacements[str_from] = str_to729 730 def process_caption(self, subset: BaseSubset, caption):731 # caption に prefix/suffix を付ける732 if subset.caption_prefix:733 caption = subset.caption_prefix + " " + caption734 if subset.caption_suffix:735 caption = caption + " " + subset.caption_suffix736 737 # dropoutの決定：tag dropがこのメソッド内にあるのでここで行うのが良い738 is_drop_out = subset.caption_dropout_rate > 0 and random.random() < subset.caption_dropout_rate739 is_drop_out = (740 is_drop_out741 or subset.caption_dropout_every_n_epochs > 0742 and self.current_epoch % subset.caption_dropout_every_n_epochs == 0743 )744 745 if is_drop_out:746 caption = ""747 else:748 # process wildcards749 if subset.enable_wildcard:750 # if caption is multiline, random choice one line751 if "\n" in caption:752 caption = random.choice(caption.split("\n"))753 754 # wildcard is like '{aaa|bbb|ccc...}'755 # escape the curly braces like {{ or }}756 replacer1 = "⦅"757 replacer2 = "⦆"758 while replacer1 in caption or replacer2 in caption:759 replacer1 += "⦅"760 replacer2 += "⦆"761 762 caption = caption.replace("{{", replacer1).replace("}}", replacer2)763 764 # replace the wildcard765 def replace_wildcard(match):766 return random.choice(match.group(1).split("|"))767 768 caption = re.sub(r"\{([^}]+)\}", replace_wildcard, caption)769 770 # unescape the curly braces771 caption = caption.replace(replacer1, "{").replace(replacer2, "}")772 else:773 # if caption is multiline, use the first line774 caption = caption.split("\n")[0]775 776 if subset.shuffle_caption or subset.token_warmup_step > 0 or subset.caption_tag_dropout_rate > 0:777 fixed_tokens = []778 flex_tokens = []779 fixed_suffix_tokens = []780 if (781 hasattr(subset, "keep_tokens_separator")782 and subset.keep_tokens_separator783 and subset.keep_tokens_separator in caption784 ):785 fixed_part, flex_part = caption.split(subset.keep_tokens_separator, 1)786 if subset.keep_tokens_separator in flex_part:787 flex_part, fixed_suffix_part = flex_part.split(subset.keep_tokens_separator, 1)788 fixed_suffix_tokens = [t.strip() for t in fixed_suffix_part.split(subset.caption_separator) if t.strip()]789 790 fixed_tokens = [t.strip() for t in fixed_part.split(subset.caption_separator) if t.strip()]791 flex_tokens = [t.strip() for t in flex_part.split(subset.caption_separator) if t.strip()]792 else:793 tokens = [t.strip() for t in caption.strip().split(subset.caption_separator)]794 flex_tokens = tokens[:]795 if subset.keep_tokens > 0:796 fixed_tokens = flex_tokens[: subset.keep_tokens]797 flex_tokens = tokens[subset.keep_tokens :]798 799 if subset.token_warmup_step < 1: # 初回に上書きする800 subset.token_warmup_step = math.floor(subset.token_warmup_step * self.max_train_steps)801 if subset.token_warmup_step and self.current_step < subset.token_warmup_step:802 tokens_len = (803 math.floor(804 (self.current_step) * ((len(flex_tokens) - subset.token_warmup_min) / (subset.token_warmup_step))805 )806 + subset.token_warmup_min807 )808 flex_tokens = flex_tokens[:tokens_len]809 810 def dropout_tags(tokens):811 if subset.caption_tag_dropout_rate <= 0:812 return tokens813 l = []814 for token in tokens:815 if random.random() >= subset.caption_tag_dropout_rate:816 l.append(token)817 return l818 819 if subset.shuffle_caption:820 random.shuffle(flex_tokens)821 822 flex_tokens = dropout_tags(flex_tokens)823 824 caption = ", ".join(fixed_tokens + flex_tokens + fixed_suffix_tokens)825 826 # process secondary separator827 if subset.secondary_separator:828 caption = caption.replace(subset.secondary_separator, subset.caption_separator)829 830 # textual inversion対応831 for str_from, str_to in self.replacements.items():832 if str_from == "":833 # replace all834 if type(str_to) == list:835 caption = random.choice(str_to)836 else:837 caption = str_to838 else:839 caption = caption.replace(str_from, str_to)840 841 return caption842 843 def get_input_ids(self, caption, tokenizer=None):844 if tokenizer is None:845 tokenizer = self.tokenizers[0]846 847 input_ids = tokenizer(848 caption, padding="max_length", truncation=True, max_length=self.tokenizer_max_length, return_tensors="pt"849 ).input_ids850 851 if self.tokenizer_max_length > tokenizer.model_max_length:852 input_ids = input_ids.squeeze(0)853 iids_list = []854 if tokenizer.pad_token_id == tokenizer.eos_token_id:855 # v1856 # 77以上の時は "<BOS> .... <EOS> <EOS> <EOS>" でトータル227とかになっているので、"<BOS>...<EOS>"の三連に変換する857 # 1111氏のやつは , で区切る、とかしているようだが とりあえず単純に858 for i in range(859 1, self.tokenizer_max_length - tokenizer.model_max_length + 2, tokenizer.model_max_length - 2860 ): # (1, 152, 75)861 ids_chunk = (862 input_ids[0].unsqueeze(0),863 input_ids[i : i + tokenizer.model_max_length - 2],864 input_ids[-1].unsqueeze(0),865 )866 ids_chunk = torch.cat(ids_chunk)867 iids_list.append(ids_chunk)868 else:869 # v2 or SDXL870 # 77以上の時は "<BOS> .... <EOS> <PAD> <PAD>..." でトータル227とかになっているので、"<BOS>...<EOS> <PAD> <PAD> ..."の三連に変換する871 for i in range(1, self.tokenizer_max_length - tokenizer.model_max_length + 2, tokenizer.model_max_length - 2):872 ids_chunk = (873 input_ids[0].unsqueeze(0), # BOS874 input_ids[i : i + tokenizer.model_max_length - 2],875 input_ids[-1].unsqueeze(0),876 ) # PAD or EOS877 ids_chunk = torch.cat(ids_chunk)878 879 # 末尾が <EOS> <PAD> または <PAD> <PAD> の場合は、何もしなくてよい880 # 末尾が x <PAD/EOS> の場合は末尾を <EOS> に変える（x <EOS> なら結果的に変化なし）881 if ids_chunk[-2] != tokenizer.eos_token_id and ids_chunk[-2] != tokenizer.pad_token_id:882 ids_chunk[-1] = tokenizer.eos_token_id883 # 先頭が <BOS> <PAD> ... の場合は <BOS> <EOS> <PAD> ... に変える884 if ids_chunk[1] == tokenizer.pad_token_id:885 ids_chunk[1] = tokenizer.eos_token_id886 887 iids_list.append(ids_chunk)888 889 input_ids = torch.stack(iids_list) # 3,77890 return input_ids891 892 def register_image(self, info: ImageInfo, subset: BaseSubset):893 self.image_data[info.image_key] = info894 self.image_to_subset[info.image_key] = subset895 896 def make_buckets(self):897 """898 bucketingを行わない場合も呼び出し必須（ひとつだけbucketを作る）899 min_size and max_size are ignored when enable_bucket is False900 """901 logger.info("loading image sizes.")902 for info in tqdm(self.image_data.values()):903 if info.image_size is None:904 info.image_size = self.get_image_size(info.absolute_path)905 906 if self.enable_bucket:907 logger.info("make buckets")908 else:909 logger.info("prepare dataset")910 911 # bucketを作成し、画像をbucketに振り分ける912 if self.enable_bucket:913 if self.bucket_manager is None: # fine tuningの場合でmetadataに定義がある場合は、すでに初期化済み914 self.bucket_manager = BucketManager(915 self.bucket_no_upscale,916 (self.width, self.height),917 self.min_bucket_reso,918 self.max_bucket_reso,919 self.bucket_reso_steps,920 )921 if not self.bucket_no_upscale:922 self.bucket_manager.make_buckets()923 else:924 logger.warning(925 "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は無視されます"926 )927 928 img_ar_errors = []929 for image_info in self.image_data.values():930 image_width, image_height = image_info.image_size931 image_info.bucket_reso, image_info.resized_size, ar_error = self.bucket_manager.select_bucket(932 image_width, image_height933 )934 935 # logger.info(image_info.image_key, image_info.bucket_reso)936 img_ar_errors.append(abs(ar_error))937 938 self.bucket_manager.sort()939 else:940 self.bucket_manager = BucketManager(False, (self.width, self.height), None, None, None)941 self.bucket_manager.set_predefined_resos([(self.width, self.height)]) # ひとつの固定サイズbucketのみ942 for image_info in self.image_data.values():943 image_width, image_height = image_info.image_size944 image_info.bucket_reso, image_info.resized_size, _ = self.bucket_manager.select_bucket(image_width, image_height)945 946 for image_info in self.image_data.values():947 for _ in range(image_info.num_repeats):948 self.bucket_manager.add_image(image_info.bucket_reso, image_info.image_key)949 950 # bucket情報を表示、格納する951 if self.enable_bucket:952 self.bucket_info = {"buckets": {}}953 logger.info("number of images (including repeats) / 各bucketの画像枚数（繰り返し回数を含む）")954 for i, (reso, bucket) in enumerate(zip(self.bucket_manager.resos, self.bucket_manager.buckets)):955 count = len(bucket)956 if count > 0:957 self.bucket_info["buckets"][i] = {"resolution": reso, "count": len(bucket)}958 logger.info(f"bucket {i}: resolution {reso}, count: {len(bucket)}")959 960 if len(img_ar_errors) == 0:961 mean_img_ar_error = 0 # avoid NaN962 else:963 img_ar_errors = np.array(img_ar_errors)964 mean_img_ar_error = np.mean(np.abs(img_ar_errors))965 self.bucket_info["mean_img_ar_error"] = mean_img_ar_error966 logger.info(f"mean ar error (without repeats): {mean_img_ar_error}")967 968 # データ参照用indexを作る。このindexはdatasetのshuffleに用いられる969 self.buckets_indices: List[BucketBatchIndex] = []970 for bucket_index, bucket in enumerate(self.bucket_manager.buckets):971 batch_count = int(math.ceil(len(bucket) / self.batch_size))972 for batch_index in range(batch_count):973 self.buckets_indices.append(BucketBatchIndex(bucket_index, self.batch_size, batch_index))974 975 # ↓以下はbucketごとのbatch件数があまりにも増えて混乱を招くので元に戻す976 # 学習時はステップ数がランダムなので、同一画像が同一batch内にあってもそれほど悪影響はないであろう、と考えられる977 #978 # # bucketが細分化されることにより、ひとつのbucketに一種類の画像のみというケースが増え、つまりそれは979 # # ひとつのbatchが同じ画像で占められることになるので、さすがに良くないであろう980 # # そのためバッチサイズを画像種類までに制限する981 # # ただそれでも同一画像が同一バッチに含まれる可能性はあるので、繰り返し回数が少ないほうがshuffleの品質は良くなることは間違いない？982 # # TO DO 正則化画像をepochまたがりで利用する仕組み983 # num_of_image_types = len(set(bucket))984 # bucket_batch_size = min(self.batch_size, num_of_image_types)985 # batch_count = int(math.ceil(len(bucket) / bucket_batch_size))986 # # logger.info(bucket_index, num_of_image_types, bucket_batch_size, batch_count)987 # for batch_index in range(batch_count):988 # self.buckets_indices.append(BucketBatchIndex(bucket_index, bucket_batch_size, batch_index))989 # ↑ここまで990 991 self.shuffle_buckets()992 self._length = len(self.buckets_indices)993 994 def shuffle_buckets(self):995 # set random seed for this epoch996 random.seed(self.seed + self.current_epoch)997 998 random.shuffle(self.buckets_indices)999 self.bucket_manager.shuffle()1000 1001 def verify_bucket_reso_steps(self, min_steps: int):1002 assert self.bucket_reso_steps is None or self.bucket_reso_steps % min_steps == 0, (1003 f"bucket_reso_steps is {self.bucket_reso_steps}. it must be divisible by {min_steps}.\n"1004 + f"bucket_reso_stepsが{self.bucket_reso_steps}です。{min_steps}で割り切れる必要があります"1005 )1006 1007 def is_latent_cacheable(self):1008 return all([not subset.color_aug and not subset.random_crop for subset in self.subsets])1009 1010 def is_text_encoder_output_cacheable(self):1011 return all(1012 [1013 not (1014 subset.caption_dropout_rate > 01015 or subset.shuffle_caption1016 or subset.token_warmup_step > 01017 or subset.caption_tag_dropout_rate > 01018 )1019 for subset in self.subsets1020 ]1021 )1022 1023 def cache_latents(self, vae, vae_batch_size=1, cache_to_disk=False, is_main_process=True):1024 # マルチGPUには対応していないので、そちらはtools/cache_latents.pyを使うこと1025 logger.info("caching latents.")1026 1027 image_infos = list(self.image_data.values())1028 1029 # sort by resolution1030 image_infos.sort(key=lambda info: info.bucket_reso[0] * info.bucket_reso[1])1031 1032 # split by resolution and some conditions1033 class Condition:1034 def __init__(self, reso, flip_aug, alpha_mask, random_crop):1035 self.reso = reso1036 self.flip_aug = flip_aug1037 self.alpha_mask = alpha_mask1038 self.random_crop = random_crop1039 1040 def __eq__(self, other):1041 return (1042 self.reso == other.reso1043 and self.flip_aug == other.flip_aug1044 and self.alpha_mask == other.alpha_mask1045 and self.random_crop == other.random_crop1046 )1047 1048 batches: List[Tuple[Condition, List[ImageInfo]]] = []1049 batch: List[ImageInfo] = []1050 current_condition = None1051 1052 logger.info("checking cache validity...")1053 for info in tqdm(image_infos):1054 subset = self.image_to_subset[info.image_key]1055 1056 if info.latents_npz is not None: # fine tuning dataset1057 continue1058 1059 # check disk cache exists and size of latents1060 if cache_to_disk:1061 info.latents_npz = os.path.splitext(info.absolute_path)[0] + ".npz"1062 if not is_main_process: # store to info only1063 continue1064 1065 cache_available = is_disk_cached_latents_is_expected(1066 info.bucket_reso, info.latents_npz, subset.flip_aug, subset.alpha_mask1067 )1068 1069 if cache_available: # do not add to batch1070 continue1071 1072 # if batch is not empty and condition is changed, flush the batch. Note that current_condition is not None if batch is not empty1073 condition = Condition(info.bucket_reso, subset.flip_aug, subset.alpha_mask, subset.random_crop)1074 if len(batch) > 0 and current_condition != condition:1075 batches.append((current_condition, batch))1076 batch = []1077 1078 batch.append(info)1079 current_condition = condition1080 1081 # if number of data in batch is enough, flush the batch1082 if len(batch) >= vae_batch_size:1083 batches.append((current_condition, batch))1084 batch = []1085 current_condition = None1086 1087 if len(batch) > 0:1088 batches.append((current_condition, batch))1089 1090 if cache_to_disk and not is_main_process: # if cache to disk, don't cache latents in non-main process, set to info only1091 return1092 1093 # iterate batches: batch doesn't have image, image will be loaded in cache_batch_latents and discarded1094 logger.info("caching latents...")1095 for condition, batch in tqdm(batches, smoothing=1, total=len(batches)):1096 cache_batch_latents(vae, cache_to_disk, batch, condition.flip_aug, condition.alpha_mask, condition.random_crop)1097 1098 # weight_dtypeを指定するとText Encoderそのもの、およひ出力がweight_dtypeになる1099 # SDXLでのみ有効だが、datasetのメソッドとする必要があるので、sdxl_train_util.pyではなくこちらに実装する1100 # SD1/2に対応するにはv2のフラグを持つ必要があるので後回し1101 def cache_text_encoder_outputs(1102 self, tokenizers, text_encoders, device, weight_dtype, cache_to_disk=False, is_main_process=True1103 ):1104 assert len(tokenizers) == 2, "only support SDXL"1105 1106 # latentsのキャッシュと同様に、ディスクへのキャッシュに対応する1107 # またマルチGPUには対応していないので、そちらはtools/cache_latents.pyを使うこと1108 logger.info("caching text encoder outputs.")1109 image_infos = list(self.image_data.values())1110 1111 logger.info("checking cache existence...")1112 image_infos_to_cache = []1113 for info in tqdm(image_infos):1114 # subset = self.image_to_subset[info.image_key]1115 if cache_to_disk:1116 te_out_npz = os.path.splitext(info.absolute_path)[0] + TEXT_ENCODER_OUTPUTS_CACHE_SUFFIX1117 info.text_encoder_outputs_npz = te_out_npz1118 1119 if not is_main_process: # store to info only1120 continue1121 1122 if os.path.exists(te_out_npz):1123 continue1124 1125 image_infos_to_cache.append(info)1126 1127 if cache_to_disk and not is_main_process: # if cache to disk, don't cache latents in non-main process, set to info only1128 return1129 1130 # prepare tokenizers and text encoders1131 for text_encoder in text_encoders:1132 text_encoder.to(device)1133 if weight_dtype is not None:1134 text_encoder.to(dtype=weight_dtype)1135 1136 # create batch1137 batch = []1138 batches = []1139 for info in image_infos_to_cache:1140 input_ids1 = self.get_input_ids(info.caption, tokenizers[0])1141 input_ids2 = self.get_input_ids(info.caption, tokenizers[1])1142 batch.append((info, input_ids1, input_ids2))1143 1144 if len(batch) >= self.batch_size:1145 batches.append(batch)1146 batch = []1147 1148 if len(batch) > 0:1149 batches.append(batch)1150 1151 # iterate batches: call text encoder and cache outputs for memory or disk1152 logger.info("caching text encoder outputs...")1153 for batch in tqdm(batches):1154 infos, input_ids1, input_ids2 = zip(*batch)1155 input_ids1 = torch.stack(input_ids1, dim=0)1156 input_ids2 = torch.stack(input_ids2, dim=0)1157 cache_batch_text_encoder_outputs(1158 infos, tokenizers, text_encoders, self.max_token_length, cache_to_disk, input_ids1, input_ids2, weight_dtype1159 )1160 1161 def get_image_size(self, image_path):1162 return imagesize.get(image_path)1163 1164 def load_image_with_face_info(self, subset: BaseSubset, image_path: str, alpha_mask=False):1165 img = load_image(image_path, alpha_mask)1166 1167 face_cx = face_cy = face_w = face_h = 01168 if subset.face_crop_aug_range is not None:1169 tokens = os.path.splitext(os.path.basename(image_path))[0].split("_")1170 if len(tokens) >= 5:1171 face_cx = int(tokens[-4])1172 face_cy = int(tokens[-3])1173 face_w = int(tokens[-2])1174 face_h = int(tokens[-1])1175 1176 return img, face_cx, face_cy, face_w, face_h1177 1178 # いい感じに切り出す1179 def crop_target(self, subset: BaseSubset, image, face_cx, face_cy, face_w, face_h):1180 height, width = image.shape[0:2]1181 if height == self.height and width == self.width:1182 return image1183 1184 # 画像サイズはsizeより大きいのでリサイズする1185 face_size = max(face_w, face_h)1186 size = min(self.height, self.width) # 短いほう1187 min_scale = max(self.height / height, self.width / width) # 画像がモデル入力サイズぴったりになる倍率（最小の倍率）1188 min_scale = min(1.0, max(min_scale, size / (face_size * subset.face_crop_aug_range[1]))) # 指定した顔最小サイズ1189 max_scale = min(1.0, max(min_scale, size / (face_size * subset.face_crop_aug_range[0]))) # 指定した顔最大サイズ1190 if min_scale >= max_scale: # range指定がmin==max1191 scale = min_scale1192 else:1193 scale = random.uniform(min_scale, max_scale)1194 1195 nh = int(height * scale + 0.5)1196 nw = int(width * scale + 0.5)1197 assert nh >= self.height and nw >= self.width, f"internal error. small scale {scale}, {width}*{height}"1198 image = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_AREA)1199 face_cx = int(face_cx * scale + 0.5)1200 face_cy = int(face_cy * scale + 0.5)