Add v8 of train_db_fixed.py

Add diffusers_fine_tuning

.gitignore

@@ -1,2 +1,3 @@
 venv
 mytraining.ps
+__pycache__

README.md

@@ -122,4 +122,5 @@ Refer to this url for more details about finetuning: https://note.com/kohya_ss/n
 
 ## Change history
 
-* 11/7 (v7): Text Encoder supports checkpoint files in different storage formats (it is converted at the time of import, so export will be in normal format). Changed the average value of EPOCH loss to output to the screen. Added a function to save epoch and global step in checkpoint in SD format (add values if there is existing data). The reg_data_dir option is enabled during fine tuning (fine tuning while mixing regularized images). Added dataset_repeats option that is valid for fine tuning (specified when the number of teacher images is small and the epoch is extremely short).
+* 11/7 (v7): Text Encoder supports checkpoint files in different storage formats (it is converted at the time of import, so export will be in normal format). Changed the average value of EPOCH loss to output to the screen. Added a function to save epoch and global step in checkpoint in SD format (add values if there is existing data). The reg_data_dir option is enabled during fine tuning (fine tuning while mixing regularized images). Added dataset_repeats option that is valid for fine tuning (specified when the number of teacher images is small and the epoch is extremely short).
+* 11/9 (v8): supports Diffusers 0.7.2. To upgrade diffusers run `pip install --upgrade diffusers[torch]`

diffusers_fine_tuning/README.md

@@ -0,0 +1,3 @@
+# Diffusers Fine Tuning
+
+This subfolder provide all the required toold  to run the diffusers fine tuning version found in this note: https://note.com/kohya_ss/n/nbf7ce8d80f29

diffusers_fine_tuning/clean_captions_and_tags.py

@@ -0,0 +1,122 @@
+# このスクリプトのライセンスは、Apache License 2.0とします
+# (c) 2022 Kohya S. @kohya_ss
+
+import argparse
+import glob
+import os
+import json
+
+from tqdm import tqdm
+
+
+def clean_tags(image_key, tags):
+  # replace '_' to ' '
+  tags = tags.replace('_', ' ')
+
+  # remove rating
+  tokens = tags.split(", rating")
+  if len(tokens) == 1:
+    print("no rating:")
+    print(f"{image_key} {tags}")
+  else:
+    if len(tokens) > 2:
+      print("multiple ratings:")
+      print(f"{image_key} {tags}")
+    tags = tokens[0]
+
+  return tags
+
+
+# 上から順に検索、置換される
+# ('置換元文字列', '置換後文字列')
+CAPTION_REPLACEMENTS = [
+    ('anime anime', 'anime'),
+    ('young ', ''),
+    ('anime girl', 'girl'),
+    ('cartoon female', 'girl'),
+    ('cartoon lady', 'girl'),
+    ('cartoon character', 'girl'),      # a or ~s
+    ('cartoon woman', 'girl'),
+    ('cartoon women', 'girls'),
+    ('cartoon girl', 'girl'),
+    ('anime female', 'girl'),
+    ('anime lady', 'girl'),
+    ('anime character', 'girl'),      # a or ~s
+    ('anime woman', 'girl'),
+    ('anime women', 'girls'),
+    ('lady', 'girl'),
+    ('female', 'girl'),
+    ('woman', 'girl'),
+    ('women', 'girls'),
+    ('people', 'girls'),
+    ('person', 'girl'),
+    ('a cartoon figure', 'a figure'),
+    ('a cartoon image', 'an image'),
+    ('a cartoon picture', 'a picture'),
+    ('an anime cartoon image', 'an image'),
+    ('a cartoon anime drawing', 'a drawing'),
+    ('a cartoon drawing', 'a drawing'),
+    ('girl girl', 'girl'),
+]
+
+
+def clean_caption(caption):
+  for rf, rt in CAPTION_REPLACEMENTS:
+    replaced = True
+    while replaced:
+      bef = caption
+      caption = caption.replace(rf, rt)
+      replaced = bef != caption
+  return caption
+
+def main(args):
+  image_paths = glob.glob(os.path.join(args.train_data_dir, "*.jpg")) + glob.glob(os.path.join(args.train_data_dir, "*.png"))
+  print(f"found {len(image_paths)} images.")
+
+  if os.path.exists(args.in_json):
+    print(f"loading existing metadata: {args.in_json}")
+    with open(args.in_json, "rt", encoding='utf-8') as f:
+      metadata = json.load(f)
+  else:
+    print("no metadata / メタデータファイルがありません")
+    return
+
+  print("cleaning captions and tags.")
+  for image_path in tqdm(image_paths):
+    tags_path = os.path.splitext(image_path)[0] + '.txt'
+    with open(tags_path, "rt", encoding='utf-8') as f:
+      tags = f.readlines()[0].strip()
+
+    image_key = os.path.splitext(os.path.basename(image_path))[0]
+    if image_key not in metadata:
+      print(f"image not in metadata / メタデータに画像がありません: {image_path}")
+      return
+
+    tags = metadata[image_key].get('tags')
+    caption = metadata[image_key].get('caption')
+    if tags is None:
+      print(f"image does not have tags / メタデータにタグがありません: {image_path}")
+      return
+    if caption is None:
+      print(f"image does not have caption / メタデータにキャプションがありません: {image_path}")
+      return
+
+    metadata[image_key]['tags'] = clean_tags(image_key, tags)
+    metadata[image_key]['caption'] = clean_caption(caption)
+
+  # metadataを書き出して終わり
+  print(f"writing metadata: {args.out_json}")
+  with open(args.out_json, "wt", encoding='utf-8') as f:
+    json.dump(metadata, f, indent=2)
+  print("done!")
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("train_data_dir", type=str, help="directory for train images / 学習画像データのディレクトリ")
+  parser.add_argument("in_json", type=str, help="metadata file to input / 読み込むメタデータファイル")
+  parser.add_argument("out_json", type=str, help="metadata file to output / メタデータファイル書き出し先")
+  # parser.add_argument("--debug", action="store_true", help="debug mode")
+
+  args = parser.parse_args()
+  main(args)

diffusers_fine_tuning/fine_tune_v1-ber.py

@@ -0,0 +1,774 @@
+# このスクリプトのライセンスは、train_dreambooth.pyと同じくApache License 2.0とします
+# (c) 2022 Kohya S. @kohya_ss
+
+import argparse
+import math
+import os
+import random
+import json
+import importlib
+
+from tqdm import tqdm
+import torch
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+from transformers import CLIPTextModel, CLIPTokenizer
+import diffusers
+from diffusers import DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel
+import numpy as np
+from einops import rearrange
+from torch import einsum
+
+import fine_tuning_utils_ber as fine_tuning_utils
+
+# Tokenizer: checkpointから読み込むのではなくあらかじめ提供されているものを使う
+TOKENIZER_PATH = "openai/clip-vit-large-patch14"
+
+# checkpointファイル名
+LAST_CHECKPOINT_NAME = "last.ckpt"
+EPOCH_CHECKPOINT_NAME = "epoch-{:06d}.ckpt"
+
+
+def collate_fn(examples):
+  return examples[0]
+
+
+class FineTuningDataset(torch.utils.data.Dataset):
+  def __init__(self, metadata, train_data_dir, batch_size, tokenizer, max_token_length, shuffle_caption, dataset_repeats, debug) -> None:
+    super().__init__()
+
+    self.metadata = metadata
+    self.train_data_dir = train_data_dir
+    self.batch_size = batch_size
+    self.tokenizer = tokenizer
+    self.max_token_length = max_token_length
+    self.shuffle_caption = shuffle_caption
+    self.debug = debug
+
+    self.tokenizer_max_length = self.tokenizer.model_max_length if max_token_length is None else max_token_length + 2
+
+    print("make buckets")
+
+    # 最初に数を数える
+    self.bucket_resos = set()
+    for img_md in metadata.values():
+      if 'train_resolution' in img_md:
+        self.bucket_resos.add(tuple(img_md['train_resolution']))
+    self.bucket_resos = list(self.bucket_resos)
+    self.bucket_resos.sort()
+    print(f"number of buckets: {len(self.bucket_resos)}")
+
+    reso_to_index = {}
+    for i, reso in enumerate(self.bucket_resos):
+      reso_to_index[reso] = i
+
+    # bucketに割り当てていく
+    self.buckets = [[] for _ in range(len(self.bucket_resos))]
+    n = 1 if dataset_repeats is None else dataset_repeats
+    images_count = 0
+    for image_key, img_md in metadata.items():
+      if 'train_resolution' not in img_md:
+        continue
+      if not os.path.exists(os.path.join(self.train_data_dir, image_key + '.npz')):
+        continue
+
+      reso = tuple(img_md['train_resolution'])
+      for _ in range(n):
+        self.buckets[reso_to_index[reso]].append(image_key)
+      images_count += n
+
+    # 参照用indexを作る
+    self.buckets_indices = []
+    for bucket_index, bucket in enumerate(self.buckets):
+      batch_count = int(math.ceil(len(bucket) / self.batch_size))
+      for batch_index in range(batch_count):
+        self.buckets_indices.append((bucket_index, batch_index))
+
+    self.shuffle_buckets()
+    self._length = len(self.buckets_indices)
+    self.images_count = images_count
+
+  def show_buckets(self):
+    for i, (reso, bucket) in enumerate(zip(self.bucket_resos, self.buckets)):
+      print(f"bucket {i}: resolution {reso}, count: {len(bucket)}")
+
+  def shuffle_buckets(self):
+    random.shuffle(self.buckets_indices)
+    for bucket in self.buckets:
+      random.shuffle(bucket)
+
+  def load_latent(self, image_key):
+    return np.load(os.path.join(self.train_data_dir, image_key + '.npz'))['arr_0']
+
+  def __len__(self):
+    return self._length
+
+  def __getitem__(self, index):
+    if index == 0:
+      self.shuffle_buckets()
+
+    bucket = self.buckets[self.buckets_indices[index][0]]
+    image_index = self.buckets_indices[index][1] * self.batch_size
+
+    input_ids_list = []
+    latents_list = []
+    captions = []
+    for image_key in bucket[image_index:image_index + self.batch_size]:
+      img_md = self.metadata[image_key]
+      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
+      assert caption is not None and len(caption) > 0, f"caption or tag is required / キャプションまたはタグは必須です:{image_key}"
+
+      latents = self.load_latent(image_key)
+
+      if self.shuffle_caption:
+        tokens = caption.strip().split(",")
+        random.shuffle(tokens)
+        caption = ",".join(tokens).strip()
+
+      captions.append(caption)
+
+      input_ids = self.tokenizer(caption, padding="max_length", truncation=True,
+                                 max_length=self.tokenizer_max_length, return_tensors="pt").input_ids
+
+      # 77以上の時は "<CLS> .... <EOS> <EOS> <EOS>" でトータル227とかになっているので、"<CLS>...<EOS>"の三連に変換する
+      # 1111氏のやつは , で区切る、とかしているようだが　とりあえず単純に
+      if self.tokenizer_max_length > self.tokenizer.model_max_length:
+        input_ids = input_ids.squeeze(0)
+        iids_list = []
+        for i in range(1, self.tokenizer_max_length - self.tokenizer.model_max_length + 2, self.tokenizer.model_max_length - 2):
+          iid = (input_ids[0].unsqueeze(0),
+                 input_ids[i:i + self.tokenizer.model_max_length - 2],
+                 input_ids[-1].unsqueeze(0))
+          iid = torch.cat(iid)
+          iids_list.append(iid)
+        input_ids = torch.stack(iids_list)      # 3,77
+
+      input_ids_list.append(input_ids)
+      latents_list.append(torch.FloatTensor(latents))
+
+    example = {}
+    example['input_ids'] = torch.stack(input_ids_list)
+    example['latents'] = torch.stack(latents_list)
+    if self.debug:
+      example['image_keys'] = bucket[image_index:image_index + self.batch_size]
+      example['captions'] = captions
+    return example
+
+
+def save_hypernetwork(output_file, hypernetwork):
+  state_dict = hypernetwork.get_state_dict()
+  torch.save(state_dict, output_file)
+
+
+def train(args):
+  fine_tuning = args.hypernetwork_module is None            # fine tuning or hypernetwork training
+
+  # モデル形式のオプション設定を確認する
+  use_stable_diffusion_format = os.path.isfile(args.pretrained_model_name_or_path)
+  if not use_stable_diffusion_format:
+    assert os.path.exists(
+        args.pretrained_model_name_or_path), f"no pretrained model / 学習元モデルがありません : {args.pretrained_model_name_or_path}"
+
+  assert not fine_tuning or (
+      args.save_every_n_epochs is None or use_stable_diffusion_format), "when loading Diffusers model, save_every_n_epochs does not work / Diffusersのモデルを読み込むときにはsave_every_n_epochsオプションは無効になります"
+
+  if args.seed is not None:
+    set_seed(args.seed)
+
+  # メタデータを読み込む
+  if os.path.exists(args.in_json):
+    print(f"loading existing metadata: {args.in_json}")
+    with open(args.in_json, "rt", encoding='utf-8') as f:
+      metadata = json.load(f)
+  else:
+    print(f"no metadata / メタデータファイルがありません: {args.in_json}")
+    return
+
+  # tokenizerを読み込む
+  print("prepare tokenizer")
+  tokenizer = CLIPTokenizer.from_pretrained(TOKENIZER_PATH)
+  if args.max_token_length is not None:
+    print(f"update token length in tokenizer: {args.max_token_length}")
+
+  # datasetを用意する
+  print("prepare dataset")
+  train_dataset = FineTuningDataset(metadata, args.train_data_dir, args.train_batch_size,
+                                    tokenizer, args.max_token_length, args.shuffle_caption, args.dataset_repeats, args.debug_dataset)
+
+  if args.debug_dataset:
+    print(f"Total dataset length / データセットの長さ: {len(train_dataset)}")
+    print(f"Total images / 画像数: {train_dataset.images_count}")
+    train_dataset.show_buckets()
+    i = 0
+    for example in train_dataset:
+      print(f"image: {example['image_keys']}")
+      print(f"captions: {example['captions']}")
+      print(f"latents: {example['latents'].shape}")
+      print(f"input_ids: {example['input_ids'].shape}")
+      print(example['input_ids'])
+      i += 1
+      if i >= 8:
+        break
+    return
+
+  # acceleratorを準備する
+  print("prepare accelerator")
+  accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps, mixed_precision=args.mixed_precision)
+
+  # モデルを読み込む
+  if use_stable_diffusion_format:
+    print("load StableDiffusion checkpoint")
+    text_encoder, _, unet = fine_tuning_utils.load_models_from_stable_diffusion_checkpoint(args.pretrained_model_name_or_path)
+  else:
+    print("load Diffusers pretrained models")
+    text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="text_encoder")
+    unet = UNet2DConditionModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="unet")
+
+  # モデルに xformers とか memory efficient attention を組み込む
+  replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
+
+  if not fine_tuning:
+    # Hypernetwork
+    print("import hypernetwork module:", args.hypernetwork_module)
+    hyp_module = importlib.import_module(args.hypernetwork_module)
+
+    hypernetwork = hyp_module.Hypernetwork()
+
+    if args.hypernetwork_weights is not None:
+      print("load hypernetwork weights from:", args.hypernetwork_weights)
+      hyp_sd = torch.load(args.hypernetwork_weights, map_location='cpu')
+      success = hypernetwork.load_from_state_dict(hyp_sd)
+      assert success, "hypernetwork weights loading failed."
+
+    print("apply hypernetwork")
+    hypernetwork.apply_to_diffusers(None, text_encoder, unet)
+
+  # mixed precisionに対応した型を用意しておき適宜castする
+  weight_dtype = torch.float32
+  if args.mixed_precision == "fp16":
+    weight_dtype = torch.float16
+  elif args.mixed_precision == "bf16":
+    weight_dtype = torch.bfloat16
+
+  # 学習を準備する
+  if fine_tuning:
+    if args.gradient_checkpointing:
+      unet.enable_gradient_checkpointing()
+    unet.requires_grad_(True)             # unetは学習しない
+    net = unet
+  else:
+    unet.requires_grad_(False)             # unetは学習しない
+    unet.eval()
+
+    hypernetwork.requires_grad_(True)
+    net = hypernetwork
+
+  # 学習に必要なクラスを準備する
+  print("prepare optimizer, data loader etc.")
+
+  # 8-bit Adamを使う
+  if args.use_8bit_adam:
+    try:
+      import bitsandbytes as bnb
+    except ImportError:
+      raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
+    print("use 8-bit Adam optimizer")
+    optimizer_class = bnb.optim.AdamW8bit
+  else:
+    optimizer_class = torch.optim.AdamW
+
+  # betaやweight decayはdiffusers DreamBoothもDreamBooth SDもデフォルト値のようなのでオプションはとりあえず省略
+  optimizer = optimizer_class(net.parameters(), lr=args.learning_rate)
+
+  # dataloaderを準備する
+  # DataLoaderのプロセス数：0はメインプロセスになる
+  n_workers = min(8, os.cpu_count() - 1)      # cpu_count-1 ただし最大8
+  train_dataloader = torch.utils.data.DataLoader(
+      train_dataset, batch_size=1, shuffle=True, collate_fn=collate_fn, num_workers=n_workers)
+
+  # lr schedulerを用意する
+  lr_scheduler = diffusers.optimization.get_scheduler(
+      "constant", optimizer, num_training_steps=args.max_train_steps * args.gradient_accumulation_steps)
+
+  # acceleratorがなんかよろしくやってくれるらしい
+  if fine_tuning:
+    unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(unet, optimizer, train_dataloader, lr_scheduler)
+    net = unet
+  else:
+    unet, hypernetwork, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+        unet, hypernetwork, optimizer, train_dataloader, lr_scheduler)
+    net = hypernetwork
+
+  text_encoder.to(accelerator.device, dtype=weight_dtype)
+  text_encoder.requires_grad_(False)             # text encoderは学習しない
+
+  # epoch数を計算する
+  num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+  num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+  # 学習する
+  total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+  print("running training / 学習開始")
+  print(f"  num examples / サンプル数: {train_dataset.images_count}")
+  print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
+  print(f"  num epochs / epoch数: {num_train_epochs}")
+  print(f"  batch size per device / バッチサイズ: {args.train_batch_size}")
+  print(f"  total train batch size (with parallel & distributed) / 総バッチサイズ（並列学習含む）: {total_batch_size}")
+  print(f"  gradient ccumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
+  print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")
+
+  progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process, desc="steps")
+  global_step = 0
+
+  noise_scheduler = DDPMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
+
+  if accelerator.is_main_process:
+    accelerator.init_trackers("finetuning" if fine_tuning else "hypernetwork")
+
+  # 以下 train_dreambooth.py からほぼコピペ
+  for epoch in range(num_train_epochs):
+    print(f"epoch {epoch+1}/{num_train_epochs}")
+    net.train()
+
+    loss_total = 0
+    for step, batch in enumerate(train_dataloader):
+      with accelerator.accumulate(unet):
+        latents = batch["latents"].to(accelerator.device)
+        latents = latents * 0.18215
+        b_size = latents.shape[0]
+
+        with torch.no_grad():
+          # Get the text embedding for conditioning
+          input_ids = batch["input_ids"].to(accelerator.device)
+          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)
+
+          encoder_hidden_states = encoder_hidden_states.reshape((b_size, -1, encoder_hidden_states.shape[-1]))
+
+          if args.max_token_length is not None:
+            # <CLS>...<EOS> の三連を <CLS>...<EOS> へ戻す
+            sts_list = [encoder_hidden_states[:, 0].unsqueeze(1)]
+            for i in range(1, args.max_token_length, tokenizer.model_max_length):
+              sts_list.append(encoder_hidden_states[:, i:i + tokenizer.model_max_length - 2])
+            sts_list.append(encoder_hidden_states[:, -1].unsqueeze(1))
+            encoder_hidden_states = torch.cat(sts_list, dim=1)
+
+        # Sample noise that we'll add to the latents
+        noise = torch.randn_like(latents, device=latents.device)
+
+        # Sample a random timestep for each image
+        timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (b_size,), device=latents.device)
+        timesteps = timesteps.long()
+
+        # Add noise to the latents according to the noise magnitude at each timestep
+        # (this is the forward diffusion process)
+        noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+
+        # Predict the noise residual
+        noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
+
+        loss = torch.nn.functional.mse_loss(noise_pred.float(), noise.float(), reduction="mean")
+
+        accelerator.backward(loss)
+        if accelerator.sync_gradients:
+          accelerator.clip_grad_norm_(net.parameters(), 1.0)  # args.max_grad_norm)
+
+        optimizer.step()
+        lr_scheduler.step()
+        optimizer.zero_grad(set_to_none=True)
+
+      # Checks if the accelerator has performed an optimization step behind the scenes
+      if accelerator.sync_gradients:
+        progress_bar.update(1)
+        global_step += 1
+
+      current_loss = loss.detach().item() * b_size
+      loss_total += current_loss
+      avr_loss = loss_total / (step+1)
+      logs = {"loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
+      progress_bar.set_postfix(**logs)
+      # accelerator.log(logs, step=global_step)
+
+      if global_step >= args.max_train_steps:
+        break
+
+    accelerator.wait_for_everyone()
+
+    if args.save_every_n_epochs is not None:
+      if (epoch + 1) % args.save_every_n_epochs == 0 and (epoch + 1) < num_train_epochs:
+        print("saving check point.")
+        os.makedirs(args.output_dir, exist_ok=True)
+        ckpt_file = os.path.join(args.output_dir, EPOCH_CHECKPOINT_NAME.format(epoch + 1))
+
+        if fine_tuning:
+          fine_tuning_utils.save_stable_diffusion_checkpoint(
+              ckpt_file, text_encoder, accelerator.unwrap_model(net), args.pretrained_model_name_or_path, epoch + 1, global_step)
+        else:
+          save_hypernetwork(ckpt_file, accelerator.unwrap_model(net))
+
+  is_main_process = accelerator.is_main_process
+  if is_main_process:
+    net = accelerator.unwrap_model(net)
+
+  accelerator.end_training()
+  del accelerator                         # この後メモリを使うのでこれは消す
+
+  if is_main_process:
+    os.makedirs(args.output_dir, exist_ok=True)
+    if fine_tuning:
+      if use_stable_diffusion_format:
+        ckpt_file = os.path.join(args.output_dir, LAST_CHECKPOINT_NAME)
+        print(f"save trained model as StableDiffusion checkpoint to {ckpt_file}")
+        fine_tuning_utils.save_stable_diffusion_checkpoint(
+            ckpt_file, text_encoder, unet, args.pretrained_model_name_or_path, epoch, global_step)
+      else:
+        # Create the pipeline using using the trained modules and save it.
+        print(f"save trained model as Diffusers to {args.output_dir}")
+        pipeline = StableDiffusionPipeline.from_pretrained(
+            args.pretrained_model_name_or_path,
+            unet=unet,
+            text_encoder=text_encoder,
+        )
+        pipeline.save_pretrained(args.output_dir)
+    else:
+      ckpt_file = os.path.join(args.output_dir, LAST_CHECKPOINT_NAME)
+      print(f"save trained model to {ckpt_file}")
+      save_hypernetwork(ckpt_file, net)
+    print("model saved.")
+
+
+# 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
+
+# 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.)
+
+        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.)
+
+        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")
+  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.6.0
+    if type(self.to_out) is torch.nn.Sequential:
+      return self.to_out(out)
+
+    # 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.6.0
+    if type(self.to_out) is torch.nn.Sequential:
+      return self.to_out(out)
+
+    # 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
+
+
+if __name__ == '__main__':
+  # torch.cuda.set_per_process_memory_fraction(0.48)
+  parser = argparse.ArgumentParser()
+  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("--in_json", type=str, default=None, help="metadata file to input / 読みこむメタデータファイル")
+  parser.add_argument("--shuffle_caption", action="store_true",
+                      help="shuffle comma-separated caption when fine tuning / fine tuning時にコンマで区切られたcaptionの各要素をshuffleする")
+  parser.add_argument("--train_data_dir", type=str, default=None, help="directory for train images / 学習画像データのディレクトリ")
+  parser.add_argument("--dataset_repeats", type=int, default=None, help="num times to repeat dataset / 学習にデータセットを繰り返す回数")
+  parser.add_argument("--output_dir", type=str, default=None,
+                      help="directory to output trained model, save as same format as input / 学習後のモデル出力先ディレクトリ（入力と同じ形式で保存）")
+  parser.add_argument("--hypernetwork_module", type=str, default=None,
+                      help='train hypernetwork instead of fine tuning, module to use / fine tuningの代わりにHypernetworkの学習をする場合、そのモジュール')
+  parser.add_argument("--hypernetwork_weights", type=str, default=None,
+                      help='hypernetwork weights to initialize for additional training / Hypernetworkの学習時に読み込む重み（Hypernetworkの追加学習）')
+  parser.add_argument("--save_every_n_epochs", type=int, default=None,
+                      help="save checkpoint every N epochs (only supports in StableDiffusion checkpoint) / 学習中のモデルを指定エポックごとに保存する（StableDiffusion形式のモデルを読み込んだ場合のみ有効）")
+  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("--train_batch_size", type=int, default=1,
+                      help="batch size for training / 学習時のバッチサイズ")
+  parser.add_argument("--use_8bit_adam", action="store_true",
+                      help="use 8bit Adam optimizer (requires bitsandbytes) / 8bit Adamオプティマイザを使う（bitsandbytesのインストールが必要）")
+  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("--learning_rate", type=float, default=2.0e-6, help="learning rate / 学習率")
+  parser.add_argument("--max_train_steps", type=int, default=1600, help="training steps / 学習ステップ数")
+  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("--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("--debug_dataset", action="store_true",
+                      help="show images for debugging (do not train) / デバッグ用に学習データを画面表示する（学習は行わない）")
+  parser.add_argument("--save_half", action="store_true",
+                      help="save ckpt model with fp16 precision")
+
+  args = parser.parse_args()
+  train(args)

diffusers_fine_tuning/fine_tune_v1.py

@@ -0,0 +1,772 @@
+# このスクリプトのライセンスは、train_dreambooth.pyと同じくApache License 2.0とします
+# (c) 2022 Kohya S. @kohya_ss
+
+import argparse
+import math
+import os
+import random
+import json
+import importlib
+
+from tqdm import tqdm
+import torch
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+from transformers import CLIPTextModel, CLIPTokenizer
+import diffusers
+from diffusers import DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel
+import numpy as np
+from einops import rearrange
+from torch import einsum
+
+import fine_tuning_utils
+
+# Tokenizer: checkpointから読み込むのではなくあらかじめ提供されているものを使う
+TOKENIZER_PATH = "openai/clip-vit-large-patch14"
+
+# checkpointファイル名
+LAST_CHECKPOINT_NAME = "last.ckpt"
+EPOCH_CHECKPOINT_NAME = "epoch-{:06d}.ckpt"
+
+
+def collate_fn(examples):
+  return examples[0]
+
+
+class FineTuningDataset(torch.utils.data.Dataset):
+  def __init__(self, metadata, train_data_dir, batch_size, tokenizer, max_token_length, shuffle_caption, dataset_repeats, debug) -> None:
+    super().__init__()
+
+    self.metadata = metadata
+    self.train_data_dir = train_data_dir
+    self.batch_size = batch_size
+    self.tokenizer = tokenizer
+    self.max_token_length = max_token_length
+    self.shuffle_caption = shuffle_caption
+    self.debug = debug
+
+    self.tokenizer_max_length = self.tokenizer.model_max_length if max_token_length is None else max_token_length + 2
+
+    print("make buckets")
+
+    # 最初に数を数える
+    self.bucket_resos = set()
+    for img_md in metadata.values():
+      if 'train_resolution' in img_md:
+        self.bucket_resos.add(tuple(img_md['train_resolution']))
+    self.bucket_resos = list(self.bucket_resos)
+    self.bucket_resos.sort()
+    print(f"number of buckets: {len(self.bucket_resos)}")
+
+    reso_to_index = {}
+    for i, reso in enumerate(self.bucket_resos):
+      reso_to_index[reso] = i
+
+    # bucketに割り当てていく
+    self.buckets = [[] for _ in range(len(self.bucket_resos))]
+    n = 1 if dataset_repeats is None else dataset_repeats
+    images_count = 0
+    for image_key, img_md in metadata.items():
+      if 'train_resolution' not in img_md:
+        continue
+      if not os.path.exists(os.path.join(self.train_data_dir, image_key + '.npz')):
+        continue
+
+      reso = tuple(img_md['train_resolution'])
+      for _ in range(n):
+        self.buckets[reso_to_index[reso]].append(image_key)
+      images_count += n
+
+    # 参照用indexを作る
+    self.buckets_indices = []
+    for bucket_index, bucket in enumerate(self.buckets):
+      batch_count = int(math.ceil(len(bucket) / self.batch_size))
+      for batch_index in range(batch_count):
+        self.buckets_indices.append((bucket_index, batch_index))
+
+    self.shuffle_buckets()
+    self._length = len(self.buckets_indices)
+    self.images_count = images_count
+
+  def show_buckets(self):
+    for i, (reso, bucket) in enumerate(zip(self.bucket_resos, self.buckets)):
+      print(f"bucket {i}: resolution {reso}, count: {len(bucket)}")
+
+  def shuffle_buckets(self):
+    random.shuffle(self.buckets_indices)
+    for bucket in self.buckets:
+      random.shuffle(bucket)
+
+  def load_latent(self, image_key):
+    return np.load(os.path.join(self.train_data_dir, image_key + '.npz'))['arr_0']
+
+  def __len__(self):
+    return self._length
+
+  def __getitem__(self, index):
+    if index == 0:
+      self.shuffle_buckets()
+
+    bucket = self.buckets[self.buckets_indices[index][0]]
+    image_index = self.buckets_indices[index][1] * self.batch_size
+
+    input_ids_list = []
+    latents_list = []
+    captions = []
+    for image_key in bucket[image_index:image_index + self.batch_size]:
+      img_md = self.metadata[image_key]
+      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
+      assert caption is not None and len(caption) > 0, f"caption or tag is required / キャプションまたはタグは必須です:{image_key}"
+
+      latents = self.load_latent(image_key)
+
+      if self.shuffle_caption:
+        tokens = caption.strip().split(",")
+        random.shuffle(tokens)
+        caption = ",".join(tokens).strip()
+
+      captions.append(caption)
+
+      input_ids = self.tokenizer(caption, padding="max_length", truncation=True,
+                                 max_length=self.tokenizer_max_length, return_tensors="pt").input_ids
+
+      # 77以上の時は "<CLS> .... <EOS> <EOS> <EOS>" でトータル227とかになっているので、"<CLS>...<EOS>"の三連に変換する
+      # 1111氏のやつは , で区切る、とかしているようだが　とりあえず単純に
+      if self.tokenizer_max_length > self.tokenizer.model_max_length:
+        input_ids = input_ids.squeeze(0)
+        iids_list = []
+        for i in range(1, self.tokenizer_max_length - self.tokenizer.model_max_length + 2, self.tokenizer.model_max_length - 2):
+          iid = (input_ids[0].unsqueeze(0),
+                 input_ids[i:i + self.tokenizer.model_max_length - 2],
+                 input_ids[-1].unsqueeze(0))
+          iid = torch.cat(iid)
+          iids_list.append(iid)
+        input_ids = torch.stack(iids_list)      # 3,77
+
+      input_ids_list.append(input_ids)
+      latents_list.append(torch.FloatTensor(latents))
+
+    example = {}
+    example['input_ids'] = torch.stack(input_ids_list)
+    example['latents'] = torch.stack(latents_list)
+    if self.debug:
+      example['image_keys'] = bucket[image_index:image_index + self.batch_size]
+      example['captions'] = captions
+    return example
+
+
+def save_hypernetwork(output_file, hypernetwork):
+  state_dict = hypernetwork.get_state_dict()
+  torch.save(state_dict, output_file)
+
+
+def train(args):
+  fine_tuning = args.hypernetwork_module is None            # fine tuning or hypernetwork training
+
+  # モデル形式のオプション設定を確認する
+  use_stable_diffusion_format = os.path.isfile(args.pretrained_model_name_or_path)
+  if not use_stable_diffusion_format:
+    assert os.path.exists(
+        args.pretrained_model_name_or_path), f"no pretrained model / 学習元モデルがありません : {args.pretrained_model_name_or_path}"
+
+  assert not fine_tuning or (
+      args.save_every_n_epochs is None or use_stable_diffusion_format), "when loading Diffusers model, save_every_n_epochs does not work / Diffusersのモデルを読み込むときにはsave_every_n_epochsオプションは無効になります"
+
+  if args.seed is not None:
+    set_seed(args.seed)
+
+  # メタデータを読み込む
+  if os.path.exists(args.in_json):
+    print(f"loading existing metadata: {args.in_json}")
+    with open(args.in_json, "rt", encoding='utf-8') as f:
+      metadata = json.load(f)
+  else:
+    print(f"no metadata / メタデータファイルがありません: {args.in_json}")
+    return
+
+  # tokenizerを読み込む
+  print("prepare tokenizer")
+  tokenizer = CLIPTokenizer.from_pretrained(TOKENIZER_PATH)
+  if args.max_token_length is not None:
+    print(f"update token length in tokenizer: {args.max_token_length}")
+
+  # datasetを用意する
+  print("prepare dataset")
+  train_dataset = FineTuningDataset(metadata, args.train_data_dir, args.train_batch_size,
+                                    tokenizer, args.max_token_length, args.shuffle_caption, args.dataset_repeats, args.debug_dataset)
+
+  if args.debug_dataset:
+    print(f"Total dataset length / データセットの長さ: {len(train_dataset)}")
+    print(f"Total images / 画像数: {train_dataset.images_count}")
+    train_dataset.show_buckets()
+    i = 0
+    for example in train_dataset:
+      print(f"image: {example['image_keys']}")
+      print(f"captions: {example['captions']}")
+      print(f"latents: {example['latents'].shape}")
+      print(f"input_ids: {example['input_ids'].shape}")
+      print(example['input_ids'])
+      i += 1
+      if i >= 8:
+        break
+    return
+
+  # acceleratorを準備する
+  print("prepare accelerator")
+  accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps, mixed_precision=args.mixed_precision)
+
+  # モデルを読み込む
+  if use_stable_diffusion_format:
+    print("load StableDiffusion checkpoint")
+    text_encoder, _, unet = fine_tuning_utils.load_models_from_stable_diffusion_checkpoint(args.pretrained_model_name_or_path)
+  else:
+    print("load Diffusers pretrained models")
+    text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="text_encoder")
+    unet = UNet2DConditionModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="unet")
+
+  # モデルに xformers とか memory efficient attention を組み込む
+  replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
+
+  if not fine_tuning:
+    # Hypernetwork
+    print("import hypernetwork module:", args.hypernetwork_module)
+    hyp_module = importlib.import_module(args.hypernetwork_module)
+
+    hypernetwork = hyp_module.Hypernetwork()
+
+    if args.hypernetwork_weights is not None:
+      print("load hypernetwork weights from:", args.hypernetwork_weights)
+      hyp_sd = torch.load(args.hypernetwork_weights, map_location='cpu')
+      success = hypernetwork.load_from_state_dict(hyp_sd)
+      assert success, "hypernetwork weights loading failed."
+
+    print("apply hypernetwork")
+    hypernetwork.apply_to_diffusers(None, text_encoder, unet)
+
+  # mixed precisionに対応した型を用意しておき適宜castする
+  weight_dtype = torch.float32
+  if args.mixed_precision == "fp16":
+    weight_dtype = torch.float16
+  elif args.mixed_precision == "bf16":
+    weight_dtype = torch.bfloat16
+
+  # 学習を準備する
+  if fine_tuning:
+    if args.gradient_checkpointing:
+      unet.enable_gradient_checkpointing()
+    unet.requires_grad_(True)             # unetは学習しない
+    net = unet
+  else:
+    unet.requires_grad_(False)             # unetは学習しない
+    unet.eval()
+
+    hypernetwork.requires_grad_(True)
+    net = hypernetwork
+
+  # 学習に必要なクラスを準備する
+  print("prepare optimizer, data loader etc.")
+
+  # 8-bit Adamを使う
+  if args.use_8bit_adam:
+    try:
+      import bitsandbytes as bnb
+    except ImportError:
+      raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
+    print("use 8-bit Adam optimizer")
+    optimizer_class = bnb.optim.AdamW8bit
+  else:
+    optimizer_class = torch.optim.AdamW
+
+  # betaやweight decayはdiffusers DreamBoothもDreamBooth SDもデフォルト値のようなのでオプションはとりあえず省略
+  optimizer = optimizer_class(net.parameters(), lr=args.learning_rate)
+
+  # dataloaderを準備する
+  # DataLoaderのプロセス数：0はメインプロセスになる
+  n_workers = min(8, os.cpu_count() - 1)      # cpu_count-1 ただし最大8
+  train_dataloader = torch.utils.data.DataLoader(
+      train_dataset, batch_size=1, shuffle=True, collate_fn=collate_fn, num_workers=n_workers)
+
+  # lr schedulerを用意する
+  lr_scheduler = diffusers.optimization.get_scheduler(
+      "constant", optimizer, num_training_steps=args.max_train_steps * args.gradient_accumulation_steps)
+
+  # acceleratorがなんかよろしくやってくれるらしい
+  if fine_tuning:
+    unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(unet, optimizer, train_dataloader, lr_scheduler)
+    net = unet
+  else:
+    unet, hypernetwork, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+        unet, hypernetwork, optimizer, train_dataloader, lr_scheduler)
+    net = hypernetwork
+
+  text_encoder.to(accelerator.device, dtype=weight_dtype)
+  text_encoder.requires_grad_(False)             # text encoderは学習しない
+
+  # epoch数を計算する
+  num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+  num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+  # 学習する
+  total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+  print("running training / 学習開始")
+  print(f"  num examples / サンプル数: {train_dataset.images_count}")
+  print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
+  print(f"  num epochs / epoch数: {num_train_epochs}")
+  print(f"  batch size per device / バッチサイズ: {args.train_batch_size}")
+  print(f"  total train batch size (with parallel & distributed) / 総バッチサイズ（並列学習含む）: {total_batch_size}")
+  print(f"  gradient ccumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
+  print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")
+
+  progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process, desc="steps")
+  global_step = 0
+
+  noise_scheduler = DDPMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
+
+  if accelerator.is_main_process:
+    accelerator.init_trackers("finetuning" if fine_tuning else "hypernetwork")
+
+  # 以下 train_dreambooth.py からほぼコピペ
+  for epoch in range(num_train_epochs):
+    print(f"epoch {epoch+1}/{num_train_epochs}")
+    net.train()
+
+    loss_total = 0
+    for step, batch in enumerate(train_dataloader):
+      with accelerator.accumulate(unet):
+        latents = batch["latents"].to(accelerator.device)
+        latents = latents * 0.18215
+        b_size = latents.shape[0]
+
+        with torch.no_grad():
+          # Get the text embedding for conditioning
+          input_ids = batch["input_ids"].to(accelerator.device)
+          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)
+
+          encoder_hidden_states = encoder_hidden_states.reshape((b_size, -1, encoder_hidden_states.shape[-1]))
+
+          if args.max_token_length is not None:
+            # <CLS>...<EOS> の三連を <CLS>...<EOS> へ戻す
+            sts_list = [encoder_hidden_states[:, 0].unsqueeze(1)]
+            for i in range(1, args.max_token_length, tokenizer.model_max_length):
+              sts_list.append(encoder_hidden_states[:, i:i + tokenizer.model_max_length - 2])
+            sts_list.append(encoder_hidden_states[:, -1].unsqueeze(1))
+            encoder_hidden_states = torch.cat(sts_list, dim=1)
+
+        # Sample noise that we'll add to the latents
+        noise = torch.randn_like(latents, device=latents.device)
+
+        # Sample a random timestep for each image
+        timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (b_size,), device=latents.device)
+        timesteps = timesteps.long()
+
+        # Add noise to the latents according to the noise magnitude at each timestep
+        # (this is the forward diffusion process)
+        noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+
+        # Predict the noise residual
+        noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
+
+        loss = torch.nn.functional.mse_loss(noise_pred.float(), noise.float(), reduction="mean")
+
+        accelerator.backward(loss)
+        if accelerator.sync_gradients:
+          accelerator.clip_grad_norm_(net.parameters(), 1.0)  # args.max_grad_norm)
+
+        optimizer.step()
+        lr_scheduler.step()
+        optimizer.zero_grad(set_to_none=True)
+
+      # Checks if the accelerator has performed an optimization step behind the scenes
+      if accelerator.sync_gradients:
+        progress_bar.update(1)
+        global_step += 1
+
+      current_loss = loss.detach().item() * b_size
+      loss_total += current_loss
+      avr_loss = loss_total / (step+1)
+      logs = {"loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
+      progress_bar.set_postfix(**logs)
+      # accelerator.log(logs, step=global_step)
+
+      if global_step >= args.max_train_steps:
+        break
+
+    accelerator.wait_for_everyone()
+
+    if args.save_every_n_epochs is not None:
+      if (epoch + 1) % args.save_every_n_epochs == 0 and (epoch + 1) < num_train_epochs:
+        print("saving check point.")
+        os.makedirs(args.output_dir, exist_ok=True)
+        ckpt_file = os.path.join(args.output_dir, EPOCH_CHECKPOINT_NAME.format(epoch + 1))
+
+        if fine_tuning:
+          fine_tuning_utils.save_stable_diffusion_checkpoint(
+              ckpt_file, text_encoder, accelerator.unwrap_model(net), args.pretrained_model_name_or_path, epoch + 1, global_step)
+        else:
+          save_hypernetwork(ckpt_file, accelerator.unwrap_model(net))
+
+  is_main_process = accelerator.is_main_process
+  if is_main_process:
+    net = accelerator.unwrap_model(net)
+
+  accelerator.end_training()
+  del accelerator                         # この後メモリを使うのでこれは消す
+
+  if is_main_process:
+    os.makedirs(args.output_dir, exist_ok=True)
+    if fine_tuning:
+      if use_stable_diffusion_format:
+        ckpt_file = os.path.join(args.output_dir, LAST_CHECKPOINT_NAME)
+        print(f"save trained model as StableDiffusion checkpoint to {ckpt_file}")
+        fine_tuning_utils.save_stable_diffusion_checkpoint(
+            ckpt_file, text_encoder, unet, args.pretrained_model_name_or_path, epoch, global_step)
+      else:
+        # Create the pipeline using using the trained modules and save it.
+        print(f"save trained model as Diffusers to {args.output_dir}")
+        pipeline = StableDiffusionPipeline.from_pretrained(
+            args.pretrained_model_name_or_path,
+            unet=unet,
+            text_encoder=text_encoder,
+        )
+        pipeline.save_pretrained(args.output_dir)
+    else:
+      ckpt_file = os.path.join(args.output_dir, LAST_CHECKPOINT_NAME)
+      print(f"save trained model to {ckpt_file}")
+      save_hypernetwork(ckpt_file, net)
+    print("model saved.")
+
+
+# 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
+
+# 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.)
+
+        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.)
+
+        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")
+  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.6.0
+    if type(self.to_out) is torch.nn.Sequential:
+      return self.to_out(out)
+
+    # 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.6.0
+    if type(self.to_out) is torch.nn.Sequential:
+      return self.to_out(out)
+
+    # 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
+
+
+if __name__ == '__main__':
+  # torch.cuda.set_per_process_memory_fraction(0.48)
+  parser = argparse.ArgumentParser()
+  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("--in_json", type=str, default=None, help="metadata file to input / 読みこむメタデータファイル")
+  parser.add_argument("--shuffle_caption", action="store_true",
+                      help="shuffle comma-separated caption when fine tuning / fine tuning時にコンマで区切られたcaptionの各要素をshuffleする")
+  parser.add_argument("--train_data_dir", type=str, default=None, help="directory for train images / 学習画像データのディレクトリ")
+  parser.add_argument("--dataset_repeats", type=int, default=None, help="num times to repeat dataset / 学習にデータセットを繰り返す回数")
+  parser.add_argument("--output_dir", type=str, default=None,
+                      help="directory to output trained model, save as same format as input / 学習後のモデル出力先ディレクトリ（入力と同じ形式で保存）")
+  parser.add_argument("--hypernetwork_module", type=str, default=None,
+                      help='train hypernetwork instead of fine tuning, module to use / fine tuningの代わりにHypernetworkの学習をする場合、そのモジュール')
+  parser.add_argument("--hypernetwork_weights", type=str, default=None,
+                      help='hypernetwork weights to initialize for additional training / Hypernetworkの学習時に読み込む重み（Hypernetworkの追加学習）')
+  parser.add_argument("--save_every_n_epochs", type=int, default=None,
+                      help="save checkpoint every N epochs (only supports in StableDiffusion checkpoint) / 学習中のモデルを指定エポックごとに保存する（StableDiffusion形式のモデルを読み込んだ場合のみ有効）")
+  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("--train_batch_size", type=int, default=1,
+                      help="batch size for training / 学習時のバッチサイズ")
+  parser.add_argument("--use_8bit_adam", action="store_true",
+                      help="use 8bit Adam optimizer (requires bitsandbytes) / 8bit Adamオプティマイザを使う（bitsandbytesのインストールが必要）")
+  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("--learning_rate", type=float, default=2.0e-6, help="learning rate / 学習率")
+  parser.add_argument("--max_train_steps", type=int, default=1600, help="training steps / 学習ステップ数")
+  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("--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("--debug_dataset", action="store_true",
+                      help="show images for debugging (do not train) / デバッグ用に学習データを画面表示する（学習は行わない）")
+
+  args = parser.parse_args()
+  train(args)

diffusers_fine_tuning/fine_tuning_utils.py

@@ -0,0 +1,763 @@
+import math
+import torch
+from transformers import CLIPTextModel
+from diffusers import AutoencoderKL, UNet2DConditionModel
+
+# Tokenizer: checkpointから読み込むのではなくあらかじめ提供されているものを使う
+TOKENIZER_PATH = "openai/clip-vit-large-patch14"
+
+# StableDiffusionのモデルパラメータ
+NUM_TRAIN_TIMESTEPS = 1000
+BETA_START = 0.00085
+BETA_END = 0.0120
+
+UNET_PARAMS_MODEL_CHANNELS = 320
+UNET_PARAMS_CHANNEL_MULT = [1, 2, 4, 4]
+UNET_PARAMS_ATTENTION_RESOLUTIONS = [4, 2, 1]
+UNET_PARAMS_IMAGE_SIZE = 32  # unused
+UNET_PARAMS_IN_CHANNELS = 4
+UNET_PARAMS_OUT_CHANNELS = 4
+UNET_PARAMS_NUM_RES_BLOCKS = 2
+UNET_PARAMS_CONTEXT_DIM = 768
+UNET_PARAMS_NUM_HEADS = 8
+
+VAE_PARAMS_Z_CHANNELS = 4
+VAE_PARAMS_RESOLUTION = 256
+VAE_PARAMS_IN_CHANNELS = 3
+VAE_PARAMS_OUT_CH = 3
+VAE_PARAMS_CH = 128
+VAE_PARAMS_CH_MULT = [1, 2, 4, 4]
+VAE_PARAMS_NUM_RES_BLOCKS = 2
+
+
+# region conversion
+# checkpoint変換など ###############################
+
+# region StableDiffusion->Diffusersの変換コード
+# convert_original_stable_diffusion_to_diffusers をコピーしている（ASL 2.0）
+
+def shave_segments(path, n_shave_prefix_segments=1):
+  """
+  Removes segments. Positive values shave the first segments, negative shave the last segments.
+  """
+  if n_shave_prefix_segments >= 0:
+    return ".".join(path.split(".")[n_shave_prefix_segments:])
+  else:
+    return ".".join(path.split(".")[:n_shave_prefix_segments])
+
+
+def renew_resnet_paths(old_list, n_shave_prefix_segments=0):
+  """
+  Updates paths inside resnets to the new naming scheme (local renaming)
+  """
+  mapping = []
+  for old_item in old_list:
+    new_item = old_item.replace("in_layers.0", "norm1")
+    new_item = new_item.replace("in_layers.2", "conv1")
+
+    new_item = new_item.replace("out_layers.0", "norm2")
+    new_item = new_item.replace("out_layers.3", "conv2")
+
+    new_item = new_item.replace("emb_layers.1", "time_emb_proj")
+    new_item = new_item.replace("skip_connection", "conv_shortcut")
+
+    new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+    mapping.append({"old": old_item, "new": new_item})
+
+  return mapping
+
+
+def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0):
+  """
+  Updates paths inside resnets to the new naming scheme (local renaming)
+  """
+  mapping = []
+  for old_item in old_list:
+    new_item = old_item
+
+    new_item = new_item.replace("nin_shortcut", "conv_shortcut")
+    new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+    mapping.append({"old": old_item, "new": new_item})
+
+  return mapping
+
+
+def renew_attention_paths(old_list, n_shave_prefix_segments=0):
+  """
+  Updates paths inside attentions to the new naming scheme (local renaming)
+  """
+  mapping = []
+  for old_item in old_list:
+    new_item = old_item
+
+    #         new_item = new_item.replace('norm.weight', 'group_norm.weight')
+    #         new_item = new_item.replace('norm.bias', 'group_norm.bias')
+
+    #         new_item = new_item.replace('proj_out.weight', 'proj_attn.weight')
+    #         new_item = new_item.replace('proj_out.bias', 'proj_attn.bias')
+
+    #         new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+    mapping.append({"old": old_item, "new": new_item})
+
+  return mapping
+
+
+def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0):
+  """
+  Updates paths inside attentions to the new naming scheme (local renaming)
+  """
+  mapping = []
+  for old_item in old_list:
+    new_item = old_item
+
+    new_item = new_item.replace("norm.weight", "group_norm.weight")
+    new_item = new_item.replace("norm.bias", "group_norm.bias")
+
+    new_item = new_item.replace("q.weight", "query.weight")
+    new_item = new_item.replace("q.bias", "query.bias")
+
+    new_item = new_item.replace("k.weight", "key.weight")
+    new_item = new_item.replace("k.bias", "key.bias")
+
+    new_item = new_item.replace("v.weight", "value.weight")
+    new_item = new_item.replace("v.bias", "value.bias")
+
+    new_item = new_item.replace("proj_out.weight", "proj_attn.weight")
+    new_item = new_item.replace("proj_out.bias", "proj_attn.bias")
+
+    new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+    mapping.append({"old": old_item, "new": new_item})
+
+  return mapping
+
+
+def assign_to_checkpoint(
+    paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None
+):
+  """
+  This does the final conversion step: take locally converted weights and apply a global renaming
+  to them. It splits attention layers, and takes into account additional replacements
+  that may arise.
+
+  Assigns the weights to the new checkpoint.
+  """
+  assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys."
+
+  # Splits the attention layers into three variables.
+  if attention_paths_to_split is not None:
+    for path, path_map in attention_paths_to_split.items():
+      old_tensor = old_checkpoint[path]
+      channels = old_tensor.shape[0] // 3
+
+      target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1)
+
+      num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3
+
+      old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:])
+      query, key, value = old_tensor.split(channels // num_heads, dim=1)
+
+      checkpoint[path_map["query"]] = query.reshape(target_shape)
+      checkpoint[path_map["key"]] = key.reshape(target_shape)
+      checkpoint[path_map["value"]] = value.reshape(target_shape)
+
+  for path in paths:
+    new_path = path["new"]
+
+    # These have already been assigned
+    if attention_paths_to_split is not None and new_path in attention_paths_to_split:
+      continue
+
+    # Global renaming happens here
+    new_path = new_path.replace("middle_block.0", "mid_block.resnets.0")
+    new_path = new_path.replace("middle_block.1", "mid_block.attentions.0")
+    new_path = new_path.replace("middle_block.2", "mid_block.resnets.1")
+
+    if additional_replacements is not None:
+      for replacement in additional_replacements:
+        new_path = new_path.replace(replacement["old"], replacement["new"])
+
+    # proj_attn.weight has to be converted from conv 1D to linear
+    if "proj_attn.weight" in new_path:
+      checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0]
+    else:
+      checkpoint[new_path] = old_checkpoint[path["old"]]
+
+
+def conv_attn_to_linear(checkpoint):
+  keys = list(checkpoint.keys())
+  attn_keys = ["query.weight", "key.weight", "value.weight"]
+  for key in keys:
+    if ".".join(key.split(".")[-2:]) in attn_keys:
+      if checkpoint[key].ndim > 2:
+        checkpoint[key] = checkpoint[key][:, :, 0, 0]
+    elif "proj_attn.weight" in key:
+      if checkpoint[key].ndim > 2:
+        checkpoint[key] = checkpoint[key][:, :, 0]
+
+
+def convert_ldm_unet_checkpoint(checkpoint, config):
+  """
+  Takes a state dict and a config, and returns a converted checkpoint.
+  """
+
+  # extract state_dict for UNet
+  unet_state_dict = {}
+  unet_key = "model.diffusion_model."
+  keys = list(checkpoint.keys())
+  for key in keys:
+    if key.startswith(unet_key):
+      unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key)
+
+  new_checkpoint = {}
+
+  new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"]
+  new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"]
+  new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"]
+  new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"]
+
+  new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"]
+  new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"]
+
+  new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"]
+  new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"]
+  new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"]
+  new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"]
+
+  # Retrieves the keys for the input blocks only
+  num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer})
+  input_blocks = {
+      layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}" in key]
+      for layer_id in range(num_input_blocks)
+  }
+
+  # Retrieves the keys for the middle blocks only
+  num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer})
+  middle_blocks = {
+      layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key]
+      for layer_id in range(num_middle_blocks)
+  }
+
+  # Retrieves the keys for the output blocks only
+  num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer})
+  output_blocks = {
+      layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}" in key]
+      for layer_id in range(num_output_blocks)
+  }
+
+  for i in range(1, num_input_blocks):
+    block_id = (i - 1) // (config["layers_per_block"] + 1)
+    layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1)
+
+    resnets = [
+        key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key
+    ]
+    attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key]
+
+    if f"input_blocks.{i}.0.op.weight" in unet_state_dict:
+      new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop(
+          f"input_blocks.{i}.0.op.weight"
+      )
+      new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop(
+          f"input_blocks.{i}.0.op.bias"
+      )
+
+    paths = renew_resnet_paths(resnets)
+    meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"}
+    assign_to_checkpoint(
+        paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+    )
+
+    if len(attentions):
+      paths = renew_attention_paths(attentions)
+      meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"}
+      assign_to_checkpoint(
+          paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+      )
+
+  resnet_0 = middle_blocks[0]
+  attentions = middle_blocks[1]
+  resnet_1 = middle_blocks[2]
+
+  resnet_0_paths = renew_resnet_paths(resnet_0)
+  assign_to_checkpoint(resnet_0_paths, new_checkpoint, unet_state_dict, config=config)
+
+  resnet_1_paths = renew_resnet_paths(resnet_1)
+  assign_to_checkpoint(resnet_1_paths, new_checkpoint, unet_state_dict, config=config)
+
+  attentions_paths = renew_attention_paths(attentions)
+  meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"}
+  assign_to_checkpoint(
+      attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+  )
+
+  for i in range(num_output_blocks):
+    block_id = i // (config["layers_per_block"] + 1)
+    layer_in_block_id = i % (config["layers_per_block"] + 1)
+    output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]]
+    output_block_list = {}
+
+    for layer in output_block_layers:
+      layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1)
+      if layer_id in output_block_list:
+        output_block_list[layer_id].append(layer_name)
+      else:
+        output_block_list[layer_id] = [layer_name]
+
+    if len(output_block_list) > 1:
+      resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key]
+      attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key]
+
+      resnet_0_paths = renew_resnet_paths(resnets)
+      paths = renew_resnet_paths(resnets)
+
+      meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"}
+      assign_to_checkpoint(
+          paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+      )
+
+      if ["conv.weight", "conv.bias"] in output_block_list.values():
+        index = list(output_block_list.values()).index(["conv.weight", "conv.bias"])
+        new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[
+            f"output_blocks.{i}.{index}.conv.weight"
+        ]
+        new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[
+            f"output_blocks.{i}.{index}.conv.bias"
+        ]
+
+        # Clear attentions as they have been attributed above.
+        if len(attentions) == 2:
+          attentions = []
+
+      if len(attentions):
+        paths = renew_attention_paths(attentions)
+        meta_path = {
+            "old": f"output_blocks.{i}.1",
+            "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}",
+        }
+        assign_to_checkpoint(
+            paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+        )
+    else:
+      resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1)
+      for path in resnet_0_paths:
+        old_path = ".".join(["output_blocks", str(i), path["old"]])
+        new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]])
+
+        new_checkpoint[new_path] = unet_state_dict[old_path]
+
+  return new_checkpoint
+
+
+def convert_ldm_vae_checkpoint(checkpoint, config):
+  # extract state dict for VAE
+  vae_state_dict = {}
+  vae_key = "first_stage_model."
+  keys = list(checkpoint.keys())
+  for key in keys:
+    if key.startswith(vae_key):
+      vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key)
+
+  new_checkpoint = {}
+
+  new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"]
+  new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"]
+  new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"]
+  new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"]
+  new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"]
+  new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"]
+
+  new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"]
+  new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"]
+  new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"]
+  new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"]
+  new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"]
+  new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"]
+
+  new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"]
+  new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"]
+  new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"]
+  new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"]
+
+  # Retrieves the keys for the encoder down blocks only
+  num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer})
+  down_blocks = {
+      layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks)
+  }
+
+  # Retrieves the keys for the decoder up blocks only
+  num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer})
+  up_blocks = {
+      layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks)
+  }
+
+  for i in range(num_down_blocks):
+    resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key]
+
+    if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict:
+      new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop(
+          f"encoder.down.{i}.downsample.conv.weight"
+      )
+      new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop(
+          f"encoder.down.{i}.downsample.conv.bias"
+      )
+
+    paths = renew_vae_resnet_paths(resnets)
+    meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+  mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key]
+  num_mid_res_blocks = 2
+  for i in range(1, num_mid_res_blocks + 1):
+    resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key]
+
+    paths = renew_vae_resnet_paths(resnets)
+    meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+  mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key]
+  paths = renew_vae_attention_paths(mid_attentions)
+  meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
+  assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+  conv_attn_to_linear(new_checkpoint)
+
+  for i in range(num_up_blocks):
+    block_id = num_up_blocks - 1 - i
+    resnets = [
+        key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key
+    ]
+
+    if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict:
+      new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[
+          f"decoder.up.{block_id}.upsample.conv.weight"
+      ]
+      new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[
+          f"decoder.up.{block_id}.upsample.conv.bias"
+      ]
+
+    paths = renew_vae_resnet_paths(resnets)
+    meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+  mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key]
+  num_mid_res_blocks = 2
+  for i in range(1, num_mid_res_blocks + 1):
+    resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key]
+
+    paths = renew_vae_resnet_paths(resnets)
+    meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+  mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key]
+  paths = renew_vae_attention_paths(mid_attentions)
+  meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
+  assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+  conv_attn_to_linear(new_checkpoint)
+  return new_checkpoint
+
+
+def create_unet_diffusers_config():
+  """
+  Creates a config for the diffusers based on the config of the LDM model.
+  """
+  # unet_params = original_config.model.params.unet_config.params
+
+  block_out_channels = [UNET_PARAMS_MODEL_CHANNELS * mult for mult in UNET_PARAMS_CHANNEL_MULT]
+
+  down_block_types = []
+  resolution = 1
+  for i in range(len(block_out_channels)):
+    block_type = "CrossAttnDownBlock2D" if resolution in UNET_PARAMS_ATTENTION_RESOLUTIONS else "DownBlock2D"
+    down_block_types.append(block_type)
+    if i != len(block_out_channels) - 1:
+      resolution *= 2
+
+  up_block_types = []
+  for i in range(len(block_out_channels)):
+    block_type = "CrossAttnUpBlock2D" if resolution in UNET_PARAMS_ATTENTION_RESOLUTIONS else "UpBlock2D"
+    up_block_types.append(block_type)
+    resolution //= 2
+
+  config = dict(
+      sample_size=UNET_PARAMS_IMAGE_SIZE,
+      in_channels=UNET_PARAMS_IN_CHANNELS,
+      out_channels=UNET_PARAMS_OUT_CHANNELS,
+      down_block_types=tuple(down_block_types),
+      up_block_types=tuple(up_block_types),
+      block_out_channels=tuple(block_out_channels),
+      layers_per_block=UNET_PARAMS_NUM_RES_BLOCKS,
+      cross_attention_dim=UNET_PARAMS_CONTEXT_DIM,
+      attention_head_dim=UNET_PARAMS_NUM_HEADS,
+  )
+
+  return config
+
+
+def create_vae_diffusers_config():
+  """
+  Creates a config for the diffusers based on the config of the LDM model.
+  """
+  # vae_params = original_config.model.params.first_stage_config.params.ddconfig
+  # _ = original_config.model.params.first_stage_config.params.embed_dim
+  block_out_channels = [VAE_PARAMS_CH * mult for mult in VAE_PARAMS_CH_MULT]
+  down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels)
+  up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels)
+
+  config = dict(
+      sample_size=VAE_PARAMS_RESOLUTION,
+      in_channels=VAE_PARAMS_IN_CHANNELS,
+      out_channels=VAE_PARAMS_OUT_CH,
+      down_block_types=tuple(down_block_types),
+      up_block_types=tuple(up_block_types),
+      block_out_channels=tuple(block_out_channels),
+      latent_channels=VAE_PARAMS_Z_CHANNELS,
+      layers_per_block=VAE_PARAMS_NUM_RES_BLOCKS,
+  )
+  return config
+
+
+def convert_ldm_clip_checkpoint(checkpoint):
+  text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+
+  keys = list(checkpoint.keys())
+
+  text_model_dict = {}
+
+  for key in keys:
+    if key.startswith("cond_stage_model.transformer"):
+      text_model_dict[key[len("cond_stage_model.transformer."):]] = checkpoint[key]
+
+  text_model.load_state_dict(text_model_dict)
+
+  return text_model
+
+# endregion
+
+
+# region Diffusers->StableDiffusion の変換コード
+# convert_diffusers_to_original_stable_diffusion をコピーしている（ASL 2.0）
+
+def convert_unet_state_dict(unet_state_dict):
+  unet_conversion_map = [
+      # (stable-diffusion, HF Diffusers)
+      ("time_embed.0.weight", "time_embedding.linear_1.weight"),
+      ("time_embed.0.bias", "time_embedding.linear_1.bias"),
+      ("time_embed.2.weight", "time_embedding.linear_2.weight"),
+      ("time_embed.2.bias", "time_embedding.linear_2.bias"),
+      ("input_blocks.0.0.weight", "conv_in.weight"),
+      ("input_blocks.0.0.bias", "conv_in.bias"),
+      ("out.0.weight", "conv_norm_out.weight"),
+      ("out.0.bias", "conv_norm_out.bias"),
+      ("out.2.weight", "conv_out.weight"),
+      ("out.2.bias", "conv_out.bias"),
+  ]
+
+  unet_conversion_map_resnet = [
+      # (stable-diffusion, HF Diffusers)
+      ("in_layers.0", "norm1"),
+      ("in_layers.2", "conv1"),
+      ("out_layers.0", "norm2"),
+      ("out_layers.3", "conv2"),
+      ("emb_layers.1", "time_emb_proj"),
+      ("skip_connection", "conv_shortcut"),
+  ]
+
+  unet_conversion_map_layer = []
+  for i in range(4):
+      # loop over downblocks/upblocks
+
+    for j in range(2):
+        # loop over resnets/attentions for downblocks
+      hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
+      sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
+      unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
+
+      if i < 3:
+        # no attention layers in down_blocks.3
+        hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
+        sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
+        unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
+
+    for j in range(3):
+      # loop over resnets/attentions for upblocks
+      hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
+      sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
+      unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
+
+      if i > 0:
+        # no attention layers in up_blocks.0
+        hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
+        sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
+        unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
+
+    if i < 3:
+      # no downsample in down_blocks.3
+      hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
+      sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
+      unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
+
+      # no upsample in up_blocks.3
+      hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+      sd_upsample_prefix = f"output_blocks.{3*i + 2}.{1 if i == 0 else 2}."
+      unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
+
+  hf_mid_atn_prefix = "mid_block.attentions.0."
+  sd_mid_atn_prefix = "middle_block.1."
+  unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
+
+  for j in range(2):
+    hf_mid_res_prefix = f"mid_block.resnets.{j}."
+    sd_mid_res_prefix = f"middle_block.{2*j}."
+    unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+  # buyer beware: this is a *brittle* function,
+  # and correct output requires that all of these pieces interact in
+  # the exact order in which I have arranged them.
+  mapping = {k: k for k in unet_state_dict.keys()}
+  for sd_name, hf_name in unet_conversion_map:
+    mapping[hf_name] = sd_name
+  for k, v in mapping.items():
+    if "resnets" in k:
+      for sd_part, hf_part in unet_conversion_map_resnet:
+        v = v.replace(hf_part, sd_part)
+      mapping[k] = v
+  for k, v in mapping.items():
+    for sd_part, hf_part in unet_conversion_map_layer:
+      v = v.replace(hf_part, sd_part)
+    mapping[k] = v
+  new_state_dict = {v: unet_state_dict[k] for k, v in mapping.items()}
+  return new_state_dict
+
+# endregion
+
+
+def load_checkpoint_with_conversion(ckpt_path):
+  # text encoderの格納形式が違うモデルに対応する ('text_model'がない)
+  TEXT_ENCODER_KEY_REPLACEMENTS = [
+      ('cond_stage_model.transformer.embeddings.', 'cond_stage_model.transformer.text_model.embeddings.'),
+      ('cond_stage_model.transformer.encoder.', 'cond_stage_model.transformer.text_model.encoder.'),
+      ('cond_stage_model.transformer.final_layer_norm.', 'cond_stage_model.transformer.text_model.final_layer_norm.')
+  ]
+
+  checkpoint = torch.load(ckpt_path, map_location="cpu")
+  state_dict = checkpoint["state_dict"]
+
+  key_reps = []
+  for rep_from, rep_to in TEXT_ENCODER_KEY_REPLACEMENTS:
+    for key in state_dict.keys():
+      if key.startswith(rep_from):
+        new_key = rep_to + key[len(rep_from):]
+        key_reps.append((key, new_key))
+
+  for key, new_key in key_reps:
+    state_dict[new_key] = state_dict[key]
+    del state_dict[key]
+
+  return checkpoint
+
+
+def load_models_from_stable_diffusion_checkpoint(ckpt_path):
+  checkpoint = load_checkpoint_with_conversion(ckpt_path)
+  state_dict = checkpoint["state_dict"]
+
+  # Convert the UNet2DConditionModel model.
+  unet_config = create_unet_diffusers_config()
+  converted_unet_checkpoint = convert_ldm_unet_checkpoint(state_dict, unet_config)
+
+  unet = UNet2DConditionModel(**unet_config)
+  unet.load_state_dict(converted_unet_checkpoint)
+
+  # Convert the VAE model.
+  vae_config = create_vae_diffusers_config()
+  converted_vae_checkpoint = convert_ldm_vae_checkpoint(state_dict, vae_config)
+
+  vae = AutoencoderKL(**vae_config)
+  vae.load_state_dict(converted_vae_checkpoint)
+
+  # convert text_model
+  text_model = convert_ldm_clip_checkpoint(state_dict)
+
+  return text_model, vae, unet
+
+
+def save_stable_diffusion_checkpoint(output_file, text_encoder, unet, ckpt_path, epochs, steps):
+  # VAEがメモリ上にないので、もう一度VAEを含めて読み込む
+  checkpoint = load_checkpoint_with_conversion(ckpt_path)
+  state_dict = checkpoint["state_dict"]
+
+  # Convert the UNet model
+  unet_state_dict = convert_unet_state_dict(unet.state_dict())
+  for k, v in unet_state_dict.items():
+    key = "model.diffusion_model." + k
+    assert key in state_dict, f"Illegal key in save SD: {key}"
+    state_dict[key] = v
+
+  # Convert the text encoder model
+  text_enc_dict = text_encoder.state_dict()             # 変換不要
+  for k, v in text_enc_dict.items():
+    key = "cond_stage_model.transformer." + k
+    assert key in state_dict, f"Illegal key in save SD: {key}"
+    state_dict[key] = v
+
+  # Put together new checkpoint
+  new_ckpt = {'state_dict': state_dict}
+
+  if 'epoch' in checkpoint:
+    epochs += checkpoint['epoch']
+  if 'global_step' in checkpoint:
+    steps += checkpoint['global_step']
+
+  new_ckpt['epoch'] = epochs
+  new_ckpt['global_step'] = steps
+
+  torch.save(new_ckpt, output_file)
+# endregion
+
+
+def make_bucket_resolutions(max_reso, min_size=256, max_size=1024, divisible=64):
+  max_width, max_height = max_reso
+  max_area = (max_width // divisible) * (max_height // divisible)
+
+  resos = set()
+
+  size = int(math.sqrt(max_area)) * divisible
+  resos.add((size, size))
+
+  size = min_size
+  while size <= max_size:
+    width = size
+    height = min(max_size, (max_area // (width // divisible)) * divisible)
+    resos.add((width, height))
+    resos.add((height, width))
+
+    # # make additional resos
+    # if width >= height and width - divisible >= min_size:
+    #   resos.add((width - divisible, height))
+    #   resos.add((height, width - divisible))
+    # if height >= width and height - divisible >= min_size:
+    #   resos.add((width, height - divisible))
+    #   resos.add((height - divisible, width))
+
+    size += divisible
+
+  resos = list(resos)
+  resos.sort()
+
+  aspect_ratios = [w / h for w, h in resos]
+  return resos, aspect_ratios
+
+
+if __name__ == '__main__':
+  resos, aspect_ratios = make_bucket_resolutions((512, 768))
+  print(len(resos))
+  print(resos)
+  print(aspect_ratios)
+
+  ars = set()
+  for ar in aspect_ratios:
+    if ar in ars:
+      print("error! duplicate ar:", ar)
+    ars.add(ar)

diffusers_fine_tuning/fine_tuning_utils_ber.py

@@ -0,0 +1,771 @@
+import math
+import torch
+from transformers import CLIPTextModel
+from diffusers import AutoencoderKL, UNet2DConditionModel
+
+# Tokenizer: checkpointから読み込むのではなくあらかじめ提供されているものを使う
+TOKENIZER_PATH = "openai/clip-vit-large-patch14"
+
+# StableDiffusionのモデルパラメータ
+NUM_TRAIN_TIMESTEPS = 1000
+BETA_START = 0.00085
+BETA_END = 0.0120
+
+UNET_PARAMS_MODEL_CHANNELS = 320
+UNET_PARAMS_CHANNEL_MULT = [1, 2, 4, 4]
+UNET_PARAMS_ATTENTION_RESOLUTIONS = [4, 2, 1]
+UNET_PARAMS_IMAGE_SIZE = 32  # unused
+UNET_PARAMS_IN_CHANNELS = 4
+UNET_PARAMS_OUT_CHANNELS = 4
+UNET_PARAMS_NUM_RES_BLOCKS = 2
+UNET_PARAMS_CONTEXT_DIM = 768
+UNET_PARAMS_NUM_HEADS = 8
+
+VAE_PARAMS_Z_CHANNELS = 4
+VAE_PARAMS_RESOLUTION = 256
+VAE_PARAMS_IN_CHANNELS = 3
+VAE_PARAMS_OUT_CH = 3
+VAE_PARAMS_CH = 128
+VAE_PARAMS_CH_MULT = [1, 2, 4, 4]
+VAE_PARAMS_NUM_RES_BLOCKS = 2
+
+
+# region conversion
+# checkpoint変換など ###############################
+
+# region StableDiffusion->Diffusersの変換コード
+# convert_original_stable_diffusion_to_diffusers をコピーしている（ASL 2.0）
+
+def shave_segments(path, n_shave_prefix_segments=1):
+  """
+  Removes segments. Positive values shave the first segments, negative shave the last segments.
+  """
+  if n_shave_prefix_segments >= 0:
+    return ".".join(path.split(".")[n_shave_prefix_segments:])
+  else:
+    return ".".join(path.split(".")[:n_shave_prefix_segments])
+
+
+def renew_resnet_paths(old_list, n_shave_prefix_segments=0):
+  """
+  Updates paths inside resnets to the new naming scheme (local renaming)
+  """
+  mapping = []
+  for old_item in old_list:
+    new_item = old_item.replace("in_layers.0", "norm1")
+    new_item = new_item.replace("in_layers.2", "conv1")
+
+    new_item = new_item.replace("out_layers.0", "norm2")
+    new_item = new_item.replace("out_layers.3", "conv2")
+
+    new_item = new_item.replace("emb_layers.1", "time_emb_proj")
+    new_item = new_item.replace("skip_connection", "conv_shortcut")
+
+    new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+    mapping.append({"old": old_item, "new": new_item})
+
+  return mapping
+
+
+def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0):
+  """
+  Updates paths inside resnets to the new naming scheme (local renaming)
+  """
+  mapping = []
+  for old_item in old_list:
+    new_item = old_item
+
+    new_item = new_item.replace("nin_shortcut", "conv_shortcut")
+    new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+    mapping.append({"old": old_item, "new": new_item})
+
+  return mapping
+
+
+def renew_attention_paths(old_list, n_shave_prefix_segments=0):
+  """
+  Updates paths inside attentions to the new naming scheme (local renaming)
+  """
+  mapping = []
+  for old_item in old_list:
+    new_item = old_item
+
+    #         new_item = new_item.replace('norm.weight', 'group_norm.weight')
+    #         new_item = new_item.replace('norm.bias', 'group_norm.bias')
+
+    #         new_item = new_item.replace('proj_out.weight', 'proj_attn.weight')
+    #         new_item = new_item.replace('proj_out.bias', 'proj_attn.bias')
+
+    #         new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+    mapping.append({"old": old_item, "new": new_item})
+
+  return mapping
+
+
+def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0):
+  """
+  Updates paths inside attentions to the new naming scheme (local renaming)
+  """
+  mapping = []
+  for old_item in old_list:
+    new_item = old_item
+
+    new_item = new_item.replace("norm.weight", "group_norm.weight")
+    new_item = new_item.replace("norm.bias", "group_norm.bias")
+
+    new_item = new_item.replace("q.weight", "query.weight")
+    new_item = new_item.replace("q.bias", "query.bias")
+
+    new_item = new_item.replace("k.weight", "key.weight")
+    new_item = new_item.replace("k.bias", "key.bias")
+
+    new_item = new_item.replace("v.weight", "value.weight")
+    new_item = new_item.replace("v.bias", "value.bias")
+
+    new_item = new_item.replace("proj_out.weight", "proj_attn.weight")
+    new_item = new_item.replace("proj_out.bias", "proj_attn.bias")
+
+    new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+    mapping.append({"old": old_item, "new": new_item})
+
+  return mapping
+
+
+def assign_to_checkpoint(
+    paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None
+):
+  """
+  This does the final conversion step: take locally converted weights and apply a global renaming
+  to them. It splits attention layers, and takes into account additional replacements
+  that may arise.
+
+  Assigns the weights to the new checkpoint.
+  """
+  assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys."
+
+  # Splits the attention layers into three variables.
+  if attention_paths_to_split is not None:
+    for path, path_map in attention_paths_to_split.items():
+      old_tensor = old_checkpoint[path]
+      channels = old_tensor.shape[0] // 3
+
+      target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1)
+
+      num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3
+
+      old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:])
+      query, key, value = old_tensor.split(channels // num_heads, dim=1)
+
+      checkpoint[path_map["query"]] = query.reshape(target_shape)
+      checkpoint[path_map["key"]] = key.reshape(target_shape)
+      checkpoint[path_map["value"]] = value.reshape(target_shape)
+
+  for path in paths:
+    new_path = path["new"]
+
+    # These have already been assigned
+    if attention_paths_to_split is not None and new_path in attention_paths_to_split:
+      continue
+
+    # Global renaming happens here
+    new_path = new_path.replace("middle_block.0", "mid_block.resnets.0")
+    new_path = new_path.replace("middle_block.1", "mid_block.attentions.0")
+    new_path = new_path.replace("middle_block.2", "mid_block.resnets.1")
+
+    if additional_replacements is not None:
+      for replacement in additional_replacements:
+        new_path = new_path.replace(replacement["old"], replacement["new"])
+
+    # proj_attn.weight has to be converted from conv 1D to linear
+    if "proj_attn.weight" in new_path:
+      checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0]
+    else:
+      checkpoint[new_path] = old_checkpoint[path["old"]]
+
+
+def conv_attn_to_linear(checkpoint):
+  keys = list(checkpoint.keys())
+  attn_keys = ["query.weight", "key.weight", "value.weight"]
+  for key in keys:
+    if ".".join(key.split(".")[-2:]) in attn_keys:
+      if checkpoint[key].ndim > 2:
+        checkpoint[key] = checkpoint[key][:, :, 0, 0]
+    elif "proj_attn.weight" in key:
+      if checkpoint[key].ndim > 2:
+        checkpoint[key] = checkpoint[key][:, :, 0]
+
+
+def convert_ldm_unet_checkpoint(checkpoint, config):
+  """
+  Takes a state dict and a config, and returns a converted checkpoint.
+  """
+
+  # extract state_dict for UNet
+  unet_state_dict = {}
+  unet_key = "model.diffusion_model."
+  keys = list(checkpoint.keys())
+  for key in keys:
+    if key.startswith(unet_key):
+      unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key)
+
+  new_checkpoint = {}
+
+  new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"]
+  new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"]
+  new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"]
+  new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"]
+
+  new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"]
+  new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"]
+
+  new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"]
+  new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"]
+  new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"]
+  new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"]
+
+  # Retrieves the keys for the input blocks only
+  num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer})
+  input_blocks = {
+      layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}" in key]
+      for layer_id in range(num_input_blocks)
+  }
+
+  # Retrieves the keys for the middle blocks only
+  num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer})
+  middle_blocks = {
+      layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key]
+      for layer_id in range(num_middle_blocks)
+  }
+
+  # Retrieves the keys for the output blocks only
+  num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer})
+  output_blocks = {
+      layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}" in key]
+      for layer_id in range(num_output_blocks)
+  }
+
+  for i in range(1, num_input_blocks):
+    block_id = (i - 1) // (config["layers_per_block"] + 1)
+    layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1)
+
+    resnets = [
+        key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key
+    ]
+    attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key]
+
+    if f"input_blocks.{i}.0.op.weight" in unet_state_dict:
+      new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop(
+          f"input_blocks.{i}.0.op.weight"
+      )
+      new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop(
+          f"input_blocks.{i}.0.op.bias"
+      )
+
+    paths = renew_resnet_paths(resnets)
+    meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"}
+    assign_to_checkpoint(
+        paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+    )
+
+    if len(attentions):
+      paths = renew_attention_paths(attentions)
+      meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"}
+      assign_to_checkpoint(
+          paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+      )
+
+  resnet_0 = middle_blocks[0]
+  attentions = middle_blocks[1]
+  resnet_1 = middle_blocks[2]
+
+  resnet_0_paths = renew_resnet_paths(resnet_0)
+  assign_to_checkpoint(resnet_0_paths, new_checkpoint, unet_state_dict, config=config)
+
+  resnet_1_paths = renew_resnet_paths(resnet_1)
+  assign_to_checkpoint(resnet_1_paths, new_checkpoint, unet_state_dict, config=config)
+
+  attentions_paths = renew_attention_paths(attentions)
+  meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"}
+  assign_to_checkpoint(
+      attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+  )
+
+  for i in range(num_output_blocks):
+    block_id = i // (config["layers_per_block"] + 1)
+    layer_in_block_id = i % (config["layers_per_block"] + 1)
+    output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]]
+    output_block_list = {}
+
+    for layer in output_block_layers:
+      layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1)
+      if layer_id in output_block_list:
+        output_block_list[layer_id].append(layer_name)
+      else:
+        output_block_list[layer_id] = [layer_name]
+
+    if len(output_block_list) > 1:
+      resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key]
+      attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key]
+
+      resnet_0_paths = renew_resnet_paths(resnets)
+      paths = renew_resnet_paths(resnets)
+
+      meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"}
+      assign_to_checkpoint(
+          paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+      )
+
+      if ["conv.weight", "conv.bias"] in output_block_list.values():
+        index = list(output_block_list.values()).index(["conv.weight", "conv.bias"])
+        new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[
+            f"output_blocks.{i}.{index}.conv.weight"
+        ]
+        new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[
+            f"output_blocks.{i}.{index}.conv.bias"
+        ]
+
+        # Clear attentions as they have been attributed above.
+        if len(attentions) == 2:
+          attentions = []
+
+      if len(attentions):
+        paths = renew_attention_paths(attentions)
+        meta_path = {
+            "old": f"output_blocks.{i}.1",
+            "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}",
+        }
+        assign_to_checkpoint(
+            paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+        )
+    else:
+      resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1)
+      for path in resnet_0_paths:
+        old_path = ".".join(["output_blocks", str(i), path["old"]])
+        new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]])
+
+        new_checkpoint[new_path] = unet_state_dict[old_path]
+
+  return new_checkpoint
+
+
+def convert_ldm_vae_checkpoint(checkpoint, config):
+  # extract state dict for VAE
+  vae_state_dict = {}
+  vae_key = "first_stage_model."
+  keys = list(checkpoint.keys())
+  for key in keys:
+    if key.startswith(vae_key):
+      vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key)
+
+  new_checkpoint = {}
+
+  new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"]
+  new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"]
+  new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"]
+  new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"]
+  new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"]
+  new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"]
+
+  new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"]
+  new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"]
+  new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"]
+  new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"]
+  new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"]
+  new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"]
+
+  new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"]
+  new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"]
+  new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"]
+  new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"]
+
+  # Retrieves the keys for the encoder down blocks only
+  num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer})
+  down_blocks = {
+      layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks)
+  }
+
+  # Retrieves the keys for the decoder up blocks only
+  num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer})
+  up_blocks = {
+      layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks)
+  }
+
+  for i in range(num_down_blocks):
+    resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key]
+
+    if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict:
+      new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop(
+          f"encoder.down.{i}.downsample.conv.weight"
+      )
+      new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop(
+          f"encoder.down.{i}.downsample.conv.bias"
+      )
+
+    paths = renew_vae_resnet_paths(resnets)
+    meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+  mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key]
+  num_mid_res_blocks = 2
+  for i in range(1, num_mid_res_blocks + 1):
+    resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key]
+
+    paths = renew_vae_resnet_paths(resnets)
+    meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+  mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key]
+  paths = renew_vae_attention_paths(mid_attentions)
+  meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
+  assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+  conv_attn_to_linear(new_checkpoint)
+
+  for i in range(num_up_blocks):
+    block_id = num_up_blocks - 1 - i
+    resnets = [
+        key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key
+    ]
+
+    if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict:
+      new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[
+          f"decoder.up.{block_id}.upsample.conv.weight"
+      ]
+      new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[
+          f"decoder.up.{block_id}.upsample.conv.bias"
+      ]
+
+    paths = renew_vae_resnet_paths(resnets)
+    meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+  mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key]
+  num_mid_res_blocks = 2
+  for i in range(1, num_mid_res_blocks + 1):
+    resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key]
+
+    paths = renew_vae_resnet_paths(resnets)
+    meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+  mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key]
+  paths = renew_vae_attention_paths(mid_attentions)
+  meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
+  assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+  conv_attn_to_linear(new_checkpoint)
+  return new_checkpoint
+
+
+def create_unet_diffusers_config():
+  """
+  Creates a config for the diffusers based on the config of the LDM model.
+  """
+  # unet_params = original_config.model.params.unet_config.params
+
+  block_out_channels = [UNET_PARAMS_MODEL_CHANNELS * mult for mult in UNET_PARAMS_CHANNEL_MULT]
+
+  down_block_types = []
+  resolution = 1
+  for i in range(len(block_out_channels)):
+    block_type = "CrossAttnDownBlock2D" if resolution in UNET_PARAMS_ATTENTION_RESOLUTIONS else "DownBlock2D"
+    down_block_types.append(block_type)
+    if i != len(block_out_channels) - 1:
+      resolution *= 2
+
+  up_block_types = []
+  for i in range(len(block_out_channels)):
+    block_type = "CrossAttnUpBlock2D" if resolution in UNET_PARAMS_ATTENTION_RESOLUTIONS else "UpBlock2D"
+    up_block_types.append(block_type)
+    resolution //= 2
+
+  config = dict(
+      sample_size=UNET_PARAMS_IMAGE_SIZE,
+      in_channels=UNET_PARAMS_IN_CHANNELS,
+      out_channels=UNET_PARAMS_OUT_CHANNELS,
+      down_block_types=tuple(down_block_types),
+      up_block_types=tuple(up_block_types),
+      block_out_channels=tuple(block_out_channels),
+      layers_per_block=UNET_PARAMS_NUM_RES_BLOCKS,
+      cross_attention_dim=UNET_PARAMS_CONTEXT_DIM,
+      attention_head_dim=UNET_PARAMS_NUM_HEADS,
+  )
+
+  return config
+
+
+def create_vae_diffusers_config():
+  """
+  Creates a config for the diffusers based on the config of the LDM model.
+  """
+  # vae_params = original_config.model.params.first_stage_config.params.ddconfig
+  # _ = original_config.model.params.first_stage_config.params.embed_dim
+  block_out_channels = [VAE_PARAMS_CH * mult for mult in VAE_PARAMS_CH_MULT]
+  down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels)
+  up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels)
+
+  config = dict(
+      sample_size=VAE_PARAMS_RESOLUTION,
+      in_channels=VAE_PARAMS_IN_CHANNELS,
+      out_channels=VAE_PARAMS_OUT_CH,
+      down_block_types=tuple(down_block_types),
+      up_block_types=tuple(up_block_types),
+      block_out_channels=tuple(block_out_channels),
+      latent_channels=VAE_PARAMS_Z_CHANNELS,
+      layers_per_block=VAE_PARAMS_NUM_RES_BLOCKS,
+  )
+  return config
+
+
+def convert_ldm_clip_checkpoint(checkpoint):
+  text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+
+  keys = list(checkpoint.keys())
+
+  text_model_dict = {}
+
+  for key in keys:
+    if key.startswith("cond_stage_model.transformer"):
+      text_model_dict[key[len("cond_stage_model.transformer."):]] = checkpoint[key]
+
+  text_model.load_state_dict(text_model_dict)
+
+  return text_model
+
+# endregion
+
+
+# region Diffusers->StableDiffusion の変換コード
+# convert_diffusers_to_original_stable_diffusion をコピーしている（ASL 2.0）
+
+def convert_unet_state_dict(unet_state_dict):
+  unet_conversion_map = [
+      # (stable-diffusion, HF Diffusers)
+      ("time_embed.0.weight", "time_embedding.linear_1.weight"),
+      ("time_embed.0.bias", "time_embedding.linear_1.bias"),
+      ("time_embed.2.weight", "time_embedding.linear_2.weight"),
+      ("time_embed.2.bias", "time_embedding.linear_2.bias"),
+      ("input_blocks.0.0.weight", "conv_in.weight"),
+      ("input_blocks.0.0.bias", "conv_in.bias"),
+      ("out.0.weight", "conv_norm_out.weight"),
+      ("out.0.bias", "conv_norm_out.bias"),
+      ("out.2.weight", "conv_out.weight"),
+      ("out.2.bias", "conv_out.bias"),
+  ]
+
+  unet_conversion_map_resnet = [
+      # (stable-diffusion, HF Diffusers)
+      ("in_layers.0", "norm1"),
+      ("in_layers.2", "conv1"),
+      ("out_layers.0", "norm2"),
+      ("out_layers.3", "conv2"),
+      ("emb_layers.1", "time_emb_proj"),
+      ("skip_connection", "conv_shortcut"),
+  ]
+
+  unet_conversion_map_layer = []
+  for i in range(4):
+      # loop over downblocks/upblocks
+
+    for j in range(2):
+        # loop over resnets/attentions for downblocks
+      hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
+      sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
+      unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
+
+      if i < 3:
+        # no attention layers in down_blocks.3
+        hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
+        sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
+        unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
+
+    for j in range(3):
+      # loop over resnets/attentions for upblocks
+      hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
+      sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
+      unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
+
+      if i > 0:
+        # no attention layers in up_blocks.0
+        hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
+        sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
+        unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
+
+    if i < 3:
+      # no downsample in down_blocks.3
+      hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
+      sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
+      unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
+
+      # no upsample in up_blocks.3
+      hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+      sd_upsample_prefix = f"output_blocks.{3*i + 2}.{1 if i == 0 else 2}."
+      unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
+
+  hf_mid_atn_prefix = "mid_block.attentions.0."
+  sd_mid_atn_prefix = "middle_block.1."
+  unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
+
+  for j in range(2):
+    hf_mid_res_prefix = f"mid_block.resnets.{j}."
+    sd_mid_res_prefix = f"middle_block.{2*j}."
+    unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+  # buyer beware: this is a *brittle* function,
+  # and correct output requires that all of these pieces interact in
+  # the exact order in which I have arranged them.
+  mapping = {k: k for k in unet_state_dict.keys()}
+  for sd_name, hf_name in unet_conversion_map:
+    mapping[hf_name] = sd_name
+  for k, v in mapping.items():
+    if "resnets" in k:
+      for sd_part, hf_part in unet_conversion_map_resnet:
+        v = v.replace(hf_part, sd_part)
+      mapping[k] = v
+  for k, v in mapping.items():
+    for sd_part, hf_part in unet_conversion_map_layer:
+      v = v.replace(hf_part, sd_part)
+    mapping[k] = v
+  new_state_dict = {v: unet_state_dict[k] for k, v in mapping.items()}
+  return new_state_dict
+
+# endregion
+
+
+def load_checkpoint_with_conversion(ckpt_path):
+  # text encoderの格納形式が違うモデルに対応する ('text_model'がない)
+  TEXT_ENCODER_KEY_REPLACEMENTS = [
+      ('cond_stage_model.transformer.embeddings.', 'cond_stage_model.transformer.text_model.embeddings.'),
+      ('cond_stage_model.transformer.encoder.', 'cond_stage_model.transformer.text_model.encoder.'),
+      ('cond_stage_model.transformer.final_layer_norm.', 'cond_stage_model.transformer.text_model.final_layer_norm.')
+  ]
+
+  checkpoint = torch.load(ckpt_path, map_location="cpu")
+  state_dict = checkpoint["state_dict"]
+
+  key_reps = []
+  for rep_from, rep_to in TEXT_ENCODER_KEY_REPLACEMENTS:
+    for key in state_dict.keys():
+      if key.startswith(rep_from):
+        new_key = rep_to + key[len(rep_from):]
+        key_reps.append((key, new_key))
+
+  for key, new_key in key_reps:
+    state_dict[new_key] = state_dict[key]
+    del state_dict[key]
+
+  return checkpoint
+
+
+def load_models_from_stable_diffusion_checkpoint(ckpt_path):
+  checkpoint = load_checkpoint_with_conversion(ckpt_path)
+  state_dict = checkpoint["state_dict"]
+
+  # Convert the UNet2DConditionModel model.
+  unet_config = create_unet_diffusers_config()
+  converted_unet_checkpoint = convert_ldm_unet_checkpoint(state_dict, unet_config)
+
+  unet = UNet2DConditionModel(**unet_config)
+  unet.load_state_dict(converted_unet_checkpoint)
+
+  # Convert the VAE model.
+  vae_config = create_vae_diffusers_config()
+  converted_vae_checkpoint = convert_ldm_vae_checkpoint(state_dict, vae_config)
+
+  vae = AutoencoderKL(**vae_config)
+  vae.load_state_dict(converted_vae_checkpoint)
+
+  # convert text_model
+  text_model = convert_ldm_clip_checkpoint(state_dict)
+
+  return text_model, vae, unet
+
+
+def save_stable_diffusion_checkpoint(output_file, text_encoder, unet, ckpt_path, epochs, steps):
+  # VAEがメモリ上にないので、もう一度VAEを含めて読み込む
+  checkpoint = load_checkpoint_with_conversion(ckpt_path)
+  state_dict = checkpoint["state_dict"]
+
+  # Convert the UNet model
+  unet_state_dict = convert_unet_state_dict(unet.state_dict())
+  for k, v in unet_state_dict.items():
+    key = "model.diffusion_model." + k
+    assert key in state_dict, f"Illegal key in save SD: {key}"
+    # if args.save_half:
+    #   state_dict[key] = v.half() # save to fp16
+    # else:
+    #   state_dict[key] = v
+    state_dict[key] = v.half() # save to fp16
+
+  # Convert the text encoder model
+  text_enc_dict = text_encoder.state_dict()             # 変換不要
+  for k, v in text_enc_dict.items():
+    key = "cond_stage_model.transformer." + k
+    assert key in state_dict, f"Illegal key in save SD: {key}"
+    # if args.save_half:
+    #   state_dict[key] = v.half() # save to fp16
+    # else:
+    #   state_dict[key] = v
+    state_dict[key] = v.half() # save to fp16
+
+  # Put together new checkpoint
+  new_ckpt = {'state_dict': state_dict}
+
+  if 'epoch' in checkpoint:
+    epochs += checkpoint['epoch']
+  if 'global_step' in checkpoint:
+    steps += checkpoint['global_step']
+
+  new_ckpt['epoch'] = epochs
+  new_ckpt['global_step'] = steps
+
+  torch.save(new_ckpt, output_file)
+# endregion
+
+
+def make_bucket_resolutions(max_reso, min_size=256, max_size=1024, divisible=64):
+  max_width, max_height = max_reso
+  max_area = (max_width // divisible) * (max_height // divisible)
+
+  resos = set()
+
+  size = int(math.sqrt(max_area)) * divisible
+  resos.add((size, size))
+
+  size = min_size
+  while size <= max_size:
+    width = size
+    height = min(max_size, (max_area // (width // divisible)) * divisible)
+    resos.add((width, height))
+    resos.add((height, width))
+
+    # # make additional resos
+    # if width >= height and width - divisible >= min_size:
+    #   resos.add((width - divisible, height))
+    #   resos.add((height, width - divisible))
+    # if height >= width and height - divisible >= min_size:
+    #   resos.add((width, height - divisible))
+    #   resos.add((height - divisible, width))
+
+    size += divisible
+
+  resos = list(resos)
+  resos.sort()
+
+  aspect_ratios = [w / h for w, h in resos]
+  return resos, aspect_ratios
+
+
+if __name__ == '__main__':
+  resos, aspect_ratios = make_bucket_resolutions((512, 768))
+  print(len(resos))
+  print(resos)
+  print(aspect_ratios)
+
+  ars = set()
+  for ar in aspect_ratios:
+    if ar in ars:
+      print("error! duplicate ar:", ar)
+    ars.add(ar)

diffusers_fine_tuning/hypernetwork_nai.py

@@ -0,0 +1,96 @@
+# NAI compatible
+
+import torch
+
+
+class HypernetworkModule(torch.nn.Module):
+  def __init__(self, dim, multiplier=1.0):
+    super().__init__()
+
+    linear1 = torch.nn.Linear(dim, dim * 2)
+    linear2 = torch.nn.Linear(dim * 2, dim)
+    linear1.weight.data.normal_(mean=0.0, std=0.01)
+    linear1.bias.data.zero_()
+    linear2.weight.data.normal_(mean=0.0, std=0.01)
+    linear2.bias.data.zero_()
+    linears = [linear1, linear2]
+
+    self.linear = torch.nn.Sequential(*linears)
+    self.multiplier = multiplier
+
+  def forward(self, x):
+    return x + self.linear(x) * self.multiplier
+
+
+class Hypernetwork(torch.nn.Module):
+  enable_sizes = [320, 640, 768, 1280]
+  # return self.modules[Hypernetwork.enable_sizes.index(size)]
+
+  def __init__(self, multiplier=1.0) -> None:
+    super().__init__()
+    self.modules = []
+    for size in Hypernetwork.enable_sizes:
+      self.modules.append((HypernetworkModule(size, multiplier), HypernetworkModule(size, multiplier)))
+      self.register_module(f"{size}_0", self.modules[-1][0])
+      self.register_module(f"{size}_1", self.modules[-1][1])
+
+  def apply_to_stable_diffusion(self, text_encoder, vae, unet):
+    blocks = unet.input_blocks + [unet.middle_block] + unet.output_blocks
+    for block in blocks:
+      for subblk in block:
+        if 'SpatialTransformer' in str(type(subblk)):
+          for tf_block in subblk.transformer_blocks:
+            for attn in [tf_block.attn1, tf_block.attn2]:
+              size = attn.context_dim
+              if size in Hypernetwork.enable_sizes:
+                attn.hypernetwork = self
+              else:
+                attn.hypernetwork = None
+
+  def apply_to_diffusers(self, text_encoder, vae, unet):
+    blocks = unet.down_blocks + [unet.mid_block] + unet.up_blocks
+    for block in blocks:
+      if hasattr(block, 'attentions'):
+        for subblk in block.attentions:
+          if 'SpatialTransformer' in str(type(subblk)) or 'Transformer2DModel' in str(type(subblk)):      # 0.6.0 and 0.7~
+            for tf_block in subblk.transformer_blocks:
+              for attn in [tf_block.attn1, tf_block.attn2]:
+                size = attn.to_k.in_features
+                if size in Hypernetwork.enable_sizes:
+                  attn.hypernetwork = self
+                else:
+                  attn.hypernetwork = None
+    return True       # TODO error checking
+
+  def forward(self, x, context):
+    size = context.shape[-1]
+    assert size in Hypernetwork.enable_sizes
+    module = self.modules[Hypernetwork.enable_sizes.index(size)]
+    return module[0].forward(context), module[1].forward(context)
+
+  def load_from_state_dict(self, state_dict):
+    # old ver to new ver
+    changes = {
+        'linear1.bias': 'linear.0.bias',
+        'linear1.weight': 'linear.0.weight',
+        'linear2.bias': 'linear.1.bias',
+        'linear2.weight': 'linear.1.weight',
+    }
+    for key_from, key_to in changes.items():
+      if key_from in state_dict:
+        state_dict[key_to] = state_dict[key_from]
+        del state_dict[key_from]
+
+    for size, sd in state_dict.items():
+      if type(size) == int:
+        self.modules[Hypernetwork.enable_sizes.index(size)][0].load_state_dict(sd[0], strict=True)
+        self.modules[Hypernetwork.enable_sizes.index(size)][1].load_state_dict(sd[1], strict=True)
+    return True
+
+  def get_state_dict(self):
+    state_dict = {}
+    for i, size in enumerate(Hypernetwork.enable_sizes):
+      sd0 = self.modules[i][0].state_dict()
+      sd1 = self.modules[i][1].state_dict()
+      state_dict[size] = [sd0, sd1]
+    return state_dict

diffusers_fine_tuning/make_captions.py

@@ -0,0 +1,90 @@
+# このスクリプトのライセンスは、Apache License 2.0とします
+# (c) 2022 Kohya S. @kohya_ss
+
+import argparse
+import glob
+import os
+import json
+
+from PIL import Image
+from tqdm import tqdm
+import numpy as np
+import torch
+from torchvision import transforms
+from torchvision.transforms.functional import InterpolationMode
+from models.blip import blip_decoder
+# from Salesforce_BLIP.models.blip import blip_decoder
+
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def main(args):
+  image_paths = glob.glob(os.path.join(args.train_data_dir, "*.jpg")) + glob.glob(os.path.join(args.train_data_dir, "*.png"))
+  print(f"found {len(image_paths)} images.")
+
+  print(f"loading BLIP caption: {args.caption_weights}")
+  image_size = 384
+  model = blip_decoder(pretrained=args.caption_weights, image_size=image_size, vit='large')
+  model.eval()
+  model = model.to(DEVICE)
+  print("BLIP loaded")
+
+  # 正方形でいいのか？　という気がするがソースがそうなので
+  transform = transforms.Compose([
+      transforms.Resize((image_size, image_size), interpolation=InterpolationMode.BICUBIC),
+      transforms.ToTensor(),
+      transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
+  ])
+
+  # captioningする
+  def run_batch(path_imgs):
+    imgs = torch.stack([im for _, im in path_imgs]).to(DEVICE)
+
+    with torch.no_grad():
+      if args.beam_search:
+        captions = model.generate(imgs, sample=False, num_beams=args.num_beams,
+                                  max_length=args.max_length, min_length=args.min_length)
+      else:
+        captions = model.generate(imgs, sample=True, top_p=args.top_p, max_length=args.max_length, min_length=args.min_length)
+
+    for (image_path, _), caption in zip(path_imgs, captions):
+      with open(os.path.splitext(image_path)[0] + args.caption_extention, "wt", encoding='utf-8') as f:
+        f.write(caption + "\n")
+        if args.debug:
+          print(image_path, caption)
+
+  b_imgs = []
+  for image_path in tqdm(image_paths):
+    raw_image = Image.open(image_path)
+    if raw_image.mode != "RGB":
+      print(f"convert image mode {raw_image.mode} to RGB: {image_path}")
+      raw_image = raw_image.convert("RGB")
+
+    image = transform(raw_image)
+    b_imgs.append((image_path, image))
+    if len(b_imgs) >= args.batch_size:
+      run_batch(b_imgs)
+      b_imgs.clear()
+  if len(b_imgs) > 0:
+    run_batch(b_imgs)
+
+  print("done!")
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("train_data_dir", type=str, help="directory for train images / 学習画像データのディレクトリ")
+  parser.add_argument("caption_weights", type=str,
+                      help="BLIP caption weights (model_large_caption.pth) / BLIP captionの重みファイル(model_large_caption.pth)")
+  parser.add_argument("--caption_extention", type=str, default=".caption", help="extention of caption file / 出力されるキャプションファイルの拡張子")
+  parser.add_argument("--beam_search", action="store_true",
+                      help="use beam search (default Nucleus sampling) / beam searchを使う（このオプション未指定時はNucleus sampling）")
+  parser.add_argument("--batch_size", type=int, default=1, help="batch size in inference / 推論時のバッチサイズ")
+  parser.add_argument("--num_beams", type=int, default=1, help="num of beams in beam search /beam search時のビーム数（多いと精度が上がるが時間がかかる）")
+  parser.add_argument("--top_p", type=float, default=0.9, help="top_p in Nucleus sampling / Nucleus sampling時のtop_p")
+  parser.add_argument("--max_length", type=int, default=75, help="max length of caption / captionの最大長")
+  parser.add_argument("--min_length", type=int, default=5, help="min length of caption / captionの最小長")
+  parser.add_argument("--debug", action="store_true", help="debug mode")
+
+  args = parser.parse_args()
+  main(args)

diffusers_fine_tuning/merge_captions_to_metadata.py

@@ -0,0 +1,61 @@
+# このスクリプトのライセンスは、Apache License 2.0とします
+# (c) 2022 Kohya S. @kohya_ss
+
+import argparse
+import glob
+import os
+import json
+
+from tqdm import tqdm
+
+
+def main(args):
+  image_paths = glob.glob(os.path.join(args.train_data_dir, "*.jpg")) + glob.glob(os.path.join(args.train_data_dir, "*.png"))
+  print(f"found {len(image_paths)} images.")
+
+  if args.in_json is not None:
+    print(f"loading existing metadata: {args.in_json}")
+    with open(args.in_json, "rt", encoding='utf-8') as f:
+      metadata = json.load(f)
+    print("captions for existing images will be overwritten / 既存の画像のキャプションは上書きされます")
+  else:
+    print("new metadata will be created / 新しいメタデータファイルが作成されます")
+    metadata = {}
+
+  print("merge caption texts to metadata json.")
+  for image_path in tqdm(image_paths):
+    caption_path = os.path.splitext(image_path)[0] + args.caption_extention
+    with open(caption_path, "rt", encoding='utf-8') as f:
+      caption = f.readlines()[0].strip()
+
+    image_key = os.path.splitext(os.path.basename(image_path))[0]
+    if image_key not in metadata:
+      # if args.verify_caption:
+      #   print(f"image not in metadata / メタデータに画像がありません: {image_path}")
+      #   return
+      metadata[image_key] = {}
+    # elif args.verify_caption and 'caption' not in metadata[image_key]:
+    #   print(f"no caption in metadata / メタデータにcaptionがありません: {image_path}")
+    #   return
+
+    metadata[image_key]['caption'] = caption
+    if args.debug:
+      print(image_key, caption)
+
+  # metadataを書き出して終わり
+  print(f"writing metadata: {args.out_json}")
+  with open(args.out_json, "wt", encoding='utf-8') as f:
+    json.dump(metadata, f, indent=2)
+  print("done!")
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("train_data_dir", type=str, help="directory for train images / 学習画像データのディレクトリ")
+  parser.add_argument("out_json", type=str, help="metadata file to output / メタデータファイル書き出し先")
+  parser.add_argument("--in_json", type=str, help="metadata file to input / 読み込むメタデータファイル")
+  parser.add_argument("--caption_extention", type=str, default=".caption", help="extention of caption file / 読み込むキャプションファイルの拡張子")
+  parser.add_argument("--debug", action="store_true", help="debug mode")
+
+  args = parser.parse_args()
+  main(args)

diffusers_fine_tuning/merge_dd_tags_to_metadata.py

@@ -0,0 +1,61 @@
+# このスクリプトのライセンスは、Apache License 2.0とします
+# (c) 2022 Kohya S. @kohya_ss
+
+import argparse
+import glob
+import os
+import json
+
+from tqdm import tqdm
+
+
+def main(args):
+  image_paths = glob.glob(os.path.join(args.train_data_dir, "*.jpg")) + glob.glob(os.path.join(args.train_data_dir, "*.png"))
+  print(f"found {len(image_paths)} images.")
+
+  if args.in_json is not None:
+    print(f"loading existing metadata: {args.in_json}")
+    with open(args.in_json, "rt", encoding='utf-8') as f:
+      metadata = json.load(f)
+    print("tags data for existing images will be overwritten / 既存の画像のタグは上書きされます")
+  else:
+    print("new metadata will be created / 新しいメタデータファイルが作成されます")
+    metadata = {}
+
+  print("merge tags to metadata json.")
+  for image_path in tqdm(image_paths):
+    tags_path = os.path.splitext(image_path)[0] + '.txt'
+    with open(tags_path, "rt", encoding='utf-8') as f:
+      tags = f.readlines()[0].strip()
+
+    image_key = os.path.splitext(os.path.basename(image_path))[0]
+    if image_key not in metadata:
+      # if args.verify_caption:
+      #   print(f"image not in metadata / メタデータに画像がありません: {image_path}")
+      #   return
+      metadata[image_key] = {}
+    # elif args.verify_caption and 'caption' not in metadata[image_key]:
+    #   print(f"no caption in metadata / メタデータにcaptionがありません: {image_path}")
+    #   return
+
+    metadata[image_key]['tags'] = tags
+    if args.debug:
+      print(image_key, tags)
+
+  # metadataを書き出して終わり
+  print(f"writing metadata: {args.out_json}")
+  with open(args.out_json, "wt", encoding='utf-8') as f:
+    json.dump(metadata, f, indent=2)
+  print("done!")
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("train_data_dir", type=str, help="directory for train images / 学習画像データのディレクトリ")
+  parser.add_argument("out_json", type=str, help="metadata file to output / メタデータファイル書き出し先")
+  parser.add_argument("--in_json", type=str, help="metadata file to input / 読み込むメタデータファイル")
+  # parser.add_argument("--verify_caption", action="store_true", help="verify caption exists / メタデータにすでにcaptionが存在することを確認する")
+  parser.add_argument("--debug", action="store_true", help="debug mode")
+
+  args = parser.parse_args()
+  main(args)

diffusers_fine_tuning/prepare_buckets_latents.py

@@ -0,0 +1,172 @@
+# このスクリプトのライセンスは、Apache License 2.0とします
+# (c) 2022 Kohya S. @kohya_ss
+
+import argparse
+import glob
+import os
+import json
+
+from tqdm import tqdm
+import numpy as np
+from diffusers import AutoencoderKL
+from PIL import Image
+import cv2
+import torch
+from torchvision import transforms
+
+import fine_tuning_utils
+
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+IMAGE_TRANSFORMS = transforms.Compose(
+    [
+        transforms.ToTensor(),
+        transforms.Normalize([0.5], [0.5]),
+    ]
+)
+
+
+def get_latents(vae, images, weight_dtype):
+  img_tensors = [IMAGE_TRANSFORMS(image) for image in images]
+  img_tensors = torch.stack(img_tensors)
+  img_tensors = img_tensors.to(DEVICE, weight_dtype)
+  with torch.no_grad():
+    latents = vae.encode(img_tensors).latent_dist.sample().float().to("cpu").numpy()
+  return latents
+
+
+def main(args):
+  image_paths = glob.glob(os.path.join(args.train_data_dir, "*.jpg")) + glob.glob(os.path.join(args.train_data_dir, "*.png"))
+  print(f"found {len(image_paths)} images.")
+
+  if os.path.exists(args.in_json):
+    print(f"loading existing metadata: {args.in_json}")
+    with open(args.in_json, "rt", encoding='utf-8') as f:
+      metadata = json.load(f)
+  else:
+    print(f"no metadata / メタデータファイルがありません: {args.in_json}")
+    return
+
+  # モデル形式のオプション設定を確認する
+  use_stable_diffusion_format = os.path.isfile(args.model_name_or_path)
+  if not use_stable_diffusion_format:
+    assert os.path.exists(args.model_name_or_path), f"no model / モデルがありません : {args.model_name_or_path}"
+
+  # モデルを読み込む
+  if use_stable_diffusion_format:
+    print("load StableDiffusion checkpoint")
+    _, vae, _ = fine_tuning_utils.load_models_from_stable_diffusion_checkpoint(args.model_name_or_path)
+  else:
+    print("load Diffusers pretrained models")
+    vae = AutoencoderKL.from_pretrained(args.model_name_or_path, subfolder="vae")
+
+  weight_dtype = torch.float32
+  if args.mixed_precision == "fp16":
+    weight_dtype = torch.float16
+  elif args.mixed_precision == "bf16":
+    weight_dtype = torch.bfloat16
+
+  vae.eval()
+  vae.to(DEVICE, dtype=weight_dtype)
+
+  # bucketのサイズを計算する
+  max_reso = tuple([int(t) for t in args.max_resolution.split(',')])
+  assert len(max_reso) == 2, f"illegal resolution (not 'width,height') / 画像サイズに誤りがあります。'幅,高さ'で指定してください: {args.max_resolution}"
+
+  bucket_resos, bucket_aspect_ratios = fine_tuning_utils.make_bucket_resolutions(max_reso)
+
+  # 画像をひとつずつ適切なbucketに割り当てながらlatentを計算する
+  bucket_aspect_ratios = np.array(bucket_aspect_ratios)
+  buckets_imgs = [[] for _ in range(len(bucket_resos))]
+  bucket_counts = [0 for _ in range(len(bucket_resos))]
+  img_ar_errors = []
+  for i, image_path in enumerate(tqdm(image_paths)):
+    image_key = os.path.splitext(os.path.basename(image_path))[0]
+    if image_key not in metadata:
+      metadata[image_key] = {}
+
+    image = Image.open(image_path)
+    if image.mode != 'RGB':
+      image = image.convert("RGB")
+
+    aspect_ratio = image.width / image.height
+    ar_errors = bucket_aspect_ratios - aspect_ratio
+    bucket_id = np.abs(ar_errors).argmin()
+    reso = bucket_resos[bucket_id]
+    ar_error = ar_errors[bucket_id]
+    img_ar_errors.append(abs(ar_error))
+
+    # どのサイズにリサイズするか→トリミングする方向で
+    if ar_error <= 0:                   # 横が長い→縦を合わせる
+      scale = reso[1] / image.height
+    else:
+      scale = reso[0] / image.width
+
+    resized_size = (int(image.width * scale + .5), int(image.height * scale + .5))
+
+    # print(image.width, image.height, bucket_id, bucket_resos[bucket_id], ar_errors[bucket_id], resized_size,
+    #       bucket_resos[bucket_id][0] - resized_size[0], bucket_resos[bucket_id][1] - resized_size[1])
+
+    assert resized_size[0] == reso[0] or resized_size[1] == reso[
+        1], f"internal error, resized size not match: {reso}, {resized_size}, {image.width}, {image.height}"
+    assert resized_size[0] >= reso[0] and resized_size[1] >= reso[
+        1], f"internal error, resized size too small: {reso}, {resized_size}, {image.width}, {image.height}"
+
+    # 画像をリサイズしてトリミングする
+    # PILにinter_areaがないのでcv2で……
+    image = np.array(image)
+    image = cv2.resize(image, resized_size, interpolation=cv2.INTER_AREA)
+    if resized_size[0] > reso[0]:
+      trim_size = resized_size[0] - reso[0]
+      image = image[:, trim_size//2:trim_size//2 + reso[0]]
+    elif resized_size[1] > reso[1]:
+      trim_size = resized_size[1] - reso[1]
+      image = image[trim_size//2:trim_size//2 + reso[1]]
+    assert image.shape[0] == reso[1] and image.shape[1] == reso[0], f"internal error, illegal trimmed size: {image.shape}, {reso}"
+
+    # # debug
+    # cv2.imwrite(f"r:\\test\\img_{i:05d}.jpg", image[:, :, ::-1])
+
+    # バッチへ追加
+    buckets_imgs[bucket_id].append((image_key, reso, image))
+    bucket_counts[bucket_id] += 1
+    metadata[image_key]['train_resolution'] = reso
+
+    # バッチを推論するか判定して推論する
+    is_last = i == len(image_paths) - 1
+    for j in range(len(buckets_imgs)):
+      bucket = buckets_imgs[j]
+      if (is_last and len(bucket) > 0) or len(bucket) >= args.batch_size:
+        latents = get_latents(vae, [img for _, _, img in bucket], weight_dtype)
+
+        for (image_key, reso, _), latent in zip(bucket, latents):
+          np.savez(os.path.join(args.train_data_dir, image_key), latent)
+
+        bucket.clear()
+
+  for i, (reso, count) in enumerate(zip(bucket_resos, bucket_counts)):
+    print(f"bucket {i} {reso}: {count}")
+  img_ar_errors = np.array(img_ar_errors)
+  print(f"mean ar error: {np.mean(img_ar_errors)}")
+
+  # metadataを書き出して終わり
+  print(f"writing metadata: {args.out_json}")
+  with open(args.out_json, "wt", encoding='utf-8') as f:
+    json.dump(metadata, f, indent=2)
+  print("done!")
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("train_data_dir", type=str, help="directory for train images / 学習画像データのディレクトリ")
+  parser.add_argument("in_json", type=str, help="metadata file to input / 読み込むメタデータファイル")
+  parser.add_argument("out_json", type=str, help="metadata file to output / メタデータファイル書き出し先")
+  parser.add_argument("model_name_or_path", type=str, help="model name or path to encode latents / latentを取得するためのモデル")
+  parser.add_argument("--batch_size", type=int, default=1, help="batch size in inference / 推論時のバッチサイズ")
+  parser.add_argument("--max_resolution", type=str, default="512,512",
+                      help="max resolution in fine tuning (width,height) / fine tuning時の最大画像サイズ 「幅,高さ」（使用メモリ量に関係します）")
+  parser.add_argument("--mixed_precision", type=str, default="no",
+                      choices=["no", "fp16", "bf16"], help="use mixed precision / 混合精度を使う場合、その精度")
+
+  args = parser.parse_args()
+  main(args)

diffusers_fine_tuning/requirements.txt

@@ -0,0 +1,6 @@
+transformers>=4.21.0
+ftfy
+albumentations
+opencv-python
+einops
+pytorch_lightning

train_db_fixed-ber.py

@@ -3,6 +3,7 @@
 
 # v7: another text encoder ckpt format, average loss, save epochs/global steps, show num of train/reg images, 
 #     enable reg images in fine-tuning, add dataset_repeats option
+# v8: supports Diffusers 0.7.2
 
 from torch.autograd.function import Function
 import argparse
@@ -1555,7 +1556,14 @@ def replace_unet_cross_attn_to_xformers():
     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)
-    return self.to_out(out)
+    # diffusers 0.6.0
+    if type(self.to_out) is torch.nn.Sequential:
+      return self.to_out(out)
+
+    # 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

train_db_fixed.py

@@ -1,8 +1,9 @@
 # このスクリプトのライセンスは、train_dreambooth.pyと同じくApache License 2.0とします
 # (c) 2022 Kohya S. @kohya_ss
 
-# v7: another text encoder ckpt format, average loss, save epochs/global steps, show num of train/reg images, 
+# v7: another text encoder ckpt format, average loss, save epochs/global steps, show num of train/reg images,
 #     enable reg images in fine-tuning, add dataset_repeats option
+# v8: supports Diffusers 0.7.2
 
 from torch.autograd.function import Function
 import argparse
@@ -1522,7 +1523,15 @@ def replace_unet_cross_attn_to_memory_efficient():
     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)')
-    return self.to_out(out)
+
+    # diffusers 0.6.0
+    if type(self.to_out) is torch.nn.Sequential:
+      return self.to_out(out)
+
+    # 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
 
@@ -1549,7 +1558,15 @@ def replace_unet_cross_attn_to_xformers():
     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)
-    return self.to_out(out)
+
+    # diffusers 0.6.0
+    if type(self.to_out) is torch.nn.Sequential:
+      return self.to_out(out)
+
+    # 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