# common functions for training import argparse import ast import importlib import json import pathlib import re import shutil import time from typing import ( Dict, List, NamedTuple, Optional, Sequence, Tuple, Union, ) from accelerate import Accelerator import glob import math import os import random import hashlib import subprocess from io import BytesIO import toml from tqdm import tqdm import torch from torch.optim import Optimizer from torchvision import transforms from transformers import CLIPTokenizer import transformers import diffusers from diffusers.optimization import SchedulerType, TYPE_TO_SCHEDULER_FUNCTION from diffusers import ( StableDiffusionPipeline, DDPMScheduler, EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler, LMSDiscreteScheduler, PNDMScheduler, DDIMScheduler, EulerDiscreteScheduler, HeunDiscreteScheduler, KDPM2DiscreteScheduler, KDPM2AncestralDiscreteScheduler, ) import albumentations as albu import numpy as np from PIL import Image import cv2 from einops import rearrange from torch import einsum import safetensors.torch from library.lpw_stable_diffusion import StableDiffusionLongPromptWeightingPipeline import library.model_util as model_util # Tokenizer: checkpointから読み込むのではなくあらかじめ提供されているものを使う TOKENIZER_PATH = "openai/clip-vit-large-patch14" V2_STABLE_DIFFUSION_PATH = "stabilityai/stable-diffusion-2" # ここからtokenizerだけ使う v2とv2.1はtokenizer仕様は同じ # checkpointファイル名 EPOCH_STATE_NAME = "{}-{:06d}-state" EPOCH_FILE_NAME = "{}-{:06d}" EPOCH_DIFFUSERS_DIR_NAME = "{}-{:06d}" LAST_STATE_NAME = "{}-state" DEFAULT_EPOCH_NAME = "epoch" DEFAULT_LAST_OUTPUT_NAME = "last" # region dataset IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".webp", ".bmp", ".PNG", ".JPG", ".JPEG", ".WEBP", ".BMP"] class ImageInfo: def __init__(self, image_key: str, num_repeats: int, caption: str, is_reg: bool, absolute_path: str) -> 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.image_size: Tuple[int, int] = None self.resized_size: Tuple[int, int] = None self.bucket_reso: Tuple[int, int] = None self.latents: torch.Tensor = None self.latents_flipped: torch.Tensor = None self.latents_npz: str = None self.latents_npz_flipped: str = None class BucketManager: def __init__(self, no_upscale, max_reso, min_size, max_size, reso_steps) -> None: 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)、学習時は image_key def add_image(self, reso, image): bucket_id = self.reso_to_id[reso] self.buckets[bucket_id].append(image) 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([]) # print(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)) # print("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 # print(b_width_rounded, b_height_in_wr, ar_width_rounded) # print(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) # print(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 # print("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 class BucketBatchIndex(NamedTuple): bucket_index: int bucket_batch_size: int batch_index: int class AugHelper: def __init__(self): # prepare all possible augmentators color_aug_method = albu.OneOf( [ albu.HueSaturationValue(8, 0, 0, p=0.5), albu.RandomGamma((95, 105), p=0.5), ], p=0.33, ) flip_aug_method = albu.HorizontalFlip(p=0.5) # key: (use_color_aug, use_flip_aug) self.augmentors = { (True, True): albu.Compose( [ color_aug_method, flip_aug_method, ], p=1.0, ), (True, False): albu.Compose( [ color_aug_method, ], p=1.0, ), (False, True): albu.Compose( [ flip_aug_method, ], p=1.0, ), (False, False): None, } def get_augmentor(self, use_color_aug: bool, use_flip_aug: bool) -> Optional[albu.Compose]: return self.augmentors[(use_color_aug, use_flip_aug)] class BaseSubset: def __init__( self, image_dir: Optional[str], num_repeats: int, shuffle_caption: bool, keep_tokens: int, color_aug: bool, flip_aug: bool, face_crop_aug_range: Optional[Tuple[float, float]], random_crop: bool, caption_dropout_rate: float, caption_dropout_every_n_epochs: int, caption_tag_dropout_rate: float, token_warmup_min: int, token_warmup_step: Union[float, int], ) -> None: self.image_dir = image_dir self.num_repeats = num_repeats self.shuffle_caption = shuffle_caption self.keep_tokens = keep_tokens self.color_aug = color_aug self.flip_aug = flip_aug self.face_crop_aug_range = face_crop_aug_range self.random_crop = random_crop self.caption_dropout_rate = caption_dropout_rate self.caption_dropout_every_n_epochs = caption_dropout_every_n_epochs self.caption_tag_dropout_rate = caption_tag_dropout_rate self.token_warmup_min = token_warmup_min # step=0におけるタグの数 self.token_warmup_step = token_warmup_step # N（N<1ならN*max_train_steps）ステップ目でタグの数が最大になる self.img_count = 0 class DreamBoothSubset(BaseSubset): def __init__( self, image_dir: str, is_reg: bool, class_tokens: Optional[str], caption_extension: str, num_repeats, shuffle_caption, keep_tokens, color_aug, flip_aug, face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate, token_warmup_min, token_warmup_step, ) -> None: assert image_dir is not None, "image_dir must be specified / image_dirは指定が必須です" super().__init__( image_dir, num_repeats, shuffle_caption, keep_tokens, color_aug, flip_aug, face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate, token_warmup_min, token_warmup_step, ) self.is_reg = is_reg self.class_tokens = class_tokens self.caption_extension = caption_extension def __eq__(self, other) -> bool: if not isinstance(other, DreamBoothSubset): return NotImplemented return self.image_dir == other.image_dir class FineTuningSubset(BaseSubset): def __init__( self, image_dir, metadata_file: str, num_repeats, shuffle_caption, keep_tokens, color_aug, flip_aug, face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate, token_warmup_min, token_warmup_step, ) -> None: assert metadata_file is not None, "metadata_file must be specified / metadata_fileは指定が必須です" super().__init__( image_dir, num_repeats, shuffle_caption, keep_tokens, color_aug, flip_aug, face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate, token_warmup_min, token_warmup_step, ) self.metadata_file = metadata_file def __eq__(self, other) -> bool: if not isinstance(other, FineTuningSubset): return NotImplemented return self.metadata_file == other.metadata_file class BaseDataset(torch.utils.data.Dataset): def __init__( self, tokenizer: CLIPTokenizer, max_token_length: int, resolution: Optional[Tuple[int, int]], debug_dataset: bool ) -> None: super().__init__() self.tokenizer = tokenizer self.max_token_length = max_token_length # width/height is used when enable_bucket==False self.width, self.height = (None, None) if resolution is None else resolution 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.tokenizer_max_length = self.tokenizer.model_max_length if max_token_length is None else max_token_length + 2 self.current_epoch: int = 0 # インスタンスがepochごとに新しく作られるようなので外側から渡さないとダメ self.current_step: int = 0 self.max_train_steps: int = 0 self.seed: int = 0 # augmentation self.aug_helper = AugHelper() self.image_transforms = transforms.Compose( [ transforms.ToTensor(), transforms.Normalize([0.5], [0.5]), ] ) self.image_data: Dict[str, ImageInfo] = {} self.image_to_subset: Dict[str, Union[DreamBoothSubset, FineTuningSubset]] = {} self.replacements = {} def set_seed(self, seed): self.seed = seed def set_current_epoch(self, epoch): if not self.current_epoch == epoch: # epochが切り替わったらバケツをシャッフルする self.shuffle_buckets() 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): # 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: if subset.shuffle_caption or subset.token_warmup_step > 0 or subset.caption_tag_dropout_rate > 0: tokens = [t.strip() for t in caption.strip().split(",")] 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(tokens) - subset.token_warmup_min) / (subset.token_warmup_step))) + subset.token_warmup_min ) tokens = 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 fixed_tokens = [] flex_tokens = tokens[:] if subset.keep_tokens > 0: fixed_tokens = flex_tokens[: subset.keep_tokens] flex_tokens = tokens[subset.keep_tokens :] if subset.shuffle_caption: random.shuffle(flex_tokens) flex_tokens = dropout_tags(flex_tokens) caption = ", ".join(fixed_tokens + flex_tokens) # 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 get_input_ids(self, caption): input_ids = self.tokenizer( caption, padding="max_length", truncation=True, max_length=self.tokenizer_max_length, return_tensors="pt" ).input_ids if self.tokenizer_max_length > self.tokenizer.model_max_length: input_ids = input_ids.squeeze(0) iids_list = [] if self.tokenizer.pad_token_id == self.tokenizer.eos_token_id: # v1 # 77以上の時は " .... " でトータル227とかになっているので、"..."の三連に変換する # 1111氏のやつは , で区切る、とかしているようだが　とりあえず単純に for i in range( 1, self.tokenizer_max_length - self.tokenizer.model_max_length + 2, self.tokenizer.model_max_length - 2 ): # (1, 152, 75) ids_chunk = ( input_ids[0].unsqueeze(0), input_ids[i : i + self.tokenizer.model_max_length - 2], input_ids[-1].unsqueeze(0), ) ids_chunk = torch.cat(ids_chunk) iids_list.append(ids_chunk) else: # v2 # 77以上の時は " .... ..." でトータル227とかになっているので、"... ..."の三連に変換する for i in range( 1, self.tokenizer_max_length - self.tokenizer.model_max_length + 2, self.tokenizer.model_max_length - 2 ): ids_chunk = ( input_ids[0].unsqueeze(0), # BOS input_ids[i : i + self.tokenizer.model_max_length - 2], input_ids[-1].unsqueeze(0), ) # PAD or EOS ids_chunk = torch.cat(ids_chunk) # 末尾が または の場合は、何もしなくてよい # 末尾が x の場合は末尾を に変える（x なら結果的に変化なし） if ids_chunk[-2] != self.tokenizer.eos_token_id and ids_chunk[-2] != self.tokenizer.pad_token_id: ids_chunk[-1] = self.tokenizer.eos_token_id # 先頭が ... の場合は ... に変える if ids_chunk[1] == self.tokenizer.pad_token_id: ids_chunk[1] = self.tokenizer.eos_token_id iids_list.append(ids_chunk) input_ids = torch.stack(iids_list) # 3,77 return input_ids 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 """ print("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) if self.enable_bucket: print("make buckets") else: print("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: print( "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 ) # print(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": {}} print("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)} print(f"bucket {i}: resolution {reso}, count: {len(bucket)}") 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 print(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)) # ↓以下はbucketごとのbatch件数があまりにも増えて混乱を招くので元に戻す # 　学習時はステップ数がランダムなので、同一画像が同一batch内にあってもそれほど悪影響はないであろう、と考えられる # # # bucketが細分化されることにより、ひとつのbucketに一種類の画像のみというケースが増え、つまりそれは # # ひとつのbatchが同じ画像で占められることになるので、さすがに良くないであろう # # そのためバッチサイズを画像種類までに制限する # # ただそれでも同一画像が同一バッチに含まれる可能性はあるので、繰り返し回数が少ないほうがshuffleの品質は良くなることは間違いない？ # # TO DO 正則化画像をepochまたがりで利用する仕組み # num_of_image_types = len(set(bucket)) # bucket_batch_size = min(self.batch_size, num_of_image_types) # batch_count = int(math.ceil(len(bucket) / bucket_batch_size)) # # print(bucket_index, num_of_image_types, bucket_batch_size, batch_count) # for batch_index in range(batch_count): # self.buckets_indices.append(BucketBatchIndex(bucket_index, bucket_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 load_image(self, image_path): image = Image.open(image_path) if not image.mode == "RGB": image = image.convert("RGB") img = np.array(image, np.uint8) return img def trim_and_resize_if_required(self, subset: BaseSubset, image, reso, resized_size): image_height, image_width = image.shape[0:2] if image_width != resized_size[0] or image_height != resized_size[1]: # リサイズする image = cv2.resize(image, resized_size, interpolation=cv2.INTER_AREA) # INTER_AREAでやりたいのでcv2でリサイズ image_height, image_width = image.shape[0:2] if image_width > reso[0]: trim_size = image_width - reso[0] p = trim_size // 2 if not subset.random_crop else random.randint(0, trim_size) # print("w", trim_size, p) image = image[:, p : p + reso[0]] if image_height > reso[1]: trim_size = image_height - reso[1] p = trim_size // 2 if not subset.random_crop else random.randint(0, trim_size) # print("h", trim_size, p) image = image[p : p + reso[1]] assert ( image.shape[0] == reso[1] and image.shape[1] == reso[0] ), f"internal error, illegal trimmed size: {image.shape}, {reso}" return image def is_latent_cacheable(self): return all([not subset.color_aug and not subset.random_crop for subset in self.subsets]) def cache_latents(self, vae, vae_batch_size=1): # ちょっと速くした print("caching latents.") 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 batches = [] batch = [] for info in image_infos: subset = self.image_to_subset[info.image_key] if info.latents_npz is not None: info.latents = self.load_latents_from_npz(info, False) info.latents = torch.FloatTensor(info.latents) info.latents_flipped = self.load_latents_from_npz(info, True) # might be None if info.latents_flipped is not None: info.latents_flipped = torch.FloatTensor(info.latents_flipped) continue # if last member of batch has different resolution, flush the batch if len(batch) > 0 and batch[-1].bucket_reso != info.bucket_reso: batches.append(batch) batch = [] batch.append(info) # if number of data in batch is enough, flush the batch if len(batch) >= vae_batch_size: batches.append(batch) batch = [] if len(batch) > 0: batches.append(batch) # iterate batches for batch in tqdm(batches, smoothing=1, total=len(batches)): images = [] for info in batch: image = self.load_image(info.absolute_path) image = self.trim_and_resize_if_required(subset, image, info.bucket_reso, info.resized_size) image = self.image_transforms(image) images.append(image) img_tensors = torch.stack(images, dim=0) img_tensors = img_tensors.to(device=vae.device, dtype=vae.dtype) latents = vae.encode(img_tensors).latent_dist.sample().to("cpu") for info, latent in zip(batch, latents): info.latents = latent if subset.flip_aug: img_tensors = torch.flip(img_tensors, dims=[3]) latents = vae.encode(img_tensors).latent_dist.sample().to("cpu") for info, latent in zip(batch, latents): info.latents_flipped = latent def get_image_size(self, image_path): image = Image.open(image_path) return image.size def load_image_with_face_info(self, subset: BaseSubset, image_path: str): img = self.load_image(image_path) 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) min_scale = max(self.height / height, self.width / width) # 画像がモデル入力サイズぴったりになる倍率（最小の倍率） min_scale = min(1.0, max(min_scale, self.size / (face_size * subset.face_crop_aug_range[1]))) # 指定した顔最小サイズ max_scale = min(1.0, max(min_scale, self.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 = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_AREA) 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 > self.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 load_latents_from_npz(self, image_info: ImageInfo, flipped): npz_file = image_info.latents_npz_flipped if flipped else image_info.latents_npz if npz_file is None: return None return np.load(npz_file)["arr_0"] 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 = [] 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] loss_weights.append(self.prior_loss_weight if image_info.is_reg else 1.0) # image/latentsを処理する if image_info.latents is not None: latents = image_info.latents if not subset.flip_aug or random.random() < 0.5 else image_info.latents_flipped image = None elif image_info.latents_npz is not None: latents = self.load_latents_from_npz(image_info, subset.flip_aug and random.random() >= 0.5) latents = torch.FloatTensor(latents) image = None else: # 画像を読み込み、必要ならcropする img, face_cx, face_cy, face_w, face_h = self.load_image_with_face_info(subset, image_info.absolute_path) im_h, im_w = img.shape[0:2] if self.enable_bucket: img = self.trim_and_resize_if_required(subset, img, image_info.bucket_reso, image_info.resized_size) 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}" # augmentation aug = self.aug_helper.get_augmentor(subset.color_aug, subset.flip_aug) if aug is not None: img = aug(image=img)["image"] latents = None image = self.image_transforms(img) # -1.0~1.0のtorch.Tensorになる images.append(image) latents_list.append(latents) caption = self.process_caption(subset, image_info.caption) if self.XTI_layers: caption_layer = [] for layer in self.XTI_layers: token_strings_from = " ".join(self.token_strings) token_strings_to = " ".join([f"{x}_{layer}" for x in self.token_strings]) caption_ = caption.replace(token_strings_from, token_strings_to) caption_layer.append(caption_) captions.append(caption_layer) else: captions.append(caption) if not self.token_padding_disabled: # this option might be omitted in future if self.XTI_layers: token_caption = self.get_input_ids(caption_layer) else: token_caption = self.get_input_ids(caption) input_ids_list.append(token_caption) example = {} example["loss_weights"] = torch.FloatTensor(loss_weights) if self.token_padding_disabled: # padding=True means pad in the batch example["input_ids"] = self.tokenizer(captions, padding=True, truncation=True, return_tensors="pt").input_ids else: # batch processing seems to be good example["input_ids"] = torch.stack(input_ids_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["latents"] = torch.stack(latents_list) if latents_list[0] is not None else None if self.debug_dataset: example["image_keys"] = bucket[image_index : image_index + self.batch_size] example["captions"] = captions return example class DreamBoothDataset(BaseDataset): def __init__( self, subsets: Sequence[DreamBoothSubset], batch_size: int, tokenizer, max_token_length, resolution, enable_bucket: bool, min_bucket_reso: int, max_bucket_reso: int, bucket_reso_steps: int, bucket_no_upscale: bool, prior_loss_weight: float, debug_dataset, ) -> None: super().__init__(tokenizer, max_token_length, resolution, debug_dataset) assert resolution is not None, f"resolution is required / resolution（解像度）指定は必須です" self.batch_size = batch_size self.size = min(self.width, self.height) # 短いほう self.prior_loss_weight = prior_loss_weight self.latents_cache = None self.enable_bucket = enable_bucket if self.enable_bucket: 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を大きくしてください" self.min_bucket_reso = min_bucket_reso self.max_bucket_reso = max_bucket_reso self.bucket_reso_steps = bucket_reso_steps self.bucket_no_upscale = bucket_no_upscale else: self.min_bucket_reso = None self.max_bucket_reso = None self.bucket_reso_steps = None # この情報は使われない self.bucket_no_upscale = False def read_caption(img_path, caption_extension): # captionの候補ファイル名を作る base_name = os.path.splitext(img_path)[0] base_name_face_det = base_name tokens = base_name.split("_") if len(tokens) >= 5: base_name_face_det = "_".join(tokens[:-4]) cap_paths = [base_name + caption_extension, base_name_face_det + caption_extension] caption = None for cap_path in cap_paths: if os.path.isfile(cap_path): with open(cap_path, "rt", encoding="utf-8") as f: try: lines = f.readlines() except UnicodeDecodeError as e: print(f"illegal char in file (not UTF-8) / ファイルにUTF-8以外の文字があります: {cap_path}") raise e assert len(lines) > 0, f"caption file is empty / キャプションファイルが空です: {cap_path}" caption = lines[0].strip() break return caption def load_dreambooth_dir(subset: DreamBoothSubset): if not os.path.isdir(subset.image_dir): print(f"not directory: {subset.image_dir}") return [], [] img_paths = glob_images(subset.image_dir, "*") print(f"found directory {subset.image_dir} contains {len(img_paths)} image files") # 画像ファイルごとにプロンプトを読み込み、もしあればそちらを使う captions = [] for img_path in img_paths: cap_for_img = read_caption(img_path, subset.caption_extension) if cap_for_img is None and subset.class_tokens is None: print(f"neither caption file nor class tokens are found. use empty caption for {img_path}") captions.append("") else: captions.append(subset.class_tokens if cap_for_img is None else cap_for_img) self.set_tag_frequency(os.path.basename(subset.image_dir), captions) # タグ頻度を記録 return img_paths, captions print("prepare images.") num_train_images = 0 num_reg_images = 0 reg_infos: List[ImageInfo] = [] for subset in subsets: if subset.num_repeats < 1: print( f"ignore subset with image_dir='{subset.image_dir}': num_repeats is less than 1 / num_repeatsが1を下回っているためサブセットを無視します: {subset.num_repeats}" ) continue if subset in self.subsets: print( f"ignore duplicated subset with image_dir='{subset.image_dir}': use the first one / 既にサブセットが登録されているため、重複した後発のサブセットを無視します" ) continue img_paths, captions = load_dreambooth_dir(subset) if len(img_paths) < 1: print(f"ignore subset with image_dir='{subset.image_dir}': no images found / 画像が見つからないためサブセットを無視します") continue if subset.is_reg: num_reg_images += subset.num_repeats * len(img_paths) else: num_train_images += subset.num_repeats * len(img_paths) for img_path, caption in zip(img_paths, captions): info = ImageInfo(img_path, subset.num_repeats, caption, subset.is_reg, img_path) if subset.is_reg: reg_infos.append(info) else: self.register_image(info, subset) subset.img_count = len(img_paths) self.subsets.append(subset) print(f"{num_train_images} train images with repeating.") self.num_train_images = num_train_images print(f"{num_reg_images} reg images.") if num_train_images < num_reg_images: print("some of reg images are not used / 正則化画像の数が多いので、一部使用されない正則化画像があります") if num_reg_images == 0: print("no regularization images / 正則化画像が見つかりませんでした") else: # num_repeatsを計算する：どうせ大した数ではないのでループで処理する n = 0 first_loop = True while n < num_train_images: for info in reg_infos: if first_loop: self.register_image(info, subset) n += info.num_repeats else: info.num_repeats += 1 # rewrite registered info n += 1 if n >= num_train_images: break first_loop = False self.num_reg_images = num_reg_images class FineTuningDataset(BaseDataset): def __init__( self, subsets: Sequence[FineTuningSubset], batch_size: int, tokenizer, max_token_length, resolution, enable_bucket: bool, min_bucket_reso: int, max_bucket_reso: int, bucket_reso_steps: int, bucket_no_upscale: bool, debug_dataset, ) -> None: super().__init__(tokenizer, max_token_length, resolution, debug_dataset) self.batch_size = batch_size self.num_train_images = 0 self.num_reg_images = 0 for subset in subsets: if subset.num_repeats < 1: print( f"ignore subset with metadata_file='{subset.metadata_file}': num_repeats is less than 1 / num_repeatsが1を下回っているためサブセットを無視します: {subset.num_repeats}" ) continue if subset in self.subsets: print( f"ignore duplicated subset with metadata_file='{subset.metadata_file}': use the first one / 既にサブセットが登録されているため、重複した後発のサブセットを無視します" ) continue # メタデータを読み込む if os.path.exists(subset.metadata_file): print(f"loading existing metadata: {subset.metadata_file}") with open(subset.metadata_file, "rt", encoding="utf-8") as f: metadata = json.load(f) else: raise ValueError(f"no metadata / メタデータファイルがありません: {subset.metadata_file}") if len(metadata) < 1: print(f"ignore subset with '{subset.metadata_file}': no image entries found / 画像に関するデータが見つからないためサブセットを無視します") continue tags_list = [] for image_key, img_md in metadata.items(): # path情報を作る if os.path.exists(image_key): abs_path = image_key elif os.path.exists(os.path.splitext(image_key)[0] + ".npz"): abs_path = os.path.splitext(image_key)[0] + ".npz" else: npz_path = os.path.join(subset.image_dir, image_key + ".npz") if os.path.exists(npz_path): abs_path = npz_path else: # わりといい加減だがいい方法が思いつかん abs_path = glob_images(subset.image_dir, image_key) assert len(abs_path) >= 1, f"no image / 画像がありません: {image_key}" abs_path = abs_path[0] caption = img_md.get("caption") tags = img_md.get("tags") if caption is None: caption = tags elif tags is not None and len(tags) > 0: caption = caption + ", " + tags tags_list.append(tags) if caption is None: caption = "" image_info = ImageInfo(image_key, subset.num_repeats, caption, False, abs_path) image_info.image_size = img_md.get("train_resolution") if not subset.color_aug and not subset.random_crop: # if npz exists, use them image_info.latents_npz, image_info.latents_npz_flipped = self.image_key_to_npz_file(subset, image_key) self.register_image(image_info, subset) self.num_train_images += len(metadata) * subset.num_repeats # TODO do not record tag freq when no tag self.set_tag_frequency(os.path.basename(subset.metadata_file), tags_list) subset.img_count = len(metadata) self.subsets.append(subset) # check existence of all npz files use_npz_latents = all([not (subset.color_aug or subset.random_crop) for subset in self.subsets]) if use_npz_latents: flip_aug_in_subset = False npz_any = False npz_all = True for image_info in self.image_data.values(): subset = self.image_to_subset[image_info.image_key] has_npz = image_info.latents_npz is not None npz_any = npz_any or has_npz if subset.flip_aug: has_npz = has_npz and image_info.latents_npz_flipped is not None flip_aug_in_subset = True npz_all = npz_all and has_npz if npz_any and not npz_all: break if not npz_any: use_npz_latents = False print(f"npz file does not exist. ignore npz files / npzファイルが見つからないためnpzファイルを無視します") elif not npz_all: use_npz_latents = False print(f"some of npz file does not exist. ignore npz files / いくつかのnpzファイルが見つからないためnpzファイルを無視します") if flip_aug_in_subset: print("maybe no flipped files / 反転されたnpzファイルがないのかもしれません") # else: # print("npz files are not used with color_aug and/or random_crop / color_augまたはrandom_cropが指定されているためnpzファイルは使用されません") # check min/max bucket size sizes = set() resos = set() for image_info in self.image_data.values(): if image_info.image_size is None: sizes = None # not calculated break sizes.add(image_info.image_size[0]) sizes.add(image_info.image_size[1]) resos.add(tuple(image_info.image_size)) if sizes is None: if use_npz_latents: use_npz_latents = False print(f"npz files exist, but no bucket info in metadata. ignore npz files / メタデータにbucket情報がないためnpzファイルを無視します") assert ( resolution is not None ), "if metadata doesn't have bucket info, resolution is required / メタデータにbucket情報がない場合はresolutionを指定してください" self.enable_bucket = enable_bucket if self.enable_bucket: self.min_bucket_reso = min_bucket_reso self.max_bucket_reso = max_bucket_reso self.bucket_reso_steps = bucket_reso_steps self.bucket_no_upscale = bucket_no_upscale else: if not enable_bucket: print("metadata has bucket info, enable bucketing / メタデータにbucket情報があるためbucketを有効にします") print("using bucket info in metadata / メタデータ内のbucket情報を使います") self.enable_bucket = True assert ( not bucket_no_upscale ), "if metadata has bucket info, bucket reso is precalculated, so bucket_no_upscale cannot be used / メタデータ内にbucket情報がある場合はbucketの解像度は計算済みのため、bucket_no_upscaleは使えません" # bucket情報を初期化しておく、make_bucketsで再作成しない self.bucket_manager = BucketManager(False, None, None, None, None) self.bucket_manager.set_predefined_resos(resos) # npz情報をきれいにしておく if not use_npz_latents: for image_info in self.image_data.values(): image_info.latents_npz = image_info.latents_npz_flipped = None def image_key_to_npz_file(self, subset: FineTuningSubset, image_key): base_name = os.path.splitext(image_key)[0] npz_file_norm = base_name + ".npz" if os.path.exists(npz_file_norm): # image_key is full path npz_file_flip = base_name + "_flip.npz" if not os.path.exists(npz_file_flip): npz_file_flip = None return npz_file_norm, npz_file_flip # if not full path, check image_dir. if image_dir is None, return None if subset.image_dir is None: return None, None # image_key is relative path npz_file_norm = os.path.join(subset.image_dir, image_key + ".npz") npz_file_flip = os.path.join(subset.image_dir, image_key + "_flip.npz") if not os.path.exists(npz_file_norm): npz_file_norm = None npz_file_flip = None elif not os.path.exists(npz_file_flip): npz_file_flip = None return npz_file_norm, npz_file_flip # behave as Dataset mock class DatasetGroup(torch.utils.data.ConcatDataset): 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 enable_XTI(self, *args, **kwargs): for dataset in self.datasets: dataset.enable_XTI(*args, **kwargs) def cache_latents(self, vae, vae_batch_size=1): for i, dataset in enumerate(self.datasets): print(f"[Dataset {i}]") dataset.cache_latents(vae, vae_batch_size) def is_latent_cacheable(self) -> bool: return all([dataset.is_latent_cacheable() for dataset in self.datasets]) 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): print(f"Total dataset length (steps) / データセットの長さ（ステップ数）: {len(train_dataset)}") print("`S` for next step, `E` for next epoch no. , Escape for exit. / Sキーで次のステップ、Eキーで次のエポック、Escキーで中断、終了します") epoch = 1 while True: print(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) print(f"steps: {steps} ({i + 1}/{len(train_dataset)})") example = train_dataset[idx] if example["latents"] is not None: print(f"sample has latents from npz file: {example['latents'].size()}") for j, (ik, cap, lw, iid) in enumerate( zip(example["image_keys"], example["captions"], example["loss_weights"], example["input_ids"]) ): print(f'{ik}, size: {train_dataset.image_data[ik].image_size}, loss weight: {lw}, caption: "{cap}"') if show_input_ids: print(f"input ids: {iid}") if example["images"] is not None: im = example["images"][j] print(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 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 # endregion # region モジュール入れ替え部 """ 高速化のためのモジュール入れ替え """ # FlashAttentionを使うCrossAttention # based on https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/memory_efficient_attention_pytorch/flash_attention.py # LICENSE MIT https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/LICENSE # constants EPSILON = 1e-6 # helper functions def exists(val): return val is not None def default(val, d): return val if exists(val) else d def model_hash(filename): """Old model hash used by stable-diffusion-webui""" try: with open(filename, "rb") as file: m = hashlib.sha256() file.seek(0x100000) m.update(file.read(0x10000)) return m.hexdigest()[0:8] except FileNotFoundError: return "NOFILE" except IsADirectoryError: # Linux? return "IsADirectory" except PermissionError: # Windows return "IsADirectory" def calculate_sha256(filename): """New model hash used by stable-diffusion-webui""" try: hash_sha256 = hashlib.sha256() blksize = 1024 * 1024 with open(filename, "rb") as f: for chunk in iter(lambda: f.read(blksize), b""): hash_sha256.update(chunk) return hash_sha256.hexdigest() except FileNotFoundError: return "NOFILE" except IsADirectoryError: # Linux? return "IsADirectory" except PermissionError: # Windows return "IsADirectory" def precalculate_safetensors_hashes(tensors, metadata): """Precalculate the model hashes needed by sd-webui-additional-networks to save time on indexing the model later.""" # Because writing user metadata to the file can change the result of # sd_models.model_hash(), only retain the training metadata for purposes of # calculating the hash, as they are meant to be immutable metadata = {k: v for k, v in metadata.items() if k.startswith("ss_")} bytes = safetensors.torch.save(tensors, metadata) b = BytesIO(bytes) model_hash = addnet_hash_safetensors(b) legacy_hash = addnet_hash_legacy(b) return model_hash, legacy_hash def addnet_hash_legacy(b): """Old model hash used by sd-webui-additional-networks for .safetensors format files""" m = hashlib.sha256() b.seek(0x100000) m.update(b.read(0x10000)) return m.hexdigest()[0:8] def addnet_hash_safetensors(b): """New model hash used by sd-webui-additional-networks for .safetensors format files""" hash_sha256 = hashlib.sha256() blksize = 1024 * 1024 b.seek(0) header = b.read(8) n = int.from_bytes(header, "little") offset = n + 8 b.seek(offset) for chunk in iter(lambda: b.read(blksize), b""): hash_sha256.update(chunk) return hash_sha256.hexdigest() def get_git_revision_hash() -> str: try: return subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=os.path.dirname(__file__)).decode("ascii").strip() except: return "(unknown)" # flash attention forwards and backwards # https://arxiv.org/abs/2205.14135 class FlashAttentionFunction(torch.autograd.function.Function): @staticmethod @torch.no_grad() def forward(ctx, q, k, v, mask, causal, q_bucket_size, k_bucket_size): """Algorithm 2 in the paper""" device = q.device dtype = q.dtype max_neg_value = -torch.finfo(q.dtype).max qk_len_diff = max(k.shape[-2] - q.shape[-2], 0) o = torch.zeros_like(q) all_row_sums = torch.zeros((*q.shape[:-1], 1), dtype=dtype, device=device) all_row_maxes = torch.full((*q.shape[:-1], 1), max_neg_value, dtype=dtype, device=device) scale = q.shape[-1] ** -0.5 if not exists(mask): mask = (None,) * math.ceil(q.shape[-2] / q_bucket_size) else: mask = rearrange(mask, "b n -> b 1 1 n") mask = mask.split(q_bucket_size, dim=-1) row_splits = zip( q.split(q_bucket_size, dim=-2), o.split(q_bucket_size, dim=-2), mask, all_row_sums.split(q_bucket_size, dim=-2), all_row_maxes.split(q_bucket_size, dim=-2), ) for ind, (qc, oc, row_mask, row_sums, row_maxes) in enumerate(row_splits): q_start_index = ind * q_bucket_size - qk_len_diff col_splits = zip( k.split(k_bucket_size, dim=-2), v.split(k_bucket_size, dim=-2), ) for k_ind, (kc, vc) in enumerate(col_splits): k_start_index = k_ind * k_bucket_size attn_weights = einsum("... i d, ... j d -> ... i j", qc, kc) * scale if exists(row_mask): attn_weights.masked_fill_(~row_mask, max_neg_value) if causal and q_start_index < (k_start_index + k_bucket_size - 1): causal_mask = torch.ones((qc.shape[-2], kc.shape[-2]), dtype=torch.bool, device=device).triu( q_start_index - k_start_index + 1 ) attn_weights.masked_fill_(causal_mask, max_neg_value) block_row_maxes = attn_weights.amax(dim=-1, keepdims=True) attn_weights -= block_row_maxes exp_weights = torch.exp(attn_weights) if exists(row_mask): exp_weights.masked_fill_(~row_mask, 0.0) block_row_sums = exp_weights.sum(dim=-1, keepdims=True).clamp(min=EPSILON) new_row_maxes = torch.maximum(block_row_maxes, row_maxes) exp_values = einsum("... i j, ... j d -> ... i d", exp_weights, vc) exp_row_max_diff = torch.exp(row_maxes - new_row_maxes) exp_block_row_max_diff = torch.exp(block_row_maxes - new_row_maxes) new_row_sums = exp_row_max_diff * row_sums + exp_block_row_max_diff * block_row_sums oc.mul_((row_sums / new_row_sums) * exp_row_max_diff).add_((exp_block_row_max_diff / new_row_sums) * exp_values) row_maxes.copy_(new_row_maxes) row_sums.copy_(new_row_sums) ctx.args = (causal, scale, mask, q_bucket_size, k_bucket_size) ctx.save_for_backward(q, k, v, o, all_row_sums, all_row_maxes) return o @staticmethod @torch.no_grad() def backward(ctx, do): """Algorithm 4 in the paper""" causal, scale, mask, q_bucket_size, k_bucket_size = ctx.args q, k, v, o, l, m = ctx.saved_tensors device = q.device max_neg_value = -torch.finfo(q.dtype).max qk_len_diff = max(k.shape[-2] - q.shape[-2], 0) dq = torch.zeros_like(q) dk = torch.zeros_like(k) dv = torch.zeros_like(v) row_splits = zip( q.split(q_bucket_size, dim=-2), o.split(q_bucket_size, dim=-2), do.split(q_bucket_size, dim=-2), mask, l.split(q_bucket_size, dim=-2), m.split(q_bucket_size, dim=-2), dq.split(q_bucket_size, dim=-2), ) for ind, (qc, oc, doc, row_mask, lc, mc, dqc) in enumerate(row_splits): q_start_index = ind * q_bucket_size - qk_len_diff col_splits = zip( k.split(k_bucket_size, dim=-2), v.split(k_bucket_size, dim=-2), dk.split(k_bucket_size, dim=-2), dv.split(k_bucket_size, dim=-2), ) for k_ind, (kc, vc, dkc, dvc) in enumerate(col_splits): k_start_index = k_ind * k_bucket_size attn_weights = einsum("... i d, ... j d -> ... i j", qc, kc) * scale if causal and q_start_index < (k_start_index + k_bucket_size - 1): causal_mask = torch.ones((qc.shape[-2], kc.shape[-2]), dtype=torch.bool, device=device).triu( q_start_index - k_start_index + 1 ) attn_weights.masked_fill_(causal_mask, max_neg_value) exp_attn_weights = torch.exp(attn_weights - mc) if exists(row_mask): exp_attn_weights.masked_fill_(~row_mask, 0.0) p = exp_attn_weights / lc dv_chunk = einsum("... i j, ... i d -> ... j d", p, doc) dp = einsum("... i d, ... j d -> ... i j", doc, vc) D = (doc * oc).sum(dim=-1, keepdims=True) ds = p * scale * (dp - D) dq_chunk = einsum("... i j, ... j d -> ... i d", ds, kc) dk_chunk = einsum("... i j, ... i d -> ... j d", ds, qc) dqc.add_(dq_chunk) dkc.add_(dk_chunk) dvc.add_(dv_chunk) return dq, dk, dv, None, None, None, None def replace_unet_modules(unet: diffusers.models.unet_2d_condition.UNet2DConditionModel, mem_eff_attn, xformers): if mem_eff_attn: replace_unet_cross_attn_to_memory_efficient() elif xformers: replace_unet_cross_attn_to_xformers() def replace_unet_cross_attn_to_memory_efficient(): print("Replace CrossAttention.forward to use FlashAttention (not xformers)") flash_func = FlashAttentionFunction def forward_flash_attn(self, x, context=None, mask=None): q_bucket_size = 512 k_bucket_size = 1024 h = self.heads q = self.to_q(x) context = context if context is not None else x context = context.to(x.dtype) if hasattr(self, "hypernetwork") and self.hypernetwork is not None: context_k, context_v = self.hypernetwork.forward(x, context) context_k = context_k.to(x.dtype) context_v = context_v.to(x.dtype) else: context_k = context context_v = context k = self.to_k(context_k) v = self.to_v(context_v) del context, x q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v)) out = flash_func.apply(q, k, v, mask, False, q_bucket_size, k_bucket_size) out = rearrange(out, "b h n d -> b n (h d)") # diffusers 0.7.0~ わざわざ変えるなよ (;´Д｀) out = self.to_out[0](out) out = self.to_out[1](out) return out diffusers.models.attention.CrossAttention.forward = forward_flash_attn def replace_unet_cross_attn_to_xformers(): print("Replace CrossAttention.forward to use xformers") try: import xformers.ops except ImportError: raise ImportError("No xformers / xformersがインストールされていないようです") def forward_xformers(self, x, context=None, mask=None): h = self.heads q_in = self.to_q(x) context = default(context, x) context = context.to(x.dtype) if hasattr(self, "hypernetwork") and self.hypernetwork is not None: context_k, context_v = self.hypernetwork.forward(x, context) context_k = context_k.to(x.dtype) context_v = context_v.to(x.dtype) else: context_k = context context_v = context k_in = self.to_k(context_k) v_in = self.to_v(context_v) q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b n h d", h=h), (q_in, k_in, v_in)) del q_in, k_in, v_in q = q.contiguous() k = k.contiguous() v = v.contiguous() out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None) # 最適なのを選んでくれる out = rearrange(out, "b n h d -> b n (h d)", h=h) # diffusers 0.7.0~ out = self.to_out[0](out) out = self.to_out[1](out) return out diffusers.models.attention.CrossAttention.forward = forward_xformers # endregion # region arguments def add_sd_models_arguments(parser: argparse.ArgumentParser): # for pretrained models parser.add_argument("--v2", action="store_true", help="load Stable Diffusion v2.0 model / Stable Diffusion 2.0のモデルを読み込む") parser.add_argument( "--v_parameterization", action="store_true", help="enable v-parameterization training / v-parameterization学習を有効にする" ) parser.add_argument( "--pretrained_model_name_or_path", type=str, default=None, help="pretrained model to train, directory to Diffusers model or StableDiffusion checkpoint / 学習元モデル、Diffusers形式モデルのディレクトリまたはStableDiffusionのckptファイル", ) parser.add_argument( "--tokenizer_cache_dir", type=str, default=None, help="directory for caching Tokenizer (for offline training) / Tokenizerをキャッシュするディレクトリ（ネット接続なしでの学習のため）", ) def add_optimizer_arguments(parser: argparse.ArgumentParser): parser.add_argument( "--optimizer_type", type=str, default="", help="Optimizer to use / オプティマイザの種類: AdamW (default), AdamW8bit, Lion, SGDNesterov, SGDNesterov8bit, DAdaptation, AdaFactor", ) # backward compatibility parser.add_argument( "--use_8bit_adam", action="store_true", help="use 8bit AdamW optimizer (requires bitsandbytes) / 8bit Adamオプティマイザを使う（bitsandbytesのインストールが必要）", ) parser.add_argument( "--use_lion_optimizer", action="store_true", help="use Lion optimizer (requires lion-pytorch) / Lionオプティマイザを使う（ lion-pytorch のインストールが必要）", ) parser.add_argument("--learning_rate", type=float, default=2.0e-6, help="learning rate / 学習率") parser.add_argument( "--max_grad_norm", default=1.0, type=float, help="Max gradient norm, 0 for no clipping / 勾配正規化の最大norm、0でclippingを行わない" ) parser.add_argument( "--optimizer_args", type=str, default=None, nargs="*", help='additional arguments for optimizer (like "weight_decay=0.01 betas=0.9,0.999 ...") / オプティマイザの追加引数（例： "weight_decay=0.01 betas=0.9,0.999 ..."）', ) parser.add_argument("--lr_scheduler_type", type=str, default="", help="custom scheduler module / 使用するスケジューラ") parser.add_argument( "--lr_scheduler_args", type=str, default=None, nargs="*", help='additional arguments for scheduler (like "T_max=100") / スケジューラの追加引数（例： "T_max100"）', ) parser.add_argument( "--lr_scheduler", type=str, default="constant", help="scheduler to use for learning rate / 学習率のスケジューラ: linear, cosine, cosine_with_restarts, polynomial, constant (default), constant_with_warmup, adafactor", ) parser.add_argument( "--lr_warmup_steps", type=int, default=0, help="Number of steps for the warmup in the lr scheduler (default is 0) / 学習率のスケジューラをウォームアップするステップ数（デフォルト0）", ) parser.add_argument( "--lr_scheduler_num_cycles", type=int, default=1, help="Number of restarts for cosine scheduler with restarts / cosine with restartsスケジューラでのリスタート回数", ) parser.add_argument( "--lr_scheduler_power", type=float, default=1, help="Polynomial power for polynomial scheduler / polynomialスケジューラでのpolynomial power", ) def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth: bool): parser.add_argument("--output_dir", type=str, default=None, help="directory to output trained model / 学習後のモデル出力先ディレクトリ") parser.add_argument("--output_name", type=str, default=None, help="base name of trained model file / 学習後のモデルの拡張子を除くファイル名") parser.add_argument( "--save_precision", type=str, default=None, choices=[None, "float", "fp16", "bf16"], help="precision in saving / 保存時に精度を変更して保存する", ) parser.add_argument( "--save_every_n_epochs", type=int, default=None, help="save checkpoint every N epochs / 学習中のモデルを指定エポックごとに保存する" ) parser.add_argument( "--save_n_epoch_ratio", type=int, default=None, help="save checkpoint N epoch ratio (for example 5 means save at least 5 files total) / 学習中のモデルを指定のエポック割合で保存する（たとえば5を指定すると最低5個のファイルが保存される）", ) parser.add_argument("--save_last_n_epochs", type=int, default=None, help="save last N checkpoints / 最大Nエポック保存する") parser.add_argument( "--save_last_n_epochs_state", type=int, default=None, help="save last N checkpoints of state (overrides the value of --save_last_n_epochs)/ 最大Nエポックstateを保存する(--save_last_n_epochsの指定を上書きします)", ) parser.add_argument( "--save_state", action="store_true", help="save training state additionally (including optimizer states etc.) / optimizerなど学習状態も含めたstateを追加で保存する", ) parser.add_argument("--resume", type=str, default=None, help="saved state to resume training / 学習再開するモデルのstate") parser.add_argument("--train_batch_size", type=int, default=1, help="batch size for training / 学習時のバッチサイズ") parser.add_argument( "--max_token_length", type=int, default=None, choices=[None, 150, 225], help="max token length of text encoder (default for 75, 150 or 225) / text encoderのトークンの最大長（未指定で75、150または225が指定可）", ) parser.add_argument( "--mem_eff_attn", action="store_true", help="use memory efficient attention for CrossAttention / CrossAttentionに省メモリ版attentionを使う", ) parser.add_argument("--xformers", action="store_true", help="use xformers for CrossAttention / CrossAttentionにxformersを使う") parser.add_argument( "--vae", type=str, default=None, help="path to checkpoint of vae to replace / VAEを入れ替える場合、VAEのcheckpointファイルまたはディレクトリ" ) parser.add_argument("--max_train_steps", type=int, default=1600, help="training steps / 学習ステップ数") parser.add_argument( "--max_train_epochs", type=int, default=None, help="training epochs (overrides max_train_steps) / 学習エポック数（max_train_stepsを上書きします）", ) parser.add_argument( "--max_data_loader_n_workers", type=int, default=8, help="max num workers for DataLoader (lower is less main RAM usage, faster epoch start and slower data loading) / DataLoaderの最大プロセス数（小さい値ではメインメモリの使用量が減りエポック間の待ち時間が減りますが、データ読み込みは遅くなります）", ) parser.add_argument( "--persistent_data_loader_workers", action="store_true", help="persistent DataLoader workers (useful for reduce time gap between epoch, but may use more memory) / DataLoader のワーカーを持続させる (エポック間の時間差を少なくするのに有効だが、より多くのメモリを消費する可能性がある)", ) parser.add_argument("--seed", type=int, default=None, help="random seed for training / 学習時の乱数のseed") parser.add_argument( "--gradient_checkpointing", action="store_true", help="enable gradient checkpointing / grandient checkpointingを有効にする" ) parser.add_argument( "--gradient_accumulation_steps", type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass / 学習時に逆伝播をする前に勾配を合計するステップ数", ) parser.add_argument( "--mixed_precision", type=str, default="no", choices=["no", "fp16", "bf16"], help="use mixed precision / 混合精度を使う場合、その精度" ) parser.add_argument("--full_fp16", action="store_true", help="fp16 training including gradients / 勾配も含めてfp16で学習する") parser.add_argument( "--clip_skip", type=int, default=None, help="use output of nth layer from back of text encoder (n>=1) / text encoderの後ろからn番目の層の出力を用いる（nは1以上）", ) parser.add_argument( "--logging_dir", type=str, default=None, help="enable logging and output TensorBoard log to this directory / ログ出力を有効にしてこのディレクトリにTensorBoard用のログを出力する", ) parser.add_argument("--log_prefix", type=str, default=None, help="add prefix for each log directory / ログディレクトリ名の先頭に追加する文字列") parser.add_argument( "--noise_offset", type=float, default=None, help="enable noise offset with this value (if enabled, around 0.1 is recommended) / Noise offsetを有効にしてこの値を設定する（有効にする場合は0.1程度を推奨）", ) parser.add_argument( "--lowram", action="store_true", help="enable low RAM optimization. e.g. load models to VRAM instead of RAM (for machines which have bigger VRAM than RAM such as Colab and Kaggle) / メインメモリが少ない環境向け最適化を有効にする。たとえばVRAMにモデルを読み込むなど（ColabやKaggleなどRAMに比べてVRAMが多い環境向け）", ) parser.add_argument( "--sample_every_n_steps", type=int, default=None, help="generate sample images every N steps / 学習中のモデルで指定ステップごとにサンプル出力する" ) parser.add_argument( "--sample_every_n_epochs", type=int, default=None, help="generate sample images every N epochs (overwrites n_steps) / 学習中のモデルで指定エポックごとにサンプル出力する（ステップ数指定を上書きします）", ) parser.add_argument( "--sample_prompts", type=str, default=None, help="file for prompts to generate sample images / 学習中モデルのサンプル出力用プロンプトのファイル" ) parser.add_argument( "--sample_sampler", type=str, default="ddim", choices=[ "ddim", "pndm", "lms", "euler", "euler_a", "heun", "dpm_2", "dpm_2_a", "dpmsolver", "dpmsolver++", "dpmsingle", "k_lms", "k_euler", "k_euler_a", "k_dpm_2", "k_dpm_2_a", ], help=f"sampler (scheduler) type for sample images / サンプル出力時のサンプラー（スケジューラ）の種類", ) parser.add_argument( "--config_file", type=str, default=None, help="using .toml instead of args to pass hyperparameter / ハイパーパラメータを引数ではなく.tomlファイルで渡す", ) parser.add_argument( "--output_config", action="store_true", help="output command line args to given .toml file / 引数を.tomlファイルに出力する" ) if support_dreambooth: # DreamBooth training parser.add_argument( "--prior_loss_weight", type=float, default=1.0, help="loss weight for regularization images / 正則化画像のlossの重み" ) def verify_training_args(args: argparse.Namespace): if args.v_parameterization and not args.v2: print("v_parameterization should be with v2 / v1でv_parameterizationを使用することは想定されていません") if args.v2 and args.clip_skip is not None: print("v2 with clip_skip will be unexpected / v2でclip_skipを使用することは想定されていません") def add_dataset_arguments( parser: argparse.ArgumentParser, support_dreambooth: bool, support_caption: bool, support_caption_dropout: bool ): # dataset common parser.add_argument("--train_data_dir", type=str, default=None, help="directory for train images / 学習画像データのディレクトリ") parser.add_argument( "--shuffle_caption", action="store_true", help="shuffle comma-separated caption / コンマで区切られたcaptionの各要素をshuffleする" ) parser.add_argument( "--caption_extension", type=str, default=".caption", help="extension of caption files / 読み込むcaptionファイルの拡張子" ) parser.add_argument( "--caption_extention", type=str, default=None, help="extension of caption files (backward compatibility) / 読み込むcaptionファイルの拡張子（スペルミスを残してあります）", ) parser.add_argument( "--keep_tokens", type=int, default=0, help="keep heading N tokens when shuffling caption tokens (token means comma separated strings) / captionのシャッフル時に、先頭からこの個数のトークンをシャッフルしないで残す（トークンはカンマ区切りの各部分を意味する）", ) parser.add_argument("--color_aug", action="store_true", help="enable weak color augmentation / 学習時に色合いのaugmentationを有効にする") parser.add_argument("--flip_aug", action="store_true", help="enable horizontal flip augmentation / 学習時に左右反転のaugmentationを有効にする") parser.add_argument( "--face_crop_aug_range", type=str, default=None, help="enable face-centered crop augmentation and its range (e.g. 2.0,4.0) / 学習時に顔を中心とした切り出しaugmentationを有効にするときは倍率を指定する（例：2.0,4.0）", ) parser.add_argument( "--random_crop", action="store_true", help="enable random crop (for style training in face-centered crop augmentation) / ランダムな切り出しを有効にする（顔を中心としたaugmentationを行うときに画風の学習用に指定する）", ) parser.add_argument( "--debug_dataset", action="store_true", help="show images for debugging (do not train) / デバッグ用に学習データを画面表示する（学習は行わない）" ) parser.add_argument( "--resolution", type=str, default=None, help="resolution in training ('size' or 'width,height') / 学習時の画像解像度（'サイズ'指定、または'幅,高さ'指定）", ) parser.add_argument( "--cache_latents", action="store_true", help="cache latents to reduce memory (augmentations must be disabled) / メモリ削減のためにlatentをcacheする（augmentationは使用不可）", ) parser.add_argument("--vae_batch_size", type=int, default=1, help="batch size for caching latents / latentのcache時のバッチサイズ") parser.add_argument( "--enable_bucket", action="store_true", help="enable buckets for multi aspect ratio training / 複数解像度学習のためのbucketを有効にする" ) parser.add_argument("--min_bucket_reso", type=int, default=256, help="minimum resolution for buckets / bucketの最小解像度") parser.add_argument("--max_bucket_reso", type=int, default=1024, help="maximum resolution for buckets / bucketの最大解像度") parser.add_argument( "--bucket_reso_steps", type=int, default=64, help="steps of resolution for buckets, divisible by 8 is recommended / bucketの解像度の単位、8で割り切れる値を推奨します", ) parser.add_argument( "--bucket_no_upscale", action="store_true", help="make bucket for each image without upscaling / 画像を拡大せずbucketを作成します" ) parser.add_argument( "--token_warmup_min", type=int, default=1, help="start learning at N tags (token means comma separated strinfloatgs) / タグ数をN個から増やしながら学習する", ) parser.add_argument( "--token_warmup_step", type=float, default=0, help="tag length reaches maximum on N steps (or N*max_train_steps if N<1) / N（N<1ならN*max_train_steps）ステップでタグ長が最大になる。デフォルトは0（最初から最大）", ) if support_caption_dropout: # Textual Inversion はcaptionのdropoutをsupportしない # いわゆるtensorのDropoutと紛らわしいのでprefixにcaptionを付けておく　every_n_epochsは他と平仄を合わせてdefault Noneに parser.add_argument( "--caption_dropout_rate", type=float, default=0.0, help="Rate out dropout caption(0.0~1.0) / captionをdropoutする割合" ) parser.add_argument( "--caption_dropout_every_n_epochs", type=int, default=0, help="Dropout all captions every N epochs / captionを指定エポックごとにdropoutする", ) parser.add_argument( "--caption_tag_dropout_rate", type=float, default=0.0, help="Rate out dropout comma separated tokens(0.0~1.0) / カンマ区切りのタグをdropoutする割合", ) if support_dreambooth: # DreamBooth dataset parser.add_argument("--reg_data_dir", type=str, default=None, help="directory for regularization images / 正則化画像データのディレクトリ") if support_caption: # caption dataset parser.add_argument("--in_json", type=str, default=None, help="json metadata for dataset / データセットのmetadataのjsonファイル") parser.add_argument( "--dataset_repeats", type=int, default=1, help="repeat dataset when training with captions / キャプションでの学習時にデータセットを繰り返す回数" ) def add_sd_saving_arguments(parser: argparse.ArgumentParser): parser.add_argument( "--save_model_as", type=str, default=None, choices=[None, "ckpt", "safetensors", "diffusers", "diffusers_safetensors"], help="format to save the model (default is same to original) / モデル保存時の形式（未指定時は元モデルと同じ）", ) parser.add_argument( "--use_safetensors", action="store_true", help="use safetensors format to save (if save_model_as is not specified) / checkpoint、モデルをsafetensors形式で保存する（save_model_as未指定時）", ) def read_config_from_file(args: argparse.Namespace, parser: argparse.ArgumentParser): if not args.config_file: return args config_path = args.config_file + ".toml" if not args.config_file.endswith(".toml") else args.config_file if args.output_config: # check if config file exists if os.path.exists(config_path): print(f"Config file already exists. Aborting... / 出力先の設定ファイルが既に存在します: {config_path}") exit(1) # convert args to dictionary args_dict = vars(args) # remove unnecessary keys for key in ["config_file", "output_config"]: if key in args_dict: del args_dict[key] # get default args from parser default_args = vars(parser.parse_args([])) # remove default values: cannot use args_dict.items directly because it will be changed during iteration for key, value in list(args_dict.items()): if key in default_args and value == default_args[key]: del args_dict[key] # convert Path to str in dictionary for key, value in args_dict.items(): if isinstance(value, pathlib.Path): args_dict[key] = str(value) # convert to toml and output to file with open(config_path, "w") as f: toml.dump(args_dict, f) print(f"Saved config file / 設定ファイルを保存しました: {config_path}") exit(0) if not os.path.exists(config_path): print(f"{config_path} not found.") exit(1) print(f"Loading settings from {config_path}...") with open(config_path, "r") as f: config_dict = toml.load(f) # combine all sections into one ignore_nesting_dict = {} for section_name, section_dict in config_dict.items(): # if value is not dict, save key and value as is if not isinstance(section_dict, dict): ignore_nesting_dict[section_name] = section_dict continue # if value is dict, save all key and value into one dict for key, value in section_dict.items(): ignore_nesting_dict[key] = value config_args = argparse.Namespace(**ignore_nesting_dict) args = parser.parse_args(namespace=config_args) args.config_file = os.path.splitext(args.config_file)[0] print(args.config_file) return args # endregion # region utils def get_optimizer(args, trainable_params): # "Optimizer to use: AdamW, AdamW8bit, Lion, SGDNesterov, SGDNesterov8bit, DAdaptation, Adafactor" optimizer_type = args.optimizer_type if args.use_8bit_adam: assert ( not args.use_lion_optimizer ), "both option use_8bit_adam and use_lion_optimizer are specified / use_8bit_adamとuse_lion_optimizerの両方のオプションが指定されています" assert ( optimizer_type is None or optimizer_type == "" ), "both option use_8bit_adam and optimizer_type are specified / use_8bit_adamとoptimizer_typeの両方のオプションが指定されています" optimizer_type = "AdamW8bit" elif args.use_lion_optimizer: assert ( optimizer_type is None or optimizer_type == "" ), "both option use_lion_optimizer and optimizer_type are specified / use_lion_optimizerとoptimizer_typeの両方のオプションが指定されています" optimizer_type = "Lion" if optimizer_type is None or optimizer_type == "": optimizer_type = "AdamW" optimizer_type = optimizer_type.lower() # 引数を分解する optimizer_kwargs = {} if args.optimizer_args is not None and len(args.optimizer_args) > 0: for arg in args.optimizer_args: key, value = arg.split("=") value = ast.literal_eval(value) # value = value.split(",") # for i in range(len(value)): # if value[i].lower() == "true" or value[i].lower() == "false": # value[i] = value[i].lower() == "true" # else: # value[i] = ast.float(value[i]) # if len(value) == 1: # value = value[0] # else: # value = tuple(value) optimizer_kwargs[key] = value # print("optkwargs:", optimizer_kwargs) lr = args.learning_rate if optimizer_type == "AdamW8bit".lower(): try: import bitsandbytes as bnb except ImportError: raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです") print(f"use 8-bit AdamW optimizer | {optimizer_kwargs}") optimizer_class = bnb.optim.AdamW8bit optimizer = optimizer_class(trainable_params, lr=lr, **optimizer_kwargs) elif optimizer_type == "SGDNesterov8bit".lower(): try: import bitsandbytes as bnb except ImportError: raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです") print(f"use 8-bit SGD with Nesterov optimizer | {optimizer_kwargs}") if "momentum" not in optimizer_kwargs: print( f"8-bit SGD with Nesterov must be with momentum, set momentum to 0.9 / 8-bit SGD with Nesterovはmomentum指定が必須のため0.9に設定します" ) optimizer_kwargs["momentum"] = 0.9 optimizer_class = bnb.optim.SGD8bit optimizer = optimizer_class(trainable_params, lr=lr, nesterov=True, **optimizer_kwargs) elif optimizer_type == "Lion".lower(): try: import lion_pytorch except ImportError: raise ImportError("No lion_pytorch / lion_pytorch がインストールされていないようです") print(f"use Lion optimizer | {optimizer_kwargs}") optimizer_class = lion_pytorch.Lion optimizer = optimizer_class(trainable_params, lr=lr, **optimizer_kwargs) elif optimizer_type == "SGDNesterov".lower(): print(f"use SGD with Nesterov optimizer | {optimizer_kwargs}") if "momentum" not in optimizer_kwargs: print(f"SGD with Nesterov must be with momentum, set momentum to 0.9 / SGD with Nesterovはmomentum指定が必須のため0.9に設定します") optimizer_kwargs["momentum"] = 0.9 optimizer_class = torch.optim.SGD optimizer = optimizer_class(trainable_params, lr=lr, nesterov=True, **optimizer_kwargs) elif optimizer_type == "DAdaptation".lower(): try: import dadaptation except ImportError: raise ImportError("No dadaptation / dadaptation がインストールされていないようです") print(f"use D-Adaptation Adam optimizer | {optimizer_kwargs}") actual_lr = lr lr_count = 1 if type(trainable_params) == list and type(trainable_params[0]) == dict: lrs = set() actual_lr = trainable_params[0].get("lr", actual_lr) for group in trainable_params: lrs.add(group.get("lr", actual_lr)) lr_count = len(lrs) if actual_lr <= 0.1: print( f"learning rate is too low. If using dadaptation, set learning rate around 1.0 / 学習率が低すぎるようです。1.0前後の値を指定してください: lr={actual_lr}" ) print("recommend option: lr=1.0 / 推奨は1.0です") if lr_count > 1: print( f"when multiple learning rates are specified with dadaptation (e.g. for Text Encoder and U-Net), only the first one will take effect / D-Adaptationで複数の学習率を指定した場合（Text EncoderとU-Netなど）、最初の学習率のみが有効になります: lr={actual_lr}" ) optimizer_class = dadaptation.DAdaptAdam optimizer = optimizer_class(trainable_params, lr=lr, **optimizer_kwargs) elif optimizer_type == "Adafactor".lower(): # 引数を確認して適宜補正する if "relative_step" not in optimizer_kwargs: optimizer_kwargs["relative_step"] = True # default if not optimizer_kwargs["relative_step"] and optimizer_kwargs.get("warmup_init", False): print(f"set relative_step to True because warmup_init is True / warmup_initがTrueのためrelative_stepをTrueにします") optimizer_kwargs["relative_step"] = True print(f"use Adafactor optimizer | {optimizer_kwargs}") if optimizer_kwargs["relative_step"]: print(f"relative_step is true / relative_stepがtrueです") if lr != 0.0: print(f"learning rate is used as initial_lr / 指定したlearning rateはinitial_lrとして使用されます") args.learning_rate = None # trainable_paramsがgroupだった時の処理：lrを削除する if type(trainable_params) == list and type(trainable_params[0]) == dict: has_group_lr = False for group in trainable_params: p = group.pop("lr", None) has_group_lr = has_group_lr or (p is not None) if has_group_lr: # 一応argsを無効にしておく TODO 依存関係が逆転してるのであまり望ましくない print(f"unet_lr and text_encoder_lr are ignored / unet_lrとtext_encoder_lrは無視されます") args.unet_lr = None args.text_encoder_lr = None if args.lr_scheduler != "adafactor": print(f"use adafactor_scheduler / スケジューラにadafactor_schedulerを使用します") args.lr_scheduler = f"adafactor:{lr}" # ちょっと微妙だけど lr = None else: if args.max_grad_norm != 0.0: print( f"because max_grad_norm is set, clip_grad_norm is enabled. consider set to 0 / max_grad_normが設定されているためclip_grad_normが有効になります。0に設定して無効にしたほうがいいかもしれません" ) if args.lr_scheduler != "constant_with_warmup": print(f"constant_with_warmup will be good / スケジューラはconstant_with_warmupが良いかもしれません") if optimizer_kwargs.get("clip_threshold", 1.0) != 1.0: print(f"clip_threshold=1.0 will be good / clip_thresholdは1.0が良いかもしれません") optimizer_class = transformers.optimization.Adafactor optimizer = optimizer_class(trainable_params, lr=lr, **optimizer_kwargs) elif optimizer_type == "AdamW".lower(): print(f"use AdamW optimizer | {optimizer_kwargs}") optimizer_class = torch.optim.AdamW optimizer = optimizer_class(trainable_params, lr=lr, **optimizer_kwargs) else: # 任意のoptimizerを使う optimizer_type = args.optimizer_type # lowerでないやつ（微妙） print(f"use {optimizer_type} | {optimizer_kwargs}") if "." not in optimizer_type: optimizer_module = torch.optim else: values = optimizer_type.split(".") optimizer_module = importlib.import_module(".".join(values[:-1])) optimizer_type = values[-1] optimizer_class = getattr(optimizer_module, optimizer_type) optimizer = optimizer_class(trainable_params, lr=lr, **optimizer_kwargs) optimizer_name = optimizer_class.__module__ + "." + optimizer_class.__name__ optimizer_args = ",".join([f"{k}={v}" for k, v in optimizer_kwargs.items()]) return optimizer_name, optimizer_args, optimizer # Monkeypatch newer get_scheduler() function overridng current version of diffusers.optimizer.get_scheduler # code is taken from https://github.com/huggingface/diffusers diffusers.optimizer, commit d87cc15977b87160c30abaace3894e802ad9e1e6 # Which is a newer release of diffusers than currently packaged with sd-scripts # This code can be removed when newer diffusers version (v0.12.1 or greater) is tested and implemented to sd-scripts def get_scheduler_fix(args, optimizer: Optimizer, num_processes: int): """ Unified API to get any scheduler from its name. """ name = args.lr_scheduler num_warmup_steps = args.lr_warmup_steps num_training_steps = args.max_train_steps * num_processes * args.gradient_accumulation_steps num_cycles = args.lr_scheduler_num_cycles power = args.lr_scheduler_power lr_scheduler_kwargs = {} # get custom lr_scheduler kwargs if args.lr_scheduler_args is not None and len(args.lr_scheduler_args) > 0: for arg in args.lr_scheduler_args: key, value = arg.split("=") value = ast.literal_eval(value) # value = value.split(",") # for i in range(len(value)): # if value[i].lower() == "true" or value[i].lower() == "false": # value[i] = value[i].lower() == "true" # else: # value[i] = ast.literal_eval(value[i]) # if len(value) == 1: # value = value[0] # else: # value = list(value) # some may use list? lr_scheduler_kwargs[key] = value # using any lr_scheduler from other library if args.lr_scheduler_type: lr_scheduler_type = args.lr_scheduler_type print(f"use {lr_scheduler_type} | {lr_scheduler_kwargs} as lr_scheduler") if "." not in lr_scheduler_type: # default to use torch.optim lr_scheduler_module = torch.optim.lr_scheduler else: values = lr_scheduler_type.split(".") lr_scheduler_module = importlib.import_module(".".join(values[:-1])) lr_scheduler_type = values[-1] lr_scheduler_class = getattr(lr_scheduler_module, lr_scheduler_type) lr_scheduler = lr_scheduler_class(optimizer, **lr_scheduler_kwargs) return lr_scheduler if name.startswith("adafactor"): assert ( type(optimizer) == transformers.optimization.Adafactor ), f"adafactor scheduler must be used with Adafactor optimizer / adafactor schedulerはAdafactorオプティマイザと同時に使ってください" initial_lr = float(name.split(":")[1]) # print("adafactor scheduler init lr", initial_lr) return transformers.optimization.AdafactorSchedule(optimizer, initial_lr) name = SchedulerType(name) schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name] if name == SchedulerType.CONSTANT: return schedule_func(optimizer) # All other schedulers require `num_warmup_steps` if num_warmup_steps is None: raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.") if name == SchedulerType.CONSTANT_WITH_WARMUP: return schedule_func(optimizer, num_warmup_steps=num_warmup_steps) # All other schedulers require `num_training_steps` if num_training_steps is None: raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.") if name == SchedulerType.COSINE_WITH_RESTARTS: return schedule_func( optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, num_cycles=num_cycles ) if name == SchedulerType.POLYNOMIAL: return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, power=power) return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps) def prepare_dataset_args(args: argparse.Namespace, support_metadata: bool): # backward compatibility if args.caption_extention is not None: args.caption_extension = args.caption_extention args.caption_extention = None # assert args.resolution is not None, f"resolution is required / resolution（解像度）を指定してください" if args.resolution is not None: args.resolution = tuple([int(r) for r in args.resolution.split(",")]) if len(args.resolution) == 1: args.resolution = (args.resolution[0], args.resolution[0]) assert ( len(args.resolution) == 2 ), f"resolution must be 'size' or 'width,height' / resolution（解像度）は'サイズ'または'幅','高さ'で指定してください: {args.resolution}" if args.face_crop_aug_range is not None: args.face_crop_aug_range = tuple([float(r) for r in args.face_crop_aug_range.split(",")]) assert ( len(args.face_crop_aug_range) == 2 and args.face_crop_aug_range[0] <= args.face_crop_aug_range[1] ), f"face_crop_aug_range must be two floats / face_crop_aug_rangeは'下限,上限'で指定してください: {args.face_crop_aug_range}" else: args.face_crop_aug_range = None if support_metadata: if args.in_json is not None and (args.color_aug or args.random_crop): print( f"latents in npz is ignored when color_aug or random_crop is True / color_augまたはrandom_cropを有効にした場合、npzファイルのlatentsは無視されます" ) def load_tokenizer(args: argparse.Namespace): print("prepare tokenizer") original_path = V2_STABLE_DIFFUSION_PATH if args.v2 else TOKENIZER_PATH tokenizer: CLIPTokenizer = None if args.tokenizer_cache_dir: local_tokenizer_path = os.path.join(args.tokenizer_cache_dir, original_path.replace("/", "_")) if os.path.exists(local_tokenizer_path): print(f"load tokenizer from cache: {local_tokenizer_path}") tokenizer = CLIPTokenizer.from_pretrained(local_tokenizer_path) # same for v1 and v2 if tokenizer is None: if args.v2: tokenizer = CLIPTokenizer.from_pretrained(original_path, subfolder="tokenizer") else: tokenizer = CLIPTokenizer.from_pretrained(original_path) if hasattr(args, "max_token_length") and args.max_token_length is not None: print(f"update token length: {args.max_token_length}") if args.tokenizer_cache_dir and not os.path.exists(local_tokenizer_path): print(f"save Tokenizer to cache: {local_tokenizer_path}") tokenizer.save_pretrained(local_tokenizer_path) return tokenizer def prepare_accelerator(args: argparse.Namespace): if args.logging_dir is None: log_with = None logging_dir = None else: log_with = "tensorboard" log_prefix = "" if args.log_prefix is None else args.log_prefix logging_dir = args.logging_dir + "/" + log_prefix + time.strftime("%Y%m%d%H%M%S", time.localtime()) accelerator = Accelerator( gradient_accumulation_steps=args.gradient_accumulation_steps, mixed_precision=args.mixed_precision, log_with=log_with, logging_dir=logging_dir, ) # accelerateの互換性問題を解決する accelerator_0_15 = True try: accelerator.unwrap_model("dummy", True) print("Using accelerator 0.15.0 or above.") except TypeError: accelerator_0_15 = False def unwrap_model(model): if accelerator_0_15: return accelerator.unwrap_model(model, True) return accelerator.unwrap_model(model) return accelerator, unwrap_model def prepare_dtype(args: argparse.Namespace): weight_dtype = torch.float32 if args.mixed_precision == "fp16": weight_dtype = torch.float16 elif args.mixed_precision == "bf16": weight_dtype = torch.bfloat16 save_dtype = None if args.save_precision == "fp16": save_dtype = torch.float16 elif args.save_precision == "bf16": save_dtype = torch.bfloat16 elif args.save_precision == "float": save_dtype = torch.float32 return weight_dtype, save_dtype def load_target_model(args: argparse.Namespace, weight_dtype, device='cpu'): name_or_path = args.pretrained_model_name_or_path name_or_path = os.readlink(name_or_path) if os.path.islink(name_or_path) else name_or_path load_stable_diffusion_format = os.path.isfile(name_or_path) # determine SD or Diffusers if load_stable_diffusion_format: print("load StableDiffusion checkpoint") text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, name_or_path, device) else: # Diffusers model is loaded to CPU print("load Diffusers pretrained models") try: pipe = StableDiffusionPipeline.from_pretrained(name_or_path, tokenizer=None, safety_checker=None) except EnvironmentError as ex: print( f"model is not found as a file or in Hugging Face, perhaps file name is wrong? / 指定したモデル名のファイル、またはHugging Faceのモデルが見つかりません。ファイル名が誤っているかもしれません: {name_or_path}" ) text_encoder = pipe.text_encoder vae = pipe.vae unet = pipe.unet del pipe # VAEを読み込む if args.vae is not None: vae = model_util.load_vae(args.vae, weight_dtype) print("additional VAE loaded") return text_encoder, vae, unet, load_stable_diffusion_format def patch_accelerator_for_fp16_training(accelerator): org_unscale_grads = accelerator.scaler._unscale_grads_ def _unscale_grads_replacer(optimizer, inv_scale, found_inf, allow_fp16): return org_unscale_grads(optimizer, inv_scale, found_inf, True) accelerator.scaler._unscale_grads_ = _unscale_grads_replacer def get_hidden_states(args: argparse.Namespace, input_ids, tokenizer, text_encoder, weight_dtype=None): # with no_token_padding, the length is not max length, return result immediately if input_ids.size()[-1] != tokenizer.model_max_length: return text_encoder(input_ids)[0] b_size = input_ids.size()[0] input_ids = input_ids.reshape((-1, tokenizer.model_max_length)) # batch_size*3, 77 if args.clip_skip is None: encoder_hidden_states = text_encoder(input_ids)[0] else: enc_out = text_encoder(input_ids, output_hidden_states=True, return_dict=True) encoder_hidden_states = enc_out["hidden_states"][-args.clip_skip] encoder_hidden_states = text_encoder.text_model.final_layer_norm(encoder_hidden_states) # bs*3, 77, 768 or 1024 encoder_hidden_states = encoder_hidden_states.reshape((b_size, -1, encoder_hidden_states.shape[-1])) if args.max_token_length is not None: if args.v2: # v2: ... ... の三連を ... ... へ戻す　正直この実装でいいのかわからん states_list = [encoder_hidden_states[:, 0].unsqueeze(1)] # for i in range(1, args.max_token_length, tokenizer.model_max_length): chunk = encoder_hidden_states[:, i : i + tokenizer.model_max_length - 2] # の後から 最後の前まで if i > 0: for j in range(len(chunk)): if input_ids[j, 1] == tokenizer.eos_token: # 空、つまり ...のパターン chunk[j, 0] = chunk[j, 1] # 次の の値をコピーする states_list.append(chunk) # の後から の前まで states_list.append(encoder_hidden_states[:, -1].unsqueeze(1)) # か のどちらか encoder_hidden_states = torch.cat(states_list, dim=1) else: # v1: ... の三連を ... へ戻す states_list = [encoder_hidden_states[:, 0].unsqueeze(1)] # for i in range(1, args.max_token_length, tokenizer.model_max_length): states_list.append(encoder_hidden_states[:, i : i + tokenizer.model_max_length - 2]) # の後から の前まで states_list.append(encoder_hidden_states[:, -1].unsqueeze(1)) # encoder_hidden_states = torch.cat(states_list, dim=1) if weight_dtype is not None: # this is required for additional network training encoder_hidden_states = encoder_hidden_states.to(weight_dtype) return encoder_hidden_states def get_epoch_ckpt_name(args: argparse.Namespace, use_safetensors, epoch): model_name = DEFAULT_EPOCH_NAME if args.output_name is None else args.output_name ckpt_name = EPOCH_FILE_NAME.format(model_name, epoch) + (".safetensors" if use_safetensors else ".ckpt") return model_name, ckpt_name def save_on_epoch_end(args: argparse.Namespace, save_func, remove_old_func, epoch_no: int, num_train_epochs: int): saving = epoch_no % args.save_every_n_epochs == 0 and epoch_no < num_train_epochs if saving: os.makedirs(args.output_dir, exist_ok=True) save_func() if args.save_last_n_epochs is not None: remove_epoch_no = epoch_no - args.save_every_n_epochs * args.save_last_n_epochs remove_old_func(remove_epoch_no) return saving def save_sd_model_on_epoch_end( args: argparse.Namespace, accelerator, src_path: str, save_stable_diffusion_format: bool, use_safetensors: bool, save_dtype: torch.dtype, epoch: int, num_train_epochs: int, global_step: int, text_encoder, unet, vae, ): epoch_no = epoch + 1 model_name, ckpt_name = get_epoch_ckpt_name(args, use_safetensors, epoch_no) if save_stable_diffusion_format: def save_sd(): ckpt_file = os.path.join(args.output_dir, ckpt_name) print(f"saving checkpoint: {ckpt_file}") model_util.save_stable_diffusion_checkpoint( args.v2, ckpt_file, text_encoder, unet, src_path, epoch_no, global_step, save_dtype, vae ) def remove_sd(old_epoch_no): _, old_ckpt_name = get_epoch_ckpt_name(args, use_safetensors, old_epoch_no) old_ckpt_file = os.path.join(args.output_dir, old_ckpt_name) if os.path.exists(old_ckpt_file): print(f"removing old checkpoint: {old_ckpt_file}") os.remove(old_ckpt_file) save_func = save_sd remove_old_func = remove_sd else: def save_du(): out_dir = os.path.join(args.output_dir, EPOCH_DIFFUSERS_DIR_NAME.format(model_name, epoch_no)) print(f"saving model: {out_dir}") os.makedirs(out_dir, exist_ok=True) model_util.save_diffusers_checkpoint( args.v2, out_dir, text_encoder, unet, src_path, vae=vae, use_safetensors=use_safetensors ) def remove_du(old_epoch_no): out_dir_old = os.path.join(args.output_dir, EPOCH_DIFFUSERS_DIR_NAME.format(model_name, old_epoch_no)) if os.path.exists(out_dir_old): print(f"removing old model: {out_dir_old}") shutil.rmtree(out_dir_old) save_func = save_du remove_old_func = remove_du saving = save_on_epoch_end(args, save_func, remove_old_func, epoch_no, num_train_epochs) if saving and args.save_state: save_state_on_epoch_end(args, accelerator, model_name, epoch_no) def save_state_on_epoch_end(args: argparse.Namespace, accelerator, model_name, epoch_no): print("saving state.") accelerator.save_state(os.path.join(args.output_dir, EPOCH_STATE_NAME.format(model_name, epoch_no))) last_n_epochs = args.save_last_n_epochs_state if args.save_last_n_epochs_state else args.save_last_n_epochs if last_n_epochs is not None: remove_epoch_no = epoch_no - args.save_every_n_epochs * last_n_epochs state_dir_old = os.path.join(args.output_dir, EPOCH_STATE_NAME.format(model_name, remove_epoch_no)) if os.path.exists(state_dir_old): print(f"removing old state: {state_dir_old}") shutil.rmtree(state_dir_old) def save_sd_model_on_train_end( args: argparse.Namespace, src_path: str, save_stable_diffusion_format: bool, use_safetensors: bool, save_dtype: torch.dtype, epoch: int, global_step: int, text_encoder, unet, vae, ): model_name = DEFAULT_LAST_OUTPUT_NAME if args.output_name is None else args.output_name if save_stable_diffusion_format: os.makedirs(args.output_dir, exist_ok=True) ckpt_name = model_name + (".safetensors" if use_safetensors else ".ckpt") ckpt_file = os.path.join(args.output_dir, ckpt_name) print(f"save trained model as StableDiffusion checkpoint to {ckpt_file}") model_util.save_stable_diffusion_checkpoint( args.v2, ckpt_file, text_encoder, unet, src_path, epoch, global_step, save_dtype, vae ) else: out_dir = os.path.join(args.output_dir, model_name) os.makedirs(out_dir, exist_ok=True) print(f"save trained model as Diffusers to {out_dir}") model_util.save_diffusers_checkpoint( args.v2, out_dir, text_encoder, unet, src_path, vae=vae, use_safetensors=use_safetensors ) def save_state_on_train_end(args: argparse.Namespace, accelerator): print("saving last state.") os.makedirs(args.output_dir, exist_ok=True) model_name = DEFAULT_LAST_OUTPUT_NAME if args.output_name is None else args.output_name accelerator.save_state(os.path.join(args.output_dir, LAST_STATE_NAME.format(model_name))) # scheduler: SCHEDULER_LINEAR_START = 0.00085 SCHEDULER_LINEAR_END = 0.0120 SCHEDULER_TIMESTEPS = 1000 SCHEDLER_SCHEDULE = "scaled_linear" def sample_images( accelerator, args: argparse.Namespace, epoch, steps, device, vae, tokenizer, text_encoder, unet, prompt_replacement=None ): """ StableDiffusionLongPromptWeightingPipelineの改造版を使うようにしたので、clip skipおよびプロンプトの重みづけに対応した """ if args.sample_every_n_steps is None and args.sample_every_n_epochs is None: return if args.sample_every_n_epochs is not None: # sample_every_n_steps は無視する if epoch is None or epoch % args.sample_every_n_epochs != 0: return else: if steps % args.sample_every_n_steps != 0 or epoch is not None: # steps is not divisible or end of epoch return print(f"generating sample images at step / サンプル画像生成 ステップ: {steps}") if not os.path.isfile(args.sample_prompts): print(f"No prompt file / プロンプトファイルがありません: {args.sample_prompts}") return org_vae_device = vae.device # CPUにいるはず vae.to(device) # read prompts with open(args.sample_prompts, "rt", encoding="utf-8") as f: prompts = f.readlines() # schedulerを用意する sched_init_args = {} if args.sample_sampler == "ddim": scheduler_cls = DDIMScheduler elif args.sample_sampler == "ddpm": # ddpmはおかしくなるのでoptionから外してある scheduler_cls = DDPMScheduler elif args.sample_sampler == "pndm": scheduler_cls = PNDMScheduler elif args.sample_sampler == "lms" or args.sample_sampler == "k_lms": scheduler_cls = LMSDiscreteScheduler elif args.sample_sampler == "euler" or args.sample_sampler == "k_euler": scheduler_cls = EulerDiscreteScheduler elif args.sample_sampler == "euler_a" or args.sample_sampler == "k_euler_a": scheduler_cls = EulerAncestralDiscreteScheduler elif args.sample_sampler == "dpmsolver" or args.sample_sampler == "dpmsolver++": scheduler_cls = DPMSolverMultistepScheduler sched_init_args["algorithm_type"] = args.sample_sampler elif args.sample_sampler == "dpmsingle": scheduler_cls = DPMSolverSinglestepScheduler elif args.sample_sampler == "heun": scheduler_cls = HeunDiscreteScheduler elif args.sample_sampler == "dpm_2" or args.sample_sampler == "k_dpm_2": scheduler_cls = KDPM2DiscreteScheduler elif args.sample_sampler == "dpm_2_a" or args.sample_sampler == "k_dpm_2_a": scheduler_cls = KDPM2AncestralDiscreteScheduler else: scheduler_cls = DDIMScheduler if args.v_parameterization: sched_init_args["prediction_type"] = "v_prediction" scheduler = scheduler_cls( num_train_timesteps=SCHEDULER_TIMESTEPS, beta_start=SCHEDULER_LINEAR_START, beta_end=SCHEDULER_LINEAR_END, beta_schedule=SCHEDLER_SCHEDULE, **sched_init_args, ) # clip_sample=Trueにする if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is False: # print("set clip_sample to True") scheduler.config.clip_sample = True pipeline = StableDiffusionLongPromptWeightingPipeline( text_encoder=text_encoder, vae=vae, unet=unet, tokenizer=tokenizer, scheduler=scheduler, clip_skip=args.clip_skip, safety_checker=None, feature_extractor=None, requires_safety_checker=False, ) pipeline.to(device) save_dir = args.output_dir + "/sample" os.makedirs(save_dir, exist_ok=True) rng_state = torch.get_rng_state() cuda_rng_state = torch.cuda.get_rng_state() with torch.no_grad(): with accelerator.autocast(): for i, prompt in enumerate(prompts): if not accelerator.is_main_process: continue prompt = prompt.strip() if len(prompt) == 0 or prompt[0] == "#": continue # subset of gen_img_diffusers prompt_args = prompt.split(" --") prompt = prompt_args[0] negative_prompt = None sample_steps = 30 width = height = 512 scale = 7.5 seed = None for parg in prompt_args: try: m = re.match(r"w (\d+)", parg, re.IGNORECASE) if m: width = int(m.group(1)) continue m = re.match(r"h (\d+)", parg, re.IGNORECASE) if m: height = int(m.group(1)) continue m = re.match(r"d (\d+)", parg, re.IGNORECASE) if m: seed = int(m.group(1)) continue m = re.match(r"s (\d+)", parg, re.IGNORECASE) if m: # steps sample_steps = max(1, min(1000, int(m.group(1)))) continue m = re.match(r"l ([\d\.]+)", parg, re.IGNORECASE) if m: # scale scale = float(m.group(1)) continue m = re.match(r"n (.+)", parg, re.IGNORECASE) if m: # negative prompt negative_prompt = m.group(1) continue except ValueError as ex: print(f"Exception in parsing / 解析エラー: {parg}") print(ex) if seed is not None: torch.manual_seed(seed) torch.cuda.manual_seed(seed) if prompt_replacement is not None: prompt = prompt.replace(prompt_replacement[0], prompt_replacement[1]) if negative_prompt is not None: negative_prompt = negative_prompt.replace(prompt_replacement[0], prompt_replacement[1]) height = max(64, height - height % 8) # round to divisible by 8 width = max(64, width - width % 8) # round to divisible by 8 print(f"prompt: {prompt}") print(f"negative_prompt: {negative_prompt}") print(f"height: {height}") print(f"width: {width}") print(f"sample_steps: {sample_steps}") print(f"scale: {scale}") image = pipeline( prompt=prompt, height=height, width=width, num_inference_steps=sample_steps, guidance_scale=scale, negative_prompt=negative_prompt, ).images[0] ts_str = time.strftime("%Y%m%d%H%M%S", time.localtime()) num_suffix = f"e{epoch:06d}" if epoch is not None else f"{steps:06d}" seed_suffix = "" if seed is None else f"_{seed}" img_filename = ( f"{'' if args.output_name is None else args.output_name + '_'}{ts_str}_{num_suffix}_{i:02d}{seed_suffix}.png" ) image.save(os.path.join(save_dir, img_filename)) # clear pipeline and cache to reduce vram usage del pipeline torch.cuda.empty_cache() torch.set_rng_state(rng_state) torch.cuda.set_rng_state(cuda_rng_state) vae.to(org_vae_device) # endregion # region 前処理用 class ImageLoadingDataset(torch.utils.data.Dataset): def __init__(self, image_paths): self.images = image_paths def __len__(self): return len(self.images) def __getitem__(self, idx): img_path = self.images[idx] try: image = Image.open(img_path).convert("RGB") # convert to tensor temporarily so dataloader will accept it tensor_pil = transforms.functional.pil_to_tensor(image) except Exception as e: print(f"Could not load image path / 画像を読み込めません: {img_path}, error: {e}") return None return (tensor_pil, img_path) # endregion # collate_fn用 epoch,stepはmultiprocessing.Value class collater_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]