Upload 54 files

.github/workflows/typos.yml

@@ -0,0 +1,21 @@
+---
+# yamllint disable rule:line-length
+name: Typos
+
+on:  # yamllint disable-line rule:truthy
+  push:
+  pull_request:
+    types:
+      - opened
+      - synchronize
+      - reopened
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v3
+
+      - name: typos-action
+        uses: crate-ci/[email protected]

.gitignore

@@ -0,0 +1,7 @@
+logs
+__pycache__
+wd14_tagger_model
+venv
+*.egg-info
+build
+.vscode

LICENSE.md

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [2022] [kohya-ss]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

_typos.toml

@@ -0,0 +1,15 @@
+# Files for typos
+# Instruction:  https://github.com/marketplace/actions/typos-action#getting-started
+
+[default.extend-identifiers]
+
+[default.extend-words]
+NIN="NIN"
+parms="parms"
+nin="nin"
+extention="extention" # Intentionally left
+nd="nd"
+
+
+[files]
+extend-exclude = ["_typos.toml"]

adastand.py

@@ -0,0 +1,291 @@
+import torch
+import math
+def __version__():
+    return 0.5
+#######################################################################################
+#NTT提案のAdastand Optimizer NTTを信用できるならAdamより少し性能高い(2019)
+#参考コード：https://github.com/bunag-public/adastand_pack/
+#似た計算式のAdaBeliefがAdamWと同じweight_decayの計算を導入していたのでAdamWのweight_decay式を使えるように
+class Adastand(torch.optim.Optimizer):
+    """Implements Adastand algorithm.
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        weight_decouple (bool, optional): if True is weight decay as in AdamW
+    """
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
+                 weight_decay=0, weight_decouple=False, fixed_decay=False, amsgrad=False):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        defaults = dict(lr=lr, betas=betas, eps=eps,
+                        weight_decay=weight_decay, weight_decouple=weight_decouple, fixed_decay=fixed_decay, amsgrad=amsgrad)
+        super(Adastand, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(Adastand, self).__setstate__(state)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
+                weight_decouple = group['weight_decouple']
+                fixed_decay = group['fixed_decay']
+                amsgrad = group['amsgrad']
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p.data, memory_format=torch.preserve_format)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p.data, memory_format=torch.preserve_format)
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of
+                        # sq. grad. values
+                        state['exp_avg_sqs'] = torch.zeros_like(
+                            p.data, memory_format=torch.preserve_format
+                        )
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                beta1, beta2 = group['betas']
+
+                state['step'] += 1
+
+                if weight_decouple:
+                    if not fixed_decay:
+                        p.data.mul_(1.0 - group['lr'] * group['weight_decay'])
+                    else:
+                        p.data.mul_(1.0 - group['weight_decay'])
+                else:
+                    if group['weight_decay'] != 0:
+                        grad.add_(p.data, alpha=group['weight_decay'])
+
+                # Decay the first and second moment running average coefficient
+                grad_residual = grad - exp_avg
+                exp_avg_sq.mul_(beta2).addcmul_(grad_residual, grad_residual, value=beta2 * (1 - beta2))
+                exp_avg.mul_(2 * beta1 - 1).add_(grad, alpha=1 - beta1)
+
+                bias_correction1 = 1 - beta1 ** state['step']
+                bias_correction2 = 1 - beta2 ** state['step']
+                if amsgrad:
+                    exp_avg_sqs = state['exp_avg_sqs']
+                    torch.max(exp_avg_sqs, exp_avg_sq, out=exp_avg_sqs)
+                    denom = exp_avg_sqs.sqrt().add_(group['eps']/ math.sqrt(bias_correction2))
+                else:
+                    denom = exp_avg_sq.sqrt().add_(group['eps'])
+                
+                step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1
+
+                p.data.addcdiv_(exp_avg, denom, value=-step_size)
+                #p.data.addcdiv_(-step_size, exp_avg, denom)
+
+        return loss
+#######################################################################################
+class Adastand_b(torch.optim.Optimizer):
+    """Implements Adastand algorithm.
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        weight_decouple (bool, optional): if True is weight decay as in AdamW
+    """
+
+    def __init__(
+        self,
+        params,
+        lr: float = 1e-3,
+        betas = (0.9, 0.999),
+        eps: float = 1e-8,
+        weight_decay: float = 0,
+        amsgrad: bool = False,
+        weight_decouple: bool = False,
+        fixed_decay: bool = False,
+        rectify: bool = False,
+    ) -> None:
+        if lr <= 0.0:
+            raise ValueError('Invalid learning rate: {}'.format(lr))
+        if eps < 0.0:
+            raise ValueError('Invalid epsilon value: {}'.format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(
+                'Invalid beta parameter at index 0: {}'.format(betas[0])
+            )
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(
+                'Invalid beta parameter at index 1: {}'.format(betas[1])
+            )
+        if weight_decay < 0:
+            raise ValueError(
+                'Invalid weight_decay value: {}'.format(weight_decay)
+            )
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            amsgrad=amsgrad,
+        )
+        super(Adastand_b, self).__init__(params, defaults)
+
+        self._weight_decouple = weight_decouple
+        self._rectify = rectify
+        self._fixed_decay = fixed_decay
+
+    def __setstate__(self, state):
+        super(Adastand_b, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('amsgrad', False)
+
+    def step(self, closure=None):
+        r"""Performs a single optimization step.
+
+        Arguments:
+            closure: A closure that reevaluates the model and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError(
+                        'AdaBelief does not support sparse gradients, '
+                        'please consider SparseAdam instead'
+                    )
+                amsgrad = group['amsgrad']
+
+                state = self.state[p]
+
+                beta1, beta2 = group['betas']
+
+                # State initialization
+                if len(state) == 0:
+                    state['rho_inf'] = 2.0 / (1.0 - beta2) - 1.0
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(
+                        p.data, memory_format=torch.preserve_format
+                    )
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_var'] = torch.zeros_like(
+                        p.data, memory_format=torch.preserve_format
+                    )
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of
+                        # sq. grad. values
+                        state['max_exp_avg_var'] = torch.zeros_like(
+                            p.data, memory_format=torch.preserve_format
+                        )
+
+                # get current state variable
+                exp_avg, exp_avg_var = state['exp_avg'], state['exp_avg_var']
+
+                state['step'] += 1
+                bias_correction1 = 1 - beta1 ** state['step']
+                bias_correction2 = 1 - beta2 ** state['step']
+
+                # perform weight decay, check if decoupled weight decay
+                if self._weight_decouple:
+                    if not self._fixed_decay:
+                        p.data.mul_(1.0 - group['lr'] * group['weight_decay'])
+                    else:
+                        p.data.mul_(1.0 - group['weight_decay'])
+                else:
+                    if group['weight_decay'] != 0:
+                        grad.add_(p.data, alpha=group['weight_decay'])
+
+                # Update first and second moment running average
+                exp_avg.mul_(2*beta1-1).add_(grad, alpha=1 - beta1)
+                grad_residual = grad - exp_avg
+                exp_avg_var.mul_(beta2).addcmul_(
+                    grad_residual, grad_residual, value=beta2*(1 - beta2)
+                )
+
+                if amsgrad:
+                    max_exp_avg_var = state['max_exp_avg_var']
+                    # Maintains the maximum of all 2nd moment running
+                    # avg. till now
+                    torch.max(
+                        max_exp_avg_var, exp_avg_var, out=max_exp_avg_var
+                    )
+
+                    # Use the max. for normalizing running avg. of gradient
+                    denom = (
+                        max_exp_avg_var.add_(group['eps']).sqrt()
+                        / math.sqrt(bias_correction2)
+                    ).add_(group['eps'])
+                else:
+                    denom = (
+                        exp_avg_var.add_(group['eps']).sqrt()
+                        / math.sqrt(bias_correction2)
+                    ).add_(group['eps'])
+
+                if not self._rectify:
+                    # Default update
+                    step_size = group['lr']* math.sqrt(bias_correction2) / bias_correction1
+                    p.data.addcdiv_(exp_avg, denom, value=-step_size)
+
+                else:  # Rectified update
+                    # calculate rho_t
+                    state['rho_t'] = state['rho_inf'] - 2 * state[
+                        'step'
+                    ] * beta2 ** state['step'] / (1.0 - beta2 ** state['step'])
+
+                    if (
+                        state['rho_t'] > 4
+                    ):  # perform Adam style update if variance is small
+                        rho_inf, rho_t = state['rho_inf'], state['rho_t']
+                        rt = (
+                            (rho_t - 4.0)
+                            * (rho_t - 2.0)
+                            * rho_inf
+                            / (rho_inf - 4.0)
+                            / (rho_inf - 2.0)
+                            / rho_t
+                        )
+                        rt = math.sqrt(rt)
+
+                        step_size = rt * group['lr'] / bias_correction1
+
+                        p.data.addcdiv_(-step_size, exp_avg, denom)
+
+                    else:  # perform SGD style update
+                        p.data.add_(-group['lr'], exp_avg)
+
+        return loss

append_module.py

@@ -0,0 +1,504 @@
+import argparse
+import json
+import shutil
+import time
+from typing import Dict, List, NamedTuple, Tuple
+from accelerate import Accelerator
+from torch.autograd.function import Function
+import glob
+import math
+import os
+import random
+import hashlib
+from io import BytesIO
+
+from tqdm import tqdm
+import torch
+from torchvision import transforms
+from transformers import CLIPTokenizer
+import diffusers
+from diffusers import DDPMScheduler, StableDiffusionPipeline
+import albumentations as albu
+import numpy as np
+from PIL import Image
+import cv2
+from einops import rearrange
+from torch import einsum
+import safetensors.torch
+
+import library.model_util as model_util
+import library.train_util as train_util
+
+#============================================================================================================
+#AdafactorScheduleに暫定的にinitial_lrを層別に適用できるようにしたもの
+#============================================================================================================
+from torch.optim.lr_scheduler import LambdaLR
+class AdafactorSchedule_append(LambdaLR):
+    """
+    Since [`~optimization.Adafactor`] performs its own scheduling, if the training loop relies on a scheduler (e.g.,
+    for logging), this class creates a proxy object that retrieves the current lr values from the optimizer.
+
+    It returns `initial_lr` during startup and the actual `lr` during stepping.
+    """
+
+    def __init__(self, optimizer, initial_lr=0.0):
+        def lr_lambda(_):
+            return initial_lr
+
+        for group in optimizer.param_groups:
+            if not type(initial_lr)==list:
+                group["initial_lr"] = initial_lr
+            else:
+                group["initial_lr"] = initial_lr.pop(0)
+        super().__init__(optimizer, lr_lambda)
+        for group in optimizer.param_groups:
+            del group["initial_lr"]
+
+    def get_lr(self):
+        opt = self.optimizer
+        lrs = [
+            opt._get_lr(group, opt.state[group["params"][0]])
+            for group in opt.param_groups
+            if group["params"][0].grad is not None
+        ]
+        if len(lrs) == 0:
+            lrs = self.base_lrs  # if called before stepping
+        return lrs
+
+#============================================================================================================
+#model_util 内より
+#============================================================================================================
+def make_bucket_resolutions_fix(max_reso, min_reso, min_size=256, max_size=1024, divisible=64, step=1):
+  max_width, max_height = max_reso
+  max_area = (max_width // divisible) * (max_height // divisible)
+
+  min_widht, min_height = min_reso
+  min_area = (min_widht // divisible) * (min_height // divisible)
+  
+  area_size_list = []
+  area_size_resos_list = []
+  _max_area = max_area
+
+  while True:
+    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]
+    area_size_list.append(_max_area)
+    area_size_resos_list.append(resos)
+    #area_size_ratio_list.append(aspect_ratios)
+
+    _max_area -= step
+    if _max_area < min_area:
+      break
+  return area_size_resos_list, area_size_list
+
+#============================================================================================================
+#train_util 内より
+#============================================================================================================
+class BucketManager_append(train_util.BucketManager):
+  def __init__(self, no_upscale, max_reso, min_size, max_size, reso_steps, min_reso=None, area_step=None) -> None:
+    super().__init__(no_upscale, max_reso, min_size, max_size, reso_steps)
+    print("BucketManager_appendを作成しました")
+    if min_reso is None:
+      self.min_reso = None
+      self.min_area = None
+    else:
+      self.min_reso = min_reso
+      self.min_area = min_reso[0] * min_reso[1]
+    self.area_step = area_step
+    self.area_sizes_flag = False
+
+  def make_buckets(self):
+    if self.min_reso:
+      print(f"make_resolution append")
+      resos, area_sizes = make_bucket_resolutions_fix(self.max_reso, self.min_reso, self.min_size, self.max_size, self.reso_steps, self.area_step)
+      self.set_predefined_resos(resos, area_sizes)
+    else:
+      resos = model_util.make_bucket_resolutions(self.max_reso, self.min_size, self.max_size, self.reso_steps)
+      self.set_predefined_resos(resos)
+
+  def set_predefined_resos(self, resos, area_sizes=None):
+    # 規定サイズから選ぶ場合の解像度、aspect ratioの情報を格納しておく
+    if area_sizes:
+      self.area_sizes_flag = True
+      self.predefined_area_sizes = np.array(area_sizes.copy())
+      self.predefined_resos_list = resos.copy()
+      self.predefined_resos_set_list = [set(reso) for reso in resos]
+      self.predefined_aspect_ratios_list = [np.array([w/h for w,h in reso]) for reso in resos]
+      self.predefined_resos = None
+      self.predefined_resos_set = None
+      self.predefined_aspect_ratios = None
+    else:
+      self.area_sizes_flag = False
+      self.predefined_area_sizes = None
+      self.predefined_resos = resos.copy()
+      self.predefined_resos_set = set(resos)
+      self.predefined_aspect_ratios = np.array([w / h for w, h in resos])
+
+  def select_bucket(self, image_width, image_height):
+    # 画像サイズを算出する
+    area_size = (image_width//64) * (image_height//64)
+    aspect_ratio = image_width / image_height
+    bucket_size_id = None
+    # 拡張したバケットサイズを使うために画像サイズのエリアを決定する
+    if self.area_sizes_flag:
+        size_errors = self.predefined_area_sizes - area_size
+        bucket_size_id = np.abs(size_errors).argmin()
+        #一定の範囲を探索して使用する画像サイズを確定する
+        serch_size_range = 1
+        bucket_size_id_list = [bucket_size_id]
+        for i in range(serch_size_range):
+          if bucket_size_id - i <0:
+            bucket_size_id_list.append(bucket_size_id + i + 1)
+          elif bucket_size_id + 1 + i >= len(self.predefined_resos_list):
+            bucket_size_id_list.append(bucket_size_id - i - 1)
+          else:
+            bucket_size_id_list.append(bucket_size_id - i - 1)
+            bucket_size_id_list.append(bucket_size_id + i + 1)
+        _min_error = 1000.
+        _min_id = bucket_size_id
+        for now_size_id in bucket_size_id:
+          self.predefined_aspect_ratios = self.predefined_aspect_ratios_list[now_size_id]
+          ar_errors = self.predefined_aspect_ratios - aspect_ratio
+          ar_error = np.abs(ar_errors).min()
+          if _min_error > ar_error:
+            _min_error = ar_error
+            _min_id = now_size_id
+          if _min_error == 0.:
+            break
+        bucket_size_id = _min_id
+        del _min_error, _min_id, ar_error #余計なものは掃除
+        self.predefined_resos = self.predefined_resos_list[bucket_size_id]
+        self.predefined_resos_set = self.predefined_resos_set_list[bucket_size_id]
+        self.predefined_aspect_ratios = self.predefined_aspect_ratios_list[bucket_size_id]
+    # --ここから処理はそのまま
+    if not self.no_upscale:
+      # 同じaspect ratioがあるかもしれないので（fine tuningで、no_upscale=Trueで前処理した場合）、解像度が同じものを優先する
+      reso = (image_width, image_height)
+      if reso in self.predefined_resos_set:
+        pass
+      else:
+        ar_errors = self.predefined_aspect_ratios - aspect_ratio
+        predefined_bucket_id = np.abs(ar_errors).argmin()          # 当該解像度以外でaspect ratio errorが最も少ないもの
+        reso = self.predefined_resos[predefined_bucket_id]
+
+      ar_reso = reso[0] / reso[1]
+      if aspect_ratio > ar_reso:                   # 横が長い→縦を合わせる
+        scale = reso[1] / image_height
+      else:
+        scale = reso[0] / image_width
+
+      resized_size = (int(image_width * scale + .5), int(image_height * scale + .5))
+      # print("use predef", image_width, image_height, reso, resized_size)
+    else:
+      if image_width * image_height > self.max_area:
+        # 画像が大きすぎるのでアスペクト比を保ったまま縮小することを前提にbucketを決める
+        resized_width = math.sqrt(self.max_area * aspect_ratio)
+        resized_height = self.max_area / resized_width
+        assert abs(resized_width / resized_height - aspect_ratio) < 1e-2, "aspect is illegal"
+
+        # リサイズ後の短辺または長辺をreso_steps単位にする：aspect ratioの差が少ないほうを選ぶ
+        # 元のbucketingと同じロジック
+        b_width_rounded = self.round_to_steps(resized_width)
+        b_height_in_wr = self.round_to_steps(b_width_rounded / aspect_ratio)
+        ar_width_rounded = b_width_rounded / b_height_in_wr
+
+        b_height_rounded = self.round_to_steps(resized_height)
+        b_width_in_hr = self.round_to_steps(b_height_rounded * aspect_ratio)
+        ar_height_rounded = b_width_in_hr / b_height_rounded
+
+        # print(b_width_rounded, b_height_in_wr, ar_width_rounded)
+        # print(b_width_in_hr, b_height_rounded, ar_height_rounded)
+
+        if abs(ar_width_rounded - aspect_ratio) < abs(ar_height_rounded - aspect_ratio):
+          resized_size = (b_width_rounded, int(b_width_rounded / aspect_ratio + .5))
+        else:
+          resized_size = (int(b_height_rounded * aspect_ratio + .5), b_height_rounded)
+        # print(resized_size)
+      else:
+        resized_size = (image_width, image_height)              # リサイズは不要
+
+      # 画像のサイズ未満をbucketのサイズとする（paddingせずにcroppingする）
+      bucket_width = resized_size[0] - resized_size[0] % self.reso_steps
+      bucket_height = resized_size[1] - resized_size[1] % self.reso_steps
+      # print("use arbitrary", image_width, image_height, resized_size, bucket_width, bucket_height)
+
+      reso = (bucket_width, bucket_height)
+
+    self.add_if_new_reso(reso)
+
+    ar_error = (reso[0] / reso[1]) - aspect_ratio
+    return reso, resized_size, ar_error
+
+class DreamBoothDataset(train_util.DreamBoothDataset):
+  def __init__(self, batch_size, train_data_dir, reg_data_dir, tokenizer, max_token_length, caption_extension, shuffle_caption, shuffle_keep_tokens, resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, prior_loss_weight, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset, min_resolution=None, area_step=None) -> None:
+    print("use append DreamBoothDataset")
+    self.min_resolution = min_resolution
+    self.area_step = area_step
+    super().__init__(batch_size, train_data_dir, reg_data_dir, tokenizer, max_token_length, caption_extension, shuffle_caption, shuffle_keep_tokens,
+                      resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, prior_loss_weight,
+                      flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset)
+  def make_buckets(self):
+    '''
+    bucketingを行わない場合も呼び出し必須（ひとつだけbucketを作る）
+    min_size and max_size are ignored when enable_bucket is False
+    '''
+    print("loading image sizes.")
+    for info in tqdm(self.image_data.values()):
+      if info.image_size is None:
+        info.image_size = self.get_image_size(info.absolute_path)
+
+    if self.enable_bucket:
+      print("make buckets")
+    else:
+      print("prepare dataset")
+
+    # bucketを作成し、画像をbucketに振り分ける
+    if self.enable_bucket:
+      if self.bucket_manager is None:                         # fine tuningの場合でmetadataに定義がある場合は、すでに初期化済み
+        #======================================================================change
+        if self.min_resolution:
+          self.bucket_manager = BucketManager_append(self.bucket_no_upscale, (self.width, self.height),
+                                              self.min_bucket_reso, self.max_bucket_reso, self.bucket_reso_steps, self.min_resolution, self.area_step)
+        else:
+          self.bucket_manager = train_util.BucketManager(self.bucket_no_upscale, (self.width, self.height),
+                                              self.min_bucket_reso, self.max_bucket_reso, self.bucket_reso_steps)
+        #======================================================================change
+        if not self.bucket_no_upscale:
+          self.bucket_manager.make_buckets()
+        else:
+          print("min_bucket_reso and max_bucket_reso are ignored if bucket_no_upscale is set, because bucket reso is defined by image size automatically / bucket_no_upscaleが指定された場合は、bucketの解像度は画像サイズから自動計算されるため、min_bucket_resoとmax_bucket_resoは無視されます")
+
+      img_ar_errors = []
+      for image_info in self.image_data.values():
+        image_width, image_height = image_info.image_size
+        image_info.bucket_reso, image_info.resized_size, ar_error = self.bucket_manager.select_bucket(image_width, image_height)
+
+        # print(image_info.image_key, image_info.bucket_reso)
+        img_ar_errors.append(abs(ar_error))
+
+      self.bucket_manager.sort()
+    else:
+      self.bucket_manager = train_util.BucketManager(False, (self.width, self.height), None, None, None)
+      self.bucket_manager.set_predefined_resos([(self.width, self.height)])  # ひとつの固定サイズbucketのみ
+      for image_info in self.image_data.values():
+        image_width, image_height = image_info.image_size
+        image_info.bucket_reso, image_info.resized_size, _ = self.bucket_manager.select_bucket(image_width, image_height)
+
+    for image_info in self.image_data.values():
+      for _ in range(image_info.num_repeats):
+        self.bucket_manager.add_image(image_info.bucket_reso, image_info.image_key)
+
+    # bucket情報を表示、格納する
+    if self.enable_bucket:
+      self.bucket_info = {"buckets": {}}
+      print("number of images (including repeats) / 各bucketの画像枚数（繰り返し回数を含む）")
+      for i, (reso, bucket) in enumerate(zip(self.bucket_manager.resos, self.bucket_manager.buckets)):
+        count = len(bucket)
+        if count > 0:
+          self.bucket_info["buckets"][i] = {"resolution": reso, "count": len(bucket)}
+          print(f"bucket {i}: resolution {reso}, count: {len(bucket)}")
+
+      img_ar_errors = np.array(img_ar_errors)
+      mean_img_ar_error = np.mean(np.abs(img_ar_errors))
+      self.bucket_info["mean_img_ar_error"] = mean_img_ar_error
+      print(f"mean ar error (without repeats): {mean_img_ar_error}")
+
+    # データ参照用indexを作る。このindexはdatasetのshuffleに用いられる
+    self.buckets_indices: List(train_util.BucketBatchIndex) = []
+    for bucket_index, bucket in enumerate(self.bucket_manager.buckets):
+      batch_count = int(math.ceil(len(bucket) / self.batch_size))
+      for batch_index in range(batch_count):
+        self.buckets_indices.append(train_util.BucketBatchIndex(bucket_index, self.batch_size, batch_index))
+
+      # ↓以下はbucketごとのbatch件数があまりにも増えて混乱を招くので元に戻す
+      # 　学習時はステップ数がランダムなので、同一画像が同一batch内にあってもそれほど悪影響はないであろう、と考えられる
+      #
+      # # bucketが細分化されることにより、ひとつのbucketに一種類の画像のみというケースが増え、つまりそれは
+      # # ひとつのbatchが同じ画像で占められることになるので、さすがに良くないであろう
+      # # そのためバッチサイズを画像種類までに制限する
+      # # ただそれでも同一画像が同一バッチに含まれる可能性はあるので、繰り返し回数が少ないほうがshuffleの品質は良くなることは間違いない？
+      # # TO DO 正則化画像をepochまたがりで利用する仕組み
+      # num_of_image_types = len(set(bucket))
+      # bucket_batch_size = min(self.batch_size, num_of_image_types)
+      # batch_count = int(math.ceil(len(bucket) / bucket_batch_size))
+      # # print(bucket_index, num_of_image_types, bucket_batch_size, batch_count)
+      # for batch_index in range(batch_count):
+      #   self.buckets_indices.append(BucketBatchIndex(bucket_index, bucket_batch_size, batch_index))
+      # ↑ここまで
+
+    self.shuffle_buckets()
+    self._length = len(self.buckets_indices)
+
+class FineTuningDataset(train_util.FineTuningDataset):
+  def __init__(self, json_file_name, batch_size, train_data_dir, tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens, resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, flip_aug, color_aug, face_crop_aug_range, random_crop, dataset_repeats, debug_dataset) -> None:
+    train_util.glob_images = glob_images
+    super().__init__( json_file_name, batch_size, train_data_dir, tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens,
+                      resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, flip_aug, color_aug, face_crop_aug_range,
+                      random_crop, dataset_repeats, debug_dataset)
+
+def glob_images(directory, base="*", npz_flag=True):
+  img_paths = []
+  dots = []
+  for ext in train_util.IMAGE_EXTENSIONS:
+    dots.append(ext)
+  if npz_flag:
+    dots.append(".npz")
+  for ext in dots:
+    if base == '*':
+      img_paths.extend(glob.glob(os.path.join(glob.escape(directory), base + ext)))
+    else:
+      img_paths.extend(glob.glob(glob.escape(os.path.join(directory, base + ext))))
+  return img_paths
+
+#============================================================================================================
+#networks.lora
+#============================================================================================================
+from networks.lora import LoRANetwork
+def replace_prepare_optimizer_params(networks):
+  def prepare_optimizer_params(self, text_encoder_lr, unet_lr, scheduler_lr=None, loranames=None):
+    
+    def enumerate_params(loras, lora_name=None):
+      params = []
+      for lora in loras:
+        if lora_name is not None:
+          if lora_name in lora.lora_name:
+            params.extend(lora.parameters())
+        else:
+          params.extend(lora.parameters())
+      return params
+
+    self.requires_grad_(True)
+    all_params = []
+    ret_scheduler_lr = []
+    
+    if loranames is not None:
+      textencoder_names = [None]
+      unet_names = [None]
+      if "text_encoder" in loranames:
+        textencoder_names = loranames["text_encoder"]
+      if "unet" in loranames:
+        unet_names = loranames["unet"]
+    
+    if self.text_encoder_loras:
+      for textencoder_name in textencoder_names:
+        param_data = {'params': enumerate_params(self.text_encoder_loras, lora_name=textencoder_name)}
+        if text_encoder_lr is not None:
+          param_data['lr'] = text_encoder_lr
+        if scheduler_lr is not None:
+          ret_scheduler_lr.append(scheduler_lr[0])
+        all_params.append(param_data)
+
+    if self.unet_loras:
+      for unet_name in unet_names:
+        param_data = {'params': enumerate_params(self.unet_loras, lora_name=unet_name)}
+        if unet_lr is not None:
+          param_data['lr'] = unet_lr
+        if scheduler_lr is not None:
+          ret_scheduler_lr.append(scheduler_lr[1])
+        all_params.append(param_data)
+
+    return all_params, ret_scheduler_lr
+  
+  LoRANetwork.prepare_optimizer_params = prepare_optimizer_params
+
+#============================================================================================================
+#新規追加
+#============================================================================================================
+def add_append_arguments(parser: argparse.ArgumentParser):
+  # for train_network_opt.py
+  parser.add_argument("--optimizer", type=str, default="AdamW", choices=["AdamW", "RAdam", "AdaBound", "AdaBelief", "AggMo", "AdamP", "Adastand", "Adastand_belief", "Apollo", "Lamb", "Ranger", "RangerVA", "Lookahead_Adam", "Lookahead_DiffGrad", "Yogi", "NovoGrad", "QHAdam", "DiffGrad", "MADGRAD", "Adafactor"], help="使用するoptimizerを指定する")
+  parser.add_argument("--optimizer_arg", type=str, default=None, nargs='*')
+  parser.add_argument("--split_lora_networks", action="store_true")
+  parser.add_argument("--split_lora_level", type=int, default=0, help="どれくらい細分化するかの設定 0がunetのみを層別に 1がunetを大枠で分割 2がtextencoder含めて層別")
+  parser.add_argument("--min_resolution", type=str, default=None)
+  parser.add_argument("--area_step", type=int, default=1)
+  parser.add_argument("--config", type=str, default=None)
+
+def create_split_names(split_flag, split_level):
+  split_names = None
+  if split_flag:
+    split_names = {}
+    text_encoder_names = [None]
+    unet_names = ["lora_unet_mid_block"]
+    if split_level==1:
+      unet_names.append(f"lora_unet_down_blocks_")
+      unet_names.append(f"lora_unet_up_blocks_")
+    elif split_level==2 or split_level==0:
+      if split_level==2:
+        text_encoder_names = []
+        for i in range(12):
+          text_encoder_names.append(f"lora_te_text_model_encoder_layers_{i}_")
+      for i in range(3):
+        unet_names.append(f"lora_unet_down_blocks_{i}")
+        unet_names.append(f"lora_unet_up_blocks_{i+1}")
+    split_names["text_encoder"] = text_encoder_names
+    split_names["unet"] = unet_names
+  return split_names
+
+def get_config(parser):
+  args = parser.parse_args()
+  if args.config is not None and (not args.config==""):
+    import yaml
+    import datetime
+    if os.path.splitext(args.config)[-1] == ".yaml":
+      args.config = os.path.splitext(args.config)[0]
+    config_path = f"./{args.config}.yaml"
+    if os.path.exists(config_path):
+      print(f"{config_path} から設定を読み込み中...")
+      margs, rest = parser.parse_known_args()
+      with open(config_path, mode="r") as f:
+          configs = yaml.unsafe_load(f)
+      #変数でのやり取りをするためargparserからDict型を取り出す
+      args_dic = vars(args)
+      #デフォから引数指定で変更があるものを確認
+      change_def_dic = {}
+      args_type_dic = {}
+      for key, v in args_dic.items():
+        if not parser.get_default(key) == v:
+          change_def_dic[key] = v
+      #デフォ指定されてるデータ型を取得する
+      for key, act in parser._option_string_actions.items():
+        if key=="-h": continue
+        key = key[2:]
+        args_type_dic[key] = act.type
+      #データタイプの確認とargsにkeyの内容を代入していく
+      for key, v in configs.items():
+        if key in args_dic:
+          if args_dic[key] is not None:
+            new_type = type(args_dic[key])
+            if (not type(v) == new_type) and (not new_type==list):
+              v = new_type(v)
+          else:
+            if v is not None:
+              if not type(v) == args_type_dic[key]:
+                v = args_type_dic[key](v)
+          args_dic[key] = v
+      #最後にデフォから指定が変わってるものを変更する
+      for key, v in change_def_dic.items():
+        args_dic[key] = v
+    else:
+      print(f"{config_path} が見つかりませんでした")
+  return args

build/lib/library/model_util.py

@@ -0,0 +1,1180 @@
+# v1: split from train_db_fixed.py.
+# v2: support safetensors
+
+import math
+import os
+import torch
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig
+from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNet2DConditionModel
+from safetensors.torch import load_file, save_file
+
+# DiffUsers版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 = 64  # fixed from old invalid value `32`
+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
+
+# V2
+V2_UNET_PARAMS_ATTENTION_HEAD_DIM = [5, 10, 20, 20]
+V2_UNET_PARAMS_CONTEXT_DIM = 1024
+
+# Diffusersの設定を読み込むための参照モデル
+DIFFUSERS_REF_MODEL_ID_V1 = "runwayml/stable-diffusion-v1-5"
+DIFFUSERS_REF_MODEL_ID_V2 = "stabilityai/stable-diffusion-2-1"
+
+
+# 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 linear_transformer_to_conv(checkpoint):
+  keys = list(checkpoint.keys())
+  tf_keys = ["proj_in.weight", "proj_out.weight"]
+  for key in keys:
+    if ".".join(key.split(".")[-2:]) in tf_keys:
+      if checkpoint[key].ndim == 2:
+        checkpoint[key] = checkpoint[key].unsqueeze(2).unsqueeze(2)
+
+
+def convert_ldm_unet_checkpoint(v2, 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"])
+
+      # biasとweightの順番に依存しないようにする：もっといいやり方がありそうだが
+      for l in output_block_list.values():
+        l.sort()
+
+      if ["conv.bias", "conv.weight"] in output_block_list.values():
+        index = list(output_block_list.values()).index(["conv.bias", "conv.weight"])
+        new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[
+            f"output_blocks.{i}.{index}.conv.bias"
+        ]
+        new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[
+            f"output_blocks.{i}.{index}.conv.weight"
+        ]
+
+        # 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]
+
+  # SDのv2では1*1のconv2dがlinearに変わっているので、linear->convに変換する
+  if v2:
+    linear_transformer_to_conv(new_checkpoint)
+
+  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)
+  # if len(vae_state_dict) == 0:
+  #   # 渡されたcheckpointは.ckptから読み込んだcheckpointではなくvaeのstate_dict
+  #   vae_state_dict = checkpoint
+
+  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(v2):
+  """
+  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 if not v2 else V2_UNET_PARAMS_CONTEXT_DIM,
+      attention_head_dim=UNET_PARAMS_NUM_HEADS if not v2 else V2_UNET_PARAMS_ATTENTION_HEAD_DIM,
+  )
+
+  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_v1(checkpoint):
+  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]
+  return text_model_dict
+
+
+def convert_ldm_clip_checkpoint_v2(checkpoint, max_length):
+  # 嫌になるくらい違うぞ！
+  def convert_key(key):
+    if not key.startswith("cond_stage_model"):
+      return None
+
+    # common conversion
+    key = key.replace("cond_stage_model.model.transformer.", "text_model.encoder.")
+    key = key.replace("cond_stage_model.model.", "text_model.")
+
+    if "resblocks" in key:
+      # resblocks conversion
+      key = key.replace(".resblocks.", ".layers.")
+      if ".ln_" in key:
+        key = key.replace(".ln_", ".layer_norm")
+      elif ".mlp." in key:
+        key = key.replace(".c_fc.", ".fc1.")
+        key = key.replace(".c_proj.", ".fc2.")
+      elif '.attn.out_proj' in key:
+        key = key.replace(".attn.out_proj.", ".self_attn.out_proj.")
+      elif '.attn.in_proj' in key:
+        key = None                  # 特殊なので後で処理する
+      else:
+        raise ValueError(f"unexpected key in SD: {key}")
+    elif '.positional_embedding' in key:
+      key = key.replace(".positional_embedding", ".embeddings.position_embedding.weight")
+    elif '.text_projection' in key:
+      key = None    # 使われない???
+    elif '.logit_scale' in key:
+      key = None    # 使われない???
+    elif '.token_embedding' in key:
+      key = key.replace(".token_embedding.weight", ".embeddings.token_embedding.weight")
+    elif '.ln_final' in key:
+      key = key.replace(".ln_final", ".final_layer_norm")
+    return key
+
+  keys = list(checkpoint.keys())
+  new_sd = {}
+  for key in keys:
+    # remove resblocks 23
+    if '.resblocks.23.' in key:
+      continue
+    new_key = convert_key(key)
+    if new_key is None:
+      continue
+    new_sd[new_key] = checkpoint[key]
+
+  # attnの変換
+  for key in keys:
+    if '.resblocks.23.' in key:
+      continue
+    if '.resblocks' in key and '.attn.in_proj_' in key:
+      # 三つに分割
+      values = torch.chunk(checkpoint[key], 3)
+
+      key_suffix = ".weight" if "weight" in key else ".bias"
+      key_pfx = key.replace("cond_stage_model.model.transformer.resblocks.", "text_model.encoder.layers.")
+      key_pfx = key_pfx.replace("_weight", "")
+      key_pfx = key_pfx.replace("_bias", "")
+      key_pfx = key_pfx.replace(".attn.in_proj", ".self_attn.")
+      new_sd[key_pfx + "q_proj" + key_suffix] = values[0]
+      new_sd[key_pfx + "k_proj" + key_suffix] = values[1]
+      new_sd[key_pfx + "v_proj" + key_suffix] = values[2]
+
+  # rename or add position_ids
+  ANOTHER_POSITION_IDS_KEY = "text_model.encoder.text_model.embeddings.position_ids"
+  if ANOTHER_POSITION_IDS_KEY in new_sd:
+    # waifu diffusion v1.4
+    position_ids = new_sd[ANOTHER_POSITION_IDS_KEY]
+    del new_sd[ANOTHER_POSITION_IDS_KEY]
+  else:
+    position_ids = torch.Tensor([list(range(max_length))]).to(torch.int64)
+
+  new_sd["text_model.embeddings.position_ids"] = position_ids
+  return new_sd
+
+# endregion
+
+
+# region Diffusers->StableDiffusion の変換コード
+# convert_diffusers_to_original_stable_diffusion をコピーして修正している（ASL 2.0）
+
+def conv_transformer_to_linear(checkpoint):
+  keys = list(checkpoint.keys())
+  tf_keys = ["proj_in.weight", "proj_out.weight"]
+  for key in keys:
+    if ".".join(key.split(".")[-2:]) in tf_keys:
+      if checkpoint[key].ndim > 2:
+        checkpoint[key] = checkpoint[key][:, :, 0, 0]
+
+
+def convert_unet_state_dict_to_sd(v2, 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()}
+
+  if v2:
+    conv_transformer_to_linear(new_state_dict)
+
+  return new_state_dict
+
+
+# ================#
+# VAE Conversion #
+# ================#
+
+def reshape_weight_for_sd(w):
+    # convert HF linear weights to SD conv2d weights
+  return w.reshape(*w.shape, 1, 1)
+
+
+def convert_vae_state_dict(vae_state_dict):
+  vae_conversion_map = [
+      # (stable-diffusion, HF Diffusers)
+      ("nin_shortcut", "conv_shortcut"),
+      ("norm_out", "conv_norm_out"),
+      ("mid.attn_1.", "mid_block.attentions.0."),
+  ]
+
+  for i in range(4):
+    # down_blocks have two resnets
+    for j in range(2):
+      hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}."
+      sd_down_prefix = f"encoder.down.{i}.block.{j}."
+      vae_conversion_map.append((sd_down_prefix, hf_down_prefix))
+
+    if i < 3:
+      hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0."
+      sd_downsample_prefix = f"down.{i}.downsample."
+      vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))
+
+      hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+      sd_upsample_prefix = f"up.{3-i}.upsample."
+      vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))
+
+    # up_blocks have three resnets
+    # also, up blocks in hf are numbered in reverse from sd
+    for j in range(3):
+      hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}."
+      sd_up_prefix = f"decoder.up.{3-i}.block.{j}."
+      vae_conversion_map.append((sd_up_prefix, hf_up_prefix))
+
+  # this part accounts for mid blocks in both the encoder and the decoder
+  for i in range(2):
+    hf_mid_res_prefix = f"mid_block.resnets.{i}."
+    sd_mid_res_prefix = f"mid.block_{i+1}."
+    vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+  vae_conversion_map_attn = [
+      # (stable-diffusion, HF Diffusers)
+      ("norm.", "group_norm."),
+      ("q.", "query."),
+      ("k.", "key."),
+      ("v.", "value."),
+      ("proj_out.", "proj_attn."),
+  ]
+
+  mapping = {k: k for k in vae_state_dict.keys()}
+  for k, v in mapping.items():
+    for sd_part, hf_part in vae_conversion_map:
+      v = v.replace(hf_part, sd_part)
+    mapping[k] = v
+  for k, v in mapping.items():
+    if "attentions" in k:
+      for sd_part, hf_part in vae_conversion_map_attn:
+        v = v.replace(hf_part, sd_part)
+      mapping[k] = v
+  new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()}
+  weights_to_convert = ["q", "k", "v", "proj_out"]
+  for k, v in new_state_dict.items():
+    for weight_name in weights_to_convert:
+      if f"mid.attn_1.{weight_name}.weight" in k:
+        # print(f"Reshaping {k} for SD format")
+        new_state_dict[k] = reshape_weight_for_sd(v)
+
+  return new_state_dict
+
+
+# endregion
+
+# region 自作のモデル読み書きなど
+
+def is_safetensors(path):
+  return os.path.splitext(path)[1].lower() == '.safetensors'
+
+
+def load_checkpoint_with_text_encoder_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.')
+  ]
+
+  if is_safetensors(ckpt_path):
+    checkpoint = None
+    state_dict = load_file(ckpt_path, "cpu")
+  else:
+    checkpoint = torch.load(ckpt_path, map_location="cpu")
+    if "state_dict" in checkpoint:
+      state_dict = checkpoint["state_dict"]
+    else:
+      state_dict = checkpoint
+      checkpoint = None
+
+  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, state_dict
+
+
+# TODO dtype指定の動作が怪しいので確認する text_encoderを指定形式で作れるか未確認
+def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, dtype=None):
+  _, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path)
+  if dtype is not None:
+    for k, v in state_dict.items():
+      if type(v) is torch.Tensor:
+        state_dict[k] = v.to(dtype)
+
+  # Convert the UNet2DConditionModel model.
+  unet_config = create_unet_diffusers_config(v2)
+  converted_unet_checkpoint = convert_ldm_unet_checkpoint(v2, state_dict, unet_config)
+
+  unet = UNet2DConditionModel(**unet_config)
+  info = unet.load_state_dict(converted_unet_checkpoint)
+  print("loading u-net:", info)
+
+  # 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)
+  info = vae.load_state_dict(converted_vae_checkpoint)
+  print("loading vae:", info)
+
+  # convert text_model
+  if v2:
+    converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v2(state_dict, 77)
+    cfg = CLIPTextConfig(
+        vocab_size=49408,
+        hidden_size=1024,
+        intermediate_size=4096,
+        num_hidden_layers=23,
+        num_attention_heads=16,
+        max_position_embeddings=77,
+        hidden_act="gelu",
+        layer_norm_eps=1e-05,
+        dropout=0.0,
+        attention_dropout=0.0,
+        initializer_range=0.02,
+        initializer_factor=1.0,
+        pad_token_id=1,
+        bos_token_id=0,
+        eos_token_id=2,
+        model_type="clip_text_model",
+        projection_dim=512,
+        torch_dtype="float32",
+        transformers_version="4.25.0.dev0",
+    )
+    text_model = CLIPTextModel._from_config(cfg)
+    info = text_model.load_state_dict(converted_text_encoder_checkpoint)
+  else:
+    converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v1(state_dict)
+    text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+    info = text_model.load_state_dict(converted_text_encoder_checkpoint)
+  print("loading text encoder:", info)
+
+  return text_model, vae, unet
+
+
+def convert_text_encoder_state_dict_to_sd_v2(checkpoint, make_dummy_weights=False):
+  def convert_key(key):
+    # position_idsの除去
+    if ".position_ids" in key:
+      return None
+
+    # common
+    key = key.replace("text_model.encoder.", "transformer.")
+    key = key.replace("text_model.", "")
+    if "layers" in key:
+      # resblocks conversion
+      key = key.replace(".layers.", ".resblocks.")
+      if ".layer_norm" in key:
+        key = key.replace(".layer_norm", ".ln_")
+      elif ".mlp." in key:
+        key = key.replace(".fc1.", ".c_fc.")
+        key = key.replace(".fc2.", ".c_proj.")
+      elif '.self_attn.out_proj' in key:
+        key = key.replace(".self_attn.out_proj.", ".attn.out_proj.")
+      elif '.self_attn.' in key:
+        key = None                  # 特殊なので後で処理する
+      else:
+        raise ValueError(f"unexpected key in DiffUsers model: {key}")
+    elif '.position_embedding' in key:
+      key = key.replace("embeddings.position_embedding.weight", "positional_embedding")
+    elif '.token_embedding' in key:
+      key = key.replace("embeddings.token_embedding.weight", "token_embedding.weight")
+    elif 'final_layer_norm' in key:
+      key = key.replace("final_layer_norm", "ln_final")
+    return key
+
+  keys = list(checkpoint.keys())
+  new_sd = {}
+  for key in keys:
+    new_key = convert_key(key)
+    if new_key is None:
+      continue
+    new_sd[new_key] = checkpoint[key]
+
+  # attnの変換
+  for key in keys:
+    if 'layers' in key and 'q_proj' in key:
+      # 三つを結合
+      key_q = key
+      key_k = key.replace("q_proj", "k_proj")
+      key_v = key.replace("q_proj", "v_proj")
+
+      value_q = checkpoint[key_q]
+      value_k = checkpoint[key_k]
+      value_v = checkpoint[key_v]
+      value = torch.cat([value_q, value_k, value_v])
+
+      new_key = key.replace("text_model.encoder.layers.", "transformer.resblocks.")
+      new_key = new_key.replace(".self_attn.q_proj.", ".attn.in_proj_")
+      new_sd[new_key] = value
+
+  # 最後の層などを捏造するか
+  if make_dummy_weights:
+    print("make dummy weights for resblock.23, text_projection and logit scale.")
+    keys = list(new_sd.keys())
+    for key in keys:
+      if key.startswith("transformer.resblocks.22."):
+        new_sd[key.replace(".22.", ".23.")] = new_sd[key].clone()          # copyしないとsafetensorsの保存で落ちる
+
+    # Diffusersに含まれない重みを作っておく
+    new_sd['text_projection'] = torch.ones((1024, 1024), dtype=new_sd[keys[0]].dtype, device=new_sd[keys[0]].device)
+    new_sd['logit_scale'] = torch.tensor(1)
+
+  return new_sd
+
+
+def save_stable_diffusion_checkpoint(v2, output_file, text_encoder, unet, ckpt_path, epochs, steps, save_dtype=None, vae=None):
+  if ckpt_path is not None:
+    # epoch/stepを参照する。またVAEがメモリ上にないときなど、もう一度VAEを含めて読み込む
+    checkpoint, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path)
+    if checkpoint is None:                # safetensors または state_dictのckpt
+      checkpoint = {}
+      strict = False
+    else:
+      strict = True
+    if "state_dict" in state_dict:
+      del state_dict["state_dict"]
+  else:
+    # 新しく作る
+    assert vae is not None, "VAE is required to save a checkpoint without a given checkpoint"
+    checkpoint = {}
+    state_dict = {}
+    strict = False
+
+  def update_sd(prefix, sd):
+    for k, v in sd.items():
+      key = prefix + k
+      assert not strict or key in state_dict, f"Illegal key in save SD: {key}"
+      if save_dtype is not None:
+        v = v.detach().clone().to("cpu").to(save_dtype)
+      state_dict[key] = v
+
+  # Convert the UNet model
+  unet_state_dict = convert_unet_state_dict_to_sd(v2, unet.state_dict())
+  update_sd("model.diffusion_model.", unet_state_dict)
+
+  # Convert the text encoder model
+  if v2:
+    make_dummy = ckpt_path is None                 # 参照元のcheckpointがない場合は最後の層を前の層から複製��て作るなどダミーの重みを入れる
+    text_enc_dict = convert_text_encoder_state_dict_to_sd_v2(text_encoder.state_dict(), make_dummy)
+    update_sd("cond_stage_model.model.", text_enc_dict)
+  else:
+    text_enc_dict = text_encoder.state_dict()
+    update_sd("cond_stage_model.transformer.", text_enc_dict)
+
+  # Convert the VAE
+  if vae is not None:
+    vae_dict = convert_vae_state_dict(vae.state_dict())
+    update_sd("first_stage_model.", vae_dict)
+
+  # Put together new checkpoint
+  key_count = len(state_dict.keys())
+  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
+
+  if is_safetensors(output_file):
+    # TODO Tensor以外のdictの値を削除したほうがいいか
+    save_file(state_dict, output_file)
+  else:
+    torch.save(new_ckpt, output_file)
+
+  return key_count
+
+
+def save_diffusers_checkpoint(v2, output_dir, text_encoder, unet, pretrained_model_name_or_path, vae=None, use_safetensors=False):
+  if pretrained_model_name_or_path is None:
+    # load default settings for v1/v2
+    if v2:
+      pretrained_model_name_or_path = DIFFUSERS_REF_MODEL_ID_V2
+    else:
+      pretrained_model_name_or_path = DIFFUSERS_REF_MODEL_ID_V1
+
+  scheduler = DDIMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
+  tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer")
+  if vae is None:
+    vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae")
+
+  pipeline = StableDiffusionPipeline(
+      unet=unet,
+      text_encoder=text_encoder,
+      vae=vae,
+      scheduler=scheduler,
+      tokenizer=tokenizer,
+      safety_checker=None,
+      feature_extractor=None,
+      requires_safety_checker=None,
+  )
+  pipeline.save_pretrained(output_dir, safe_serialization=use_safetensors)
+
+
+VAE_PREFIX = "first_stage_model."
+
+
+def load_vae(vae_id, dtype):
+  print(f"load VAE: {vae_id}")
+  if os.path.isdir(vae_id) or not os.path.isfile(vae_id):
+    # Diffusers local/remote
+    try:
+      vae = AutoencoderKL.from_pretrained(vae_id, subfolder=None, torch_dtype=dtype)
+    except EnvironmentError as e:
+      print(f"exception occurs in loading vae: {e}")
+      print("retry with subfolder='vae'")
+      vae = AutoencoderKL.from_pretrained(vae_id, subfolder="vae", torch_dtype=dtype)
+    return vae
+
+  # local
+  vae_config = create_vae_diffusers_config()
+
+  if vae_id.endswith(".bin"):
+    # SD 1.5 VAE on Huggingface
+    converted_vae_checkpoint = torch.load(vae_id, map_location="cpu")
+  else:
+    # StableDiffusion
+    vae_model = (load_file(vae_id, "cpu") if is_safetensors(vae_id)
+                 else torch.load(vae_id, map_location="cpu"))
+    vae_sd = vae_model['state_dict'] if 'state_dict' in vae_model else vae_model
+
+    # vae only or full model
+    full_model = False
+    for vae_key in vae_sd:
+      if vae_key.startswith(VAE_PREFIX):
+        full_model = True
+        break
+    if not full_model:
+      sd = {}
+      for key, value in vae_sd.items():
+        sd[VAE_PREFIX + key] = value
+      vae_sd = sd
+      del sd
+
+    # Convert the VAE model.
+    converted_vae_checkpoint = convert_ldm_vae_checkpoint(vae_sd, vae_config)
+
+  vae = AutoencoderKL(**vae_config)
+  vae.load_state_dict(converted_vae_checkpoint)
+  return vae
+
+# 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()
+  return resos
+
+
+if __name__ == '__main__':
+  resos = make_bucket_resolutions((512, 768))
+  print(len(resos))
+  print(resos)
+  aspect_ratios = [w / h for w, h in resos]
+  print(aspect_ratios)
+
+  ars = set()
+  for ar in aspect_ratios:
+    if ar in ars:
+      print("error! duplicate ar:", ar)
+    ars.add(ar)

build/lib/library/train_util.py

@@ -0,0 +1,1796 @@
+# common functions for training
+
+import argparse
+import json
+import shutil
+import time
+from typing import Dict, List, NamedTuple, Tuple
+from accelerate import Accelerator
+from torch.autograd.function import Function
+import glob
+import math
+import os
+import random
+import hashlib
+import subprocess
+from io import BytesIO
+
+from tqdm import tqdm
+import torch
+from torchvision import transforms
+from transformers import CLIPTokenizer
+import diffusers
+from diffusers import DDPMScheduler, StableDiffusionPipeline
+import albumentations as albu
+import numpy as np
+from PIL import Image
+import cv2
+from einops import rearrange
+from torch import einsum
+import safetensors.torch
+
+import library.model_util as model_util
+
+# Tokenizer: checkpointから読み込むのではなくあらかじめ提供されているものを使う
+TOKENIZER_PATH = "openai/clip-vit-large-patch14"
+V2_STABLE_DIFFUSION_PATH = "stabilityai/stable-diffusion-2"     # ここからtokenizerだけ使う v2とv2.1はtokenizer仕様は同じ
+
+# checkpointファイル名
+EPOCH_STATE_NAME = "{}-{:06d}-state"
+EPOCH_FILE_NAME = "{}-{:06d}"
+EPOCH_DIFFUSERS_DIR_NAME = "{}-{:06d}"
+LAST_STATE_NAME = "{}-state"
+DEFAULT_EPOCH_NAME = "epoch"
+DEFAULT_LAST_OUTPUT_NAME = "last"
+
+# region dataset
+
+IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".webp", ".bmp"]
+# , ".PNG", ".JPG", ".JPEG", ".WEBP", ".BMP"]         # Linux?
+
+
+class ImageInfo():
+  def __init__(self, image_key: str, num_repeats: int, caption: str, is_reg: bool, absolute_path: str) -> None:
+    self.image_key: str = image_key
+    self.num_repeats: int = num_repeats
+    self.caption: str = caption
+    self.is_reg: bool = is_reg
+    self.absolute_path: str = absolute_path
+    self.image_size: Tuple[int, int] = None
+    self.resized_size: Tuple[int, int] = None
+    self.bucket_reso: Tuple[int, int] = None
+    self.latents: torch.Tensor = None
+    self.latents_flipped: torch.Tensor = None
+    self.latents_npz: str = None
+    self.latents_npz_flipped: str = None
+
+
+class BucketManager():
+  def __init__(self, no_upscale, max_reso, min_size, max_size, reso_steps) -> None:
+    self.no_upscale = no_upscale
+    if max_reso is None:
+      self.max_reso = None
+      self.max_area = None
+    else:
+      self.max_reso = max_reso
+      self.max_area = max_reso[0] * max_reso[1]
+    self.min_size = min_size
+    self.max_size = max_size
+    self.reso_steps = reso_steps
+
+    self.resos = []
+    self.reso_to_id = {}
+    self.buckets = []                     # 前処理時は (image_key, image)、学習時は image_key
+
+  def add_image(self, reso, image):
+    bucket_id = self.reso_to_id[reso]
+    self.buckets[bucket_id].append(image)
+
+  def shuffle(self):
+    for bucket in self.buckets:
+      random.shuffle(bucket)
+
+  def sort(self):
+    # 解像度順にソートする（表示時、メタデータ格納時の見栄えをよくするためだけ）。bucketsも入れ替えてreso_to_idも振り直す
+    sorted_resos = self.resos.copy()
+    sorted_resos.sort()
+
+    sorted_buckets = []
+    sorted_reso_to_id = {}
+    for i, reso in enumerate(sorted_resos):
+      bucket_id = self.reso_to_id[reso]
+      sorted_buckets.append(self.buckets[bucket_id])
+      sorted_reso_to_id[reso] = i
+
+    self.resos = sorted_resos
+    self.buckets = sorted_buckets
+    self.reso_to_id = sorted_reso_to_id
+
+  def make_buckets(self):
+    resos = model_util.make_bucket_resolutions(self.max_reso, self.min_size, self.max_size, self.reso_steps)
+    self.set_predefined_resos(resos)
+
+  def set_predefined_resos(self, resos):
+    # 規定サイズから選ぶ場合の解像度、aspect ratioの情報を格納しておく
+    self.predefined_resos = resos.copy()
+    self.predefined_resos_set = set(resos)
+    self.predefined_aspect_ratios = np.array([w / h for w, h in resos])
+
+  def add_if_new_reso(self, reso):
+    if reso not in self.reso_to_id:
+      bucket_id = len(self.resos)
+      self.reso_to_id[reso] = bucket_id
+      self.resos.append(reso)
+      self.buckets.append([])
+      # print(reso, bucket_id, len(self.buckets))
+
+  def round_to_steps(self, x):
+    x = int(x + .5)
+    return x - x % self.reso_steps
+
+  def select_bucket(self, image_width, image_height):
+    aspect_ratio = image_width / image_height
+    if not self.no_upscale:
+      # 同じaspect ratioがあるかもしれないので（fine tuningで、no_upscale=Trueで前処理した場合）、解像度が同じものを優先する
+      reso = (image_width, image_height)
+      if reso in self.predefined_resos_set:
+        pass
+      else:
+        ar_errors = self.predefined_aspect_ratios - aspect_ratio
+        predefined_bucket_id = np.abs(ar_errors).argmin()          # 当該解像度以外でaspect ratio errorが最も少ないもの
+        reso = self.predefined_resos[predefined_bucket_id]
+
+      ar_reso = reso[0] / reso[1]
+      if aspect_ratio > ar_reso:                   # 横が長い→縦を合わせる
+        scale = reso[1] / image_height
+      else:
+        scale = reso[0] / image_width
+
+      resized_size = (int(image_width * scale + .5), int(image_height * scale + .5))
+      # print("use predef", image_width, image_height, reso, resized_size)
+    else:
+      if image_width * image_height > self.max_area:
+        # 画像が大きすぎるのでアスペクト比を保ったまま縮小することを前提にbucketを決める
+        resized_width = math.sqrt(self.max_area * aspect_ratio)
+        resized_height = self.max_area / resized_width
+        assert abs(resized_width / resized_height - aspect_ratio) < 1e-2, "aspect is illegal"
+
+        # リサイズ後の短辺または長辺をreso_steps単位にする：aspect ratioの差が少ないほうを選ぶ
+        # 元のbucketingと同じロジック
+        b_width_rounded = self.round_to_steps(resized_width)
+        b_height_in_wr = self.round_to_steps(b_width_rounded / aspect_ratio)
+        ar_width_rounded = b_width_rounded / b_height_in_wr
+
+        b_height_rounded = self.round_to_steps(resized_height)
+        b_width_in_hr = self.round_to_steps(b_height_rounded * aspect_ratio)
+        ar_height_rounded = b_width_in_hr / b_height_rounded
+
+        # print(b_width_rounded, b_height_in_wr, ar_width_rounded)
+        # print(b_width_in_hr, b_height_rounded, ar_height_rounded)
+
+        if abs(ar_width_rounded - aspect_ratio) < abs(ar_height_rounded - aspect_ratio):
+          resized_size = (b_width_rounded, int(b_width_rounded / aspect_ratio + .5))
+        else:
+          resized_size = (int(b_height_rounded * aspect_ratio + .5), b_height_rounded)
+        # print(resized_size)
+      else:
+        resized_size = (image_width, image_height)              # リサイズは不要
+
+      # 画像のサイズ未満をbucketのサイズとする（paddingせずにcroppingする）
+      bucket_width = resized_size[0] - resized_size[0] % self.reso_steps
+      bucket_height = resized_size[1] - resized_size[1] % self.reso_steps
+      # print("use arbitrary", image_width, image_height, resized_size, bucket_width, bucket_height)
+
+      reso = (bucket_width, bucket_height)
+
+    self.add_if_new_reso(reso)
+
+    ar_error = (reso[0] / reso[1]) - aspect_ratio
+    return reso, resized_size, ar_error
+
+
+class BucketBatchIndex(NamedTuple):
+  bucket_index: int
+  bucket_batch_size: int
+  batch_index: int
+
+
+class BaseDataset(torch.utils.data.Dataset):
+  def __init__(self, tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens, resolution, flip_aug: bool, color_aug: bool, face_crop_aug_range, random_crop, debug_dataset: bool) -> None:
+    super().__init__()
+    self.tokenizer: CLIPTokenizer = tokenizer
+    self.max_token_length = max_token_length
+    self.shuffle_caption = shuffle_caption
+    self.shuffle_keep_tokens = shuffle_keep_tokens
+    # width/height is used when enable_bucket==False
+    self.width, self.height = (None, None) if resolution is None else resolution
+    self.face_crop_aug_range = face_crop_aug_range
+    self.flip_aug = flip_aug
+    self.color_aug = color_aug
+    self.debug_dataset = debug_dataset
+    self.random_crop = random_crop
+    self.token_padding_disabled = False
+    self.dataset_dirs_info = {}
+    self.reg_dataset_dirs_info = {}
+    self.tag_frequency = {}
+
+    self.enable_bucket = False
+    self.bucket_manager: BucketManager = None                         # not initialized
+    self.min_bucket_reso = None
+    self.max_bucket_reso = None
+    self.bucket_reso_steps = None
+    self.bucket_no_upscale = None
+    self.bucket_info = None                                           # for metadata
+
+    self.tokenizer_max_length = self.tokenizer.model_max_length if max_token_length is None else max_token_length + 2
+
+    self.current_epoch: int = 0            # インスタンスがepochごとに新しく作られるようなので外側から渡さないとダメ
+    self.dropout_rate: float = 0
+    self.dropout_every_n_epochs: int = None
+    self.tag_dropout_rate: float = 0
+
+    # augmentation
+    flip_p = 0.5 if flip_aug else 0.0
+    if color_aug:
+      # わりと弱めの色合いaugmentation：brightness/contrastあたりは画像のpixel valueの最大値・最小値を変えてしまうのでよくないのではという想定でgamma/hueあたりを触る
+      self.aug = albu.Compose([
+          albu.OneOf([
+              albu.HueSaturationValue(8, 0, 0, p=.5),
+              albu.RandomGamma((95, 105), p=.5),
+          ], p=.33),
+          albu.HorizontalFlip(p=flip_p)
+      ], p=1.)
+    elif flip_aug:
+      self.aug = albu.Compose([
+          albu.HorizontalFlip(p=flip_p)
+      ], p=1.)
+    else:
+      self.aug = None
+
+    self.image_transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5]), ])
+
+    self.image_data: Dict[str, ImageInfo] = {}
+
+    self.replacements = {}
+
+  def set_current_epoch(self, epoch):
+    self.current_epoch = epoch
+
+  def set_caption_dropout(self, dropout_rate, dropout_every_n_epochs, tag_dropout_rate):
+    # コンストラクタで渡さないのはTextual Inversionで意識したくないから（ということにしておく）
+    self.dropout_rate = dropout_rate
+    self.dropout_every_n_epochs = dropout_every_n_epochs
+    self.tag_dropout_rate = tag_dropout_rate
+
+  def set_tag_frequency(self, dir_name, captions):
+    frequency_for_dir = self.tag_frequency.get(dir_name, {})
+    self.tag_frequency[dir_name] = frequency_for_dir
+    for caption in captions:
+      for tag in caption.split(","):
+        if tag and not tag.isspace():
+          tag = tag.lower()
+          frequency = frequency_for_dir.get(tag, 0)
+          frequency_for_dir[tag] = frequency + 1
+
+  def disable_token_padding(self):
+    self.token_padding_disabled = True
+
+  def add_replacement(self, str_from, str_to):
+    self.replacements[str_from] = str_to
+
+  def process_caption(self, caption):
+    # dropoutの決定：tag dropがこのメソッド内にあるのでここで行うのが良い
+    is_drop_out = self.dropout_rate > 0 and random.random() < self.dropout_rate
+    is_drop_out = is_drop_out or self.dropout_every_n_epochs and self.current_epoch % self.dropout_every_n_epochs == 0
+
+    if is_drop_out:
+      caption = ""
+    else:
+      if self.shuffle_caption or self.tag_dropout_rate > 0:
+        def dropout_tags(tokens):
+          if self.tag_dropout_rate <= 0:
+            return tokens
+          l = []
+          for token in tokens:
+            if random.random() >= self.tag_dropout_rate:
+              l.append(token)
+          return l
+
+        tokens = [t.strip() for t in caption.strip().split(",")]
+        if self.shuffle_keep_tokens is None:
+          if self.shuffle_caption:
+            random.shuffle(tokens)
+
+          tokens = dropout_tags(tokens)
+        else:
+          if len(tokens) > self.shuffle_keep_tokens:
+            keep_tokens = tokens[:self.shuffle_keep_tokens]
+            tokens = tokens[self.shuffle_keep_tokens:]
+
+            if self.shuffle_caption:
+              random.shuffle(tokens)
+
+            tokens = dropout_tags(tokens)
+
+            tokens = keep_tokens + tokens
+        caption = ", ".join(tokens)
+
+      # textual inversion対応
+      for str_from, str_to in self.replacements.items():
+        if str_from == "":
+          # replace all
+          if type(str_to) == list:
+            caption = random.choice(str_to)
+          else:
+            caption = str_to
+        else:
+          caption = caption.replace(str_from, str_to)
+
+    return caption
+
+  def get_input_ids(self, caption):
+    input_ids = self.tokenizer(caption, padding="max_length", truncation=True,
+                               max_length=self.tokenizer_max_length, return_tensors="pt").input_ids
+
+    if self.tokenizer_max_length > self.tokenizer.model_max_length:
+      input_ids = input_ids.squeeze(0)
+      iids_list = []
+      if self.tokenizer.pad_token_id == self.tokenizer.eos_token_id:
+        # v1
+        # 77以上の時は "<BOS> .... <EOS> <EOS> <EOS>" でトータル227とかになっているので、"<BOS>...<EOS>"の三連に変換する
+        # 1111氏のやつは , で区切る、とかしているようだが　とりあえず単純に
+        for i in range(1, self.tokenizer_max_length - self.tokenizer.model_max_length + 2, self.tokenizer.model_max_length - 2):  # (1, 152, 75)
+          ids_chunk = (input_ids[0].unsqueeze(0),
+                       input_ids[i:i + self.tokenizer.model_max_length - 2],
+                       input_ids[-1].unsqueeze(0))
+          ids_chunk = torch.cat(ids_chunk)
+          iids_list.append(ids_chunk)
+      else:
+        # v2
+        # 77以上の時は "<BOS> .... <EOS> <PAD> <PAD>..." でトータル227とかになっているので、"<BOS>...<EOS> <PAD> <PAD> ..."の三連に変換する
+        for i in range(1, self.tokenizer_max_length - self.tokenizer.model_max_length + 2, self.tokenizer.model_max_length - 2):
+          ids_chunk = (input_ids[0].unsqueeze(0),       # BOS
+                       input_ids[i:i + self.tokenizer.model_max_length - 2],
+                       input_ids[-1].unsqueeze(0))      # PAD or EOS
+          ids_chunk = torch.cat(ids_chunk)
+
+          # 末尾が <EOS> <PAD> または <PAD> <PAD> の場合は、何もしなくてよい
+          # 末尾が x <PAD/EOS> の場合は末尾を <EOS> に変える（x <EOS> なら結果的に変化なし）
+          if ids_chunk[-2] != self.tokenizer.eos_token_id and ids_chunk[-2] != self.tokenizer.pad_token_id:
+            ids_chunk[-1] = self.tokenizer.eos_token_id
+          # 先頭が <BOS> <PAD> ... の場合は <BOS> <EOS> <PAD> ... に変える
+          if ids_chunk[1] == self.tokenizer.pad_token_id:
+            ids_chunk[1] = self.tokenizer.eos_token_id
+
+          iids_list.append(ids_chunk)
+
+      input_ids = torch.stack(iids_list)      # 3,77
+    return input_ids
+
+  def register_image(self, info: ImageInfo):
+    self.image_data[info.image_key] = info
+
+  def make_buckets(self):
+    '''
+    bucketingを行わない場合も呼び出し必須（ひとつだけbucketを作る）
+    min_size and max_size are ignored when enable_bucket is False
+    '''
+    print("loading image sizes.")
+    for info in tqdm(self.image_data.values()):
+      if info.image_size is None:
+        info.image_size = self.get_image_size(info.absolute_path)
+
+    if self.enable_bucket:
+      print("make buckets")
+    else:
+      print("prepare dataset")
+
+    # bucketを作成し、画像をbucketに振り分ける
+    if self.enable_bucket:
+      if self.bucket_manager is None:                         # fine tuningの場合でmetadataに定義がある場合は、すでに初期化済み
+        self.bucket_manager = BucketManager(self.bucket_no_upscale, (self.width, self.height),
+                                            self.min_bucket_reso, self.max_bucket_reso, self.bucket_reso_steps)
+        if not self.bucket_no_upscale:
+          self.bucket_manager.make_buckets()
+        else:
+          print("min_bucket_reso and max_bucket_reso are ignored if bucket_no_upscale is set, because bucket reso is defined by image size automatically / bucket_no_upscaleが指定された場合は、bucketの解像度は画像サイズから自動計算されるため、min_bucket_resoとmax_bucket_resoは無視されます")
+
+      img_ar_errors = []
+      for image_info in self.image_data.values():
+        image_width, image_height = image_info.image_size
+        image_info.bucket_reso, image_info.resized_size, ar_error = self.bucket_manager.select_bucket(image_width, image_height)
+
+        # print(image_info.image_key, image_info.bucket_reso)
+        img_ar_errors.append(abs(ar_error))
+
+      self.bucket_manager.sort()
+    else:
+      self.bucket_manager = BucketManager(False, (self.width, self.height), None, None, None)
+      self.bucket_manager.set_predefined_resos([(self.width, self.height)])  # ひとつの固定サイズbucketのみ
+      for image_info in self.image_data.values():
+        image_width, image_height = image_info.image_size
+        image_info.bucket_reso, image_info.resized_size, _ = self.bucket_manager.select_bucket(image_width, image_height)
+
+    for image_info in self.image_data.values():
+      for _ in range(image_info.num_repeats):
+        self.bucket_manager.add_image(image_info.bucket_reso, image_info.image_key)
+
+    # bucket情報を表示、格納する
+    if self.enable_bucket:
+      self.bucket_info = {"buckets": {}}
+      print("number of images (including repeats) / 各bucketの画像枚数（繰り返し回数を含む）")
+      for i, (reso, bucket) in enumerate(zip(self.bucket_manager.resos, self.bucket_manager.buckets)):
+        count = len(bucket)
+        if count > 0:
+          self.bucket_info["buckets"][i] = {"resolution": reso, "count": len(bucket)}
+          print(f"bucket {i}: resolution {reso}, count: {len(bucket)}")
+
+      img_ar_errors = np.array(img_ar_errors)
+      mean_img_ar_error = np.mean(np.abs(img_ar_errors))
+      self.bucket_info["mean_img_ar_error"] = mean_img_ar_error
+      print(f"mean ar error (without repeats): {mean_img_ar_error}")
+
+    # データ参照用indexを作る。このindexはdatasetのshuffleに用いられる
+    self.buckets_indices: List(BucketBatchIndex) = []
+    for bucket_index, bucket in enumerate(self.bucket_manager.buckets):
+      batch_count = int(math.ceil(len(bucket) / self.batch_size))
+      for batch_index in range(batch_count):
+        self.buckets_indices.append(BucketBatchIndex(bucket_index, self.batch_size, batch_index))
+
+      # ↓以下はbucketごとのbatch件数があまりにも増えて混乱を招くので元に戻す
+      # 　学習時はステップ数がランダムなので、同一画像が同一batch内にあってもそれほど悪影響はないであろう、と考えられる
+      #
+      # # bucketが細分化されることにより、ひとつのbucketに一種類の画像のみというケースが増え、つまりそれは
+      # # ひとつのbatchが同じ画像で占められることになるので、さすがに良くないであろう
+      # # そのためバッチサイズを画像種類までに制限する
+      # # ただそれでも同一画像が同一バッチに含まれる可能性はあるので、繰り返し回数が少ないほうがshuffleの品質は良くなることは間違いない？
+      # # TO DO 正則化画像をepochまたがりで利用する仕組み
+      # num_of_image_types = len(set(bucket))
+      # bucket_batch_size = min(self.batch_size, num_of_image_types)
+      # batch_count = int(math.ceil(len(bucket) / bucket_batch_size))
+      # # print(bucket_index, num_of_image_types, bucket_batch_size, batch_count)
+      # for batch_index in range(batch_count):
+      #   self.buckets_indices.append(BucketBatchIndex(bucket_index, bucket_batch_size, batch_index))
+      # ↑ここまで
+
+    self.shuffle_buckets()
+    self._length = len(self.buckets_indices)
+
+  def shuffle_buckets(self):
+    random.shuffle(self.buckets_indices)
+    self.bucket_manager.shuffle()
+
+  def load_image(self, image_path):
+    image = Image.open(image_path)
+    if not image.mode == "RGB":
+      image = image.convert("RGB")
+    img = np.array(image, np.uint8)
+    return img
+
+  def trim_and_resize_if_required(self, image, reso, resized_size):
+    image_height, image_width = image.shape[0:2]
+
+    if image_width != resized_size[0] or image_height != resized_size[1]:
+      # リサイズする
+      image = cv2.resize(image, resized_size, interpolation=cv2.INTER_AREA)       # INTER_AREAでやりたいのでcv2でリサイズ
+
+    image_height, image_width = image.shape[0:2]
+    if image_width > reso[0]:
+      trim_size = image_width - reso[0]
+      p = trim_size // 2 if not self.random_crop else random.randint(0, trim_size)
+      # print("w", trim_size, p)
+      image = image[:, p:p + reso[0]]
+    if image_height > reso[1]:
+      trim_size = image_height - reso[1]
+      p = trim_size // 2 if not self.random_crop else random.randint(0, trim_size)
+      # print("h", trim_size, p)
+      image = image[p:p + reso[1]]
+
+    assert image.shape[0] == reso[1] and image.shape[1] == reso[0], f"internal error, illegal trimmed size: {image.shape}, {reso}"
+    return image
+
+  def cache_latents(self, vae):
+    # TODO ここを高速化したい
+    print("caching latents.")
+    for info in tqdm(self.image_data.values()):
+      if info.latents_npz is not None:
+        info.latents = self.load_latents_from_npz(info, False)
+        info.latents = torch.FloatTensor(info.latents)
+        info.latents_flipped = self.load_latents_from_npz(info, True)             # might be None
+        if info.latents_flipped is not None:
+          info.latents_flipped = torch.FloatTensor(info.latents_flipped)
+        continue
+
+      image = self.load_image(info.absolute_path)
+      image = self.trim_and_resize_if_required(image, info.bucket_reso, info.resized_size)
+
+      img_tensor = self.image_transforms(image)
+      img_tensor = img_tensor.unsqueeze(0).to(device=vae.device, dtype=vae.dtype)
+      info.latents = vae.encode(img_tensor).latent_dist.sample().squeeze(0).to("cpu")
+
+      if self.flip_aug:
+        image = image[:, ::-1].copy()     # cannot convert to Tensor without copy
+        img_tensor = self.image_transforms(image)
+        img_tensor = img_tensor.unsqueeze(0).to(device=vae.device, dtype=vae.dtype)
+        info.latents_flipped = vae.encode(img_tensor).latent_dist.sample().squeeze(0).to("cpu")
+
+  def get_image_size(self, image_path):
+    image = Image.open(image_path)
+    return image.size
+
+  def load_image_with_face_info(self, image_path: str):
+    img = self.load_image(image_path)
+
+    face_cx = face_cy = face_w = face_h = 0
+    if self.face_crop_aug_range is not None:
+      tokens = os.path.splitext(os.path.basename(image_path))[0].split('_')
+      if len(tokens) >= 5:
+        face_cx = int(tokens[-4])
+        face_cy = int(tokens[-3])
+        face_w = int(tokens[-2])
+        face_h = int(tokens[-1])
+
+    return img, face_cx, face_cy, face_w, face_h
+
+  # いい感じに切り出す
+  def crop_target(self, image, face_cx, face_cy, face_w, face_h):
+    height, width = image.shape[0:2]
+    if height == self.height and width == self.width:
+      return image
+
+    # 画像サイズはsizeより大きいのでリサイズする
+    face_size = max(face_w, face_h)
+    min_scale = max(self.height / height, self.width / width)        # 画像がモデル入力サイズぴったりになる倍率（最小の倍率）
+    min_scale = min(1.0, max(min_scale, self.size / (face_size * self.face_crop_aug_range[1])))             # 指定した顔最小サイズ
+    max_scale = min(1.0, max(min_scale, self.size / (face_size * self.face_crop_aug_range[0])))             # 指定した顔最大サイズ
+    if min_scale >= max_scale:          # range指定がmin==max
+      scale = min_scale
+    else:
+      scale = random.uniform(min_scale, max_scale)
+
+    nh = int(height * scale + .5)
+    nw = int(width * scale + .5)
+    assert nh >= self.height and nw >= self.width, f"internal error. small scale {scale}, {width}*{height}"
+    image = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_AREA)
+    face_cx = int(face_cx * scale + .5)
+    face_cy = int(face_cy * scale + .5)
+    height, width = nh, nw
+
+    # 顔を中心として448*640とかへ切り出す
+    for axis, (target_size, length, face_p) in enumerate(zip((self.height, self.width), (height, width), (face_cy, face_cx))):
+      p1 = face_p - target_size // 2                # 顔を中心に持ってくるための切り出し位置
+
+      if self.random_crop:
+        # 背景も含めるために顔を中心に置く確率を高めつつずらす
+        range = max(length - face_p, face_p)        # 画像の端から顔中心までの距離の長いほう
+        p1 = p1 + (random.randint(0, range) + random.randint(0, range)) - range     # -range ~ +range までのいい感じの乱数
+      else:
+        # range指定があるときのみ、すこしだけランダムに（わりと適当）
+        if self.face_crop_aug_range[0] != self.face_crop_aug_range[1]:
+          if face_size > self.size // 10 and face_size >= 40:
+            p1 = p1 + random.randint(-face_size // 20, +face_size // 20)
+
+      p1 = max(0, min(p1, length - target_size))
+
+      if axis == 0:
+        image = image[p1:p1 + target_size, :]
+      else:
+        image = image[:, p1:p1 + target_size]
+
+    return image
+
+  def load_latents_from_npz(self, image_info: ImageInfo, flipped):
+    npz_file = image_info.latents_npz_flipped if flipped else image_info.latents_npz
+    if npz_file is None:
+      return None
+    return np.load(npz_file)['arr_0']
+
+  def __len__(self):
+    return self._length
+
+  def __getitem__(self, index):
+    if index == 0:
+      self.shuffle_buckets()
+
+    bucket = self.bucket_manager.buckets[self.buckets_indices[index].bucket_index]
+    bucket_batch_size = self.buckets_indices[index].bucket_batch_size
+    image_index = self.buckets_indices[index].batch_index * bucket_batch_size
+
+    loss_weights = []
+    captions = []
+    input_ids_list = []
+    latents_list = []
+    images = []
+
+    for image_key in bucket[image_index:image_index + bucket_batch_size]:
+      image_info = self.image_data[image_key]
+      loss_weights.append(self.prior_loss_weight if image_info.is_reg else 1.0)
+
+      # image/latentsを処理する
+      if image_info.latents is not None:
+        latents = image_info.latents if not self.flip_aug or random.random() < .5 else image_info.latents_flipped
+        image = None
+      elif image_info.latents_npz is not None:
+        latents = self.load_latents_from_npz(image_info, self.flip_aug and random.random() >= .5)
+        latents = torch.FloatTensor(latents)
+        image = None
+      else:
+        # 画像を読み込み、必要ならcropする
+        img, face_cx, face_cy, face_w, face_h = self.load_image_with_face_info(image_info.absolute_path)
+        im_h, im_w = img.shape[0:2]
+
+        if self.enable_bucket:
+          img = self.trim_and_resize_if_required(img, image_info.bucket_reso, image_info.resized_size)
+        else:
+          if face_cx > 0:                   # 顔位置情報あり
+            img = self.crop_target(img, face_cx, face_cy, face_w, face_h)
+          elif im_h > self.height or im_w > self.width:
+            assert self.random_crop, f"image too large, but cropping and bucketing are disabled / 画像サイズが大きいのでface_crop_aug_rangeかrandom_crop、またはbucketを有効にしてください: {image_info.absolute_path}"
+            if im_h > self.height:
+              p = random.randint(0, im_h - self.height)
+              img = img[p:p + self.height]
+            if im_w > self.width:
+              p = random.randint(0, im_w - self.width)
+              img = img[:, p:p + self.width]
+
+          im_h, im_w = img.shape[0:2]
+          assert im_h == self.height and im_w == self.width, f"image size is small / 画像サイズが小さいようです: {image_info.absolute_path}"
+
+        # augmentation
+        if self.aug is not None:
+          img = self.aug(image=img)['image']
+
+        latents = None
+        image = self.image_transforms(img)      # -1.0~1.0のtorch.Tensorになる
+
+      images.append(image)
+      latents_list.append(latents)
+
+      caption = self.process_caption(image_info.caption)
+      captions.append(caption)
+      if not self.token_padding_disabled:                     # this option might be omitted in future
+        input_ids_list.append(self.get_input_ids(caption))
+
+    example = {}
+    example['loss_weights'] = torch.FloatTensor(loss_weights)
+
+    if self.token_padding_disabled:
+      # padding=True means pad in the batch
+      example['input_ids'] = self.tokenizer(captions, padding=True, truncation=True, return_tensors="pt").input_ids
+    else:
+      # batch processing seems to be good
+      example['input_ids'] = torch.stack(input_ids_list)
+
+    if images[0] is not None:
+      images = torch.stack(images)
+      images = images.to(memory_format=torch.contiguous_format).float()
+    else:
+      images = None
+    example['images'] = images
+
+    example['latents'] = torch.stack(latents_list) if latents_list[0] is not None else None
+
+    if self.debug_dataset:
+      example['image_keys'] = bucket[image_index:image_index + self.batch_size]
+      example['captions'] = captions
+    return example
+
+
+class DreamBoothDataset(BaseDataset):
+  def __init__(self, batch_size, train_data_dir, reg_data_dir, tokenizer, max_token_length, caption_extension, shuffle_caption, shuffle_keep_tokens, resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, prior_loss_weight, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset) -> None:
+    super().__init__(tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens,
+                     resolution, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset)
+
+    assert resolution is not None, f"resolution is required / resolution（解像度）指定は必須です"
+
+    self.batch_size = batch_size
+    self.size = min(self.width, self.height)                  # 短いほう
+    self.prior_loss_weight = prior_loss_weight
+    self.latents_cache = None
+
+    self.enable_bucket = enable_bucket
+    if self.enable_bucket:
+      assert min(resolution) >= min_bucket_reso, f"min_bucket_reso must be equal or less than resolution / min_bucket_resoは最小解像度より大きくできません。解像度を大きくするかmin_bucket_resoを小さくしてください"
+      assert max(resolution) <= max_bucket_reso, f"max_bucket_reso must be equal or greater than resolution / max_bucket_resoは最大解像度より小さくできません。解像度を小さくするかmin_bucket_resoを大きくしてください"
+      self.min_bucket_reso = min_bucket_reso
+      self.max_bucket_reso = max_bucket_reso
+      self.bucket_reso_steps = bucket_reso_steps
+      self.bucket_no_upscale = bucket_no_upscale
+    else:
+      self.min_bucket_reso = None
+      self.max_bucket_reso = None
+      self.bucket_reso_steps = None                              # この情報は使われない
+      self.bucket_no_upscale = False
+
+    def read_caption(img_path):
+      # captionの候補ファイル名を作る
+      base_name = os.path.splitext(img_path)[0]
+      base_name_face_det = base_name
+      tokens = base_name.split("_")
+      if len(tokens) >= 5:
+        base_name_face_det = "_".join(tokens[:-4])
+      cap_paths = [base_name + caption_extension, base_name_face_det + caption_extension]
+
+      caption = None
+      for cap_path in cap_paths:
+        if os.path.isfile(cap_path):
+          with open(cap_path, "rt", encoding='utf-8') as f:
+            try:
+              lines = f.readlines()
+            except UnicodeDecodeError as e:
+              print(f"illegal char in file (not UTF-8) / ファイルにUTF-8以外の文字があります: {cap_path}")
+              raise e
+            assert len(lines) > 0, f"caption file is empty / キャプションファイルが空です: {cap_path}"
+            caption = lines[0].strip()
+          break
+      return caption
+
+    def load_dreambooth_dir(dir):
+      if not os.path.isdir(dir):
+        # print(f"ignore file: {dir}")
+        return 0, [], []
+
+      tokens = os.path.basename(dir).split('_')
+      try:
+        n_repeats = int(tokens[0])
+      except ValueError as e:
+        print(f"ignore directory without repeats / 繰り返し回数のないディレクトリを無視します: {dir}")
+        return 0, [], []
+
+      caption_by_folder = '_'.join(tokens[1:])
+      img_paths = glob_images(dir, "*")
+      print(f"found directory {n_repeats}_{caption_by_folder} contains {len(img_paths)} image files")
+
+      # 画像ファイルごとにプロンプトを読み込み、もしあればそちらを使う
+      captions = []
+      for img_path in img_paths:
+        cap_for_img = read_caption(img_path)
+        captions.append(caption_by_folder if cap_for_img is None else cap_for_img)
+
+      self.set_tag_frequency(os.path.basename(dir), captions)         # タグ頻度を記録
+
+      return n_repeats, img_paths, captions
+
+    print("prepare train images.")
+    train_dirs = os.listdir(train_data_dir)
+    num_train_images = 0
+    for dir in train_dirs:
+      n_repeats, img_paths, captions = load_dreambooth_dir(os.path.join(train_data_dir, dir))
+      num_train_images += n_repeats * len(img_paths)
+
+      for img_path, caption in zip(img_paths, captions):
+        info = ImageInfo(img_path, n_repeats, caption, False, img_path)
+        self.register_image(info)
+
+      self.dataset_dirs_info[os.path.basename(dir)] = {"n_repeats": n_repeats, "img_count": len(img_paths)}
+
+    print(f"{num_train_images} train images with repeating.")
+    self.num_train_images = num_train_images
+
+    # reg imageは数を数えて学習画像と同じ枚数にする
+    num_reg_images = 0
+    if reg_data_dir:
+      print("prepare reg images.")
+      reg_infos: List[ImageInfo] = []
+
+      reg_dirs = os.listdir(reg_data_dir)
+      for dir in reg_dirs:
+        n_repeats, img_paths, captions = load_dreambooth_dir(os.path.join(reg_data_dir, dir))
+        num_reg_images += n_repeats * len(img_paths)
+
+        for img_path, caption in zip(img_paths, captions):
+          info = ImageInfo(img_path, n_repeats, caption, True, img_path)
+          reg_infos.append(info)
+
+        self.reg_dataset_dirs_info[os.path.basename(dir)] = {"n_repeats": n_repeats, "img_count": len(img_paths)}
+
+      print(f"{num_reg_images} reg images.")
+      if num_train_images < num_reg_images:
+        print("some of reg images are not used / 正則化画像の数が多いので、一部使用されない正則化画像があります")
+
+      if num_reg_images == 0:
+        print("no regularization images / 正則化画像が見つかりませんでした")
+      else:
+        # num_repeatsを計算する：どうせ大した数ではないのでループで処理する
+        n = 0
+        first_loop = True
+        while n < num_train_images:
+          for info in reg_infos:
+            if first_loop:
+              self.register_image(info)
+              n += info.num_repeats
+            else:
+              info.num_repeats += 1
+              n += 1
+            if n >= num_train_images:
+              break
+          first_loop = False
+
+    self.num_reg_images = num_reg_images
+
+
+class FineTuningDataset(BaseDataset):
+  def __init__(self, json_file_name, batch_size, train_data_dir, tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens, resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, flip_aug, color_aug, face_crop_aug_range, random_crop, dataset_repeats, debug_dataset) -> None:
+    super().__init__(tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens,
+                     resolution, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset)
+
+    # メタデータを読み込む
+    if os.path.exists(json_file_name):
+      print(f"loading existing metadata: {json_file_name}")
+      with open(json_file_name, "rt", encoding='utf-8') as f:
+        metadata = json.load(f)
+    else:
+      raise ValueError(f"no metadata / メタデータファイルがありません: {json_file_name}")
+
+    self.metadata = metadata
+    self.train_data_dir = train_data_dir
+    self.batch_size = batch_size
+
+    tags_list = []
+    for image_key, img_md in metadata.items():
+      # path情報を作る
+      if os.path.exists(image_key):
+        abs_path = image_key
+      else:
+        # わりといい加減だがいい方法が思いつかん
+        abs_path = glob_images(train_data_dir, image_key)
+        assert len(abs_path) >= 1, f"no image / 画像がありません: {image_key}"
+        abs_path = abs_path[0]
+
+      caption = img_md.get('caption')
+      tags = img_md.get('tags')
+      if caption is None:
+        caption = tags
+      elif tags is not None and len(tags) > 0:
+        caption = caption + ', ' + tags
+        tags_list.append(tags)
+      assert caption is not None and len(caption) > 0, f"caption or tag is required / キャプションまたはタグは必須です:{abs_path}"
+
+      image_info = ImageInfo(image_key, dataset_repeats, caption, False, abs_path)
+      image_info.image_size = img_md.get('train_resolution')
+
+      if not self.color_aug and not self.random_crop:
+        # if npz exists, use them
+        image_info.latents_npz, image_info.latents_npz_flipped = self.image_key_to_npz_file(image_key)
+
+      self.register_image(image_info)
+    self.num_train_images = len(metadata) * dataset_repeats
+    self.num_reg_images = 0
+
+    # TODO do not record tag freq when no tag
+    self.set_tag_frequency(os.path.basename(json_file_name), tags_list)
+    self.dataset_dirs_info[os.path.basename(json_file_name)] = {"n_repeats": dataset_repeats, "img_count": len(metadata)}
+
+    # check existence of all npz files
+    use_npz_latents = not (self.color_aug or self.random_crop)
+    if use_npz_latents:
+      npz_any = False
+      npz_all = True
+      for image_info in self.image_data.values():
+        has_npz = image_info.latents_npz is not None
+        npz_any = npz_any or has_npz
+
+        if self.flip_aug:
+          has_npz = has_npz and image_info.latents_npz_flipped is not None
+        npz_all = npz_all and has_npz
+
+        if npz_any and not npz_all:
+          break
+
+      if not npz_any:
+        use_npz_latents = False
+        print(f"npz file does not exist. ignore npz files / npzファイルが見つからないためnpzファイルを無視します")
+      elif not npz_all:
+        use_npz_latents = False
+        print(f"some of npz file does not exist. ignore npz files / いくつかのnpzファイルが見つからないためnpzファイルを無視します")
+        if self.flip_aug:
+          print("maybe no flipped files / 反転されたnpzファイルがないのかもしれません")
+    # else:
+    #   print("npz files are not used with color_aug and/or random_crop / color_augまたはrandom_cropが指定されているためnpzファイルは使用されません")
+
+    # check min/max bucket size
+    sizes = set()
+    resos = set()
+    for image_info in self.image_data.values():
+      if image_info.image_size is None:
+        sizes = None                  # not calculated
+        break
+      sizes.add(image_info.image_size[0])
+      sizes.add(image_info.image_size[1])
+      resos.add(tuple(image_info.image_size))
+
+    if sizes is None:
+      if use_npz_latents:
+        use_npz_latents = False
+        print(f"npz files exist, but no bucket info in metadata. ignore npz files / メタデータにbucket情報がないためnpzファイルを無視します")
+
+      assert resolution is not None, "if metadata doesn't have bucket info, resolution is required / メタデータにbucket情報がない場合はresolutionを指定してください"
+
+      self.enable_bucket = enable_bucket
+      if self.enable_bucket:
+        self.min_bucket_reso = min_bucket_reso
+        self.max_bucket_reso = max_bucket_reso
+        self.bucket_reso_steps = bucket_reso_steps
+        self.bucket_no_upscale = bucket_no_upscale
+    else:
+      if not enable_bucket:
+        print("metadata has bucket info, enable bucketing / メタデータにbucket情報があるためbucketを有効にします")
+      print("using bucket info in metadata / メタデータ内のbucket情報を使います")
+      self.enable_bucket = True
+
+      assert not bucket_no_upscale, "if metadata has bucket info, bucket reso is precalculated, so bucket_no_upscale cannot be used / メタデータ内にbucket情報がある場合はbucketの解像度は計算済みのため、bucket_no_upscaleは使えません"
+
+      # bucket情報を初期化しておく、make_bucketsで再作成しない
+      self.bucket_manager = BucketManager(False, None, None, None, None)
+      self.bucket_manager.set_predefined_resos(resos)
+
+    # npz情報をきれいにしておく
+    if not use_npz_latents:
+      for image_info in self.image_data.values():
+        image_info.latents_npz = image_info.latents_npz_flipped = None
+
+  def image_key_to_npz_file(self, image_key):
+    base_name = os.path.splitext(image_key)[0]
+    npz_file_norm = base_name + '.npz'
+
+    if os.path.exists(npz_file_norm):
+      # image_key is full path
+      npz_file_flip = base_name + '_flip.npz'
+      if not os.path.exists(npz_file_flip):
+        npz_file_flip = None
+      return npz_file_norm, npz_file_flip
+
+    # image_key is relative path
+    npz_file_norm = os.path.join(self.train_data_dir, image_key + '.npz')
+    npz_file_flip = os.path.join(self.train_data_dir, image_key + '_flip.npz')
+
+    if not os.path.exists(npz_file_norm):
+      npz_file_norm = None
+      npz_file_flip = None
+    elif not os.path.exists(npz_file_flip):
+      npz_file_flip = None
+
+    return npz_file_norm, npz_file_flip
+
+
+def debug_dataset(train_dataset, show_input_ids=False):
+  print(f"Total dataset length (steps) / データセットの長さ（ステップ数）: {len(train_dataset)}")
+  print("Escape for exit. / Escキーで中断、終了します")
+
+  train_dataset.set_current_epoch(1)
+  k = 0
+  for i, example in enumerate(train_dataset):
+    if example['latents'] is not None:
+      print(f"sample has latents from npz file: {example['latents'].size()}")
+    for j, (ik, cap, lw, iid) in enumerate(zip(example['image_keys'], example['captions'], example['loss_weights'], example['input_ids'])):
+      print(f'{ik}, size: {train_dataset.image_data[ik].image_size}, loss weight: {lw}, caption: "{cap}"')
+      if show_input_ids:
+        print(f"input ids: {iid}")
+      if example['images'] is not None:
+        im = example['images'][j]
+        print(f"image size: {im.size()}")
+        im = ((im.numpy() + 1.0) * 127.5).astype(np.uint8)
+        im = np.transpose(im, (1, 2, 0))                # c,H,W -> H,W,c
+        im = im[:, :, ::-1]                             # RGB -> BGR (OpenCV)
+        if os.name == 'nt':                             # only windows
+          cv2.imshow("img", im)
+        k = cv2.waitKey()
+        cv2.destroyAllWindows()
+        if k == 27:
+          break
+    if k == 27 or (example['images'] is None and i >= 8):
+      break
+
+
+def glob_images(directory, base="*"):
+  img_paths = []
+  for ext in IMAGE_EXTENSIONS:
+    if base == '*':
+      img_paths.extend(glob.glob(os.path.join(glob.escape(directory), base + ext)))
+    else:
+      img_paths.extend(glob.glob(glob.escape(os.path.join(directory, base + ext))))
+  # img_paths = list(set(img_paths))                    # 重複を排除
+  # img_paths.sort()
+  return img_paths
+
+
+def glob_images_pathlib(dir_path, recursive):
+  image_paths = []
+  if recursive:
+    for ext in IMAGE_EXTENSIONS:
+      image_paths += list(dir_path.rglob('*' + ext))
+  else:
+    for ext in IMAGE_EXTENSIONS:
+      image_paths += list(dir_path.glob('*' + ext))
+  # image_paths = list(set(image_paths))        # 重複を排除
+  # image_paths.sort()
+  return image_paths
+
+# endregion
+
+
+# region モジュール入れ替え部
+"""
+高速化のためのモジュール入れ替え
+"""
+
+# FlashAttentionを使うCrossAttention
+# based on https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/memory_efficient_attention_pytorch/flash_attention.py
+# LICENSE MIT https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/LICENSE
+
+# constants
+
+EPSILON = 1e-6
+
+# helper functions
+
+
+def exists(val):
+  return val is not None
+
+
+def default(val, d):
+  return val if exists(val) else d
+
+
+def model_hash(filename):
+  """Old model hash used by stable-diffusion-webui"""
+  try:
+    with open(filename, "rb") as file:
+      m = hashlib.sha256()
+
+      file.seek(0x100000)
+      m.update(file.read(0x10000))
+      return m.hexdigest()[0:8]
+  except FileNotFoundError:
+    return 'NOFILE'
+
+
+def calculate_sha256(filename):
+  """New model hash used by stable-diffusion-webui"""
+  hash_sha256 = hashlib.sha256()
+  blksize = 1024 * 1024
+
+  with open(filename, "rb") as f:
+    for chunk in iter(lambda: f.read(blksize), b""):
+      hash_sha256.update(chunk)
+
+  return hash_sha256.hexdigest()
+
+
+def precalculate_safetensors_hashes(tensors, metadata):
+  """Precalculate the model hashes needed by sd-webui-additional-networks to
+  save time on indexing the model later."""
+
+  # Because writing user metadata to the file can change the result of
+  # sd_models.model_hash(), only retain the training metadata for purposes of
+  # calculating the hash, as they are meant to be immutable
+  metadata = {k: v for k, v in metadata.items() if k.startswith("ss_")}
+
+  bytes = safetensors.torch.save(tensors, metadata)
+  b = BytesIO(bytes)
+
+  model_hash = addnet_hash_safetensors(b)
+  legacy_hash = addnet_hash_legacy(b)
+  return model_hash, legacy_hash
+
+
+def addnet_hash_legacy(b):
+  """Old model hash used by sd-webui-additional-networks for .safetensors format files"""
+  m = hashlib.sha256()
+
+  b.seek(0x100000)
+  m.update(b.read(0x10000))
+  return m.hexdigest()[0:8]
+
+
+def addnet_hash_safetensors(b):
+  """New model hash used by sd-webui-additional-networks for .safetensors format files"""
+  hash_sha256 = hashlib.sha256()
+  blksize = 1024 * 1024
+
+  b.seek(0)
+  header = b.read(8)
+  n = int.from_bytes(header, "little")
+
+  offset = n + 8
+  b.seek(offset)
+  for chunk in iter(lambda: b.read(blksize), b""):
+    hash_sha256.update(chunk)
+
+  return hash_sha256.hexdigest()
+
+
+def get_git_revision_hash() -> str:
+  try:
+    return subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=os.path.dirname(__file__)).decode('ascii').strip()
+  except:
+    return "(unknown)"
+
+
+# flash attention forwards and backwards
+
+# https://arxiv.org/abs/2205.14135
+
+
+class FlashAttentionFunction(torch.autograd.function.Function):
+  @ staticmethod
+  @ torch.no_grad()
+  def forward(ctx, q, k, v, mask, causal, q_bucket_size, k_bucket_size):
+    """ Algorithm 2 in the paper """
+
+    device = q.device
+    dtype = q.dtype
+    max_neg_value = -torch.finfo(q.dtype).max
+    qk_len_diff = max(k.shape[-2] - q.shape[-2], 0)
+
+    o = torch.zeros_like(q)
+    all_row_sums = torch.zeros((*q.shape[:-1], 1), dtype=dtype, device=device)
+    all_row_maxes = torch.full((*q.shape[:-1], 1), max_neg_value, dtype=dtype, device=device)
+
+    scale = (q.shape[-1] ** -0.5)
+
+    if not exists(mask):
+      mask = (None,) * math.ceil(q.shape[-2] / q_bucket_size)
+    else:
+      mask = rearrange(mask, 'b n -> b 1 1 n')
+      mask = mask.split(q_bucket_size, dim=-1)
+
+    row_splits = zip(
+        q.split(q_bucket_size, dim=-2),
+        o.split(q_bucket_size, dim=-2),
+        mask,
+        all_row_sums.split(q_bucket_size, dim=-2),
+        all_row_maxes.split(q_bucket_size, dim=-2),
+    )
+
+    for ind, (qc, oc, row_mask, row_sums, row_maxes) in enumerate(row_splits):
+      q_start_index = ind * q_bucket_size - qk_len_diff
+
+      col_splits = zip(
+          k.split(k_bucket_size, dim=-2),
+          v.split(k_bucket_size, dim=-2),
+      )
+
+      for k_ind, (kc, vc) in enumerate(col_splits):
+        k_start_index = k_ind * k_bucket_size
+
+        attn_weights = einsum('... i d, ... j d -> ... i j', qc, kc) * scale
+
+        if exists(row_mask):
+          attn_weights.masked_fill_(~row_mask, max_neg_value)
+
+        if causal and q_start_index < (k_start_index + k_bucket_size - 1):
+          causal_mask = torch.ones((qc.shape[-2], kc.shape[-2]), dtype=torch.bool,
+                                   device=device).triu(q_start_index - k_start_index + 1)
+          attn_weights.masked_fill_(causal_mask, max_neg_value)
+
+        block_row_maxes = attn_weights.amax(dim=-1, keepdims=True)
+        attn_weights -= block_row_maxes
+        exp_weights = torch.exp(attn_weights)
+
+        if exists(row_mask):
+          exp_weights.masked_fill_(~row_mask, 0.)
+
+        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 (not xformers)")
+  flash_func = FlashAttentionFunction
+
+  def forward_flash_attn(self, x, context=None, mask=None):
+    q_bucket_size = 512
+    k_bucket_size = 1024
+
+    h = self.heads
+    q = self.to_q(x)
+
+    context = context if context is not None else x
+    context = context.to(x.dtype)
+
+    if hasattr(self, 'hypernetwork') and self.hypernetwork is not None:
+      context_k, context_v = self.hypernetwork.forward(x, context)
+      context_k = context_k.to(x.dtype)
+      context_v = context_v.to(x.dtype)
+    else:
+      context_k = context
+      context_v = context
+
+    k = self.to_k(context_k)
+    v = self.to_v(context_v)
+    del context, x
+
+    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), (q, k, v))
+
+    out = flash_func.apply(q, k, v, mask, False, q_bucket_size, k_bucket_size)
+
+    out = rearrange(out, 'b h n d -> b n (h d)')
+
+    # diffusers 0.7.0~  わざわざ変えるなよ (;´Д｀)
+    out = self.to_out[0](out)
+    out = self.to_out[1](out)
+    return out
+
+  diffusers.models.attention.CrossAttention.forward = forward_flash_attn
+
+
+def replace_unet_cross_attn_to_xformers():
+  print("Replace CrossAttention.forward to use xformers")
+  try:
+    import xformers.ops
+  except ImportError:
+    raise ImportError("No xformers / xformersがインストールされていないようです")
+
+  def forward_xformers(self, x, context=None, mask=None):
+    h = self.heads
+    q_in = self.to_q(x)
+
+    context = default(context, x)
+    context = context.to(x.dtype)
+
+    if hasattr(self, 'hypernetwork') and self.hypernetwork is not None:
+      context_k, context_v = self.hypernetwork.forward(x, context)
+      context_k = context_k.to(x.dtype)
+      context_v = context_v.to(x.dtype)
+    else:
+      context_k = context
+      context_v = context
+
+    k_in = self.to_k(context_k)
+    v_in = self.to_v(context_v)
+
+    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b n h d', h=h), (q_in, k_in, v_in))
+    del q_in, k_in, v_in
+
+    q = q.contiguous()
+    k = k.contiguous()
+    v = v.contiguous()
+    out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None)        # 最適なのを選んでくれる
+
+    out = rearrange(out, 'b n h d -> b n (h d)', h=h)
+
+    # diffusers 0.7.0~
+    out = self.to_out[0](out)
+    out = self.to_out[1](out)
+    return out
+
+  diffusers.models.attention.CrossAttention.forward = forward_xformers
+# endregion
+
+
+# region arguments
+
+def add_sd_models_arguments(parser: argparse.ArgumentParser):
+  # for pretrained models
+  parser.add_argument("--v2", action='store_true',
+                      help='load Stable Diffusion v2.0 model / Stable Diffusion 2.0のモデルを読み込む')
+  parser.add_argument("--v_parameterization", action='store_true',
+                      help='enable v-parameterization training / v-parameterization学習を有効にする')
+  parser.add_argument("--pretrained_model_name_or_path", type=str, default=None,
+                      help="pretrained model to train, directory to Diffusers model or StableDiffusion checkpoint / 学習元モデル、Diffusers形式モデルのディレクトリまたはStableDiffusionのckptファイル")
+
+
+def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth: bool):
+  parser.add_argument("--output_dir", type=str, default=None,
+                      help="directory to output trained model / 学習後のモデル出力先ディレクトリ")
+  parser.add_argument("--output_name", type=str, default=None,
+                      help="base name of trained model file / 学習後のモデルの拡張子を除くファイル名")
+  parser.add_argument("--save_precision", type=str, default=None,
+                      choices=[None, "float", "fp16", "bf16"], help="precision in saving / 保存時に精度を変更して保存する")
+  parser.add_argument("--save_every_n_epochs", type=int, default=None,
+                      help="save checkpoint every N epochs / 学習中のモデルを指定エポックごとに保存する")
+  parser.add_argument("--save_n_epoch_ratio", type=int, default=None,
+                      help="save checkpoint N epoch ratio (for example 5 means save at least 5 files total) / 学習中のモデルを指定のエポック割合で保存する（たとえば5を指定すると最低5個のファイルが保存される）")
+  parser.add_argument("--save_last_n_epochs", type=int, default=None, help="save last N checkpoints / 最大Nエポック保存する")
+  parser.add_argument("--save_last_n_epochs_state", type=int, default=None,
+                      help="save last N checkpoints of state (overrides the value of --save_last_n_epochs)/ 最大Nエポックstateを保存する(--save_last_n_epochsの指定を上書きします)")
+  parser.add_argument("--save_state", action="store_true",
+                      help="save training state additionally (including optimizer states etc.) / optimizerなど学習状態も含めたstateを追加で保存する")
+  parser.add_argument("--resume", type=str, default=None, help="saved state to resume training / 学習再開するモデルのstate")
+
+  parser.add_argument("--train_batch_size", type=int, default=1, help="batch size for training / 学習時のバッチサイズ")
+  parser.add_argument("--max_token_length", type=int, default=None, choices=[None, 150, 225],
+                      help="max token length of text encoder (default for 75, 150 or 225) / text encoderのトークンの最大長（未指定で75、150または225が指定可）")
+  parser.add_argument("--use_8bit_adam", action="store_true",
+                      help="use 8bit Adam optimizer (requires bitsandbytes) / 8bit Adamオプティマイザを使う（bitsandbytesのインストールが必要）")
+  parser.add_argument("--use_lion_optimizer", action="store_true",
+                      help="use Lion optimizer (requires lion-pytorch) / Lionオプティマイザを使う（ lion-pytorch のインストールが必要）")
+  parser.add_argument("--mem_eff_attn", action="store_true",
+                      help="use memory efficient attention for CrossAttention / CrossAttentionに省メモリ版attentionを使う")
+  parser.add_argument("--xformers", action="store_true",
+                      help="use xformers for CrossAttention / CrossAttentionにxformersを使う")
+  parser.add_argument("--vae", type=str, default=None,
+                      help="path to checkpoint of vae to replace / VAEを入れ替える場合、VAEのcheckpointファイルまたはディレクトリ")
+
+  parser.add_argument("--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("--max_train_epochs", type=int, default=None,
+                      help="training epochs (overrides max_train_steps) / 学習エポック数（max_train_stepsを上書きします）")
+  parser.add_argument("--max_data_loader_n_workers", type=int, default=8,
+                      help="max num workers for DataLoader (lower is less main RAM usage, faster epoch start and slower data loading) / DataLoaderの最大プロセス数（小さい値ではメインメモリの使用量が減りエポック間の待ち時間が減りますが、データ読み込みは遅くなります）")
+  parser.add_argument("--persistent_data_loader_workers", action="store_true",
+                      help="persistent DataLoader workers (useful for reduce time gap between epoch, but may use more memory) / DataLoader のワーカーを持続させる (エポック間の時間差を少なくするのに有効だが、より多くのメモリを消費する可能性がある)")
+  parser.add_argument("--seed", type=int, default=None, help="random seed for training / 学習時の乱数のseed")
+  parser.add_argument("--gradient_checkpointing", action="store_true",
+                      help="enable gradient checkpointing / grandient checkpointingを有効にする")
+  parser.add_argument("--gradient_accumulation_steps", type=int, default=1,
+                      help="Number of updates steps to accumulate before performing a backward/update pass / 学習時に逆伝播をする前に勾配を合計するステップ数")
+  parser.add_argument("--mixed_precision", type=str, default="no",
+                      choices=["no", "fp16", "bf16"], help="use mixed precision / 混合精度を使う場合、その精度")
+  parser.add_argument("--full_fp16", action="store_true", help="fp16 training including gradients / 勾配も含めてfp16で学習する")
+  parser.add_argument("--clip_skip", type=int, default=None,
+                      help="use output of nth layer from back of text encoder (n>=1) / text encoderの後ろからn番目の層の出力を用いる（nは1以上）")
+  parser.add_argument("--logging_dir", type=str, default=None,
+                      help="enable logging and output TensorBoard log to this directory / ログ出力を有効にしてこのディレクトリにTensorBoard用のログを出力する")
+  parser.add_argument("--log_prefix", type=str, default=None, help="add prefix for each log directory / ログディレクトリ名の先頭に追加する文字列")
+  parser.add_argument("--lr_scheduler", type=str, default="constant",
+                      help="scheduler to use for learning rate / 学習率のスケジューラ: linear, cosine, cosine_with_restarts, polynomial, constant (default), constant_with_warmup")
+  parser.add_argument("--lr_warmup_steps", type=int, default=0,
+                      help="Number of steps for the warmup in the lr scheduler (default is 0) / 学習率のスケジューラをウォームアップするステップ数（デフォルト0）")
+  parser.add_argument("--noise_offset", type=float, default=None,
+                      help="enable noise offset with this value (if enabled, around 0.1 is recommended) / Noise offsetを有効にしてこの値を設定する（有効にする場合は0.1程度を推奨）")
+  parser.add_argument("--lowram", action="store_true",
+                      help="enable low RAM optimization. e.g. load models to VRAM instead of RAM (for machines which have bigger VRAM than RAM such as Colab and Kaggle) / メインメモリが少ない環境向け最適化を有効にする。たとえばVRAMにモデルを読み込むなど（ColabやKaggleなどRAMに比べてVRAMが多い環境向け）")
+
+  if support_dreambooth:
+    # DreamBooth training
+    parser.add_argument("--prior_loss_weight", type=float, default=1.0,
+                        help="loss weight for regularization images / 正則化画像のlossの重み")
+
+
+def verify_training_args(args: argparse.Namespace):
+  if args.v_parameterization and not args.v2:
+    print("v_parameterization should be with v2 / v1でv_parameterizationを使用することは想定されていません")
+  if args.v2 and args.clip_skip is not None:
+    print("v2 with clip_skip will be unexpected / v2でclip_skipを使用することは想定されていません")
+
+
+def add_dataset_arguments(parser: argparse.ArgumentParser, support_dreambooth: bool, support_caption: bool, support_caption_dropout: bool):
+  # dataset common
+  parser.add_argument("--train_data_dir", type=str, default=None, help="directory for train images / 学習画像データのディレクトリ")
+  parser.add_argument("--shuffle_caption", action="store_true",
+                      help="shuffle comma-separated caption / コンマで区切られたcaptionの各要素をshuffleする")
+  parser.add_argument("--caption_extension", type=str, default=".caption", help="extension of caption files / 読み込むcaptionファイルの拡張子")
+  parser.add_argument("--caption_extention", type=str, default=None,
+                      help="extension of caption files (backward compatibility) / 読み込むcaptionファイルの拡張子（スペルミスを残してあります）")
+  parser.add_argument("--keep_tokens", type=int, default=None,
+                      help="keep heading N tokens when shuffling caption tokens / captionのシャッフル時に、先頭からこの個数のトークンをシャッフルしないで残す")
+  parser.add_argument("--color_aug", action="store_true", help="enable weak color augmentation / 学習時に色合いのaugmentationを有効にする")
+  parser.add_argument("--flip_aug", action="store_true", help="enable horizontal flip augmentation / 学習時に左右反転のaugmentationを有効にする")
+  parser.add_argument("--face_crop_aug_range", type=str, default=None,
+                      help="enable face-centered crop augmentation and its range (e.g. 2.0,4.0) / 学習時に顔を中心とした切り出しaugmentationを有効にするときは倍率を指定する（例：2.0,4.0）")
+  parser.add_argument("--random_crop", action="store_true",
+                      help="enable random crop (for style training in face-centered crop augmentation) / ランダムな切り出しを有効にする（顔を中心としたaugmentationを行うときに画風の学習用に指定する）")
+  parser.add_argument("--debug_dataset", action="store_true",
+                      help="show images for debugging (do not train) / デバッグ用に学習データを画面表示する（学習は行わない）")
+  parser.add_argument("--resolution", type=str, default=None,
+                      help="resolution in training ('size' or 'width,height') / 学習時の画像解像度（'サイズ'指定、または'幅,高さ'指定）")
+  parser.add_argument("--cache_latents", action="store_true",
+                      help="cache latents to reduce memory (augmentations must be disabled) / メモリ削減のためにlatentをcacheする（augmentationは使用不可）")
+  parser.add_argument("--enable_bucket", action="store_true",
+                      help="enable buckets for multi aspect ratio training / 複数解像度学習のためのbucketを有効にする")
+  parser.add_argument("--min_bucket_reso", type=int, default=256, help="minimum resolution for buckets / bucketの最小解像度")
+  parser.add_argument("--max_bucket_reso", type=int, default=1024, help="maximum resolution for buckets / bucketの最大解像度")
+  parser.add_argument("--bucket_reso_steps", type=int, default=64,
+                      help="steps of resolution for buckets, divisible by 8 is recommended / bucketの解像度の単位、8で割り切れる値を推奨します")
+  parser.add_argument("--bucket_no_upscale", action="store_true",
+                      help="make bucket for each image without upscaling / 画像を拡大せずbucketを作成します")
+
+  if support_caption_dropout:
+    # Textual Inversion はcaptionのdropoutをsupportしない
+    # いわゆるtensorのDropoutと紛らわしいのでprefixにcaptionを付けておく　every_n_epochsは他と平仄を合わせてdefault Noneに
+    parser.add_argument("--caption_dropout_rate", type=float, default=0,
+                        help="Rate out dropout caption(0.0~1.0) / captionをdropoutする割合")
+    parser.add_argument("--caption_dropout_every_n_epochs", type=int, default=None,
+                        help="Dropout all captions every N epochs / captionを指定エポックごとにdropoutする")
+    parser.add_argument("--caption_tag_dropout_rate", type=float, default=0,
+                        help="Rate out dropout comma separated tokens(0.0~1.0) / カンマ区切りのタグをdropoutする割合")
+
+  if support_dreambooth:
+    # DreamBooth dataset
+    parser.add_argument("--reg_data_dir", type=str, default=None, help="directory for regularization images / 正則化画像データのディレクトリ")
+
+  if support_caption:
+    # caption dataset
+    parser.add_argument("--in_json", type=str, default=None, help="json metadata for dataset / データセットのmetadataのjsonファイル")
+    parser.add_argument("--dataset_repeats", type=int, default=1,
+                        help="repeat dataset when training with captions / キャプションでの学習時にデータセットを繰り返す回数")
+
+
+def add_sd_saving_arguments(parser: argparse.ArgumentParser):
+  parser.add_argument("--save_model_as", type=str, default=None, choices=[None, "ckpt", "safetensors", "diffusers", "diffusers_safetensors"],
+                      help="format to save the model (default is same to original) / モデル保存時の形式（未指定時は元モデルと同じ）")
+  parser.add_argument("--use_safetensors", action='store_true',
+                      help="use safetensors format to save (if save_model_as is not specified) / checkpoint、モデルをsafetensors形式で保存する（save_model_as未指定時）")
+
+# endregion
+
+# region utils
+
+
+def prepare_dataset_args(args: argparse.Namespace, support_metadata: bool):
+  # backward compatibility
+  if args.caption_extention is not None:
+    args.caption_extension = args.caption_extention
+    args.caption_extention = None
+
+  if args.cache_latents:
+    assert not args.color_aug, "when caching latents, color_aug cannot be used / latentをキャッシュするときはcolor_augは使えません"
+    assert not args.random_crop, "when caching latents, random_crop cannot be used / latentをキャッシュするときはrandom_cropは使えません"
+
+  # assert args.resolution is not None, f"resolution is required / resolution（解像度）を指定してください"
+  if args.resolution is not None:
+    args.resolution = tuple([int(r) for r in args.resolution.split(',')])
+    if len(args.resolution) == 1:
+      args.resolution = (args.resolution[0], args.resolution[0])
+    assert len(args.resolution) == 2, \
+        f"resolution must be 'size' or 'width,height' / resolution（解像度）は'サイズ'または'幅','高さ'で指定してください: {args.resolution}"
+
+  if args.face_crop_aug_range is not None:
+    args.face_crop_aug_range = tuple([float(r) for r in args.face_crop_aug_range.split(',')])
+    assert len(args.face_crop_aug_range) == 2 and args.face_crop_aug_range[0] <= args.face_crop_aug_range[1], \
+        f"face_crop_aug_range must be two floats / face_crop_aug_rangeは'下限,上限'で指定してください: {args.face_crop_aug_range}"
+  else:
+    args.face_crop_aug_range = None
+
+  if support_metadata:
+    if args.in_json is not None and (args.color_aug or args.random_crop):
+      print(f"latents in npz is ignored when color_aug or random_crop is True / color_augまたはrandom_cropを有効にした場合、npzファイルのlatentsは無視されます")
+
+
+def load_tokenizer(args: argparse.Namespace):
+  print("prepare tokenizer")
+  if args.v2:
+    tokenizer = CLIPTokenizer.from_pretrained(V2_STABLE_DIFFUSION_PATH, subfolder="tokenizer")
+  else:
+    tokenizer = CLIPTokenizer.from_pretrained(TOKENIZER_PATH)
+  if args.max_token_length is not None:
+    print(f"update token length: {args.max_token_length}")
+  return tokenizer
+
+
+def prepare_accelerator(args: argparse.Namespace):
+  if args.logging_dir is None:
+    log_with = None
+    logging_dir = None
+  else:
+    log_with = "tensorboard"
+    log_prefix = "" if args.log_prefix is None else args.log_prefix
+    logging_dir = args.logging_dir + "/" + log_prefix + time.strftime('%Y%m%d%H%M%S', time.localtime())
+
+  accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps, mixed_precision=args.mixed_precision,
+                            log_with=log_with, logging_dir=logging_dir)
+
+  # accelerateの互換性問題を解決する
+  accelerator_0_15 = True
+  try:
+    accelerator.unwrap_model("dummy", True)
+    print("Using accelerator 0.15.0 or above.")
+  except TypeError:
+    accelerator_0_15 = False
+
+  def unwrap_model(model):
+    if accelerator_0_15:
+      return accelerator.unwrap_model(model, True)
+    return accelerator.unwrap_model(model)
+
+  return accelerator, unwrap_model
+
+
+def prepare_dtype(args: argparse.Namespace):
+  weight_dtype = torch.float32
+  if args.mixed_precision == "fp16":
+    weight_dtype = torch.float16
+  elif args.mixed_precision == "bf16":
+    weight_dtype = torch.bfloat16
+
+  save_dtype = None
+  if args.save_precision == "fp16":
+    save_dtype = torch.float16
+  elif args.save_precision == "bf16":
+    save_dtype = torch.bfloat16
+  elif args.save_precision == "float":
+    save_dtype = torch.float32
+
+  return weight_dtype, save_dtype
+
+
+def load_target_model(args: argparse.Namespace, weight_dtype):
+  load_stable_diffusion_format = os.path.isfile(args.pretrained_model_name_or_path)           # determine SD or Diffusers
+  if load_stable_diffusion_format:
+    print("load StableDiffusion checkpoint")
+    text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.pretrained_model_name_or_path)
+  else:
+    print("load Diffusers pretrained models")
+    pipe = StableDiffusionPipeline.from_pretrained(args.pretrained_model_name_or_path, tokenizer=None, safety_checker=None)
+    text_encoder = pipe.text_encoder
+    vae = pipe.vae
+    unet = pipe.unet
+    del pipe
+
+  # VAEを読み込む
+  if args.vae is not None:
+    vae = model_util.load_vae(args.vae, weight_dtype)
+    print("additional VAE loaded")
+
+  return text_encoder, vae, unet, load_stable_diffusion_format
+
+
+def patch_accelerator_for_fp16_training(accelerator):
+  org_unscale_grads = accelerator.scaler._unscale_grads_
+
+  def _unscale_grads_replacer(optimizer, inv_scale, found_inf, allow_fp16):
+    return org_unscale_grads(optimizer, inv_scale, found_inf, True)
+
+  accelerator.scaler._unscale_grads_ = _unscale_grads_replacer
+
+
+def get_hidden_states(args: argparse.Namespace, input_ids, tokenizer, text_encoder, weight_dtype=None):
+  # with no_token_padding, the length is not max length, return result immediately
+  if input_ids.size()[-1] != tokenizer.model_max_length:
+    return text_encoder(input_ids)[0]
+
+  b_size = input_ids.size()[0]
+  input_ids = input_ids.reshape((-1, tokenizer.model_max_length))     # batch_size*3, 77
+
+  if args.clip_skip is None:
+    encoder_hidden_states = text_encoder(input_ids)[0]
+  else:
+    enc_out = text_encoder(input_ids, output_hidden_states=True, return_dict=True)
+    encoder_hidden_states = enc_out['hidden_states'][-args.clip_skip]
+    encoder_hidden_states = text_encoder.text_model.final_layer_norm(encoder_hidden_states)
+
+  # bs*3, 77, 768 or 1024
+  encoder_hidden_states = encoder_hidden_states.reshape((b_size, -1, encoder_hidden_states.shape[-1]))
+
+  if args.max_token_length is not None:
+    if args.v2:
+        # v2: <BOS>...<EOS> <PAD> ... の三連を <BOS>...<EOS> <PAD> ... へ戻す　正直この実装でいいのかわからん
+      states_list = [encoder_hidden_states[:, 0].unsqueeze(1)]                              # <BOS>
+      for i in range(1, args.max_token_length, tokenizer.model_max_length):
+        chunk = encoder_hidden_states[:, i:i + tokenizer.model_max_length - 2]              # <BOS> の後から 最後の前まで
+        if i > 0:
+          for j in range(len(chunk)):
+            if input_ids[j, 1] == tokenizer.eos_token:                                      # 空、つまり <BOS> <EOS> <PAD> ...のパターン
+              chunk[j, 0] = chunk[j, 1]                                                     # 次の <PAD> の値をコピーする
+        states_list.append(chunk)  # <BOS> の後から <EOS> の前まで
+      states_list.append(encoder_hidden_states[:, -1].unsqueeze(1))                         # <EOS> か <PAD> のどちらか
+      encoder_hidden_states = torch.cat(states_list, dim=1)
+    else:
+      # v1: <BOS>...<EOS> の三連を <BOS>...<EOS> へ戻す
+      states_list = [encoder_hidden_states[:, 0].unsqueeze(1)]                              # <BOS>
+      for i in range(1, args.max_token_length, tokenizer.model_max_length):
+        states_list.append(encoder_hidden_states[:, i:i + tokenizer.model_max_length - 2])  # <BOS> の後から <EOS> の前まで
+      states_list.append(encoder_hidden_states[:, -1].unsqueeze(1))                         # <EOS>
+      encoder_hidden_states = torch.cat(states_list, dim=1)
+
+  if weight_dtype is not None:
+    # this is required for additional network training
+    encoder_hidden_states = encoder_hidden_states.to(weight_dtype)
+
+  return encoder_hidden_states
+
+
+def get_epoch_ckpt_name(args: argparse.Namespace, use_safetensors, epoch):
+  model_name = DEFAULT_EPOCH_NAME if args.output_name is None else args.output_name
+  ckpt_name = EPOCH_FILE_NAME.format(model_name, epoch) + (".safetensors" if use_safetensors else ".ckpt")
+  return model_name, ckpt_name
+
+
+def save_on_epoch_end(args: argparse.Namespace, save_func, remove_old_func, epoch_no: int, num_train_epochs: int):
+  saving = epoch_no % args.save_every_n_epochs == 0 and epoch_no < num_train_epochs
+  if saving:
+    os.makedirs(args.output_dir, exist_ok=True)
+    save_func()
+
+    if args.save_last_n_epochs is not None:
+      remove_epoch_no = epoch_no - args.save_every_n_epochs * args.save_last_n_epochs
+      remove_old_func(remove_epoch_no)
+  return saving
+
+
+def save_sd_model_on_epoch_end(args: argparse.Namespace, accelerator, src_path: str, save_stable_diffusion_format: bool, use_safetensors: bool, save_dtype: torch.dtype, epoch: int, num_train_epochs: int, global_step: int, text_encoder, unet, vae):
+  epoch_no = epoch + 1
+  model_name, ckpt_name = get_epoch_ckpt_name(args, use_safetensors, epoch_no)
+
+  if save_stable_diffusion_format:
+    def save_sd():
+      ckpt_file = os.path.join(args.output_dir, ckpt_name)
+      print(f"saving checkpoint: {ckpt_file}")
+      model_util.save_stable_diffusion_checkpoint(args.v2, ckpt_file, text_encoder, unet,
+                                                  src_path, epoch_no, global_step, save_dtype, vae)
+
+    def remove_sd(old_epoch_no):
+      _, old_ckpt_name = get_epoch_ckpt_name(args,  use_safetensors, old_epoch_no)
+      old_ckpt_file = os.path.join(args.output_dir, old_ckpt_name)
+      if os.path.exists(old_ckpt_file):
+        print(f"removing old checkpoint: {old_ckpt_file}")
+        os.remove(old_ckpt_file)
+
+    save_func = save_sd
+    remove_old_func = remove_sd
+  else:
+    def save_du():
+      out_dir = os.path.join(args.output_dir, EPOCH_DIFFUSERS_DIR_NAME.format(model_name, epoch_no))
+      print(f"saving model: {out_dir}")
+      os.makedirs(out_dir, exist_ok=True)
+      model_util.save_diffusers_checkpoint(args.v2, out_dir, text_encoder, unet,
+                                           src_path, vae=vae, use_safetensors=use_safetensors)
+
+    def remove_du(old_epoch_no):
+      out_dir_old = os.path.join(args.output_dir, EPOCH_DIFFUSERS_DIR_NAME.format(model_name, old_epoch_no))
+      if os.path.exists(out_dir_old):
+        print(f"removing old model: {out_dir_old}")
+        shutil.rmtree(out_dir_old)
+
+    save_func = save_du
+    remove_old_func = remove_du
+
+  saving = save_on_epoch_end(args, save_func, remove_old_func, epoch_no, num_train_epochs)
+  if saving and args.save_state:
+    save_state_on_epoch_end(args, accelerator, model_name, epoch_no)
+
+
+def save_state_on_epoch_end(args: argparse.Namespace, accelerator, model_name, epoch_no):
+  print("saving state.")
+  accelerator.save_state(os.path.join(args.output_dir, EPOCH_STATE_NAME.format(model_name, epoch_no)))
+
+  last_n_epochs = args.save_last_n_epochs_state if args.save_last_n_epochs_state else args.save_last_n_epochs
+  if last_n_epochs is not None:
+    remove_epoch_no = epoch_no - args.save_every_n_epochs * last_n_epochs
+    state_dir_old = os.path.join(args.output_dir, EPOCH_STATE_NAME.format(model_name, remove_epoch_no))
+    if os.path.exists(state_dir_old):
+      print(f"removing old state: {state_dir_old}")
+      shutil.rmtree(state_dir_old)
+
+
+def save_sd_model_on_train_end(args: argparse.Namespace, src_path: str, save_stable_diffusion_format: bool, use_safetensors: bool, save_dtype: torch.dtype, epoch: int, global_step: int, text_encoder, unet, vae):
+  model_name = DEFAULT_LAST_OUTPUT_NAME if args.output_name is None else args.output_name
+
+  if save_stable_diffusion_format:
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    ckpt_name = model_name + (".safetensors" if use_safetensors else ".ckpt")
+    ckpt_file = os.path.join(args.output_dir, ckpt_name)
+
+    print(f"save trained model as StableDiffusion checkpoint to {ckpt_file}")
+    model_util.save_stable_diffusion_checkpoint(args.v2, ckpt_file, text_encoder, unet,
+                                                src_path, epoch, global_step, save_dtype, vae)
+  else:
+    out_dir = os.path.join(args.output_dir, model_name)
+    os.makedirs(out_dir, exist_ok=True)
+
+    print(f"save trained model as Diffusers to {out_dir}")
+    model_util.save_diffusers_checkpoint(args.v2, out_dir, text_encoder, unet,
+                                         src_path, vae=vae, use_safetensors=use_safetensors)
+
+
+def save_state_on_train_end(args: argparse.Namespace, accelerator):
+  print("saving last state.")
+  os.makedirs(args.output_dir, exist_ok=True)
+  model_name = DEFAULT_LAST_OUTPUT_NAME if args.output_name is None else args.output_name
+  accelerator.save_state(os.path.join(args.output_dir, LAST_STATE_NAME.format(model_name)))
+
+# endregion
+
+# region 前処理用
+
+
+class ImageLoadingDataset(torch.utils.data.Dataset):
+  def __init__(self, image_paths):
+    self.images = image_paths
+
+  def __len__(self):
+    return len(self.images)
+
+  def __getitem__(self, idx):
+    img_path = self.images[idx]
+
+    try:
+      image = Image.open(img_path).convert("RGB")
+      # convert to tensor temporarily so dataloader will accept it
+      tensor_pil = transforms.functional.pil_to_tensor(image)
+    except Exception as e:
+      print(f"Could not load image path / 画像を読み込めません: {img_path}, error: {e}")
+      return None
+
+    return (tensor_pil, img_path)
+
+
+# endregion

fine_tune.py

@@ -0,0 +1,360 @@
+# training with captions
+# XXX dropped option: hypernetwork training
+
+import argparse
+import gc
+import math
+import os
+
+from tqdm import tqdm
+import torch
+from accelerate.utils import set_seed
+import diffusers
+from diffusers import DDPMScheduler
+
+import library.train_util as train_util
+
+
+def collate_fn(examples):
+  return examples[0]
+
+
+def train(args):
+  train_util.verify_training_args(args)
+  train_util.prepare_dataset_args(args, True)
+
+  cache_latents = args.cache_latents
+
+  if args.seed is not None:
+    set_seed(args.seed)                           # 乱数系列を初期化する
+
+  tokenizer = train_util.load_tokenizer(args)
+
+  train_dataset = train_util.FineTuningDataset(args.in_json, args.train_batch_size, args.train_data_dir,
+                                               tokenizer, args.max_token_length, args.shuffle_caption, args.keep_tokens,
+                                               args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
+                                               args.bucket_reso_steps, args.bucket_no_upscale,
+                                               args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop,
+                                               args.dataset_repeats, args.debug_dataset)
+
+  # 学習データのdropout率を設定する
+  train_dataset.set_caption_dropout(args.caption_dropout_rate, args.caption_dropout_every_n_epochs, args.caption_tag_dropout_rate)
+
+  train_dataset.make_buckets()
+
+  if args.debug_dataset:
+    train_util.debug_dataset(train_dataset)
+    return
+  if len(train_dataset) == 0:
+    print("No data found. Please verify the metadata file and train_data_dir option. / 画像がありません。メタデータおよびtrain_data_dirオプションを確認してください。")
+    return
+
+  # acceleratorを準備する
+  print("prepare accelerator")
+  accelerator, unwrap_model = train_util.prepare_accelerator(args)
+
+  # mixed precisionに対応した型を用意しておき適宜castする
+  weight_dtype, save_dtype = train_util.prepare_dtype(args)
+
+  # モデルを読み込む
+  text_encoder, vae, unet, load_stable_diffusion_format = train_util.load_target_model(args, weight_dtype)
+
+  # verify load/save model formats
+  if load_stable_diffusion_format:
+    src_stable_diffusion_ckpt = args.pretrained_model_name_or_path
+    src_diffusers_model_path = None
+  else:
+    src_stable_diffusion_ckpt = None
+    src_diffusers_model_path = args.pretrained_model_name_or_path
+
+  if args.save_model_as is None:
+    save_stable_diffusion_format = load_stable_diffusion_format
+    use_safetensors = args.use_safetensors
+  else:
+    save_stable_diffusion_format = args.save_model_as.lower() == 'ckpt' or args.save_model_as.lower() == 'safetensors'
+    use_safetensors = args.use_safetensors or ("safetensors" in args.save_model_as.lower())
+
+  # Diffusers版のxformers使用フラグを設定する関数
+  def set_diffusers_xformers_flag(model, valid):
+    #   model.set_use_memory_efficient_attention_xformers(valid)            # 次のリリースでなくなりそう
+    # pipeが自動で再帰的にset_use_memory_efficient_attention_xformersを探すんだって(;´Д｀)
+    # U-Netだけ使う時にはどうすればいいのか……仕方ないからコピって使うか
+    # 0.10.2でなんか巻き戻って個別に指定するようになった(;^ω^)
+
+    # Recursively walk through all the children.
+    # Any children which exposes the set_use_memory_efficient_attention_xformers method
+    # gets the message
+    def fn_recursive_set_mem_eff(module: torch.nn.Module):
+      if hasattr(module, "set_use_memory_efficient_attention_xformers"):
+        module.set_use_memory_efficient_attention_xformers(valid)
+
+      for child in module.children():
+        fn_recursive_set_mem_eff(child)
+
+    fn_recursive_set_mem_eff(model)
+
+  # モデルに xformers とか memory efficient attention を組み込む
+  if args.diffusers_xformers:
+    print("Use xformers by Diffusers")
+    set_diffusers_xformers_flag(unet, True)
+  else:
+    # Windows版のxformersはfloatで学習できないのでxformersを使わない設定も可能にしておく必要がある
+    print("Disable Diffusers' xformers")
+    set_diffusers_xformers_flag(unet, False)
+    train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
+
+  # 学習を準備する
+  if cache_latents:
+    vae.to(accelerator.device, dtype=weight_dtype)
+    vae.requires_grad_(False)
+    vae.eval()
+    with torch.no_grad():
+      train_dataset.cache_latents(vae)
+    vae.to("cpu")
+    if torch.cuda.is_available():
+      torch.cuda.empty_cache()
+    gc.collect()
+
+  # 学習を準備する：モデルを適切な状態にする
+  training_models = []
+  if args.gradient_checkpointing:
+    unet.enable_gradient_checkpointing()
+  training_models.append(unet)
+
+  if args.train_text_encoder:
+    print("enable text encoder training")
+    if args.gradient_checkpointing:
+      text_encoder.gradient_checkpointing_enable()
+    training_models.append(text_encoder)
+  else:
+    text_encoder.to(accelerator.device, dtype=weight_dtype)
+    text_encoder.requires_grad_(False)             # text encoderは学習しない
+    if args.gradient_checkpointing:
+      text_encoder.gradient_checkpointing_enable()
+      text_encoder.train()                # required for gradient_checkpointing
+    else:
+      text_encoder.eval()
+
+  if not cache_latents:
+    vae.requires_grad_(False)
+    vae.eval()
+    vae.to(accelerator.device, dtype=weight_dtype)
+
+  for m in training_models:
+    m.requires_grad_(True)
+  params = []
+  for m in training_models:
+    params.extend(m.parameters())
+  params_to_optimize = params
+
+  # 学習に必要なクラスを準備する
+  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
+  elif args.use_lion_optimizer:
+    try:
+      import lion_pytorch
+    except ImportError:
+      raise ImportError("No lion_pytorch / lion_pytorch がインストールされていないようです")
+    print("use Lion optimizer")
+    optimizer_class = lion_pytorch.Lion
+  else:
+    optimizer_class = torch.optim.AdamW
+
+  # betaやweight decayはdiffusers DreamBoothもDreamBooth SDもデフォルト値のようなのでオプションはとりあえず省略
+  optimizer = optimizer_class(params_to_optimize, lr=args.learning_rate)
+
+  # dataloaderを準備する
+  # DataLoaderのプロセス数：0はメインプロセスになる
+  n_workers = min(args.max_data_loader_n_workers, os.cpu_count() - 1)      # cpu_count-1 ただし最大で指定された数まで
+  train_dataloader = torch.utils.data.DataLoader(
+      train_dataset, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=n_workers, persistent_workers=args.persistent_data_loader_workers)
+
+  # 学習ステップ数を計算する
+  if args.max_train_epochs is not None:
+    args.max_train_steps = args.max_train_epochs * len(train_dataloader)
+    print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
+
+  # lr schedulerを用意する
+  lr_scheduler = diffusers.optimization.get_scheduler(
+      args.lr_scheduler, optimizer, num_warmup_steps=args.lr_warmup_steps, num_training_steps=args.max_train_steps * args.gradient_accumulation_steps)
+
+  # 実験的機能：勾配も含めたfp16学習を行う　モデル全体をfp16にする
+  if args.full_fp16:
+    assert args.mixed_precision == "fp16", "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
+    print("enable full fp16 training.")
+    unet.to(weight_dtype)
+    text_encoder.to(weight_dtype)
+
+  # acceleratorがなんかよろしくやってくれるらしい
+  if args.train_text_encoder:
+    unet, text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+        unet, text_encoder, optimizer, train_dataloader, lr_scheduler)
+  else:
+    unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(unet, optimizer, train_dataloader, lr_scheduler)
+
+  # 実験的機能：勾配も含めたfp16学習を行う　PyTorchにパッチを当ててfp16でのgrad scaleを有効にする
+  if args.full_fp16:
+    train_util.patch_accelerator_for_fp16_training(accelerator)
+
+  # resumeする
+  if args.resume is not None:
+    print(f"resume training from state: {args.resume}")
+    accelerator.load_state(args.resume)
+
+  # 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)
+  if (args.save_n_epoch_ratio is not None) and (args.save_n_epoch_ratio > 0):
+    args.save_every_n_epochs = math.floor(num_train_epochs / args.save_n_epoch_ratio) or 1
+
+  # 学習する
+  total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+  print("running training / 学習開始")
+  print(f"  num examples / サンプル数: {train_dataset.num_train_images}")
+  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 & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {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), smoothing=0, 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, clip_sample=False)
+
+  if accelerator.is_main_process:
+    accelerator.init_trackers("finetuning")
+
+  for epoch in range(num_train_epochs):
+    print(f"epoch {epoch+1}/{num_train_epochs}")
+    train_dataset.set_current_epoch(epoch + 1)
+
+    for m in training_models:
+      m.train()
+
+    loss_total = 0
+    for step, batch in enumerate(train_dataloader):
+      with accelerator.accumulate(training_models[0]):  # 複数モデルに対応していない模様だがとりあえずこうしておく
+        with torch.no_grad():
+          if "latents" in batch and batch["latents"] is not None:
+            latents = batch["latents"].to(accelerator.device)
+          else:
+            # latentに変換
+            latents = vae.encode(batch["images"].to(dtype=weight_dtype)).latent_dist.sample()
+          latents = latents * 0.18215
+        b_size = latents.shape[0]
+
+        with torch.set_grad_enabled(args.train_text_encoder):
+          # Get the text embedding for conditioning
+          input_ids = batch["input_ids"].to(accelerator.device)
+          encoder_hidden_states = train_util.get_hidden_states(
+              args, input_ids, tokenizer, text_encoder, None if not args.full_fp16 else weight_dtype)
+
+        # Sample noise that we'll add to the latents
+        noise = torch.randn_like(latents, device=latents.device)
+        if args.noise_offset:
+          # https://www.crosslabs.org//blog/diffusion-with-offset-noise
+          noise += args.noise_offset * torch.randn((latents.shape[0], latents.shape[1], 1, 1), 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
+
+        if args.v_parameterization:
+          # v-parameterization training
+          target = noise_scheduler.get_velocity(latents, noise, timesteps)
+        else:
+          target = noise
+
+        loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="mean")
+
+        accelerator.backward(loss)
+        if accelerator.sync_gradients:
+          params_to_clip = []
+          for m in training_models:
+            params_to_clip.extend(m.parameters())
+          accelerator.clip_grad_norm_(params_to_clip, 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()        # 平均なのでbatch sizeは関係ないはず
+      if args.logging_dir is not None:
+        logs = {"loss": current_loss, "lr": lr_scheduler.get_last_lr()[0]}
+        accelerator.log(logs, step=global_step)
+
+      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)
+
+      if global_step >= args.max_train_steps:
+        break
+
+    if args.logging_dir is not None:
+      logs = {"epoch_loss": loss_total / len(train_dataloader)}
+      accelerator.log(logs, step=epoch+1)
+
+    accelerator.wait_for_everyone()
+
+    if args.save_every_n_epochs is not None:
+      src_path = src_stable_diffusion_ckpt if save_stable_diffusion_format else src_diffusers_model_path
+      train_util.save_sd_model_on_epoch_end(args, accelerator, src_path, save_stable_diffusion_format, use_safetensors,
+                                            save_dtype, epoch, num_train_epochs, global_step,  unwrap_model(text_encoder), unwrap_model(unet), vae)
+
+  is_main_process = accelerator.is_main_process
+  if is_main_process:
+    unet = unwrap_model(unet)
+    text_encoder = unwrap_model(text_encoder)
+
+  accelerator.end_training()
+
+  if args.save_state:
+    train_util.save_state_on_train_end(args, accelerator)
+
+  del accelerator                         # この後メモリを使うのでこれは消す
+
+  if is_main_process:
+    src_path = src_stable_diffusion_ckpt if save_stable_diffusion_format else src_diffusers_model_path
+    train_util.save_sd_model_on_train_end(args, src_path, save_stable_diffusion_format, use_safetensors,
+                                          save_dtype, epoch, global_step,  text_encoder, unet, vae)
+    print("model saved.")
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+
+  train_util.add_sd_models_arguments(parser)
+  train_util.add_dataset_arguments(parser, False, True, True)
+  train_util.add_training_arguments(parser, False)
+  train_util.add_sd_saving_arguments(parser)
+
+  parser.add_argument("--diffusers_xformers", action='store_true',
+                      help='use xformers by diffusers / Diffusersでxformersを使用する')
+  parser.add_argument("--train_text_encoder", action="store_true", help="train text encoder / text encoderも学習する")
+
+  args = parser.parse_args()
+  train(args)

library.egg-info/PKG-INFO

@@ -0,0 +1,4 @@
+Metadata-Version: 2.1
+Name: library
+Version: 0.0.0
+License-File: LICENSE.md

library.egg-info/SOURCES.txt

@@ -0,0 +1,10 @@
+LICENSE.md
+README.md
+setup.py
+library/__init__.py
+library/model_util.py
+library/train_util.py
+library.egg-info/PKG-INFO
+library.egg-info/SOURCES.txt
+library.egg-info/dependency_links.txt
+library.egg-info/top_level.txt

library.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

library.egg-info/top_level.txt

@@ -0,0 +1 @@
+library

library/model_util.py

@@ -0,0 +1,1180 @@
+# v1: split from train_db_fixed.py.
+# v2: support safetensors
+
+import math
+import os
+import torch
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig
+from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNet2DConditionModel
+from safetensors.torch import load_file, save_file
+
+# DiffUsers版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 = 64  # fixed from old invalid value `32`
+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
+
+# V2
+V2_UNET_PARAMS_ATTENTION_HEAD_DIM = [5, 10, 20, 20]
+V2_UNET_PARAMS_CONTEXT_DIM = 1024
+
+# Diffusersの設定を読み込むための参照モデル
+DIFFUSERS_REF_MODEL_ID_V1 = "runwayml/stable-diffusion-v1-5"
+DIFFUSERS_REF_MODEL_ID_V2 = "stabilityai/stable-diffusion-2-1"
+
+
+# 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 linear_transformer_to_conv(checkpoint):
+  keys = list(checkpoint.keys())
+  tf_keys = ["proj_in.weight", "proj_out.weight"]
+  for key in keys:
+    if ".".join(key.split(".")[-2:]) in tf_keys:
+      if checkpoint[key].ndim == 2:
+        checkpoint[key] = checkpoint[key].unsqueeze(2).unsqueeze(2)
+
+
+def convert_ldm_unet_checkpoint(v2, 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"])
+
+      # biasとweightの順番に依存しないようにする：もっといいやり方がありそうだが
+      for l in output_block_list.values():
+        l.sort()
+
+      if ["conv.bias", "conv.weight"] in output_block_list.values():
+        index = list(output_block_list.values()).index(["conv.bias", "conv.weight"])
+        new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[
+            f"output_blocks.{i}.{index}.conv.bias"
+        ]
+        new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[
+            f"output_blocks.{i}.{index}.conv.weight"
+        ]
+
+        # 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]
+
+  # SDのv2では1*1のconv2dがlinearに変わっているので、linear->convに変換する
+  if v2:
+    linear_transformer_to_conv(new_checkpoint)
+
+  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)
+  # if len(vae_state_dict) == 0:
+  #   # 渡されたcheckpointは.ckptから読み込んだcheckpointではなくvaeのstate_dict
+  #   vae_state_dict = checkpoint
+
+  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(v2):
+  """
+  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 if not v2 else V2_UNET_PARAMS_CONTEXT_DIM,
+      attention_head_dim=UNET_PARAMS_NUM_HEADS if not v2 else V2_UNET_PARAMS_ATTENTION_HEAD_DIM,
+  )
+
+  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_v1(checkpoint):
+  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]
+  return text_model_dict
+
+
+def convert_ldm_clip_checkpoint_v2(checkpoint, max_length):
+  # 嫌になるくらい違うぞ！
+  def convert_key(key):
+    if not key.startswith("cond_stage_model"):
+      return None
+
+    # common conversion
+    key = key.replace("cond_stage_model.model.transformer.", "text_model.encoder.")
+    key = key.replace("cond_stage_model.model.", "text_model.")
+
+    if "resblocks" in key:
+      # resblocks conversion
+      key = key.replace(".resblocks.", ".layers.")
+      if ".ln_" in key:
+        key = key.replace(".ln_", ".layer_norm")
+      elif ".mlp." in key:
+        key = key.replace(".c_fc.", ".fc1.")
+        key = key.replace(".c_proj.", ".fc2.")
+      elif '.attn.out_proj' in key:
+        key = key.replace(".attn.out_proj.", ".self_attn.out_proj.")
+      elif '.attn.in_proj' in key:
+        key = None                  # 特殊なので後で処理する
+      else:
+        raise ValueError(f"unexpected key in SD: {key}")
+    elif '.positional_embedding' in key:
+      key = key.replace(".positional_embedding", ".embeddings.position_embedding.weight")
+    elif '.text_projection' in key:
+      key = None    # 使われない???
+    elif '.logit_scale' in key:
+      key = None    # 使われない???
+    elif '.token_embedding' in key:
+      key = key.replace(".token_embedding.weight", ".embeddings.token_embedding.weight")
+    elif '.ln_final' in key:
+      key = key.replace(".ln_final", ".final_layer_norm")
+    return key
+
+  keys = list(checkpoint.keys())
+  new_sd = {}
+  for key in keys:
+    # remove resblocks 23
+    if '.resblocks.23.' in key:
+      continue
+    new_key = convert_key(key)
+    if new_key is None:
+      continue
+    new_sd[new_key] = checkpoint[key]
+
+  # attnの変換
+  for key in keys:
+    if '.resblocks.23.' in key:
+      continue
+    if '.resblocks' in key and '.attn.in_proj_' in key:
+      # 三つに分割
+      values = torch.chunk(checkpoint[key], 3)
+
+      key_suffix = ".weight" if "weight" in key else ".bias"
+      key_pfx = key.replace("cond_stage_model.model.transformer.resblocks.", "text_model.encoder.layers.")
+      key_pfx = key_pfx.replace("_weight", "")
+      key_pfx = key_pfx.replace("_bias", "")
+      key_pfx = key_pfx.replace(".attn.in_proj", ".self_attn.")
+      new_sd[key_pfx + "q_proj" + key_suffix] = values[0]
+      new_sd[key_pfx + "k_proj" + key_suffix] = values[1]
+      new_sd[key_pfx + "v_proj" + key_suffix] = values[2]
+
+  # rename or add position_ids
+  ANOTHER_POSITION_IDS_KEY = "text_model.encoder.text_model.embeddings.position_ids"
+  if ANOTHER_POSITION_IDS_KEY in new_sd:
+    # waifu diffusion v1.4
+    position_ids = new_sd[ANOTHER_POSITION_IDS_KEY]
+    del new_sd[ANOTHER_POSITION_IDS_KEY]
+  else:
+    position_ids = torch.Tensor([list(range(max_length))]).to(torch.int64)
+
+  new_sd["text_model.embeddings.position_ids"] = position_ids
+  return new_sd
+
+# endregion
+
+
+# region Diffusers->StableDiffusion の変換コード
+# convert_diffusers_to_original_stable_diffusion をコピーして修正している（ASL 2.0）
+
+def conv_transformer_to_linear(checkpoint):
+  keys = list(checkpoint.keys())
+  tf_keys = ["proj_in.weight", "proj_out.weight"]
+  for key in keys:
+    if ".".join(key.split(".")[-2:]) in tf_keys:
+      if checkpoint[key].ndim > 2:
+        checkpoint[key] = checkpoint[key][:, :, 0, 0]
+
+
+def convert_unet_state_dict_to_sd(v2, 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()}
+
+  if v2:
+    conv_transformer_to_linear(new_state_dict)
+
+  return new_state_dict
+
+
+# ================#
+# VAE Conversion #
+# ================#
+
+def reshape_weight_for_sd(w):
+    # convert HF linear weights to SD conv2d weights
+  return w.reshape(*w.shape, 1, 1)
+
+
+def convert_vae_state_dict(vae_state_dict):
+  vae_conversion_map = [
+      # (stable-diffusion, HF Diffusers)
+      ("nin_shortcut", "conv_shortcut"),
+      ("norm_out", "conv_norm_out"),
+      ("mid.attn_1.", "mid_block.attentions.0."),
+  ]
+
+  for i in range(4):
+    # down_blocks have two resnets
+    for j in range(2):
+      hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}."
+      sd_down_prefix = f"encoder.down.{i}.block.{j}."
+      vae_conversion_map.append((sd_down_prefix, hf_down_prefix))
+
+    if i < 3:
+      hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0."
+      sd_downsample_prefix = f"down.{i}.downsample."
+      vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))
+
+      hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+      sd_upsample_prefix = f"up.{3-i}.upsample."
+      vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))
+
+    # up_blocks have three resnets
+    # also, up blocks in hf are numbered in reverse from sd
+    for j in range(3):
+      hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}."
+      sd_up_prefix = f"decoder.up.{3-i}.block.{j}."
+      vae_conversion_map.append((sd_up_prefix, hf_up_prefix))
+
+  # this part accounts for mid blocks in both the encoder and the decoder
+  for i in range(2):
+    hf_mid_res_prefix = f"mid_block.resnets.{i}."
+    sd_mid_res_prefix = f"mid.block_{i+1}."
+    vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+  vae_conversion_map_attn = [
+      # (stable-diffusion, HF Diffusers)
+      ("norm.", "group_norm."),
+      ("q.", "query."),
+      ("k.", "key."),
+      ("v.", "value."),
+      ("proj_out.", "proj_attn."),
+  ]
+
+  mapping = {k: k for k in vae_state_dict.keys()}
+  for k, v in mapping.items():
+    for sd_part, hf_part in vae_conversion_map:
+      v = v.replace(hf_part, sd_part)
+    mapping[k] = v
+  for k, v in mapping.items():
+    if "attentions" in k:
+      for sd_part, hf_part in vae_conversion_map_attn:
+        v = v.replace(hf_part, sd_part)
+      mapping[k] = v
+  new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()}
+  weights_to_convert = ["q", "k", "v", "proj_out"]
+  for k, v in new_state_dict.items():
+    for weight_name in weights_to_convert:
+      if f"mid.attn_1.{weight_name}.weight" in k:
+        # print(f"Reshaping {k} for SD format")
+        new_state_dict[k] = reshape_weight_for_sd(v)
+
+  return new_state_dict
+
+
+# endregion
+
+# region 自作のモデル読み書きなど
+
+def is_safetensors(path):
+  return os.path.splitext(path)[1].lower() == '.safetensors'
+
+
+def load_checkpoint_with_text_encoder_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.')
+  ]
+
+  if is_safetensors(ckpt_path):
+    checkpoint = None
+    state_dict = load_file(ckpt_path, "cpu")
+  else:
+    checkpoint = torch.load(ckpt_path, map_location="cpu")
+    if "state_dict" in checkpoint:
+      state_dict = checkpoint["state_dict"]
+    else:
+      state_dict = checkpoint
+      checkpoint = None
+
+  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, state_dict
+
+
+# TODO dtype指定の動作が怪しいので確認する text_encoderを指定形式で作れるか未確認
+def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, dtype=None):
+  _, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path)
+  if dtype is not None:
+    for k, v in state_dict.items():
+      if type(v) is torch.Tensor:
+        state_dict[k] = v.to(dtype)
+
+  # Convert the UNet2DConditionModel model.
+  unet_config = create_unet_diffusers_config(v2)
+  converted_unet_checkpoint = convert_ldm_unet_checkpoint(v2, state_dict, unet_config)
+
+  unet = UNet2DConditionModel(**unet_config)
+  info = unet.load_state_dict(converted_unet_checkpoint)
+  print("loading u-net:", info)
+
+  # 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)
+  info = vae.load_state_dict(converted_vae_checkpoint)
+  print("loading vae:", info)
+
+  # convert text_model
+  if v2:
+    converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v2(state_dict, 77)
+    cfg = CLIPTextConfig(
+        vocab_size=49408,
+        hidden_size=1024,
+        intermediate_size=4096,
+        num_hidden_layers=23,
+        num_attention_heads=16,
+        max_position_embeddings=77,
+        hidden_act="gelu",
+        layer_norm_eps=1e-05,
+        dropout=0.0,
+        attention_dropout=0.0,
+        initializer_range=0.02,
+        initializer_factor=1.0,
+        pad_token_id=1,
+        bos_token_id=0,
+        eos_token_id=2,
+        model_type="clip_text_model",
+        projection_dim=512,
+        torch_dtype="float32",
+        transformers_version="4.25.0.dev0",
+    )
+    text_model = CLIPTextModel._from_config(cfg)
+    info = text_model.load_state_dict(converted_text_encoder_checkpoint)
+  else:
+    converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v1(state_dict)
+    text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+    info = text_model.load_state_dict(converted_text_encoder_checkpoint)
+  print("loading text encoder:", info)
+
+  return text_model, vae, unet
+
+
+def convert_text_encoder_state_dict_to_sd_v2(checkpoint, make_dummy_weights=False):
+  def convert_key(key):
+    # position_idsの除去
+    if ".position_ids" in key:
+      return None
+
+    # common
+    key = key.replace("text_model.encoder.", "transformer.")
+    key = key.replace("text_model.", "")
+    if "layers" in key:
+      # resblocks conversion
+      key = key.replace(".layers.", ".resblocks.")
+      if ".layer_norm" in key:
+        key = key.replace(".layer_norm", ".ln_")
+      elif ".mlp." in key:
+        key = key.replace(".fc1.", ".c_fc.")
+        key = key.replace(".fc2.", ".c_proj.")
+      elif '.self_attn.out_proj' in key:
+        key = key.replace(".self_attn.out_proj.", ".attn.out_proj.")
+      elif '.self_attn.' in key:
+        key = None                  # 特殊なので後で処理する
+      else:
+        raise ValueError(f"unexpected key in DiffUsers model: {key}")
+    elif '.position_embedding' in key:
+      key = key.replace("embeddings.position_embedding.weight", "positional_embedding")
+    elif '.token_embedding' in key:
+      key = key.replace("embeddings.token_embedding.weight", "token_embedding.weight")
+    elif 'final_layer_norm' in key:
+      key = key.replace("final_layer_norm", "ln_final")
+    return key
+
+  keys = list(checkpoint.keys())
+  new_sd = {}
+  for key in keys:
+    new_key = convert_key(key)
+    if new_key is None:
+      continue
+    new_sd[new_key] = checkpoint[key]
+
+  # attnの変換
+  for key in keys:
+    if 'layers' in key and 'q_proj' in key:
+      # 三つを結合
+      key_q = key
+      key_k = key.replace("q_proj", "k_proj")
+      key_v = key.replace("q_proj", "v_proj")
+
+      value_q = checkpoint[key_q]
+      value_k = checkpoint[key_k]
+      value_v = checkpoint[key_v]
+      value = torch.cat([value_q, value_k, value_v])
+
+      new_key = key.replace("text_model.encoder.layers.", "transformer.resblocks.")
+      new_key = new_key.replace(".self_attn.q_proj.", ".attn.in_proj_")
+      new_sd[new_key] = value
+
+  # 最後の層などを捏造するか
+  if make_dummy_weights:
+    print("make dummy weights for resblock.23, text_projection and logit scale.")
+    keys = list(new_sd.keys())
+    for key in keys:
+      if key.startswith("transformer.resblocks.22."):
+        new_sd[key.replace(".22.", ".23.")] = new_sd[key].clone()          # copyしないとsafetensorsの保存で落ちる
+
+    # Diffusersに含まれない重みを作っておく
+    new_sd['text_projection'] = torch.ones((1024, 1024), dtype=new_sd[keys[0]].dtype, device=new_sd[keys[0]].device)
+    new_sd['logit_scale'] = torch.tensor(1)
+
+  return new_sd
+
+
+def save_stable_diffusion_checkpoint(v2, output_file, text_encoder, unet, ckpt_path, epochs, steps, save_dtype=None, vae=None):
+  if ckpt_path is not None:
+    # epoch/stepを参照する。またVAEがメモリ上にないときなど、もう一度VAEを含めて読み込む
+    checkpoint, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path)
+    if checkpoint is None:                # safetensors または state_dictのckpt
+      checkpoint = {}
+      strict = False
+    else:
+      strict = True
+    if "state_dict" in state_dict:
+      del state_dict["state_dict"]
+  else:
+    # 新しく作る
+    assert vae is not None, "VAE is required to save a checkpoint without a given checkpoint"
+    checkpoint = {}
+    state_dict = {}
+    strict = False
+
+  def update_sd(prefix, sd):
+    for k, v in sd.items():
+      key = prefix + k
+      assert not strict or key in state_dict, f"Illegal key in save SD: {key}"
+      if save_dtype is not None:
+        v = v.detach().clone().to("cpu").to(save_dtype)
+      state_dict[key] = v
+
+  # Convert the UNet model
+  unet_state_dict = convert_unet_state_dict_to_sd(v2, unet.state_dict())
+  update_sd("model.diffusion_model.", unet_state_dict)
+
+  # Convert the text encoder model
+  if v2:
+    make_dummy = ckpt_path is None                 # 参照元のcheckpointがない場合は最後の層を前の層から複製��て作るなどダミーの重みを入れる
+    text_enc_dict = convert_text_encoder_state_dict_to_sd_v2(text_encoder.state_dict(), make_dummy)
+    update_sd("cond_stage_model.model.", text_enc_dict)
+  else:
+    text_enc_dict = text_encoder.state_dict()
+    update_sd("cond_stage_model.transformer.", text_enc_dict)
+
+  # Convert the VAE
+  if vae is not None:
+    vae_dict = convert_vae_state_dict(vae.state_dict())
+    update_sd("first_stage_model.", vae_dict)
+
+  # Put together new checkpoint
+  key_count = len(state_dict.keys())
+  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
+
+  if is_safetensors(output_file):
+    # TODO Tensor以外のdictの値を削除したほうがいいか
+    save_file(state_dict, output_file)
+  else:
+    torch.save(new_ckpt, output_file)
+
+  return key_count
+
+
+def save_diffusers_checkpoint(v2, output_dir, text_encoder, unet, pretrained_model_name_or_path, vae=None, use_safetensors=False):
+  if pretrained_model_name_or_path is None:
+    # load default settings for v1/v2
+    if v2:
+      pretrained_model_name_or_path = DIFFUSERS_REF_MODEL_ID_V2
+    else:
+      pretrained_model_name_or_path = DIFFUSERS_REF_MODEL_ID_V1
+
+  scheduler = DDIMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
+  tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer")
+  if vae is None:
+    vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae")
+
+  pipeline = StableDiffusionPipeline(
+      unet=unet,
+      text_encoder=text_encoder,
+      vae=vae,
+      scheduler=scheduler,
+      tokenizer=tokenizer,
+      safety_checker=None,
+      feature_extractor=None,
+      requires_safety_checker=None,
+  )
+  pipeline.save_pretrained(output_dir, safe_serialization=use_safetensors)
+
+
+VAE_PREFIX = "first_stage_model."
+
+
+def load_vae(vae_id, dtype):
+  print(f"load VAE: {vae_id}")
+  if os.path.isdir(vae_id) or not os.path.isfile(vae_id):
+    # Diffusers local/remote
+    try:
+      vae = AutoencoderKL.from_pretrained(vae_id, subfolder=None, torch_dtype=dtype)
+    except EnvironmentError as e:
+      print(f"exception occurs in loading vae: {e}")
+      print("retry with subfolder='vae'")
+      vae = AutoencoderKL.from_pretrained(vae_id, subfolder="vae", torch_dtype=dtype)
+    return vae
+
+  # local
+  vae_config = create_vae_diffusers_config()
+
+  if vae_id.endswith(".bin"):
+    # SD 1.5 VAE on Huggingface
+    converted_vae_checkpoint = torch.load(vae_id, map_location="cpu")
+  else:
+    # StableDiffusion
+    vae_model = (load_file(vae_id, "cpu") if is_safetensors(vae_id)
+                 else torch.load(vae_id, map_location="cpu"))
+    vae_sd = vae_model['state_dict'] if 'state_dict' in vae_model else vae_model
+
+    # vae only or full model
+    full_model = False
+    for vae_key in vae_sd:
+      if vae_key.startswith(VAE_PREFIX):
+        full_model = True
+        break
+    if not full_model:
+      sd = {}
+      for key, value in vae_sd.items():
+        sd[VAE_PREFIX + key] = value
+      vae_sd = sd
+      del sd
+
+    # Convert the VAE model.
+    converted_vae_checkpoint = convert_ldm_vae_checkpoint(vae_sd, vae_config)
+
+  vae = AutoencoderKL(**vae_config)
+  vae.load_state_dict(converted_vae_checkpoint)
+  return vae
+
+# 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()
+  return resos
+
+
+if __name__ == '__main__':
+  resos = make_bucket_resolutions((512, 768))
+  print(len(resos))
+  print(resos)
+  aspect_ratios = [w / h for w, h in resos]
+  print(aspect_ratios)
+
+  ars = set()
+  for ar in aspect_ratios:
+    if ar in ars:
+      print("error! duplicate ar:", ar)
+    ars.add(ar)

library/train_util.py

@@ -0,0 +1,1796 @@
+# common functions for training
+
+import argparse
+import json
+import shutil
+import time
+from typing import Dict, List, NamedTuple, Tuple
+from accelerate import Accelerator
+from torch.autograd.function import Function
+import glob
+import math
+import os
+import random
+import hashlib
+import subprocess
+from io import BytesIO
+
+from tqdm import tqdm
+import torch
+from torchvision import transforms
+from transformers import CLIPTokenizer
+import diffusers
+from diffusers import DDPMScheduler, StableDiffusionPipeline
+import albumentations as albu
+import numpy as np
+from PIL import Image
+import cv2
+from einops import rearrange
+from torch import einsum
+import safetensors.torch
+
+import library.model_util as model_util
+
+# Tokenizer: checkpointから読み込むのではなくあらかじめ提供されているものを使う
+TOKENIZER_PATH = "openai/clip-vit-large-patch14"
+V2_STABLE_DIFFUSION_PATH = "stabilityai/stable-diffusion-2"     # ここからtokenizerだけ使う v2とv2.1はtokenizer仕様は同じ
+
+# checkpointファイル名
+EPOCH_STATE_NAME = "{}-{:06d}-state"
+EPOCH_FILE_NAME = "{}-{:06d}"
+EPOCH_DIFFUSERS_DIR_NAME = "{}-{:06d}"
+LAST_STATE_NAME = "{}-state"
+DEFAULT_EPOCH_NAME = "epoch"
+DEFAULT_LAST_OUTPUT_NAME = "last"
+
+# region dataset
+
+IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".webp", ".bmp"]
+# , ".PNG", ".JPG", ".JPEG", ".WEBP", ".BMP"]         # Linux?
+
+
+class ImageInfo():
+  def __init__(self, image_key: str, num_repeats: int, caption: str, is_reg: bool, absolute_path: str) -> None:
+    self.image_key: str = image_key
+    self.num_repeats: int = num_repeats
+    self.caption: str = caption
+    self.is_reg: bool = is_reg
+    self.absolute_path: str = absolute_path
+    self.image_size: Tuple[int, int] = None
+    self.resized_size: Tuple[int, int] = None
+    self.bucket_reso: Tuple[int, int] = None
+    self.latents: torch.Tensor = None
+    self.latents_flipped: torch.Tensor = None
+    self.latents_npz: str = None
+    self.latents_npz_flipped: str = None
+
+
+class BucketManager():
+  def __init__(self, no_upscale, max_reso, min_size, max_size, reso_steps) -> None:
+    self.no_upscale = no_upscale
+    if max_reso is None:
+      self.max_reso = None
+      self.max_area = None
+    else:
+      self.max_reso = max_reso
+      self.max_area = max_reso[0] * max_reso[1]
+    self.min_size = min_size
+    self.max_size = max_size
+    self.reso_steps = reso_steps
+
+    self.resos = []
+    self.reso_to_id = {}
+    self.buckets = []                     # 前処理時は (image_key, image)、学習時は image_key
+
+  def add_image(self, reso, image):
+    bucket_id = self.reso_to_id[reso]
+    self.buckets[bucket_id].append(image)
+
+  def shuffle(self):
+    for bucket in self.buckets:
+      random.shuffle(bucket)
+
+  def sort(self):
+    # 解像度順にソートする（表示時、メタデータ格納時の見栄えをよくするためだけ）。bucketsも入れ替えてreso_to_idも振り直す
+    sorted_resos = self.resos.copy()
+    sorted_resos.sort()
+
+    sorted_buckets = []
+    sorted_reso_to_id = {}
+    for i, reso in enumerate(sorted_resos):
+      bucket_id = self.reso_to_id[reso]
+      sorted_buckets.append(self.buckets[bucket_id])
+      sorted_reso_to_id[reso] = i
+
+    self.resos = sorted_resos
+    self.buckets = sorted_buckets
+    self.reso_to_id = sorted_reso_to_id
+
+  def make_buckets(self):
+    resos = model_util.make_bucket_resolutions(self.max_reso, self.min_size, self.max_size, self.reso_steps)
+    self.set_predefined_resos(resos)
+
+  def set_predefined_resos(self, resos):
+    # 規定サイズから選ぶ場合の解像度、aspect ratioの情報を格納しておく
+    self.predefined_resos = resos.copy()
+    self.predefined_resos_set = set(resos)
+    self.predefined_aspect_ratios = np.array([w / h for w, h in resos])
+
+  def add_if_new_reso(self, reso):
+    if reso not in self.reso_to_id:
+      bucket_id = len(self.resos)
+      self.reso_to_id[reso] = bucket_id
+      self.resos.append(reso)
+      self.buckets.append([])
+      # print(reso, bucket_id, len(self.buckets))
+
+  def round_to_steps(self, x):
+    x = int(x + .5)
+    return x - x % self.reso_steps
+
+  def select_bucket(self, image_width, image_height):
+    aspect_ratio = image_width / image_height
+    if not self.no_upscale:
+      # 同じaspect ratioがあるかもしれないので（fine tuningで、no_upscale=Trueで前処理した場合）、解像度が同じものを優先する
+      reso = (image_width, image_height)
+      if reso in self.predefined_resos_set:
+        pass
+      else:
+        ar_errors = self.predefined_aspect_ratios - aspect_ratio
+        predefined_bucket_id = np.abs(ar_errors).argmin()          # 当該解像度以外でaspect ratio errorが最も少ないもの
+        reso = self.predefined_resos[predefined_bucket_id]
+
+      ar_reso = reso[0] / reso[1]
+      if aspect_ratio > ar_reso:                   # 横が長い→縦を合わせる
+        scale = reso[1] / image_height
+      else:
+        scale = reso[0] / image_width
+
+      resized_size = (int(image_width * scale + .5), int(image_height * scale + .5))
+      # print("use predef", image_width, image_height, reso, resized_size)
+    else:
+      if image_width * image_height > self.max_area:
+        # 画像が大きすぎるのでアスペクト比を保ったまま縮小することを前提にbucketを決める
+        resized_width = math.sqrt(self.max_area * aspect_ratio)
+        resized_height = self.max_area / resized_width
+        assert abs(resized_width / resized_height - aspect_ratio) < 1e-2, "aspect is illegal"
+
+        # リサイズ後の短辺または長辺をreso_steps単位にする：aspect ratioの差が少ないほうを選ぶ
+        # 元のbucketingと同じロジック
+        b_width_rounded = self.round_to_steps(resized_width)
+        b_height_in_wr = self.round_to_steps(b_width_rounded / aspect_ratio)
+        ar_width_rounded = b_width_rounded / b_height_in_wr
+
+        b_height_rounded = self.round_to_steps(resized_height)
+        b_width_in_hr = self.round_to_steps(b_height_rounded * aspect_ratio)
+        ar_height_rounded = b_width_in_hr / b_height_rounded
+
+        # print(b_width_rounded, b_height_in_wr, ar_width_rounded)
+        # print(b_width_in_hr, b_height_rounded, ar_height_rounded)
+
+        if abs(ar_width_rounded - aspect_ratio) < abs(ar_height_rounded - aspect_ratio):
+          resized_size = (b_width_rounded, int(b_width_rounded / aspect_ratio + .5))
+        else:
+          resized_size = (int(b_height_rounded * aspect_ratio + .5), b_height_rounded)
+        # print(resized_size)
+      else:
+        resized_size = (image_width, image_height)              # リサイズは不要
+
+      # 画像のサイズ未満をbucketのサイズとする（paddingせずにcroppingする）
+      bucket_width = resized_size[0] - resized_size[0] % self.reso_steps
+      bucket_height = resized_size[1] - resized_size[1] % self.reso_steps
+      # print("use arbitrary", image_width, image_height, resized_size, bucket_width, bucket_height)
+
+      reso = (bucket_width, bucket_height)
+
+    self.add_if_new_reso(reso)
+
+    ar_error = (reso[0] / reso[1]) - aspect_ratio
+    return reso, resized_size, ar_error
+
+
+class BucketBatchIndex(NamedTuple):
+  bucket_index: int
+  bucket_batch_size: int
+  batch_index: int
+
+
+class BaseDataset(torch.utils.data.Dataset):
+  def __init__(self, tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens, resolution, flip_aug: bool, color_aug: bool, face_crop_aug_range, random_crop, debug_dataset: bool) -> None:
+    super().__init__()
+    self.tokenizer: CLIPTokenizer = tokenizer
+    self.max_token_length = max_token_length
+    self.shuffle_caption = shuffle_caption
+    self.shuffle_keep_tokens = shuffle_keep_tokens
+    # width/height is used when enable_bucket==False
+    self.width, self.height = (None, None) if resolution is None else resolution
+    self.face_crop_aug_range = face_crop_aug_range
+    self.flip_aug = flip_aug
+    self.color_aug = color_aug
+    self.debug_dataset = debug_dataset
+    self.random_crop = random_crop
+    self.token_padding_disabled = False
+    self.dataset_dirs_info = {}
+    self.reg_dataset_dirs_info = {}
+    self.tag_frequency = {}
+
+    self.enable_bucket = False
+    self.bucket_manager: BucketManager = None                         # not initialized
+    self.min_bucket_reso = None
+    self.max_bucket_reso = None
+    self.bucket_reso_steps = None
+    self.bucket_no_upscale = None
+    self.bucket_info = None                                           # for metadata
+
+    self.tokenizer_max_length = self.tokenizer.model_max_length if max_token_length is None else max_token_length + 2
+
+    self.current_epoch: int = 0            # インスタンスがepochごとに新しく作られるようなので外側から渡さないとダメ
+    self.dropout_rate: float = 0
+    self.dropout_every_n_epochs: int = None
+    self.tag_dropout_rate: float = 0
+
+    # augmentation
+    flip_p = 0.5 if flip_aug else 0.0
+    if color_aug:
+      # わりと弱めの色合いaugmentation：brightness/contrastあたりは画像のpixel valueの最大値・最小値を変えてしまうのでよくないのではという想定でgamma/hueあたりを触る
+      self.aug = albu.Compose([
+          albu.OneOf([
+              albu.HueSaturationValue(8, 0, 0, p=.5),
+              albu.RandomGamma((95, 105), p=.5),
+          ], p=.33),
+          albu.HorizontalFlip(p=flip_p)
+      ], p=1.)
+    elif flip_aug:
+      self.aug = albu.Compose([
+          albu.HorizontalFlip(p=flip_p)
+      ], p=1.)
+    else:
+      self.aug = None
+
+    self.image_transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5]), ])
+
+    self.image_data: Dict[str, ImageInfo] = {}
+
+    self.replacements = {}
+
+  def set_current_epoch(self, epoch):
+    self.current_epoch = epoch
+
+  def set_caption_dropout(self, dropout_rate, dropout_every_n_epochs, tag_dropout_rate):
+    # コンストラクタで渡さないのはTextual Inversionで意識したくないから（ということにしておく）
+    self.dropout_rate = dropout_rate
+    self.dropout_every_n_epochs = dropout_every_n_epochs
+    self.tag_dropout_rate = tag_dropout_rate
+
+  def set_tag_frequency(self, dir_name, captions):
+    frequency_for_dir = self.tag_frequency.get(dir_name, {})
+    self.tag_frequency[dir_name] = frequency_for_dir
+    for caption in captions:
+      for tag in caption.split(","):
+        if tag and not tag.isspace():
+          tag = tag.lower()
+          frequency = frequency_for_dir.get(tag, 0)
+          frequency_for_dir[tag] = frequency + 1
+
+  def disable_token_padding(self):
+    self.token_padding_disabled = True
+
+  def add_replacement(self, str_from, str_to):
+    self.replacements[str_from] = str_to
+
+  def process_caption(self, caption):
+    # dropoutの決定：tag dropがこのメソッド内にあるのでここで行うのが良い
+    is_drop_out = self.dropout_rate > 0 and random.random() < self.dropout_rate
+    is_drop_out = is_drop_out or self.dropout_every_n_epochs and self.current_epoch % self.dropout_every_n_epochs == 0
+
+    if is_drop_out:
+      caption = ""
+    else:
+      if self.shuffle_caption or self.tag_dropout_rate > 0:
+        def dropout_tags(tokens):
+          if self.tag_dropout_rate <= 0:
+            return tokens
+          l = []
+          for token in tokens:
+            if random.random() >= self.tag_dropout_rate:
+              l.append(token)
+          return l
+
+        tokens = [t.strip() for t in caption.strip().split(",")]
+        if self.shuffle_keep_tokens is None:
+          if self.shuffle_caption:
+            random.shuffle(tokens)
+
+          tokens = dropout_tags(tokens)
+        else:
+          if len(tokens) > self.shuffle_keep_tokens:
+            keep_tokens = tokens[:self.shuffle_keep_tokens]
+            tokens = tokens[self.shuffle_keep_tokens:]
+
+            if self.shuffle_caption:
+              random.shuffle(tokens)
+
+            tokens = dropout_tags(tokens)
+
+            tokens = keep_tokens + tokens
+        caption = ", ".join(tokens)
+
+      # textual inversion対応
+      for str_from, str_to in self.replacements.items():
+        if str_from == "":
+          # replace all
+          if type(str_to) == list:
+            caption = random.choice(str_to)
+          else:
+            caption = str_to
+        else:
+          caption = caption.replace(str_from, str_to)
+
+    return caption
+
+  def get_input_ids(self, caption):
+    input_ids = self.tokenizer(caption, padding="max_length", truncation=True,
+                               max_length=self.tokenizer_max_length, return_tensors="pt").input_ids
+
+    if self.tokenizer_max_length > self.tokenizer.model_max_length:
+      input_ids = input_ids.squeeze(0)
+      iids_list = []
+      if self.tokenizer.pad_token_id == self.tokenizer.eos_token_id:
+        # v1
+        # 77以上の時は "<BOS> .... <EOS> <EOS> <EOS>" でトータル227とかになっているので、"<BOS>...<EOS>"の三連に変換する
+        # 1111氏のやつは , で区切る、とかしているようだが　とりあえず単純に
+        for i in range(1, self.tokenizer_max_length - self.tokenizer.model_max_length + 2, self.tokenizer.model_max_length - 2):  # (1, 152, 75)
+          ids_chunk = (input_ids[0].unsqueeze(0),
+                       input_ids[i:i + self.tokenizer.model_max_length - 2],
+                       input_ids[-1].unsqueeze(0))
+          ids_chunk = torch.cat(ids_chunk)
+          iids_list.append(ids_chunk)
+      else:
+        # v2
+        # 77以上の時は "<BOS> .... <EOS> <PAD> <PAD>..." でトータル227とかになっているので、"<BOS>...<EOS> <PAD> <PAD> ..."の三連に変換する
+        for i in range(1, self.tokenizer_max_length - self.tokenizer.model_max_length + 2, self.tokenizer.model_max_length - 2):
+          ids_chunk = (input_ids[0].unsqueeze(0),       # BOS
+                       input_ids[i:i + self.tokenizer.model_max_length - 2],
+                       input_ids[-1].unsqueeze(0))      # PAD or EOS
+          ids_chunk = torch.cat(ids_chunk)
+
+          # 末尾が <EOS> <PAD> または <PAD> <PAD> の場合は、何もしなくてよい
+          # 末尾が x <PAD/EOS> の場合は末尾を <EOS> に変える（x <EOS> なら結果的に変化なし）
+          if ids_chunk[-2] != self.tokenizer.eos_token_id and ids_chunk[-2] != self.tokenizer.pad_token_id:
+            ids_chunk[-1] = self.tokenizer.eos_token_id
+          # 先頭が <BOS> <PAD> ... の場合は <BOS> <EOS> <PAD> ... に変える
+          if ids_chunk[1] == self.tokenizer.pad_token_id:
+            ids_chunk[1] = self.tokenizer.eos_token_id
+
+          iids_list.append(ids_chunk)
+
+      input_ids = torch.stack(iids_list)      # 3,77
+    return input_ids
+
+  def register_image(self, info: ImageInfo):
+    self.image_data[info.image_key] = info
+
+  def make_buckets(self):
+    '''
+    bucketingを行わない場合も呼び出し必須（ひとつだけbucketを作る）
+    min_size and max_size are ignored when enable_bucket is False
+    '''
+    print("loading image sizes.")
+    for info in tqdm(self.image_data.values()):
+      if info.image_size is None:
+        info.image_size = self.get_image_size(info.absolute_path)
+
+    if self.enable_bucket:
+      print("make buckets")
+    else:
+      print("prepare dataset")
+
+    # bucketを作成し、画像をbucketに振り分ける
+    if self.enable_bucket:
+      if self.bucket_manager is None:                         # fine tuningの場合でmetadataに定義がある場合は、すでに初期化済み
+        self.bucket_manager = BucketManager(self.bucket_no_upscale, (self.width, self.height),
+                                            self.min_bucket_reso, self.max_bucket_reso, self.bucket_reso_steps)
+        if not self.bucket_no_upscale:
+          self.bucket_manager.make_buckets()
+        else:
+          print("min_bucket_reso and max_bucket_reso are ignored if bucket_no_upscale is set, because bucket reso is defined by image size automatically / bucket_no_upscaleが指定された場合は、bucketの解像度は画像サイズから自動計算されるため、min_bucket_resoとmax_bucket_resoは無視されます")
+
+      img_ar_errors = []
+      for image_info in self.image_data.values():
+        image_width, image_height = image_info.image_size
+        image_info.bucket_reso, image_info.resized_size, ar_error = self.bucket_manager.select_bucket(image_width, image_height)
+
+        # print(image_info.image_key, image_info.bucket_reso)
+        img_ar_errors.append(abs(ar_error))
+
+      self.bucket_manager.sort()
+    else:
+      self.bucket_manager = BucketManager(False, (self.width, self.height), None, None, None)
+      self.bucket_manager.set_predefined_resos([(self.width, self.height)])  # ひとつの固定サイズbucketのみ
+      for image_info in self.image_data.values():
+        image_width, image_height = image_info.image_size
+        image_info.bucket_reso, image_info.resized_size, _ = self.bucket_manager.select_bucket(image_width, image_height)
+
+    for image_info in self.image_data.values():
+      for _ in range(image_info.num_repeats):
+        self.bucket_manager.add_image(image_info.bucket_reso, image_info.image_key)
+
+    # bucket情報を表示、格納する
+    if self.enable_bucket:
+      self.bucket_info = {"buckets": {}}
+      print("number of images (including repeats) / 各bucketの画像枚数（繰り返し回数を含む）")
+      for i, (reso, bucket) in enumerate(zip(self.bucket_manager.resos, self.bucket_manager.buckets)):
+        count = len(bucket)
+        if count > 0:
+          self.bucket_info["buckets"][i] = {"resolution": reso, "count": len(bucket)}
+          print(f"bucket {i}: resolution {reso}, count: {len(bucket)}")
+
+      img_ar_errors = np.array(img_ar_errors)
+      mean_img_ar_error = np.mean(np.abs(img_ar_errors))
+      self.bucket_info["mean_img_ar_error"] = mean_img_ar_error
+      print(f"mean ar error (without repeats): {mean_img_ar_error}")
+
+    # データ参照用indexを作る。このindexはdatasetのshuffleに用いられる
+    self.buckets_indices: List(BucketBatchIndex) = []
+    for bucket_index, bucket in enumerate(self.bucket_manager.buckets):
+      batch_count = int(math.ceil(len(bucket) / self.batch_size))
+      for batch_index in range(batch_count):
+        self.buckets_indices.append(BucketBatchIndex(bucket_index, self.batch_size, batch_index))
+
+      # ↓以下はbucketごとのbatch件数があまりにも増えて混乱を招くので元に戻す
+      # 　学習時はステップ数がランダムなので、同一画像が同一batch内にあってもそれほど悪影響はないであろう、と考えられる
+      #
+      # # bucketが細分化されることにより、ひとつのbucketに一種類の画像のみというケースが増え、つまりそれは
+      # # ひとつのbatchが同じ画像で占められることになるので、さすがに良くないであろう
+      # # そのためバッチサイズを画像種類までに制限する
+      # # ただそれでも同一画像が同一バッチに含まれる可能性はあるので、繰り返し回数が少ないほうがshuffleの品質は良くなることは間違いない？
+      # # TO DO 正則化画像をepochまたがりで利用する仕組み
+      # num_of_image_types = len(set(bucket))
+      # bucket_batch_size = min(self.batch_size, num_of_image_types)
+      # batch_count = int(math.ceil(len(bucket) / bucket_batch_size))
+      # # print(bucket_index, num_of_image_types, bucket_batch_size, batch_count)
+      # for batch_index in range(batch_count):
+      #   self.buckets_indices.append(BucketBatchIndex(bucket_index, bucket_batch_size, batch_index))
+      # ↑ここまで
+
+    self.shuffle_buckets()
+    self._length = len(self.buckets_indices)
+
+  def shuffle_buckets(self):
+    random.shuffle(self.buckets_indices)
+    self.bucket_manager.shuffle()
+
+  def load_image(self, image_path):
+    image = Image.open(image_path)
+    if not image.mode == "RGB":
+      image = image.convert("RGB")
+    img = np.array(image, np.uint8)
+    return img
+
+  def trim_and_resize_if_required(self, image, reso, resized_size):
+    image_height, image_width = image.shape[0:2]
+
+    if image_width != resized_size[0] or image_height != resized_size[1]:
+      # リサイズする
+      image = cv2.resize(image, resized_size, interpolation=cv2.INTER_AREA)       # INTER_AREAでやりたいのでcv2でリサイズ
+
+    image_height, image_width = image.shape[0:2]
+    if image_width > reso[0]:
+      trim_size = image_width - reso[0]
+      p = trim_size // 2 if not self.random_crop else random.randint(0, trim_size)
+      # print("w", trim_size, p)
+      image = image[:, p:p + reso[0]]
+    if image_height > reso[1]:
+      trim_size = image_height - reso[1]
+      p = trim_size // 2 if not self.random_crop else random.randint(0, trim_size)
+      # print("h", trim_size, p)
+      image = image[p:p + reso[1]]
+
+    assert image.shape[0] == reso[1] and image.shape[1] == reso[0], f"internal error, illegal trimmed size: {image.shape}, {reso}"
+    return image
+
+  def cache_latents(self, vae):
+    # TODO ここを高速化したい
+    print("caching latents.")
+    for info in tqdm(self.image_data.values()):
+      if info.latents_npz is not None:
+        info.latents = self.load_latents_from_npz(info, False)
+        info.latents = torch.FloatTensor(info.latents)
+        info.latents_flipped = self.load_latents_from_npz(info, True)             # might be None
+        if info.latents_flipped is not None:
+          info.latents_flipped = torch.FloatTensor(info.latents_flipped)
+        continue
+
+      image = self.load_image(info.absolute_path)
+      image = self.trim_and_resize_if_required(image, info.bucket_reso, info.resized_size)
+
+      img_tensor = self.image_transforms(image)
+      img_tensor = img_tensor.unsqueeze(0).to(device=vae.device, dtype=vae.dtype)
+      info.latents = vae.encode(img_tensor).latent_dist.sample().squeeze(0).to("cpu")
+
+      if self.flip_aug:
+        image = image[:, ::-1].copy()     # cannot convert to Tensor without copy
+        img_tensor = self.image_transforms(image)
+        img_tensor = img_tensor.unsqueeze(0).to(device=vae.device, dtype=vae.dtype)
+        info.latents_flipped = vae.encode(img_tensor).latent_dist.sample().squeeze(0).to("cpu")
+
+  def get_image_size(self, image_path):
+    image = Image.open(image_path)
+    return image.size
+
+  def load_image_with_face_info(self, image_path: str):
+    img = self.load_image(image_path)
+
+    face_cx = face_cy = face_w = face_h = 0
+    if self.face_crop_aug_range is not None:
+      tokens = os.path.splitext(os.path.basename(image_path))[0].split('_')
+      if len(tokens) >= 5:
+        face_cx = int(tokens[-4])
+        face_cy = int(tokens[-3])
+        face_w = int(tokens[-2])
+        face_h = int(tokens[-1])
+
+    return img, face_cx, face_cy, face_w, face_h
+
+  # いい感じに切り出す
+  def crop_target(self, image, face_cx, face_cy, face_w, face_h):
+    height, width = image.shape[0:2]
+    if height == self.height and width == self.width:
+      return image
+
+    # 画像サイズはsizeより大きいのでリサイズする
+    face_size = max(face_w, face_h)
+    min_scale = max(self.height / height, self.width / width)        # 画像がモデル入力サイズぴったりになる倍率（最小の倍率）
+    min_scale = min(1.0, max(min_scale, self.size / (face_size * self.face_crop_aug_range[1])))             # 指定した顔最小サイズ
+    max_scale = min(1.0, max(min_scale, self.size / (face_size * self.face_crop_aug_range[0])))             # 指定した顔最大サイズ
+    if min_scale >= max_scale:          # range指定がmin==max
+      scale = min_scale
+    else:
+      scale = random.uniform(min_scale, max_scale)
+
+    nh = int(height * scale + .5)
+    nw = int(width * scale + .5)
+    assert nh >= self.height and nw >= self.width, f"internal error. small scale {scale}, {width}*{height}"
+    image = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_AREA)
+    face_cx = int(face_cx * scale + .5)
+    face_cy = int(face_cy * scale + .5)
+    height, width = nh, nw
+
+    # 顔を中心として448*640とかへ切り出す
+    for axis, (target_size, length, face_p) in enumerate(zip((self.height, self.width), (height, width), (face_cy, face_cx))):
+      p1 = face_p - target_size // 2                # 顔を中心に持ってくるための切り出し位置
+
+      if self.random_crop:
+        # 背景も含めるために顔を中心に置く確率を高めつつずらす
+        range = max(length - face_p, face_p)        # 画像の端から顔中心までの距離の長いほう
+        p1 = p1 + (random.randint(0, range) + random.randint(0, range)) - range     # -range ~ +range までのいい感じの乱数
+      else:
+        # range指定があるときのみ、すこしだけランダムに（わりと適当）
+        if self.face_crop_aug_range[0] != self.face_crop_aug_range[1]:
+          if face_size > self.size // 10 and face_size >= 40:
+            p1 = p1 + random.randint(-face_size // 20, +face_size // 20)
+
+      p1 = max(0, min(p1, length - target_size))
+
+      if axis == 0:
+        image = image[p1:p1 + target_size, :]
+      else:
+        image = image[:, p1:p1 + target_size]
+
+    return image
+
+  def load_latents_from_npz(self, image_info: ImageInfo, flipped):
+    npz_file = image_info.latents_npz_flipped if flipped else image_info.latents_npz
+    if npz_file is None:
+      return None
+    return np.load(npz_file)['arr_0']
+
+  def __len__(self):
+    return self._length
+
+  def __getitem__(self, index):
+    if index == 0:
+      self.shuffle_buckets()
+
+    bucket = self.bucket_manager.buckets[self.buckets_indices[index].bucket_index]
+    bucket_batch_size = self.buckets_indices[index].bucket_batch_size
+    image_index = self.buckets_indices[index].batch_index * bucket_batch_size
+
+    loss_weights = []
+    captions = []
+    input_ids_list = []
+    latents_list = []
+    images = []
+
+    for image_key in bucket[image_index:image_index + bucket_batch_size]:
+      image_info = self.image_data[image_key]
+      loss_weights.append(self.prior_loss_weight if image_info.is_reg else 1.0)
+
+      # image/latentsを処理する
+      if image_info.latents is not None:
+        latents = image_info.latents if not self.flip_aug or random.random() < .5 else image_info.latents_flipped
+        image = None
+      elif image_info.latents_npz is not None:
+        latents = self.load_latents_from_npz(image_info, self.flip_aug and random.random() >= .5)
+        latents = torch.FloatTensor(latents)
+        image = None
+      else:
+        # 画像を読み込み、必要ならcropする
+        img, face_cx, face_cy, face_w, face_h = self.load_image_with_face_info(image_info.absolute_path)
+        im_h, im_w = img.shape[0:2]
+
+        if self.enable_bucket:
+          img = self.trim_and_resize_if_required(img, image_info.bucket_reso, image_info.resized_size)
+        else:
+          if face_cx > 0:                   # 顔位置情報あり
+            img = self.crop_target(img, face_cx, face_cy, face_w, face_h)
+          elif im_h > self.height or im_w > self.width:
+            assert self.random_crop, f"image too large, but cropping and bucketing are disabled / 画像サイズが大きいのでface_crop_aug_rangeかrandom_crop、またはbucketを有効にしてください: {image_info.absolute_path}"
+            if im_h > self.height:
+              p = random.randint(0, im_h - self.height)
+              img = img[p:p + self.height]
+            if im_w > self.width:
+              p = random.randint(0, im_w - self.width)
+              img = img[:, p:p + self.width]
+
+          im_h, im_w = img.shape[0:2]
+          assert im_h == self.height and im_w == self.width, f"image size is small / 画像サイズが小さいようです: {image_info.absolute_path}"
+
+        # augmentation
+        if self.aug is not None:
+          img = self.aug(image=img)['image']
+
+        latents = None
+        image = self.image_transforms(img)      # -1.0~1.0のtorch.Tensorになる
+
+      images.append(image)
+      latents_list.append(latents)
+
+      caption = self.process_caption(image_info.caption)
+      captions.append(caption)
+      if not self.token_padding_disabled:                     # this option might be omitted in future
+        input_ids_list.append(self.get_input_ids(caption))
+
+    example = {}
+    example['loss_weights'] = torch.FloatTensor(loss_weights)
+
+    if self.token_padding_disabled:
+      # padding=True means pad in the batch
+      example['input_ids'] = self.tokenizer(captions, padding=True, truncation=True, return_tensors="pt").input_ids
+    else:
+      # batch processing seems to be good
+      example['input_ids'] = torch.stack(input_ids_list)
+
+    if images[0] is not None:
+      images = torch.stack(images)
+      images = images.to(memory_format=torch.contiguous_format).float()
+    else:
+      images = None
+    example['images'] = images
+
+    example['latents'] = torch.stack(latents_list) if latents_list[0] is not None else None
+
+    if self.debug_dataset:
+      example['image_keys'] = bucket[image_index:image_index + self.batch_size]
+      example['captions'] = captions
+    return example
+
+
+class DreamBoothDataset(BaseDataset):
+  def __init__(self, batch_size, train_data_dir, reg_data_dir, tokenizer, max_token_length, caption_extension, shuffle_caption, shuffle_keep_tokens, resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, prior_loss_weight, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset) -> None:
+    super().__init__(tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens,
+                     resolution, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset)
+
+    assert resolution is not None, f"resolution is required / resolution（解像度）指定は必須です"
+
+    self.batch_size = batch_size
+    self.size = min(self.width, self.height)                  # 短いほう
+    self.prior_loss_weight = prior_loss_weight
+    self.latents_cache = None
+
+    self.enable_bucket = enable_bucket
+    if self.enable_bucket:
+      assert min(resolution) >= min_bucket_reso, f"min_bucket_reso must be equal or less than resolution / min_bucket_resoは最小解像度より大きくできません。解像度を大きくするかmin_bucket_resoを小さくしてください"
+      assert max(resolution) <= max_bucket_reso, f"max_bucket_reso must be equal or greater than resolution / max_bucket_resoは最大解像度より小さくできません。解像度を小さくするかmin_bucket_resoを大きくしてください"
+      self.min_bucket_reso = min_bucket_reso
+      self.max_bucket_reso = max_bucket_reso
+      self.bucket_reso_steps = bucket_reso_steps
+      self.bucket_no_upscale = bucket_no_upscale
+    else:
+      self.min_bucket_reso = None
+      self.max_bucket_reso = None
+      self.bucket_reso_steps = None                              # この情報は使われない
+      self.bucket_no_upscale = False
+
+    def read_caption(img_path):
+      # captionの候補ファイル名を作る
+      base_name = os.path.splitext(img_path)[0]
+      base_name_face_det = base_name
+      tokens = base_name.split("_")
+      if len(tokens) >= 5:
+        base_name_face_det = "_".join(tokens[:-4])
+      cap_paths = [base_name + caption_extension, base_name_face_det + caption_extension]
+
+      caption = None
+      for cap_path in cap_paths:
+        if os.path.isfile(cap_path):
+          with open(cap_path, "rt", encoding='utf-8') as f:
+            try:
+              lines = f.readlines()
+            except UnicodeDecodeError as e:
+              print(f"illegal char in file (not UTF-8) / ファイルにUTF-8以外の文字があります: {cap_path}")
+              raise e
+            assert len(lines) > 0, f"caption file is empty / キャプションファイルが空です: {cap_path}"
+            caption = lines[0].strip()
+          break
+      return caption
+
+    def load_dreambooth_dir(dir):
+      if not os.path.isdir(dir):
+        # print(f"ignore file: {dir}")
+        return 0, [], []
+
+      tokens = os.path.basename(dir).split('_')
+      try:
+        n_repeats = int(tokens[0])
+      except ValueError as e:
+        print(f"ignore directory without repeats / 繰り返し回数のないディレクトリを無視します: {dir}")
+        return 0, [], []
+
+      caption_by_folder = '_'.join(tokens[1:])
+      img_paths = glob_images(dir, "*")
+      print(f"found directory {n_repeats}_{caption_by_folder} contains {len(img_paths)} image files")
+
+      # 画像ファイルごとにプロンプトを読み込み、もしあればそちらを使う
+      captions = []
+      for img_path in img_paths:
+        cap_for_img = read_caption(img_path)
+        captions.append(caption_by_folder if cap_for_img is None else cap_for_img)
+
+      self.set_tag_frequency(os.path.basename(dir), captions)         # タグ頻度を記録
+
+      return n_repeats, img_paths, captions
+
+    print("prepare train images.")
+    train_dirs = os.listdir(train_data_dir)
+    num_train_images = 0
+    for dir in train_dirs:
+      n_repeats, img_paths, captions = load_dreambooth_dir(os.path.join(train_data_dir, dir))
+      num_train_images += n_repeats * len(img_paths)
+
+      for img_path, caption in zip(img_paths, captions):
+        info = ImageInfo(img_path, n_repeats, caption, False, img_path)
+        self.register_image(info)
+
+      self.dataset_dirs_info[os.path.basename(dir)] = {"n_repeats": n_repeats, "img_count": len(img_paths)}
+
+    print(f"{num_train_images} train images with repeating.")
+    self.num_train_images = num_train_images
+
+    # reg imageは数を数えて学習画像と同じ枚数にする
+    num_reg_images = 0
+    if reg_data_dir:
+      print("prepare reg images.")
+      reg_infos: List[ImageInfo] = []
+
+      reg_dirs = os.listdir(reg_data_dir)
+      for dir in reg_dirs:
+        n_repeats, img_paths, captions = load_dreambooth_dir(os.path.join(reg_data_dir, dir))
+        num_reg_images += n_repeats * len(img_paths)
+
+        for img_path, caption in zip(img_paths, captions):
+          info = ImageInfo(img_path, n_repeats, caption, True, img_path)
+          reg_infos.append(info)
+
+        self.reg_dataset_dirs_info[os.path.basename(dir)] = {"n_repeats": n_repeats, "img_count": len(img_paths)}
+
+      print(f"{num_reg_images} reg images.")
+      if num_train_images < num_reg_images:
+        print("some of reg images are not used / 正則化画像の数が多いので、一部使用されない正則化画像があります")
+
+      if num_reg_images == 0:
+        print("no regularization images / 正則化画像が見つかりませんでした")
+      else:
+        # num_repeatsを計算する：どうせ大した数ではないのでループで処理する
+        n = 0
+        first_loop = True
+        while n < num_train_images:
+          for info in reg_infos:
+            if first_loop:
+              self.register_image(info)
+              n += info.num_repeats
+            else:
+              info.num_repeats += 1
+              n += 1
+            if n >= num_train_images:
+              break
+          first_loop = False
+
+    self.num_reg_images = num_reg_images
+
+
+class FineTuningDataset(BaseDataset):
+  def __init__(self, json_file_name, batch_size, train_data_dir, tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens, resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, flip_aug, color_aug, face_crop_aug_range, random_crop, dataset_repeats, debug_dataset) -> None:
+    super().__init__(tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens,
+                     resolution, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset)
+
+    # メタデータを読み込む
+    if os.path.exists(json_file_name):
+      print(f"loading existing metadata: {json_file_name}")
+      with open(json_file_name, "rt", encoding='utf-8') as f:
+        metadata = json.load(f)
+    else:
+      raise ValueError(f"no metadata / メタデータファイルがありません: {json_file_name}")
+
+    self.metadata = metadata
+    self.train_data_dir = train_data_dir
+    self.batch_size = batch_size
+
+    tags_list = []
+    for image_key, img_md in metadata.items():
+      # path情報を作る
+      if os.path.exists(image_key):
+        abs_path = image_key
+      else:
+        # わりといい加減だがいい方法が思いつかん
+        abs_path = glob_images(train_data_dir, image_key)
+        assert len(abs_path) >= 1, f"no image / 画像がありません: {image_key}"
+        abs_path = abs_path[0]
+
+      caption = img_md.get('caption')
+      tags = img_md.get('tags')
+      if caption is None:
+        caption = tags
+      elif tags is not None and len(tags) > 0:
+        caption = caption + ', ' + tags
+        tags_list.append(tags)
+      assert caption is not None and len(caption) > 0, f"caption or tag is required / キャプションまたはタグは必須です:{abs_path}"
+
+      image_info = ImageInfo(image_key, dataset_repeats, caption, False, abs_path)
+      image_info.image_size = img_md.get('train_resolution')
+
+      if not self.color_aug and not self.random_crop:
+        # if npz exists, use them
+        image_info.latents_npz, image_info.latents_npz_flipped = self.image_key_to_npz_file(image_key)
+
+      self.register_image(image_info)
+    self.num_train_images = len(metadata) * dataset_repeats
+    self.num_reg_images = 0
+
+    # TODO do not record tag freq when no tag
+    self.set_tag_frequency(os.path.basename(json_file_name), tags_list)
+    self.dataset_dirs_info[os.path.basename(json_file_name)] = {"n_repeats": dataset_repeats, "img_count": len(metadata)}
+
+    # check existence of all npz files
+    use_npz_latents = not (self.color_aug or self.random_crop)
+    if use_npz_latents:
+      npz_any = False
+      npz_all = True
+      for image_info in self.image_data.values():
+        has_npz = image_info.latents_npz is not None
+        npz_any = npz_any or has_npz
+
+        if self.flip_aug:
+          has_npz = has_npz and image_info.latents_npz_flipped is not None
+        npz_all = npz_all and has_npz
+
+        if npz_any and not npz_all:
+          break
+
+      if not npz_any:
+        use_npz_latents = False
+        print(f"npz file does not exist. ignore npz files / npzファイルが見つからないためnpzファイルを無視します")
+      elif not npz_all:
+        use_npz_latents = False
+        print(f"some of npz file does not exist. ignore npz files / いくつかのnpzファイルが見つからないためnpzファイルを無視します")
+        if self.flip_aug:
+          print("maybe no flipped files / 反転されたnpzファイルがないのかもしれません")
+    # else:
+    #   print("npz files are not used with color_aug and/or random_crop / color_augまたはrandom_cropが指定されているためnpzファイルは使用されません")
+
+    # check min/max bucket size
+    sizes = set()
+    resos = set()
+    for image_info in self.image_data.values():
+      if image_info.image_size is None:
+        sizes = None                  # not calculated
+        break
+      sizes.add(image_info.image_size[0])
+      sizes.add(image_info.image_size[1])
+      resos.add(tuple(image_info.image_size))
+
+    if sizes is None:
+      if use_npz_latents:
+        use_npz_latents = False
+        print(f"npz files exist, but no bucket info in metadata. ignore npz files / メタデータにbucket情報がないためnpzファイルを無視します")
+
+      assert resolution is not None, "if metadata doesn't have bucket info, resolution is required / メタデータにbucket情報がない場合はresolutionを指定してください"
+
+      self.enable_bucket = enable_bucket
+      if self.enable_bucket:
+        self.min_bucket_reso = min_bucket_reso
+        self.max_bucket_reso = max_bucket_reso
+        self.bucket_reso_steps = bucket_reso_steps
+        self.bucket_no_upscale = bucket_no_upscale
+    else:
+      if not enable_bucket:
+        print("metadata has bucket info, enable bucketing / メタデータにbucket情報があるためbucketを有効にします")
+      print("using bucket info in metadata / メタデータ内のbucket情報を使います")
+      self.enable_bucket = True
+
+      assert not bucket_no_upscale, "if metadata has bucket info, bucket reso is precalculated, so bucket_no_upscale cannot be used / メタデータ内にbucket情報がある場合はbucketの解像度は計算済みのため、bucket_no_upscaleは使えません"
+
+      # bucket情報を初期化しておく、make_bucketsで再作成しない
+      self.bucket_manager = BucketManager(False, None, None, None, None)
+      self.bucket_manager.set_predefined_resos(resos)
+
+    # npz情報をきれいにしておく
+    if not use_npz_latents:
+      for image_info in self.image_data.values():
+        image_info.latents_npz = image_info.latents_npz_flipped = None
+
+  def image_key_to_npz_file(self, image_key):
+    base_name = os.path.splitext(image_key)[0]
+    npz_file_norm = base_name + '.npz'
+
+    if os.path.exists(npz_file_norm):
+      # image_key is full path
+      npz_file_flip = base_name + '_flip.npz'
+      if not os.path.exists(npz_file_flip):
+        npz_file_flip = None
+      return npz_file_norm, npz_file_flip
+
+    # image_key is relative path
+    npz_file_norm = os.path.join(self.train_data_dir, image_key + '.npz')
+    npz_file_flip = os.path.join(self.train_data_dir, image_key + '_flip.npz')
+
+    if not os.path.exists(npz_file_norm):
+      npz_file_norm = None
+      npz_file_flip = None
+    elif not os.path.exists(npz_file_flip):
+      npz_file_flip = None
+
+    return npz_file_norm, npz_file_flip
+
+
+def debug_dataset(train_dataset, show_input_ids=False):
+  print(f"Total dataset length (steps) / データセットの長さ（ステップ数）: {len(train_dataset)}")
+  print("Escape for exit. / Escキーで中断、終了します")
+
+  train_dataset.set_current_epoch(1)
+  k = 0
+  for i, example in enumerate(train_dataset):
+    if example['latents'] is not None:
+      print(f"sample has latents from npz file: {example['latents'].size()}")
+    for j, (ik, cap, lw, iid) in enumerate(zip(example['image_keys'], example['captions'], example['loss_weights'], example['input_ids'])):
+      print(f'{ik}, size: {train_dataset.image_data[ik].image_size}, loss weight: {lw}, caption: "{cap}"')
+      if show_input_ids:
+        print(f"input ids: {iid}")
+      if example['images'] is not None:
+        im = example['images'][j]
+        print(f"image size: {im.size()}")
+        im = ((im.numpy() + 1.0) * 127.5).astype(np.uint8)
+        im = np.transpose(im, (1, 2, 0))                # c,H,W -> H,W,c
+        im = im[:, :, ::-1]                             # RGB -> BGR (OpenCV)
+        if os.name == 'nt':                             # only windows
+          cv2.imshow("img", im)
+        k = cv2.waitKey()
+        cv2.destroyAllWindows()
+        if k == 27:
+          break
+    if k == 27 or (example['images'] is None and i >= 8):
+      break
+
+
+def glob_images(directory, base="*"):
+  img_paths = []
+  for ext in IMAGE_EXTENSIONS:
+    if base == '*':
+      img_paths.extend(glob.glob(os.path.join(glob.escape(directory), base + ext)))
+    else:
+      img_paths.extend(glob.glob(glob.escape(os.path.join(directory, base + ext))))
+  # img_paths = list(set(img_paths))                    # 重複を排除
+  # img_paths.sort()
+  return img_paths
+
+
+def glob_images_pathlib(dir_path, recursive):
+  image_paths = []
+  if recursive:
+    for ext in IMAGE_EXTENSIONS:
+      image_paths += list(dir_path.rglob('*' + ext))
+  else:
+    for ext in IMAGE_EXTENSIONS:
+      image_paths += list(dir_path.glob('*' + ext))
+  # image_paths = list(set(image_paths))        # 重複を排除
+  # image_paths.sort()
+  return image_paths
+
+# endregion
+
+
+# region モジュール入れ替え部
+"""
+高速化のためのモジュール入れ替え
+"""
+
+# FlashAttentionを使うCrossAttention
+# based on https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/memory_efficient_attention_pytorch/flash_attention.py
+# LICENSE MIT https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/LICENSE
+
+# constants
+
+EPSILON = 1e-6
+
+# helper functions
+
+
+def exists(val):
+  return val is not None
+
+
+def default(val, d):
+  return val if exists(val) else d
+
+
+def model_hash(filename):
+  """Old model hash used by stable-diffusion-webui"""
+  try:
+    with open(filename, "rb") as file:
+      m = hashlib.sha256()
+
+      file.seek(0x100000)
+      m.update(file.read(0x10000))
+      return m.hexdigest()[0:8]
+  except FileNotFoundError:
+    return 'NOFILE'
+
+
+def calculate_sha256(filename):
+  """New model hash used by stable-diffusion-webui"""
+  hash_sha256 = hashlib.sha256()
+  blksize = 1024 * 1024
+
+  with open(filename, "rb") as f:
+    for chunk in iter(lambda: f.read(blksize), b""):
+      hash_sha256.update(chunk)
+
+  return hash_sha256.hexdigest()
+
+
+def precalculate_safetensors_hashes(tensors, metadata):
+  """Precalculate the model hashes needed by sd-webui-additional-networks to
+  save time on indexing the model later."""
+
+  # Because writing user metadata to the file can change the result of
+  # sd_models.model_hash(), only retain the training metadata for purposes of
+  # calculating the hash, as they are meant to be immutable
+  metadata = {k: v for k, v in metadata.items() if k.startswith("ss_")}
+
+  bytes = safetensors.torch.save(tensors, metadata)
+  b = BytesIO(bytes)
+
+  model_hash = addnet_hash_safetensors(b)
+  legacy_hash = addnet_hash_legacy(b)
+  return model_hash, legacy_hash
+
+
+def addnet_hash_legacy(b):
+  """Old model hash used by sd-webui-additional-networks for .safetensors format files"""
+  m = hashlib.sha256()
+
+  b.seek(0x100000)
+  m.update(b.read(0x10000))
+  return m.hexdigest()[0:8]
+
+
+def addnet_hash_safetensors(b):
+  """New model hash used by sd-webui-additional-networks for .safetensors format files"""
+  hash_sha256 = hashlib.sha256()
+  blksize = 1024 * 1024
+
+  b.seek(0)
+  header = b.read(8)
+  n = int.from_bytes(header, "little")
+
+  offset = n + 8
+  b.seek(offset)
+  for chunk in iter(lambda: b.read(blksize), b""):
+    hash_sha256.update(chunk)
+
+  return hash_sha256.hexdigest()
+
+
+def get_git_revision_hash() -> str:
+  try:
+    return subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=os.path.dirname(__file__)).decode('ascii').strip()
+  except:
+    return "(unknown)"
+
+
+# flash attention forwards and backwards
+
+# https://arxiv.org/abs/2205.14135
+
+
+class FlashAttentionFunction(torch.autograd.function.Function):
+  @ staticmethod
+  @ torch.no_grad()
+  def forward(ctx, q, k, v, mask, causal, q_bucket_size, k_bucket_size):
+    """ Algorithm 2 in the paper """
+
+    device = q.device
+    dtype = q.dtype
+    max_neg_value = -torch.finfo(q.dtype).max
+    qk_len_diff = max(k.shape[-2] - q.shape[-2], 0)
+
+    o = torch.zeros_like(q)
+    all_row_sums = torch.zeros((*q.shape[:-1], 1), dtype=dtype, device=device)
+    all_row_maxes = torch.full((*q.shape[:-1], 1), max_neg_value, dtype=dtype, device=device)
+
+    scale = (q.shape[-1] ** -0.5)
+
+    if not exists(mask):
+      mask = (None,) * math.ceil(q.shape[-2] / q_bucket_size)
+    else:
+      mask = rearrange(mask, 'b n -> b 1 1 n')
+      mask = mask.split(q_bucket_size, dim=-1)
+
+    row_splits = zip(
+        q.split(q_bucket_size, dim=-2),
+        o.split(q_bucket_size, dim=-2),
+        mask,
+        all_row_sums.split(q_bucket_size, dim=-2),
+        all_row_maxes.split(q_bucket_size, dim=-2),
+    )
+
+    for ind, (qc, oc, row_mask, row_sums, row_maxes) in enumerate(row_splits):
+      q_start_index = ind * q_bucket_size - qk_len_diff
+
+      col_splits = zip(
+          k.split(k_bucket_size, dim=-2),
+          v.split(k_bucket_size, dim=-2),
+      )
+
+      for k_ind, (kc, vc) in enumerate(col_splits):
+        k_start_index = k_ind * k_bucket_size
+
+        attn_weights = einsum('... i d, ... j d -> ... i j', qc, kc) * scale
+
+        if exists(row_mask):
+          attn_weights.masked_fill_(~row_mask, max_neg_value)
+
+        if causal and q_start_index < (k_start_index + k_bucket_size - 1):
+          causal_mask = torch.ones((qc.shape[-2], kc.shape[-2]), dtype=torch.bool,
+                                   device=device).triu(q_start_index - k_start_index + 1)
+          attn_weights.masked_fill_(causal_mask, max_neg_value)
+
+        block_row_maxes = attn_weights.amax(dim=-1, keepdims=True)
+        attn_weights -= block_row_maxes
+        exp_weights = torch.exp(attn_weights)
+
+        if exists(row_mask):
+          exp_weights.masked_fill_(~row_mask, 0.)
+
+        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 (not xformers)")
+  flash_func = FlashAttentionFunction
+
+  def forward_flash_attn(self, x, context=None, mask=None):
+    q_bucket_size = 512
+    k_bucket_size = 1024
+
+    h = self.heads
+    q = self.to_q(x)
+
+    context = context if context is not None else x
+    context = context.to(x.dtype)
+
+    if hasattr(self, 'hypernetwork') and self.hypernetwork is not None:
+      context_k, context_v = self.hypernetwork.forward(x, context)
+      context_k = context_k.to(x.dtype)
+      context_v = context_v.to(x.dtype)
+    else:
+      context_k = context
+      context_v = context
+
+    k = self.to_k(context_k)
+    v = self.to_v(context_v)
+    del context, x
+
+    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), (q, k, v))
+
+    out = flash_func.apply(q, k, v, mask, False, q_bucket_size, k_bucket_size)
+
+    out = rearrange(out, 'b h n d -> b n (h d)')
+
+    # diffusers 0.7.0~  わざわざ変えるなよ (;´Д｀)
+    out = self.to_out[0](out)
+    out = self.to_out[1](out)
+    return out
+
+  diffusers.models.attention.CrossAttention.forward = forward_flash_attn
+
+
+def replace_unet_cross_attn_to_xformers():
+  print("Replace CrossAttention.forward to use xformers")
+  try:
+    import xformers.ops
+  except ImportError:
+    raise ImportError("No xformers / xformersがインストールされていないようです")
+
+  def forward_xformers(self, x, context=None, mask=None):
+    h = self.heads
+    q_in = self.to_q(x)
+
+    context = default(context, x)
+    context = context.to(x.dtype)
+
+    if hasattr(self, 'hypernetwork') and self.hypernetwork is not None:
+      context_k, context_v = self.hypernetwork.forward(x, context)
+      context_k = context_k.to(x.dtype)
+      context_v = context_v.to(x.dtype)
+    else:
+      context_k = context
+      context_v = context
+
+    k_in = self.to_k(context_k)
+    v_in = self.to_v(context_v)
+
+    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b n h d', h=h), (q_in, k_in, v_in))
+    del q_in, k_in, v_in
+
+    q = q.contiguous()
+    k = k.contiguous()
+    v = v.contiguous()
+    out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None)        # 最適なのを選んでくれる
+
+    out = rearrange(out, 'b n h d -> b n (h d)', h=h)
+
+    # diffusers 0.7.0~
+    out = self.to_out[0](out)
+    out = self.to_out[1](out)
+    return out
+
+  diffusers.models.attention.CrossAttention.forward = forward_xformers
+# endregion
+
+
+# region arguments
+
+def add_sd_models_arguments(parser: argparse.ArgumentParser):
+  # for pretrained models
+  parser.add_argument("--v2", action='store_true',
+                      help='load Stable Diffusion v2.0 model / Stable Diffusion 2.0のモデルを読み込む')
+  parser.add_argument("--v_parameterization", action='store_true',
+                      help='enable v-parameterization training / v-parameterization学習を有効にする')
+  parser.add_argument("--pretrained_model_name_or_path", type=str, default=None,
+                      help="pretrained model to train, directory to Diffusers model or StableDiffusion checkpoint / 学習元モデル、Diffusers形式モデルのディレクトリまたはStableDiffusionのckptファイル")
+
+
+def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth: bool):
+  parser.add_argument("--output_dir", type=str, default=None,
+                      help="directory to output trained model / 学習後のモデル出力先ディレクトリ")
+  parser.add_argument("--output_name", type=str, default=None,
+                      help="base name of trained model file / 学習後のモデルの拡張子を除くファイル名")
+  parser.add_argument("--save_precision", type=str, default=None,
+                      choices=[None, "float", "fp16", "bf16"], help="precision in saving / 保存時に精度を変更して保存する")
+  parser.add_argument("--save_every_n_epochs", type=int, default=None,
+                      help="save checkpoint every N epochs / 学習中のモデルを指定エポックごとに保存する")
+  parser.add_argument("--save_n_epoch_ratio", type=int, default=None,
+                      help="save checkpoint N epoch ratio (for example 5 means save at least 5 files total) / 学習中のモデルを指定のエポック割合で保存する（たとえば5を指定すると最低5個のファイルが保存される）")
+  parser.add_argument("--save_last_n_epochs", type=int, default=None, help="save last N checkpoints / 最大Nエポック保存する")
+  parser.add_argument("--save_last_n_epochs_state", type=int, default=None,
+                      help="save last N checkpoints of state (overrides the value of --save_last_n_epochs)/ 最大Nエポックstateを保存する(--save_last_n_epochsの指定を上書きします)")
+  parser.add_argument("--save_state", action="store_true",
+                      help="save training state additionally (including optimizer states etc.) / optimizerなど学習状態も含めたstateを追加で保存する")
+  parser.add_argument("--resume", type=str, default=None, help="saved state to resume training / 学習再開するモデルのstate")
+
+  parser.add_argument("--train_batch_size", type=int, default=1, help="batch size for training / 学習時のバッチサイズ")
+  parser.add_argument("--max_token_length", type=int, default=None, choices=[None, 150, 225],
+                      help="max token length of text encoder (default for 75, 150 or 225) / text encoderのトークンの最大長（未指定で75、150または225が指定可）")
+  parser.add_argument("--use_8bit_adam", action="store_true",
+                      help="use 8bit Adam optimizer (requires bitsandbytes) / 8bit Adamオプティマイザを使う（bitsandbytesのインストールが必要）")
+  parser.add_argument("--use_lion_optimizer", action="store_true",
+                      help="use Lion optimizer (requires lion-pytorch) / Lionオプティマイザを使う（ lion-pytorch のインストールが必要）")
+  parser.add_argument("--mem_eff_attn", action="store_true",
+                      help="use memory efficient attention for CrossAttention / CrossAttentionに省メモリ版attentionを使う")
+  parser.add_argument("--xformers", action="store_true",
+                      help="use xformers for CrossAttention / CrossAttentionにxformersを使う")
+  parser.add_argument("--vae", type=str, default=None,
+                      help="path to checkpoint of vae to replace / VAEを入れ替える場合、VAEのcheckpointファイルまたはディレクトリ")
+
+  parser.add_argument("--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("--max_train_epochs", type=int, default=None,
+                      help="training epochs (overrides max_train_steps) / 学習エポック数（max_train_stepsを上書きします）")
+  parser.add_argument("--max_data_loader_n_workers", type=int, default=8,
+                      help="max num workers for DataLoader (lower is less main RAM usage, faster epoch start and slower data loading) / DataLoaderの最大プロセス数（小さい値ではメインメモリの使用量が減りエポック間の待ち時間が減りますが、データ読み込みは遅くなります）")
+  parser.add_argument("--persistent_data_loader_workers", action="store_true",
+                      help="persistent DataLoader workers (useful for reduce time gap between epoch, but may use more memory) / DataLoader のワーカーを持続させる (エポック間の時間差を少なくするのに有効だが、より多くのメモリを消費する可能性がある)")
+  parser.add_argument("--seed", type=int, default=None, help="random seed for training / 学習時の乱数のseed")
+  parser.add_argument("--gradient_checkpointing", action="store_true",
+                      help="enable gradient checkpointing / grandient checkpointingを有効にする")
+  parser.add_argument("--gradient_accumulation_steps", type=int, default=1,
+                      help="Number of updates steps to accumulate before performing a backward/update pass / 学習時に逆伝播をする前に勾配を合計するステップ数")
+  parser.add_argument("--mixed_precision", type=str, default="no",
+                      choices=["no", "fp16", "bf16"], help="use mixed precision / 混合精度を使う場合、その精度")
+  parser.add_argument("--full_fp16", action="store_true", help="fp16 training including gradients / 勾配も含めてfp16で学習する")
+  parser.add_argument("--clip_skip", type=int, default=None,
+                      help="use output of nth layer from back of text encoder (n>=1) / text encoderの後ろからn番目の層の出力を用いる（nは1以上）")
+  parser.add_argument("--logging_dir", type=str, default=None,
+                      help="enable logging and output TensorBoard log to this directory / ログ出力を有効にしてこのディレクトリにTensorBoard用のログを出力する")
+  parser.add_argument("--log_prefix", type=str, default=None, help="add prefix for each log directory / ログディレクトリ名の先頭に追加する文字列")
+  parser.add_argument("--lr_scheduler", type=str, default="constant",
+                      help="scheduler to use for learning rate / 学習率のスケジューラ: linear, cosine, cosine_with_restarts, polynomial, constant (default), constant_with_warmup")
+  parser.add_argument("--lr_warmup_steps", type=int, default=0,
+                      help="Number of steps for the warmup in the lr scheduler (default is 0) / 学習率のスケジューラをウォームアップするステップ数（デフォルト0）")
+  parser.add_argument("--noise_offset", type=float, default=None,
+                      help="enable noise offset with this value (if enabled, around 0.1 is recommended) / Noise offsetを有効にしてこの値を設定する（有効にする場合は0.1程度を推奨）")
+  parser.add_argument("--lowram", action="store_true",
+                      help="enable low RAM optimization. e.g. load models to VRAM instead of RAM (for machines which have bigger VRAM than RAM such as Colab and Kaggle) / メインメモリが少ない環境向け最適化を有効にする。たとえばVRAMにモデルを読み込むなど（ColabやKaggleなどRAMに比べてVRAMが多い環境向け）")
+
+  if support_dreambooth:
+    # DreamBooth training
+    parser.add_argument("--prior_loss_weight", type=float, default=1.0,
+                        help="loss weight for regularization images / 正則化画像のlossの重み")
+
+
+def verify_training_args(args: argparse.Namespace):
+  if args.v_parameterization and not args.v2:
+    print("v_parameterization should be with v2 / v1でv_parameterizationを使用することは想定されていません")
+  if args.v2 and args.clip_skip is not None:
+    print("v2 with clip_skip will be unexpected / v2でclip_skipを使用することは想定されていません")
+
+
+def add_dataset_arguments(parser: argparse.ArgumentParser, support_dreambooth: bool, support_caption: bool, support_caption_dropout: bool):
+  # dataset common
+  parser.add_argument("--train_data_dir", type=str, default=None, help="directory for train images / 学習画像データのディレクトリ")
+  parser.add_argument("--shuffle_caption", action="store_true",
+                      help="shuffle comma-separated caption / コンマで区切られたcaptionの各要素をshuffleする")
+  parser.add_argument("--caption_extension", type=str, default=".caption", help="extension of caption files / 読み込むcaptionファイルの拡張子")
+  parser.add_argument("--caption_extention", type=str, default=None,
+                      help="extension of caption files (backward compatibility) / 読み込むcaptionファイルの拡張子（スペルミスを残してあります）")
+  parser.add_argument("--keep_tokens", type=int, default=None,
+                      help="keep heading N tokens when shuffling caption tokens / captionのシャッフル時に、先頭からこの個数のトークンをシャッフルしないで残す")
+  parser.add_argument("--color_aug", action="store_true", help="enable weak color augmentation / 学習時に色合いのaugmentationを有効にする")
+  parser.add_argument("--flip_aug", action="store_true", help="enable horizontal flip augmentation / 学習時に左右反転のaugmentationを有効にする")
+  parser.add_argument("--face_crop_aug_range", type=str, default=None,
+                      help="enable face-centered crop augmentation and its range (e.g. 2.0,4.0) / 学習時に顔を中心とした切り出しaugmentationを有効にするときは倍率を指定する（例：2.0,4.0）")
+  parser.add_argument("--random_crop", action="store_true",
+                      help="enable random crop (for style training in face-centered crop augmentation) / ランダムな切り出しを有効にする（顔を中心としたaugmentationを行うときに画風の学習用に指定する）")
+  parser.add_argument("--debug_dataset", action="store_true",
+                      help="show images for debugging (do not train) / デバッグ用に学習データを画面表示する（学習は行わない）")
+  parser.add_argument("--resolution", type=str, default=None,
+                      help="resolution in training ('size' or 'width,height') / 学習時の画像解像度（'サイズ'指定、または'幅,高さ'指定）")
+  parser.add_argument("--cache_latents", action="store_true",
+                      help="cache latents to reduce memory (augmentations must be disabled) / メモリ削減のためにlatentをcacheする（augmentationは使用不可）")
+  parser.add_argument("--enable_bucket", action="store_true",
+                      help="enable buckets for multi aspect ratio training / 複数解像度学習のためのbucketを有効にする")
+  parser.add_argument("--min_bucket_reso", type=int, default=256, help="minimum resolution for buckets / bucketの最小解像度")
+  parser.add_argument("--max_bucket_reso", type=int, default=1024, help="maximum resolution for buckets / bucketの最大解像度")
+  parser.add_argument("--bucket_reso_steps", type=int, default=64,
+                      help="steps of resolution for buckets, divisible by 8 is recommended / bucketの解像度の単位、8で割り切れる値を推奨します")
+  parser.add_argument("--bucket_no_upscale", action="store_true",
+                      help="make bucket for each image without upscaling / 画像を拡大せずbucketを作成します")
+
+  if support_caption_dropout:
+    # Textual Inversion はcaptionのdropoutをsupportしない
+    # いわゆるtensorのDropoutと紛らわしいのでprefixにcaptionを付けておく　every_n_epochsは他と平仄を合わせてdefault Noneに
+    parser.add_argument("--caption_dropout_rate", type=float, default=0,
+                        help="Rate out dropout caption(0.0~1.0) / captionをdropoutする割合")
+    parser.add_argument("--caption_dropout_every_n_epochs", type=int, default=None,
+                        help="Dropout all captions every N epochs / captionを指定エポックごとにdropoutする")
+    parser.add_argument("--caption_tag_dropout_rate", type=float, default=0,
+                        help="Rate out dropout comma separated tokens(0.0~1.0) / カンマ区切りのタグをdropoutする割合")
+
+  if support_dreambooth:
+    # DreamBooth dataset
+    parser.add_argument("--reg_data_dir", type=str, default=None, help="directory for regularization images / 正則化画像データのディレクトリ")
+
+  if support_caption:
+    # caption dataset
+    parser.add_argument("--in_json", type=str, default=None, help="json metadata for dataset / データセットのmetadataのjsonファイル")
+    parser.add_argument("--dataset_repeats", type=int, default=1,
+                        help="repeat dataset when training with captions / キャプションでの学習時にデータセットを繰り返す回数")
+
+
+def add_sd_saving_arguments(parser: argparse.ArgumentParser):
+  parser.add_argument("--save_model_as", type=str, default=None, choices=[None, "ckpt", "safetensors", "diffusers", "diffusers_safetensors"],
+                      help="format to save the model (default is same to original) / モデル保存時の形式（未指定時は元モデルと同じ）")
+  parser.add_argument("--use_safetensors", action='store_true',
+                      help="use safetensors format to save (if save_model_as is not specified) / checkpoint、モデルをsafetensors形式で保存する（save_model_as未指定時）")
+
+# endregion
+
+# region utils
+
+
+def prepare_dataset_args(args: argparse.Namespace, support_metadata: bool):
+  # backward compatibility
+  if args.caption_extention is not None:
+    args.caption_extension = args.caption_extention
+    args.caption_extention = None
+
+  if args.cache_latents:
+    assert not args.color_aug, "when caching latents, color_aug cannot be used / latentをキャッシュするときはcolor_augは使えません"
+    assert not args.random_crop, "when caching latents, random_crop cannot be used / latentをキャッシュするときはrandom_cropは使えません"
+
+  # assert args.resolution is not None, f"resolution is required / resolution（解像度）を指定してください"
+  if args.resolution is not None:
+    args.resolution = tuple([int(r) for r in args.resolution.split(',')])
+    if len(args.resolution) == 1:
+      args.resolution = (args.resolution[0], args.resolution[0])
+    assert len(args.resolution) == 2, \
+        f"resolution must be 'size' or 'width,height' / resolution（解像度）は'サイズ'または'幅','高さ'で指定してください: {args.resolution}"
+
+  if args.face_crop_aug_range is not None:
+    args.face_crop_aug_range = tuple([float(r) for r in args.face_crop_aug_range.split(',')])
+    assert len(args.face_crop_aug_range) == 2 and args.face_crop_aug_range[0] <= args.face_crop_aug_range[1], \
+        f"face_crop_aug_range must be two floats / face_crop_aug_rangeは'下限,上限'で指定してください: {args.face_crop_aug_range}"
+  else:
+    args.face_crop_aug_range = None
+
+  if support_metadata:
+    if args.in_json is not None and (args.color_aug or args.random_crop):
+      print(f"latents in npz is ignored when color_aug or random_crop is True / color_augまたはrandom_cropを有効にした場合、npzファイルのlatentsは無視されます")
+
+
+def load_tokenizer(args: argparse.Namespace):
+  print("prepare tokenizer")
+  if args.v2:
+    tokenizer = CLIPTokenizer.from_pretrained(V2_STABLE_DIFFUSION_PATH, subfolder="tokenizer")
+  else:
+    tokenizer = CLIPTokenizer.from_pretrained(TOKENIZER_PATH)
+  if args.max_token_length is not None:
+    print(f"update token length: {args.max_token_length}")
+  return tokenizer
+
+
+def prepare_accelerator(args: argparse.Namespace):
+  if args.logging_dir is None:
+    log_with = None
+    logging_dir = None
+  else:
+    log_with = "tensorboard"
+    log_prefix = "" if args.log_prefix is None else args.log_prefix
+    logging_dir = args.logging_dir + "/" + log_prefix + time.strftime('%Y%m%d%H%M%S', time.localtime())
+
+  accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps, mixed_precision=args.mixed_precision,
+                            log_with=log_with, logging_dir=logging_dir)
+
+  # accelerateの互換性問題を解決する
+  accelerator_0_15 = True
+  try:
+    accelerator.unwrap_model("dummy", True)
+    print("Using accelerator 0.15.0 or above.")
+  except TypeError:
+    accelerator_0_15 = False
+
+  def unwrap_model(model):
+    if accelerator_0_15:
+      return accelerator.unwrap_model(model, True)
+    return accelerator.unwrap_model(model)
+
+  return accelerator, unwrap_model
+
+
+def prepare_dtype(args: argparse.Namespace):
+  weight_dtype = torch.float32
+  if args.mixed_precision == "fp16":
+    weight_dtype = torch.float16
+  elif args.mixed_precision == "bf16":
+    weight_dtype = torch.bfloat16
+
+  save_dtype = None
+  if args.save_precision == "fp16":
+    save_dtype = torch.float16
+  elif args.save_precision == "bf16":
+    save_dtype = torch.bfloat16
+  elif args.save_precision == "float":
+    save_dtype = torch.float32
+
+  return weight_dtype, save_dtype
+
+
+def load_target_model(args: argparse.Namespace, weight_dtype):
+  load_stable_diffusion_format = os.path.isfile(args.pretrained_model_name_or_path)           # determine SD or Diffusers
+  if load_stable_diffusion_format:
+    print("load StableDiffusion checkpoint")
+    text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.pretrained_model_name_or_path)
+  else:
+    print("load Diffusers pretrained models")
+    pipe = StableDiffusionPipeline.from_pretrained(args.pretrained_model_name_or_path, tokenizer=None, safety_checker=None)
+    text_encoder = pipe.text_encoder
+    vae = pipe.vae
+    unet = pipe.unet
+    del pipe
+
+  # VAEを読み込む
+  if args.vae is not None:
+    vae = model_util.load_vae(args.vae, weight_dtype)
+    print("additional VAE loaded")
+
+  return text_encoder, vae, unet, load_stable_diffusion_format
+
+
+def patch_accelerator_for_fp16_training(accelerator):
+  org_unscale_grads = accelerator.scaler._unscale_grads_
+
+  def _unscale_grads_replacer(optimizer, inv_scale, found_inf, allow_fp16):
+    return org_unscale_grads(optimizer, inv_scale, found_inf, True)
+
+  accelerator.scaler._unscale_grads_ = _unscale_grads_replacer
+
+
+def get_hidden_states(args: argparse.Namespace, input_ids, tokenizer, text_encoder, weight_dtype=None):
+  # with no_token_padding, the length is not max length, return result immediately
+  if input_ids.size()[-1] != tokenizer.model_max_length:
+    return text_encoder(input_ids)[0]
+
+  b_size = input_ids.size()[0]
+  input_ids = input_ids.reshape((-1, tokenizer.model_max_length))     # batch_size*3, 77
+
+  if args.clip_skip is None:
+    encoder_hidden_states = text_encoder(input_ids)[0]
+  else:
+    enc_out = text_encoder(input_ids, output_hidden_states=True, return_dict=True)
+    encoder_hidden_states = enc_out['hidden_states'][-args.clip_skip]
+    encoder_hidden_states = text_encoder.text_model.final_layer_norm(encoder_hidden_states)
+
+  # bs*3, 77, 768 or 1024
+  encoder_hidden_states = encoder_hidden_states.reshape((b_size, -1, encoder_hidden_states.shape[-1]))
+
+  if args.max_token_length is not None:
+    if args.v2:
+        # v2: <BOS>...<EOS> <PAD> ... の三連を <BOS>...<EOS> <PAD> ... へ戻す　正直この実装でいいのかわからん
+      states_list = [encoder_hidden_states[:, 0].unsqueeze(1)]                              # <BOS>
+      for i in range(1, args.max_token_length, tokenizer.model_max_length):
+        chunk = encoder_hidden_states[:, i:i + tokenizer.model_max_length - 2]              # <BOS> の後から 最後の前まで
+        if i > 0:
+          for j in range(len(chunk)):
+            if input_ids[j, 1] == tokenizer.eos_token:                                      # 空、つまり <BOS> <EOS> <PAD> ...のパターン
+              chunk[j, 0] = chunk[j, 1]                                                     # 次の <PAD> の値をコピーする
+        states_list.append(chunk)  # <BOS> の後から <EOS> の前まで
+      states_list.append(encoder_hidden_states[:, -1].unsqueeze(1))                         # <EOS> か <PAD> のどちらか
+      encoder_hidden_states = torch.cat(states_list, dim=1)
+    else:
+      # v1: <BOS>...<EOS> の三連を <BOS>...<EOS> へ戻す
+      states_list = [encoder_hidden_states[:, 0].unsqueeze(1)]                              # <BOS>
+      for i in range(1, args.max_token_length, tokenizer.model_max_length):
+        states_list.append(encoder_hidden_states[:, i:i + tokenizer.model_max_length - 2])  # <BOS> の後から <EOS> の前まで
+      states_list.append(encoder_hidden_states[:, -1].unsqueeze(1))                         # <EOS>
+      encoder_hidden_states = torch.cat(states_list, dim=1)
+
+  if weight_dtype is not None:
+    # this is required for additional network training
+    encoder_hidden_states = encoder_hidden_states.to(weight_dtype)
+
+  return encoder_hidden_states
+
+
+def get_epoch_ckpt_name(args: argparse.Namespace, use_safetensors, epoch):
+  model_name = DEFAULT_EPOCH_NAME if args.output_name is None else args.output_name
+  ckpt_name = EPOCH_FILE_NAME.format(model_name, epoch) + (".safetensors" if use_safetensors else ".ckpt")
+  return model_name, ckpt_name
+
+
+def save_on_epoch_end(args: argparse.Namespace, save_func, remove_old_func, epoch_no: int, num_train_epochs: int):
+  saving = epoch_no % args.save_every_n_epochs == 0 and epoch_no < num_train_epochs
+  if saving:
+    os.makedirs(args.output_dir, exist_ok=True)
+    save_func()
+
+    if args.save_last_n_epochs is not None:
+      remove_epoch_no = epoch_no - args.save_every_n_epochs * args.save_last_n_epochs
+      remove_old_func(remove_epoch_no)
+  return saving
+
+
+def save_sd_model_on_epoch_end(args: argparse.Namespace, accelerator, src_path: str, save_stable_diffusion_format: bool, use_safetensors: bool, save_dtype: torch.dtype, epoch: int, num_train_epochs: int, global_step: int, text_encoder, unet, vae):
+  epoch_no = epoch + 1
+  model_name, ckpt_name = get_epoch_ckpt_name(args, use_safetensors, epoch_no)
+
+  if save_stable_diffusion_format:
+    def save_sd():
+      ckpt_file = os.path.join(args.output_dir, ckpt_name)
+      print(f"saving checkpoint: {ckpt_file}")
+      model_util.save_stable_diffusion_checkpoint(args.v2, ckpt_file, text_encoder, unet,
+                                                  src_path, epoch_no, global_step, save_dtype, vae)
+
+    def remove_sd(old_epoch_no):
+      _, old_ckpt_name = get_epoch_ckpt_name(args,  use_safetensors, old_epoch_no)
+      old_ckpt_file = os.path.join(args.output_dir, old_ckpt_name)
+      if os.path.exists(old_ckpt_file):
+        print(f"removing old checkpoint: {old_ckpt_file}")
+        os.remove(old_ckpt_file)
+
+    save_func = save_sd
+    remove_old_func = remove_sd
+  else:
+    def save_du():
+      out_dir = os.path.join(args.output_dir, EPOCH_DIFFUSERS_DIR_NAME.format(model_name, epoch_no))
+      print(f"saving model: {out_dir}")
+      os.makedirs(out_dir, exist_ok=True)
+      model_util.save_diffusers_checkpoint(args.v2, out_dir, text_encoder, unet,
+                                           src_path, vae=vae, use_safetensors=use_safetensors)
+
+    def remove_du(old_epoch_no):
+      out_dir_old = os.path.join(args.output_dir, EPOCH_DIFFUSERS_DIR_NAME.format(model_name, old_epoch_no))
+      if os.path.exists(out_dir_old):
+        print(f"removing old model: {out_dir_old}")
+        shutil.rmtree(out_dir_old)
+
+    save_func = save_du
+    remove_old_func = remove_du
+
+  saving = save_on_epoch_end(args, save_func, remove_old_func, epoch_no, num_train_epochs)
+  if saving and args.save_state:
+    save_state_on_epoch_end(args, accelerator, model_name, epoch_no)
+
+
+def save_state_on_epoch_end(args: argparse.Namespace, accelerator, model_name, epoch_no):
+  print("saving state.")
+  accelerator.save_state(os.path.join(args.output_dir, EPOCH_STATE_NAME.format(model_name, epoch_no)))
+
+  last_n_epochs = args.save_last_n_epochs_state if args.save_last_n_epochs_state else args.save_last_n_epochs
+  if last_n_epochs is not None:
+    remove_epoch_no = epoch_no - args.save_every_n_epochs * last_n_epochs
+    state_dir_old = os.path.join(args.output_dir, EPOCH_STATE_NAME.format(model_name, remove_epoch_no))
+    if os.path.exists(state_dir_old):
+      print(f"removing old state: {state_dir_old}")
+      shutil.rmtree(state_dir_old)
+
+
+def save_sd_model_on_train_end(args: argparse.Namespace, src_path: str, save_stable_diffusion_format: bool, use_safetensors: bool, save_dtype: torch.dtype, epoch: int, global_step: int, text_encoder, unet, vae):
+  model_name = DEFAULT_LAST_OUTPUT_NAME if args.output_name is None else args.output_name
+
+  if save_stable_diffusion_format:
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    ckpt_name = model_name + (".safetensors" if use_safetensors else ".ckpt")
+    ckpt_file = os.path.join(args.output_dir, ckpt_name)
+
+    print(f"save trained model as StableDiffusion checkpoint to {ckpt_file}")
+    model_util.save_stable_diffusion_checkpoint(args.v2, ckpt_file, text_encoder, unet,
+                                                src_path, epoch, global_step, save_dtype, vae)
+  else:
+    out_dir = os.path.join(args.output_dir, model_name)
+    os.makedirs(out_dir, exist_ok=True)
+
+    print(f"save trained model as Diffusers to {out_dir}")
+    model_util.save_diffusers_checkpoint(args.v2, out_dir, text_encoder, unet,
+                                         src_path, vae=vae, use_safetensors=use_safetensors)
+
+
+def save_state_on_train_end(args: argparse.Namespace, accelerator):
+  print("saving last state.")
+  os.makedirs(args.output_dir, exist_ok=True)
+  model_name = DEFAULT_LAST_OUTPUT_NAME if args.output_name is None else args.output_name
+  accelerator.save_state(os.path.join(args.output_dir, LAST_STATE_NAME.format(model_name)))
+
+# endregion
+
+# region 前処理用
+
+
+class ImageLoadingDataset(torch.utils.data.Dataset):
+  def __init__(self, image_paths):
+    self.images = image_paths
+
+  def __len__(self):
+    return len(self.images)
+
+  def __getitem__(self, idx):
+    img_path = self.images[idx]
+
+    try:
+      image = Image.open(img_path).convert("RGB")
+      # convert to tensor temporarily so dataloader will accept it
+      tensor_pil = transforms.functional.pil_to_tensor(image)
+    except Exception as e:
+      print(f"Could not load image path / 画像を読み込めません: {img_path}, error: {e}")
+      return None
+
+    return (tensor_pil, img_path)
+
+
+# endregion

locon/kohya_model_utils.py

@@ -0,0 +1,1184 @@
+'''
+https://github.com/kohya-ss/sd-scripts/blob/main/library/model_util.py
+'''
+# v1: split from train_db_fixed.py.
+# v2: support safetensors
+
+import math
+import os
+import torch
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig
+from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNet2DConditionModel
+from safetensors.torch import load_file, save_file
+
+# DiffUsers版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
+
+# V2
+V2_UNET_PARAMS_ATTENTION_HEAD_DIM = [5, 10, 20, 20]
+V2_UNET_PARAMS_CONTEXT_DIM = 1024
+
+# Diffusersの設定を読み込むための参照モデル
+DIFFUSERS_REF_MODEL_ID_V1 = "runwayml/stable-diffusion-v1-5"
+DIFFUSERS_REF_MODEL_ID_V2 = "stabilityai/stable-diffusion-2-1"
+
+
+# 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 linear_transformer_to_conv(checkpoint):
+  keys = list(checkpoint.keys())
+  tf_keys = ["proj_in.weight", "proj_out.weight"]
+  for key in keys:
+    if ".".join(key.split(".")[-2:]) in tf_keys:
+      if checkpoint[key].ndim == 2:
+        checkpoint[key] = checkpoint[key].unsqueeze(2).unsqueeze(2)
+
+
+def convert_ldm_unet_checkpoint(v2, 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"])
+
+      # biasとweightの順番に依存しないようにする：もっといいやり方がありそうだが
+      for l in output_block_list.values():
+        l.sort()
+
+      if ["conv.bias", "conv.weight"] in output_block_list.values():
+        index = list(output_block_list.values()).index(["conv.bias", "conv.weight"])
+        new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[
+            f"output_blocks.{i}.{index}.conv.bias"
+        ]
+        new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[
+            f"output_blocks.{i}.{index}.conv.weight"
+        ]
+
+        # 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]
+
+  # SDのv2では1*1のconv2dがlinearに変わっているので、linear->convに変換する
+  if v2:
+    linear_transformer_to_conv(new_checkpoint)
+
+  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)
+  # if len(vae_state_dict) == 0:
+  #   # 渡されたcheckpointは.ckptから読み込んだcheckpointではなくvaeのstate_dict
+  #   vae_state_dict = checkpoint
+
+  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(v2):
+  """
+  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 if not v2 else V2_UNET_PARAMS_CONTEXT_DIM,
+      attention_head_dim=UNET_PARAMS_NUM_HEADS if not v2 else V2_UNET_PARAMS_ATTENTION_HEAD_DIM,
+  )
+
+  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_v1(checkpoint):
+  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]
+  return text_model_dict
+
+
+def convert_ldm_clip_checkpoint_v2(checkpoint, max_length):
+  # 嫌になるくらい違うぞ！
+  def convert_key(key):
+    if not key.startswith("cond_stage_model"):
+      return None
+
+    # common conversion
+    key = key.replace("cond_stage_model.model.transformer.", "text_model.encoder.")
+    key = key.replace("cond_stage_model.model.", "text_model.")
+
+    if "resblocks" in key:
+      # resblocks conversion
+      key = key.replace(".resblocks.", ".layers.")
+      if ".ln_" in key:
+        key = key.replace(".ln_", ".layer_norm")
+      elif ".mlp." in key:
+        key = key.replace(".c_fc.", ".fc1.")
+        key = key.replace(".c_proj.", ".fc2.")
+      elif '.attn.out_proj' in key:
+        key = key.replace(".attn.out_proj.", ".self_attn.out_proj.")
+      elif '.attn.in_proj' in key:
+        key = None                  # 特殊なので後で処理する
+      else:
+        raise ValueError(f"unexpected key in SD: {key}")
+    elif '.positional_embedding' in key:
+      key = key.replace(".positional_embedding", ".embeddings.position_embedding.weight")
+    elif '.text_projection' in key:
+      key = None    # 使われない???
+    elif '.logit_scale' in key:
+      key = None    # 使われない???
+    elif '.token_embedding' in key:
+      key = key.replace(".token_embedding.weight", ".embeddings.token_embedding.weight")
+    elif '.ln_final' in key:
+      key = key.replace(".ln_final", ".final_layer_norm")
+    return key
+
+  keys = list(checkpoint.keys())
+  new_sd = {}
+  for key in keys:
+    # remove resblocks 23
+    if '.resblocks.23.' in key:
+      continue
+    new_key = convert_key(key)
+    if new_key is None:
+      continue
+    new_sd[new_key] = checkpoint[key]
+
+  # attnの変換
+  for key in keys:
+    if '.resblocks.23.' in key:
+      continue
+    if '.resblocks' in key and '.attn.in_proj_' in key:
+      # 三つに分割
+      values = torch.chunk(checkpoint[key], 3)
+
+      key_suffix = ".weight" if "weight" in key else ".bias"
+      key_pfx = key.replace("cond_stage_model.model.transformer.resblocks.", "text_model.encoder.layers.")
+      key_pfx = key_pfx.replace("_weight", "")
+      key_pfx = key_pfx.replace("_bias", "")
+      key_pfx = key_pfx.replace(".attn.in_proj", ".self_attn.")
+      new_sd[key_pfx + "q_proj" + key_suffix] = values[0]
+      new_sd[key_pfx + "k_proj" + key_suffix] = values[1]
+      new_sd[key_pfx + "v_proj" + key_suffix] = values[2]
+
+  # rename or add position_ids
+  ANOTHER_POSITION_IDS_KEY = "text_model.encoder.text_model.embeddings.position_ids"
+  if ANOTHER_POSITION_IDS_KEY in new_sd:
+    # waifu diffusion v1.4
+    position_ids = new_sd[ANOTHER_POSITION_IDS_KEY]
+    del new_sd[ANOTHER_POSITION_IDS_KEY]
+  else:
+    position_ids = torch.Tensor([list(range(max_length))]).to(torch.int64)
+
+  new_sd["text_model.embeddings.position_ids"] = position_ids
+  return new_sd
+
+# endregion
+
+
+# region Diffusers->StableDiffusion の変換コード
+# convert_diffusers_to_original_stable_diffusion をコピーして修正している（ASL 2.0）
+
+def conv_transformer_to_linear(checkpoint):
+  keys = list(checkpoint.keys())
+  tf_keys = ["proj_in.weight", "proj_out.weight"]
+  for key in keys:
+    if ".".join(key.split(".")[-2:]) in tf_keys:
+      if checkpoint[key].ndim > 2:
+        checkpoint[key] = checkpoint[key][:, :, 0, 0]
+
+
+def convert_unet_state_dict_to_sd(v2, 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()}
+
+  if v2:
+    conv_transformer_to_linear(new_state_dict)
+
+  return new_state_dict
+
+
+# ================#
+# VAE Conversion #
+# ================#
+
+def reshape_weight_for_sd(w):
+    # convert HF linear weights to SD conv2d weights
+  return w.reshape(*w.shape, 1, 1)
+
+
+def convert_vae_state_dict(vae_state_dict):
+  vae_conversion_map = [
+      # (stable-diffusion, HF Diffusers)
+      ("nin_shortcut", "conv_shortcut"),
+      ("norm_out", "conv_norm_out"),
+      ("mid.attn_1.", "mid_block.attentions.0."),
+  ]
+
+  for i in range(4):
+    # down_blocks have two resnets
+    for j in range(2):
+      hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}."
+      sd_down_prefix = f"encoder.down.{i}.block.{j}."
+      vae_conversion_map.append((sd_down_prefix, hf_down_prefix))
+
+    if i < 3:
+      hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0."
+      sd_downsample_prefix = f"down.{i}.downsample."
+      vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))
+
+      hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+      sd_upsample_prefix = f"up.{3-i}.upsample."
+      vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))
+
+    # up_blocks have three resnets
+    # also, up blocks in hf are numbered in reverse from sd
+    for j in range(3):
+      hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}."
+      sd_up_prefix = f"decoder.up.{3-i}.block.{j}."
+      vae_conversion_map.append((sd_up_prefix, hf_up_prefix))
+
+  # this part accounts for mid blocks in both the encoder and the decoder
+  for i in range(2):
+    hf_mid_res_prefix = f"mid_block.resnets.{i}."
+    sd_mid_res_prefix = f"mid.block_{i+1}."
+    vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+  vae_conversion_map_attn = [
+      # (stable-diffusion, HF Diffusers)
+      ("norm.", "group_norm."),
+      ("q.", "query."),
+      ("k.", "key."),
+      ("v.", "value."),
+      ("proj_out.", "proj_attn."),
+  ]
+
+  mapping = {k: k for k in vae_state_dict.keys()}
+  for k, v in mapping.items():
+    for sd_part, hf_part in vae_conversion_map:
+      v = v.replace(hf_part, sd_part)
+    mapping[k] = v
+  for k, v in mapping.items():
+    if "attentions" in k:
+      for sd_part, hf_part in vae_conversion_map_attn:
+        v = v.replace(hf_part, sd_part)
+      mapping[k] = v
+  new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()}
+  weights_to_convert = ["q", "k", "v", "proj_out"]
+  for k, v in new_state_dict.items():
+    for weight_name in weights_to_convert:
+      if f"mid.attn_1.{weight_name}.weight" in k:
+        # print(f"Reshaping {k} for SD format")
+        new_state_dict[k] = reshape_weight_for_sd(v)
+
+  return new_state_dict
+
+
+# endregion
+
+# region 自作のモデル読み書きなど
+
+def is_safetensors(path):
+  return os.path.splitext(path)[1].lower() == '.safetensors'
+
+
+def load_checkpoint_with_text_encoder_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.')
+  ]
+
+  if is_safetensors(ckpt_path):
+    checkpoint = None
+    state_dict = load_file(ckpt_path, "cpu")
+  else:
+    checkpoint = torch.load(ckpt_path, map_location="cpu")
+    if "state_dict" in checkpoint:
+      state_dict = checkpoint["state_dict"]
+    else:
+      state_dict = checkpoint
+      checkpoint = None
+
+  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, state_dict
+
+
+# TODO dtype指定の動作が怪しいので確認する text_encoderを指定形式で作れるか未確認
+def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, dtype=None):
+  _, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path)
+  if dtype is not None:
+    for k, v in state_dict.items():
+      if type(v) is torch.Tensor:
+        state_dict[k] = v.to(dtype)
+
+  # Convert the UNet2DConditionModel model.
+  unet_config = create_unet_diffusers_config(v2)
+  converted_unet_checkpoint = convert_ldm_unet_checkpoint(v2, state_dict, unet_config)
+
+  unet = UNet2DConditionModel(**unet_config)
+  info = unet.load_state_dict(converted_unet_checkpoint)
+  print("loading u-net:", info)
+
+  # 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)
+  info = vae.load_state_dict(converted_vae_checkpoint)
+  print("loading vae:", info)
+
+  # convert text_model
+  if v2:
+    converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v2(state_dict, 77)
+    cfg = CLIPTextConfig(
+        vocab_size=49408,
+        hidden_size=1024,
+        intermediate_size=4096,
+        num_hidden_layers=23,
+        num_attention_heads=16,
+        max_position_embeddings=77,
+        hidden_act="gelu",
+        layer_norm_eps=1e-05,
+        dropout=0.0,
+        attention_dropout=0.0,
+        initializer_range=0.02,
+        initializer_factor=1.0,
+        pad_token_id=1,
+        bos_token_id=0,
+        eos_token_id=2,
+        model_type="clip_text_model",
+        projection_dim=512,
+        torch_dtype="float32",
+        transformers_version="4.25.0.dev0",
+    )
+    text_model = CLIPTextModel._from_config(cfg)
+    info = text_model.load_state_dict(converted_text_encoder_checkpoint)
+  else:
+    converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v1(state_dict)
+    text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+    info = text_model.load_state_dict(converted_text_encoder_checkpoint)
+  print("loading text encoder:", info)
+
+  return text_model, vae, unet
+
+
+def convert_text_encoder_state_dict_to_sd_v2(checkpoint, make_dummy_weights=False):
+  def convert_key(key):
+    # position_idsの除去
+    if ".position_ids" in key:
+      return None
+
+    # common
+    key = key.replace("text_model.encoder.", "transformer.")
+    key = key.replace("text_model.", "")
+    if "layers" in key:
+      # resblocks conversion
+      key = key.replace(".layers.", ".resblocks.")
+      if ".layer_norm" in key:
+        key = key.replace(".layer_norm", ".ln_")
+      elif ".mlp." in key:
+        key = key.replace(".fc1.", ".c_fc.")
+        key = key.replace(".fc2.", ".c_proj.")
+      elif '.self_attn.out_proj' in key:
+        key = key.replace(".self_attn.out_proj.", ".attn.out_proj.")
+      elif '.self_attn.' in key:
+        key = None                  # 特殊なので後で処理する
+      else:
+        raise ValueError(f"unexpected key in DiffUsers model: {key}")
+    elif '.position_embedding' in key:
+      key = key.replace("embeddings.position_embedding.weight", "positional_embedding")
+    elif '.token_embedding' in key:
+      key = key.replace("embeddings.token_embedding.weight", "token_embedding.weight")
+    elif 'final_layer_norm' in key:
+      key = key.replace("final_layer_norm", "ln_final")
+    return key
+
+  keys = list(checkpoint.keys())
+  new_sd = {}
+  for key in keys:
+    new_key = convert_key(key)
+    if new_key is None:
+      continue
+    new_sd[new_key] = checkpoint[key]
+
+  # attnの変換
+  for key in keys:
+    if 'layers' in key and 'q_proj' in key:
+      # 三つを結合
+      key_q = key
+      key_k = key.replace("q_proj", "k_proj")
+      key_v = key.replace("q_proj", "v_proj")
+
+      value_q = checkpoint[key_q]
+      value_k = checkpoint[key_k]
+      value_v = checkpoint[key_v]
+      value = torch.cat([value_q, value_k, value_v])
+
+      new_key = key.replace("text_model.encoder.layers.", "transformer.resblocks.")
+      new_key = new_key.replace(".self_attn.q_proj.", ".attn.in_proj_")
+      new_sd[new_key] = value
+
+  # 最後の層などを捏造するか
+  if make_dummy_weights:
+    print("make dummy weights for resblock.23, text_projection and logit scale.")
+    keys = list(new_sd.keys())
+    for key in keys:
+      if key.startswith("transformer.resblocks.22."):
+        new_sd[key.replace(".22.", ".23.")] = new_sd[key].clone()          # copyしないとsafetensorsの保存で落ちる
+
+    # Diffusersに含まれない重みを作っておく
+    new_sd['text_projection'] = torch.ones((1024, 1024), dtype=new_sd[keys[0]].dtype, device=new_sd[keys[0]].device)
+    new_sd['logit_scale'] = torch.tensor(1)
+
+  return new_sd
+
+
+def save_stable_diffusion_checkpoint(v2, output_file, text_encoder, unet, ckpt_path, epochs, steps, save_dtype=None, vae=None):
+  if ckpt_path is not None:
+    # epoch/stepを参照する。またVAEがメモリ上にないときなど、もう一度VAEを含めて読み込む
+    checkpoint, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path)
+    if checkpoint is None:                # safetensors または state_dictのckpt
+      checkpoint = {}
+      strict = False
+    else:
+      strict = True
+    if "state_dict" in state_dict:
+      del state_dict["state_dict"]
+  else:
+    # 新しく作る
+    assert vae is not None, "VAE is required to save a checkpoint without a given checkpoint"
+    checkpoint = {}
+    state_dict = {}
+    strict = False
+
+  def update_sd(prefix, sd):
+    for k, v in sd.items():
+      key = prefix + k
+      assert not strict or key in state_dict, f"Illegal key in save SD: {key}"
+      if save_dtype is not None:
+        v = v.detach().clone().to("cpu").to(save_dtype)
+      state_dict[key] = v
+
+  # Convert the UNet model
+  unet_state_dict = convert_unet_state_dict_to_sd(v2, unet.state_dict())
+  update_sd("model.diffusion_model.", unet_state_dict)
+
+  # Convert the text encoder model
+  if v2:
+    make_dummy = ckpt_path is None                 # 参照元のcheckpoint���ない場合は最後の層を前の層から複製して作るなどダミーの重みを入れる
+    text_enc_dict = convert_text_encoder_state_dict_to_sd_v2(text_encoder.state_dict(), make_dummy)
+    update_sd("cond_stage_model.model.", text_enc_dict)
+  else:
+    text_enc_dict = text_encoder.state_dict()
+    update_sd("cond_stage_model.transformer.", text_enc_dict)
+
+  # Convert the VAE
+  if vae is not None:
+    vae_dict = convert_vae_state_dict(vae.state_dict())
+    update_sd("first_stage_model.", vae_dict)
+
+  # Put together new checkpoint
+  key_count = len(state_dict.keys())
+  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
+
+  if is_safetensors(output_file):
+    # TODO Tensor以外のdictの値を削除したほうがいいか
+    save_file(state_dict, output_file)
+  else:
+    torch.save(new_ckpt, output_file)
+
+  return key_count
+
+
+def save_diffusers_checkpoint(v2, output_dir, text_encoder, unet, pretrained_model_name_or_path, vae=None, use_safetensors=False):
+  if pretrained_model_name_or_path is None:
+    # load default settings for v1/v2
+    if v2:
+      pretrained_model_name_or_path = DIFFUSERS_REF_MODEL_ID_V2
+    else:
+      pretrained_model_name_or_path = DIFFUSERS_REF_MODEL_ID_V1
+
+  scheduler = DDIMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
+  tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer")
+  if vae is None:
+    vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae")
+
+  pipeline = StableDiffusionPipeline(
+      unet=unet,
+      text_encoder=text_encoder,
+      vae=vae,
+      scheduler=scheduler,
+      tokenizer=tokenizer,
+      safety_checker=None,
+      feature_extractor=None,
+      requires_safety_checker=None,
+  )
+  pipeline.save_pretrained(output_dir, safe_serialization=use_safetensors)
+
+
+VAE_PREFIX = "first_stage_model."
+
+
+def load_vae(vae_id, dtype):
+  print(f"load VAE: {vae_id}")
+  if os.path.isdir(vae_id) or not os.path.isfile(vae_id):
+    # Diffusers local/remote
+    try:
+      vae = AutoencoderKL.from_pretrained(vae_id, subfolder=None, torch_dtype=dtype)
+    except EnvironmentError as e:
+      print(f"exception occurs in loading vae: {e}")
+      print("retry with subfolder='vae'")
+      vae = AutoencoderKL.from_pretrained(vae_id, subfolder="vae", torch_dtype=dtype)
+    return vae
+
+  # local
+  vae_config = create_vae_diffusers_config()
+
+  if vae_id.endswith(".bin"):
+    # SD 1.5 VAE on Huggingface
+    converted_vae_checkpoint = torch.load(vae_id, map_location="cpu")
+  else:
+    # StableDiffusion
+    vae_model = (load_file(vae_id, "cpu") if is_safetensors(vae_id)
+                 else torch.load(vae_id, map_location="cpu"))
+    vae_sd = vae_model['state_dict'] if 'state_dict' in vae_model else vae_model
+
+    # vae only or full model
+    full_model = False
+    for vae_key in vae_sd:
+      if vae_key.startswith(VAE_PREFIX):
+        full_model = True
+        break
+    if not full_model:
+      sd = {}
+      for key, value in vae_sd.items():
+        sd[VAE_PREFIX + key] = value
+      vae_sd = sd
+      del sd
+
+    # Convert the VAE model.
+    converted_vae_checkpoint = convert_ldm_vae_checkpoint(vae_sd, vae_config)
+
+  vae = AutoencoderKL(**vae_config)
+  vae.load_state_dict(converted_vae_checkpoint)
+  return vae
+
+# 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)

locon/kohya_utils.py

@@ -0,0 +1,48 @@
+# part of https://github.com/kohya-ss/sd-scripts/blob/main/library/train_util.py
+
+import hashlib
+import safetensors
+from io import BytesIO
+
+
+def addnet_hash_legacy(b):
+    """Old model hash used by sd-webui-additional-networks for .safetensors format files"""
+    m = hashlib.sha256()
+
+    b.seek(0x100000)
+    m.update(b.read(0x10000))
+    return m.hexdigest()[0:8]
+
+
+def addnet_hash_safetensors(b):
+    """New model hash used by sd-webui-additional-networks for .safetensors format files"""
+    hash_sha256 = hashlib.sha256()
+    blksize = 1024 * 1024
+
+    b.seek(0)
+    header = b.read(8)
+    n = int.from_bytes(header, "little")
+
+    offset = n + 8
+    b.seek(offset)
+    for chunk in iter(lambda: b.read(blksize), b""):
+        hash_sha256.update(chunk)
+
+    return hash_sha256.hexdigest()
+
+
+def precalculate_safetensors_hashes(tensors, metadata):
+    """Precalculate the model hashes needed by sd-webui-additional-networks to
+    save time on indexing the model later."""
+
+    # Because writing user metadata to the file can change the result of
+    # sd_models.model_hash(), only retain the training metadata for purposes of
+    # calculating the hash, as they are meant to be immutable
+    metadata = {k: v for k, v in metadata.items() if k.startswith("ss_")}
+
+    bytes = safetensors.torch.save(tensors, metadata)
+    b = BytesIO(bytes)
+
+    model_hash = addnet_hash_safetensors(b)
+    legacy_hash = addnet_hash_legacy(b)
+    return model_hash, legacy_hash

locon/locon.py

@@ -0,0 +1,53 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LoConModule(nn.Module):
+    """
+    modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
+    """
+
+    def __init__(self, lora_name, org_module: nn.Module, multiplier=1.0, lora_dim=4, alpha=1):
+        """ if alpha == 0 or None, alpha is rank (no scaling). """
+        super().__init__()
+        self.lora_name = lora_name
+        self.lora_dim = lora_dim
+
+        if org_module.__class__.__name__ == 'Conv2d':
+            # For general LoCon
+            in_dim = org_module.in_channels
+            k_size = org_module.kernel_size
+            stride = org_module.stride
+            padding = org_module.padding
+            out_dim = org_module.out_channels
+            self.lora_down = nn.Conv2d(in_dim, lora_dim, k_size, stride, padding, bias=False)
+            self.lora_up = nn.Conv2d(lora_dim, out_dim, (1, 1), bias=False)
+        else:
+            in_dim = org_module.in_features
+            out_dim = org_module.out_features
+            self.lora_down = nn.Linear(in_dim, lora_dim, bias=False)
+            self.lora_up = nn.Linear(lora_dim, out_dim, bias=False)
+
+        if type(alpha) == torch.Tensor:
+            alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
+        alpha = lora_dim if alpha is None or alpha == 0 else alpha
+        self.scale = alpha / self.lora_dim
+        self.register_buffer('alpha', torch.tensor(alpha)) # 定数として扱える
+
+        # same as microsoft's
+        torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
+        torch.nn.init.zeros_(self.lora_up.weight)
+
+        self.multiplier = multiplier
+        self.org_module = org_module # remove in applying
+
+    def apply_to(self):
+        self.org_forward = self.org_module.forward
+        self.org_module.forward = self.forward
+        del self.org_module
+
+    def forward(self, x):
+        return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale

locon/locon_kohya.py

@@ -0,0 +1,243 @@
+# LoCon network module
+# reference:
+# https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
+# https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/lora.py
+# https://github.com/kohya-ss/sd-scripts/blob/main/networks/lora.py
+
+import math
+import os
+from typing import List
+import torch
+
+from .kohya_utils import *
+from .locon import LoConModule
+
+
+def create_network(multiplier, network_dim, network_alpha, vae, text_encoder, unet, **kwargs):
+    if network_dim is None:
+        network_dim = 4                     # default
+    conv_dim = kwargs.get('conv_dim', network_dim)
+    conv_alpha = kwargs.get('conv_alpha', network_alpha)
+    network = LoRANetwork(
+        text_encoder, unet, 
+        multiplier=multiplier, 
+        lora_dim=network_dim, conv_lora_dim=conv_dim, 
+        alpha=network_alpha, conv_alpha=conv_alpha
+    )
+    return network
+
+
+def create_network_from_weights(multiplier, file, vae, text_encoder, unet, **kwargs):
+    if os.path.splitext(file)[1] == '.safetensors':
+        from safetensors.torch import load_file, safe_open
+        weights_sd = load_file(file)
+    else:
+        weights_sd = torch.load(file, map_location='cpu')
+
+    # get dim (rank)
+    network_alpha = None
+    network_dim = None
+    for key, value in weights_sd.items():
+        if network_alpha is None and 'alpha' in key:
+            network_alpha = value
+        if network_dim is None and 'lora_down' in key and len(value.size()) == 2:
+            network_dim = value.size()[0]
+
+    if network_alpha is None:
+        network_alpha = network_dim
+
+    network = LoRANetwork(text_encoder, unet, multiplier=multiplier, lora_dim=network_dim, alpha=network_alpha)
+    network.weights_sd = weights_sd
+    return network
+
+torch.nn.Conv2d
+class LoRANetwork(torch.nn.Module):
+    '''
+    LoRA + LoCon
+    '''
+    # Ignore proj_in or proj_out, their channels is only a few.
+    UNET_TARGET_REPLACE_MODULE = [
+        "Transformer2DModel", 
+        "Attention", 
+        "ResnetBlock2D", 
+        "Downsample2D", 
+        "Upsample2D"
+    ]
+    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
+    LORA_PREFIX_UNET = 'lora_unet'
+    LORA_PREFIX_TEXT_ENCODER = 'lora_te'
+
+    def __init__(
+        self, 
+        text_encoder, unet, 
+        multiplier=1.0, 
+        lora_dim=4, conv_lora_dim=4, 
+        alpha=1, conv_alpha=1
+    ) -> None:
+        super().__init__()
+        self.multiplier = multiplier
+        self.lora_dim = lora_dim
+        self.conv_lora_dim = int(conv_lora_dim)
+        if self.conv_lora_dim != self.lora_dim: 
+            print('Apply different lora dim for conv layer')
+            print(f'LoCon Dim: {conv_lora_dim}, LoRA Dim: {lora_dim}')
+        self.alpha = alpha
+        self.conv_alpha = float(conv_alpha)
+        if self.alpha != self.conv_alpha: 
+            print('Apply different alpha value for conv layer')
+            print(f'LoCon alpha: {conv_alpha}, LoRA alpha: {alpha}')
+
+        # create module instances
+        def create_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> List[LoConModule]:
+            print('Create LoCon Module')
+            loras = []
+            for name, module in root_module.named_modules():
+                if module.__class__.__name__ in target_replace_modules:
+                    for child_name, child_module in module.named_modules():
+                        lora_name = prefix + '.' + name + '.' + child_name
+                        lora_name = lora_name.replace('.', '_')
+                        if child_module.__class__.__name__ == 'Linear':
+                            lora = LoConModule(lora_name, child_module, self.multiplier, self.lora_dim, self.alpha)
+                        elif child_module.__class__.__name__ == 'Conv2d':
+                            k_size, *_ = child_module.kernel_size
+                            if k_size==1:
+                                lora = LoConModule(lora_name, child_module, self.multiplier, self.lora_dim, self.alpha)
+                            else:
+                                lora = LoConModule(lora_name, child_module, self.multiplier, self.conv_lora_dim, self.conv_alpha)
+                        else:
+                            continue
+                        loras.append(lora)
+            return loras
+
+        self.text_encoder_loras = create_modules(
+            LoRANetwork.LORA_PREFIX_TEXT_ENCODER,
+            text_encoder, 
+            LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE
+        )
+        print(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
+
+        self.unet_loras = create_modules(LoRANetwork.LORA_PREFIX_UNET, unet, LoRANetwork.UNET_TARGET_REPLACE_MODULE)
+        print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
+
+        self.weights_sd = None
+
+        # assertion
+        names = set()
+        for lora in self.text_encoder_loras + self.unet_loras:
+            assert lora.lora_name not in names, f"duplicated lora name: {lora.lora_name}"
+            names.add(lora.lora_name)
+
+    def set_multiplier(self, multiplier):
+        self.multiplier = multiplier
+        for lora in self.text_encoder_loras + self.unet_loras:
+            lora.multiplier = self.multiplier
+            
+    def load_weights(self, file):
+        if os.path.splitext(file)[1] == '.safetensors':
+            from safetensors.torch import load_file, safe_open
+            self.weights_sd = load_file(file)
+        else:
+            self.weights_sd = torch.load(file, map_location='cpu')
+
+    def apply_to(self, text_encoder, unet, apply_text_encoder=None, apply_unet=None):
+        if self.weights_sd:
+            weights_has_text_encoder = weights_has_unet = False
+            for key in self.weights_sd.keys():
+                if key.startswith(LoRANetwork.LORA_PREFIX_TEXT_ENCODER):
+                    weights_has_text_encoder = True
+                elif key.startswith(LoRANetwork.LORA_PREFIX_UNET):
+                    weights_has_unet = True
+
+            if apply_text_encoder is None:
+                apply_text_encoder = weights_has_text_encoder
+            else:
+                assert apply_text_encoder == weights_has_text_encoder, f"text encoder weights: {weights_has_text_encoder} but text encoder flag: {apply_text_encoder} / 重みとText Encoderのフラグが矛盾しています"
+
+            if apply_unet is None:
+                apply_unet = weights_has_unet
+            else:
+                assert apply_unet == weights_has_unet, f"u-net weights: {weights_has_unet} but u-net flag: {apply_unet} / 重みとU-Netのフラグが矛盾しています"
+        else:
+            assert apply_text_encoder is not None and apply_unet is not None, f"internal error: flag not set"
+
+        if apply_text_encoder:
+            print("enable LoRA for text encoder")
+        else:
+            self.text_encoder_loras = []
+
+        if apply_unet:
+            print("enable LoRA for U-Net")
+        else:
+            self.unet_loras = []
+
+        for lora in self.text_encoder_loras + self.unet_loras:
+            lora.apply_to()
+            self.add_module(lora.lora_name, lora)
+
+        if self.weights_sd:
+            # if some weights are not in state dict, it is ok because initial LoRA does nothing (lora_up is initialized by zeros)
+            info = self.load_state_dict(self.weights_sd, False)
+            print(f"weights are loaded: {info}")
+
+    def enable_gradient_checkpointing(self):
+        # not supported
+        pass
+
+    def prepare_optimizer_params(self, text_encoder_lr, unet_lr):
+        def enumerate_params(loras):
+            params = []
+            for lora in loras:
+                params.extend(lora.parameters())
+            return params
+
+        self.requires_grad_(True)
+        all_params = []
+
+        if self.text_encoder_loras:
+            param_data = {'params': enumerate_params(self.text_encoder_loras)}
+            if text_encoder_lr is not None:
+                param_data['lr'] = text_encoder_lr
+            all_params.append(param_data)
+
+        if self.unet_loras:
+            param_data = {'params': enumerate_params(self.unet_loras)}
+            if unet_lr is not None:
+                param_data['lr'] = unet_lr
+            all_params.append(param_data)
+
+        return all_params
+
+    def prepare_grad_etc(self, text_encoder, unet):
+        self.requires_grad_(True)
+
+    def on_epoch_start(self, text_encoder, unet):
+        self.train()
+
+    def get_trainable_params(self):
+        return self.parameters()
+
+    def save_weights(self, file, dtype, metadata):
+        if metadata is not None and len(metadata) == 0:
+            metadata = None
+
+        state_dict = self.state_dict()
+
+        if dtype is not None:
+            for key in list(state_dict.keys()):
+                v = state_dict[key]
+                v = v.detach().clone().to("cpu").to(dtype)
+                state_dict[key] = v
+
+        if os.path.splitext(file)[1] == '.safetensors':
+            from safetensors.torch import save_file
+
+            # Precalculate model hashes to save time on indexing
+            if metadata is None:
+                metadata = {}
+            model_hash, legacy_hash = precalculate_safetensors_hashes(state_dict, metadata)
+            metadata["sshs_model_hash"] = model_hash
+            metadata["sshs_legacy_hash"] = legacy_hash
+
+            save_file(state_dict, file, metadata)
+        else:
+            torch.save(state_dict, file)

locon/utils.py

@@ -0,0 +1,148 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import torch.linalg as linalg
+
+from tqdm import tqdm
+
+
+def extract_conv(
+    weight: nn.Parameter|torch.Tensor,
+    lora_rank = 8
+) -> tuple[nn.Parameter, nn.Parameter]:
+    out_ch, in_ch, kernel_size, _ = weight.shape
+    lora_rank = min(out_ch, in_ch, lora_rank)
+    
+    U, S, Vh = linalg.svd(weight.reshape(out_ch, -1))
+    
+    U = U[:, :lora_rank]
+    S = S[:lora_rank]
+    U = U @ torch.diag(S)
+    Vh = Vh[:lora_rank, :]
+    
+    extract_weight_A = Vh.reshape(lora_rank, in_ch, kernel_size, kernel_size).cpu()
+    extract_weight_B = U.reshape(out_ch, lora_rank, 1, 1).cpu()
+    del U, S, Vh, weight
+    return extract_weight_A, extract_weight_B
+
+
+def merge_conv(
+    weight_a: nn.Parameter|torch.Tensor,
+    weight_b: nn.Parameter|torch.Tensor,
+):
+    rank, in_ch, kernel_size, k_ = weight_a.shape
+    out_ch, rank_, _, _ = weight_b.shape
+    
+    assert rank == rank_ and kernel_size == k_
+    
+    merged = weight_b.reshape(out_ch, -1) @ weight_a.reshape(rank, -1)
+    weight = merged.reshape(out_ch, in_ch, kernel_size, kernel_size)
+    return weight
+
+
+def extract_linear(
+    weight: nn.Parameter|torch.Tensor,
+    lora_rank = 8
+) -> tuple[nn.Parameter, nn.Parameter]:
+    out_ch, in_ch = weight.shape
+    lora_rank = min(out_ch, in_ch, lora_rank)
+    
+    U, S, Vh = linalg.svd(weight)
+    
+    U = U[:, :lora_rank]
+    S = S[:lora_rank]
+    U = U @ torch.diag(S)
+    Vh = Vh[:lora_rank, :]
+    
+    extract_weight_A = Vh.reshape(lora_rank, in_ch).cpu()
+    extract_weight_B = U.reshape(out_ch, lora_rank).cpu()
+    del U, S, Vh, weight
+    return extract_weight_A, extract_weight_B
+
+
+def merge_linear(
+    weight_a: nn.Parameter|torch.Tensor,
+    weight_b: nn.Parameter|torch.Tensor,
+):
+    rank, in_ch = weight_a.shape
+    out_ch, rank_ = weight_b.shape
+    
+    assert rank == rank_
+    
+    weight = weight_b @ weight_a
+    return weight
+
+
+def extract_diff(
+    base_model,
+    db_model,
+    lora_dim=4, 
+    conv_lora_dim=4,
+    extract_device = 'cuda',
+):
+    UNET_TARGET_REPLACE_MODULE = [
+        "Transformer2DModel", 
+        "Attention", 
+        "ResnetBlock2D", 
+        "Downsample2D", 
+        "Upsample2D"
+    ]
+    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
+    LORA_PREFIX_UNET = 'lora_unet'
+    LORA_PREFIX_TEXT_ENCODER = 'lora_te'
+    def make_state_dict(
+        prefix, 
+        root_module: torch.nn.Module,
+        target_module: torch.nn.Module,
+        target_replace_modules
+    ):
+        loras = {}
+        temp = {}
+        
+        for name, module in root_module.named_modules():
+            if module.__class__.__name__ in target_replace_modules:
+                temp[name] = {}
+                for child_name, child_module in module.named_modules():
+                    if child_module.__class__.__name__ not in {'Linear', 'Conv2d'}:
+                        continue
+                    temp[name][child_name] = child_module.weight
+        
+        for name, module in tqdm(list(target_module.named_modules())):
+            if name in temp:
+                weights = temp[name]
+                for child_name, child_module in module.named_modules():
+                    lora_name = prefix + '.' + name + '.' + child_name
+                    lora_name = lora_name.replace('.', '_')
+                    if child_module.__class__.__name__ == 'Linear':
+                        extract_a, extract_b = extract_linear(
+                            (child_module.weight - weights[child_name]),
+                            lora_dim
+                        )
+                    elif child_module.__class__.__name__ == 'Conv2d':
+                        extract_a, extract_b = extract_conv(
+                            (child_module.weight - weights[child_name]), 
+                            conv_lora_dim
+                        )
+                    else:
+                        continue
+                    
+                    loras[f'{lora_name}.lora_down.weight'] = extract_a.detach().cpu().half()
+                    loras[f'{lora_name}.lora_up.weight'] = extract_b.detach().cpu().half()
+                    loras[f'{lora_name}.alpha'] = torch.Tensor([int(extract_a.shape[0])]).detach().cpu().half()
+                    del extract_a, extract_b
+        return loras
+    
+    text_encoder_loras = make_state_dict(
+        LORA_PREFIX_TEXT_ENCODER, 
+        base_model[0], db_model[0], 
+        TEXT_ENCODER_TARGET_REPLACE_MODULE
+    )
+    
+    unet_loras = make_state_dict(
+        LORA_PREFIX_UNET,
+        base_model[2], db_model[2], 
+        UNET_TARGET_REPLACE_MODULE
+    )
+    print(len(text_encoder_loras), len(unet_loras))
+    return text_encoder_loras|unet_loras

lora_train_popup.py

@@ -0,0 +1,862 @@
+import gc
+import json
+import time
+from functools import partial
+from typing import Union
+import os
+import tkinter as tk
+from tkinter import filedialog as fd, ttk
+from tkinter import simpledialog as sd
+from tkinter import messagebox as mb
+
+import torch.cuda
+import train_network
+import library.train_util as util
+import argparse
+
+
+class ArgStore:
+    # Represents the entirety of all possible inputs for sd-scripts. they are ordered from most important to least
+    def __init__(self):
+        # Important, these are the most likely things you will modify
+        self.base_model: str = r""  # example path, r"E:\sd\stable-diffusion-webui\models\Stable-diffusion\nai.ckpt"
+        self.img_folder: str = r""  # is the folder path to your img folder, make sure to follow the guide here for folder setup: https://rentry.org/2chAI_LoRA_Dreambooth_guide_english#for-kohyas-script
+        self.output_folder: str = r""  # just the folder all epochs/safetensors are output
+        self.change_output_name: Union[str, None] = None  # changes the output name of the epochs
+        self.save_json_folder: Union[str, None] = None  # OPTIONAL, saves a json folder of your config to whatever location you set here.
+        self.load_json_path: Union[str, None] = None  # OPTIONAL, loads a json file partially changes the config to match. things like folder paths do not get modified.
+        self.json_load_skip_list: Union[list[str], None] = ["save_json_folder", "reg_img_folder",
+                                                            "lora_model_for_resume", "change_output_name",
+                                                            "training_comment",
+                                                            "json_load_skip_list"]  # OPTIONAL, allows the user to define what they skip when loading a json, by default it loads everything, including all paths, set it up like this ["base_model", "img_folder", "output_folder"]
+        self.caption_dropout_rate: Union[float, None] = None  # The rate at which captions for files get dropped.
+        self.caption_dropout_every_n_epochs: Union[int, None] = None  # Defines how often an epoch will completely ignore
+        # captions, EX. 3 means it will ignore captions at epochs 3, 6, and 9
+        self.caption_tag_dropout_rate: Union[float, None] = None  # Defines the rate at which a tag would be dropped, rather than the entire caption file
+        self.noise_offset: Union[float, None] = None  # OPTIONAL, seems to help allow SD to gen better blacks and whites
+                                                      # Kohya recommends, if you have it set, to use 0.1, not sure how
+                                                      # high the value can be, I'm going to assume maximum of 1
+
+        self.net_dim: int = 128  # network dimension, 128 is the most common, however you might be able to get lesser to work
+        self.alpha: float = 128  # represents the scalar for training. the lower the alpha, the less gets learned per step. if you want the older way of training, set this to dim
+        # list of schedulers: linear, cosine, cosine_with_restarts, polynomial, constant, constant_with_warmup
+        self.scheduler: str = "cosine_with_restarts"  # the scheduler for learning rate. Each does something specific
+        self.cosine_restarts: Union[int, None] = 1  # OPTIONAL, represents the number of times it restarts. Only matters if you are using cosine_with_restarts
+        self.scheduler_power: Union[float, None] = 1  # OPTIONAL, represents the power of the polynomial. Only matters if you are using polynomial
+        self.warmup_lr_ratio: Union[float, None] = None  # OPTIONAL, Calculates the number of warmup steps based on the ratio given. Make sure to set this if you are using constant_with_warmup, None to ignore
+        self.learning_rate: Union[float, None] = 1e-4  # OPTIONAL, when not set, lr gets set to 1e-3 as per adamW. Personally, I suggest actually setting this as lower lr seems to be a small bit better.
+        self.text_encoder_lr: Union[float, None] = None  # OPTIONAL, Sets a specific lr for the text encoder, this overwrites the base lr I believe, None to ignore
+        self.unet_lr: Union[float, None] = None  # OPTIONAL, Sets a specific lr for the unet, this overwrites the base lr I believe, None to ignore
+        self.num_workers: int = 1  # The number of threads that are being used to load images, lower speeds up the start of epochs, but slows down the loading of data. The assumption here is that it increases the training time as you reduce this value
+        self.persistent_workers: bool = True  # makes workers persistent, further reduces/eliminates the lag in between epochs. however it may increase memory usage
+
+        self.batch_size: int = 1  # The number of images that get processed at one time, this is directly proportional to your vram and resolution. with 12gb of vram, at 512 reso, you can get a maximum of 6 batch size
+        self.num_epochs: int = 1  # The number of epochs, if you set max steps this value is ignored as it doesn't calculate steps.
+        self.save_every_n_epochs: Union[int, None] = None  # OPTIONAL, how often to save epochs, None to ignore
+        self.shuffle_captions: bool = False  # OPTIONAL, False to ignore
+        self.keep_tokens: Union[int, None] = None  # OPTIONAL, None to ignore
+        self.max_steps: Union[int, None] = None  # OPTIONAL, if you have specific steps you want to hit, this allows you to set it directly. None to ignore
+        self.tag_occurrence_txt_file: bool = False  # OPTIONAL, creates a txt file that has the entire occurrence of all tags in your dataset
+                                                    # the metadata will also have this so long as you have metadata on, so no reason to have this on by default
+                                                    # will automatically output to the same folder as your output checkpoints
+        self.sort_tag_occurrence_alphabetically: bool = False  # OPTIONAL, only applies if tag_occurrence_txt_file is also true
+                                                               # Will change the output to be alphabetically vs being occurrence based
+
+        # These are the second most likely things you will modify
+        self.train_resolution: int = 512
+        self.min_bucket_resolution: int = 320
+        self.max_bucket_resolution: int = 960
+        self.lora_model_for_resume: Union[str, None] = None  # OPTIONAL, takes an input lora to continue training from, not exactly the way it *should* be, but it works, None to ignore
+        self.save_state: bool = False  # OPTIONAL, is the intended way to save a training state to use for continuing training, False to ignore
+        self.load_previous_save_state: Union[str, None] = None  # OPTIONAL, is the intended way to load a training state to use for continuing training, None to ignore
+        self.training_comment: Union[str, None] = None  # OPTIONAL, great way to put in things like activation tokens right into the metadata. seems to not work at this point and time
+        self.unet_only: bool = False  # OPTIONAL, set it to only train the unet
+        self.text_only: bool = False  # OPTIONAL, set it to only train the text encoder
+
+        # These are the least likely things you will modify
+        self.reg_img_folder: Union[str, None] = None  # OPTIONAL, None to ignore
+        self.clip_skip: int = 2  # If you are training on a model that is anime based, keep this at 2 as most models are designed for that
+        self.test_seed: int = 23  # this is the "reproducable seed", basically if you set the seed to this, you should be able to input a prompt from one of your training images and get a close representation of it
+        self.prior_loss_weight: float = 1  # is the loss weight much like Dreambooth, is required for LoRA training
+        self.gradient_checkpointing: bool = False  # OPTIONAL, enables gradient checkpointing
+        self.gradient_acc_steps: Union[int, None] = None  # OPTIONAL, not sure exactly what this means
+        self.mixed_precision: str = "fp16"  # If you have the ability to use bf16, do it, it's better
+        self.save_precision: str = "fp16"  # You can also save in bf16, but because it's not universally supported, I suggest you keep saving at fp16
+        self.save_as: str = "safetensors"  # list is pt, ckpt, safetensors
+        self.caption_extension: str = ".txt"  # the other option is .captions, but since wd1.4 tagger outputs as txt files, this is the default
+        self.max_clip_token_length = 150  # can be 75, 150, or 225 I believe, there is no reason to go higher than 150 though
+        self.buckets: bool = True
+        self.xformers: bool = True
+        self.use_8bit_adam: bool = True
+        self.cache_latents: bool = True
+        self.color_aug: bool = False  # IMPORTANT: Clashes with cache_latents, only have one of the two on!
+        self.flip_aug: bool = False
+        self.vae: Union[str, None] = None  # Seems to only make results worse when not using that specific vae, should probably not use
+        self.no_meta: bool = False  # This removes the metadata that now gets saved into safetensors, (you should keep this on)
+        self.log_dir: Union[str, None] = None  # output of logs, not useful to most people.
+        self.v2: bool = False  # Sets up training for SD2.1
+        self.v_parameterization: bool = False  # Only is used when v2 is also set and you are using the 768x version of v2
+
+    # Creates the dict that is used for the rest of the code, to facilitate easier json saving and loading
+    @staticmethod
+    def convert_args_to_dict():
+        return ArgStore().__dict__
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    setup_args(parser)
+    pre_args = parser.parse_args()
+    queues = 0
+    args_queue = []
+    cont = True
+    while cont:
+        arg_dict = ArgStore.convert_args_to_dict()
+        ret = mb.askyesno(message="Do you want to load a json config file?")
+        if ret:
+            load_json(ask_file("select json to load from", {"json"}), arg_dict)
+            arg_dict = ask_elements_trunc(arg_dict)
+        else:
+            arg_dict = ask_elements(arg_dict)
+        if pre_args.save_json_path or arg_dict["save_json_folder"]:
+            save_json(pre_args.save_json_path if pre_args.save_json_path else arg_dict['save_json_folder'], arg_dict)
+        args = create_arg_space(arg_dict)
+        args = parser.parse_args(args)
+        queues += 1
+        args_queue.append(args)
+        if arg_dict['tag_occurrence_txt_file']:
+            get_occurrence_of_tags(arg_dict)
+        ret = mb.askyesno(message="Do you want to queue another training?")
+        if not ret:
+            cont = False
+    for args in args_queue:
+        try:
+            train_network.train(args)
+        except Exception as e:
+            print(f"Failed to train this set of args.\nSkipping this training session.\nError is: {e}")
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+def create_arg_space(args: dict) -> [str]:
+    # This is the list of args that are to be used regardless of setup
+    output = ["--network_module=networks.lora", f"--pretrained_model_name_or_path={args['base_model']}",
+              f"--train_data_dir={args['img_folder']}", f"--output_dir={args['output_folder']}",
+              f"--prior_loss_weight={args['prior_loss_weight']}", f"--caption_extension=" + args['caption_extension'],
+              f"--resolution={args['train_resolution']}", f"--train_batch_size={args['batch_size']}",
+              f"--mixed_precision={args['mixed_precision']}", f"--save_precision={args['save_precision']}",
+              f"--network_dim={args['net_dim']}", f"--save_model_as={args['save_as']}",
+              f"--clip_skip={args['clip_skip']}", f"--seed={args['test_seed']}",
+              f"--max_token_length={args['max_clip_token_length']}", f"--lr_scheduler={args['scheduler']}",
+              f"--network_alpha={args['alpha']}", f"--max_data_loader_n_workers={args['num_workers']}"]
+    if not args['max_steps']:
+        output.append(f"--max_train_epochs={args['num_epochs']}")
+        output += create_optional_args(args, find_max_steps(args))
+    else:
+        output.append(f"--max_train_steps={args['max_steps']}")
+        output += create_optional_args(args, args['max_steps'])
+    return output
+
+
+def create_optional_args(args: dict, steps):
+    output = []
+    if args["reg_img_folder"]:
+        output.append(f"--reg_data_dir={args['reg_img_folder']}")
+
+    if args['lora_model_for_resume']:
+        output.append(f"--network_weights={args['lora_model_for_resume']}")
+
+    if args['save_every_n_epochs']:
+        output.append(f"--save_every_n_epochs={args['save_every_n_epochs']}")
+    else:
+        output.append("--save_every_n_epochs=999999")
+
+    if args['shuffle_captions']:
+        output.append("--shuffle_caption")
+
+    if args['keep_tokens'] and args['keep_tokens'] > 0:
+        output.append(f"--keep_tokens={args['keep_tokens']}")
+
+    if args['buckets']:
+        output.append("--enable_bucket")
+        output.append(f"--min_bucket_reso={args['min_bucket_resolution']}")
+        output.append(f"--max_bucket_reso={args['max_bucket_resolution']}")
+
+    if args['use_8bit_adam']:
+        output.append("--use_8bit_adam")
+
+    if args['xformers']:
+        output.append("--xformers")
+
+    if args['color_aug']:
+        if args['cache_latents']:
+            print("color_aug and cache_latents conflict with one another. Please select only one")
+            quit(1)
+        output.append("--color_aug")
+
+    if args['flip_aug']:
+        output.append("--flip_aug")
+
+    if args['cache_latents']:
+        output.append("--cache_latents")
+
+    if args['warmup_lr_ratio'] and args['warmup_lr_ratio'] > 0:
+        warmup_steps = int(steps * args['warmup_lr_ratio'])
+        output.append(f"--lr_warmup_steps={warmup_steps}")
+
+    if args['gradient_checkpointing']:
+        output.append("--gradient_checkpointing")
+
+    if args['gradient_acc_steps'] and args['gradient_acc_steps'] > 0 and args['gradient_checkpointing']:
+        output.append(f"--gradient_accumulation_steps={args['gradient_acc_steps']}")
+
+    if args['learning_rate'] and args['learning_rate'] > 0:
+        output.append(f"--learning_rate={args['learning_rate']}")
+
+    if args['text_encoder_lr'] and args['text_encoder_lr'] > 0:
+        output.append(f"--text_encoder_lr={args['text_encoder_lr']}")
+
+    if args['unet_lr'] and args['unet_lr'] > 0:
+        output.append(f"--unet_lr={args['unet_lr']}")
+
+    if args['vae']:
+        output.append(f"--vae={args['vae']}")
+
+    if args['no_meta']:
+        output.append("--no_metadata")
+
+    if args['save_state']:
+        output.append("--save_state")
+
+    if args['load_previous_save_state']:
+        output.append(f"--resume={args['load_previous_save_state']}")
+
+    if args['change_output_name']:
+        output.append(f"--output_name={args['change_output_name']}")
+
+    if args['training_comment']:
+        output.append(f"--training_comment={args['training_comment']}")
+
+    if args['cosine_restarts'] and args['scheduler'] == "cosine_with_restarts":
+        output.append(f"--lr_scheduler_num_cycles={args['cosine_restarts']}")
+
+    if args['scheduler_power'] and args['scheduler'] == "polynomial":
+        output.append(f"--lr_scheduler_power={args['scheduler_power']}")
+
+    if args['persistent_workers']:
+        output.append(f"--persistent_data_loader_workers")
+
+    if args['unet_only']:
+        output.append("--network_train_unet_only")
+
+    if args['text_only'] and not args['unet_only']:
+        output.append("--network_train_text_encoder_only")
+
+    if args["log_dir"]:
+        output.append(f"--logging_dir={args['log_dir']}")
+
+    if args['caption_dropout_rate']:
+        output.append(f"--caption_dropout_rate={args['caption_dropout_rate']}")
+
+    if args['caption_dropout_every_n_epochs']:
+        output.append(f"--caption_dropout_every_n_epochs={args['caption_dropout_every_n_epochs']}")
+
+    if args['caption_tag_dropout_rate']:
+        output.append(f"--caption_tag_dropout_rate={args['caption_tag_dropout_rate']}")
+
+    if args['v2']:
+        output.append("--v2")
+
+    if args['v2'] and args['v_parameterization']:
+        output.append("--v_parameterization")
+
+    if args['noise_offset']:
+        output.append(f"--noise_offset={args['noise_offset']}")
+    return output
+
+
+def find_max_steps(args: dict) -> int:
+    total_steps = 0
+    folders = os.listdir(args["img_folder"])
+    for folder in folders:
+        if not os.path.isdir(os.path.join(args["img_folder"], folder)):
+            continue
+        num_repeats = folder.split("_")
+        if len(num_repeats) < 2:
+            print(f"folder {folder} is not in the correct format. Format is x_name. skipping")
+            continue
+        try:
+            num_repeats = int(num_repeats[0])
+        except ValueError:
+            print(f"folder {folder} is not in the correct format. Format is x_name. skipping")
+            continue
+        imgs = 0
+        for file in os.listdir(os.path.join(args["img_folder"], folder)):
+            if os.path.isdir(file):
+                continue
+            ext = file.split(".")
+            if ext[-1].lower() in {"png", "bmp", "gif", "jpeg", "jpg", "webp"}:
+                imgs += 1
+        total_steps += (num_repeats * imgs)
+    total_steps = int((total_steps / args["batch_size"]) * args["num_epochs"])
+    return total_steps
+
+
+def add_misc_args(parser):
+    parser.add_argument("--save_json_path", type=str, default=None,
+                        help="Path to save a configuration json file to")
+    parser.add_argument("--load_json_path", type=str, default=None,
+                        help="Path to a json file to configure things from")
+    parser.add_argument("--no_metadata", action='store_true',
+                        help="do not save metadata in output model / メタデータを出力先モデルに保存しない")
+    parser.add_argument("--save_model_as", type=str, default="safetensors", choices=[None, "ckpt", "pt", "safetensors"],
+                        help="format to save the model (default is .safetensors) / モデル保存時の形式（デフォルトはsafetensors）")
+
+    parser.add_argument("--unet_lr", type=float, default=None, help="learning rate for U-Net / U-Netの学習率")
+    parser.add_argument("--text_encoder_lr", type=float, default=None,
+                        help="learning rate for Text Encoder / Text Encoderの学習率")
+    parser.add_argument("--lr_scheduler_num_cycles", type=int, default=1,
+                        help="Number of restarts for cosine scheduler with restarts / cosine with restartsスケジューラでのリスタート回数")
+    parser.add_argument("--lr_scheduler_power", type=float, default=1,
+                        help="Polynomial power for polynomial scheduler / polynomialスケジューラでのpolynomial power")
+
+    parser.add_argument("--network_weights", type=str, default=None,
+                        help="pretrained weights for network / 学習するネットワークの初期重み")
+    parser.add_argument("--network_module", type=str, default=None,
+                        help='network module to train / 学習対象のネットワークのモジュール')
+    parser.add_argument("--network_dim", type=int, default=None,
+                        help='network dimensions (depends on each network) / モジュールの次元数（ネットワークにより定義は異なります）')
+    parser.add_argument("--network_alpha", type=float, default=1,
+                        help='alpha for LoRA weight scaling, default 1 (same as network_dim for same behavior as old version) / LoRaの重み調整のalpha値、デフォルト1（旧バージョンと同じ動作をするにはnetwork_dimと同じ値を指定）')
+    parser.add_argument("--network_args", type=str, default=None, nargs='*',
+                        help='additional argmuments for network (key=value) / ネットワークへの追加の引数')
+    parser.add_argument("--network_train_unet_only", action="store_true",
+                        help="only training U-Net part / U-Net関連部分のみ学習する")
+    parser.add_argument("--network_train_text_encoder_only", action="store_true",
+                        help="only training Text Encoder part / Text Encoder関連部分のみ学習する")
+    parser.add_argument("--training_comment", type=str, default=None,
+                        help="arbitrary comment string stored in metadata / メタデータに記録する任意のコメント文字列")
+
+
+def setup_args(parser):
+    util.add_sd_models_arguments(parser)
+    util.add_dataset_arguments(parser, True, True, True)
+    util.add_training_arguments(parser, True)
+    add_misc_args(parser)
+
+
+def get_occurrence_of_tags(args):
+    extension = args['caption_extension']
+    img_folder = args['img_folder']
+    output_folder = args['output_folder']
+    occurrence_dict = {}
+    print(img_folder)
+    for folder in os.listdir(img_folder):
+        print(folder)
+        if not os.path.isdir(os.path.join(img_folder, folder)):
+            continue
+        for file in os.listdir(os.path.join(img_folder, folder)):
+            if not os.path.isfile(os.path.join(img_folder, folder, file)):
+                continue
+            ext = os.path.splitext(file)[1]
+            if ext != extension:
+                continue
+            get_tags_from_file(os.path.join(img_folder, folder, file), occurrence_dict)
+    if not args['sort_tag_occurrence_alphabetically']:
+        output_list = {k: v for k, v in sorted(occurrence_dict.items(), key=lambda item: item[1], reverse=True)}
+    else:
+        output_list = {k: v for k, v in sorted(occurrence_dict.items(), key=lambda item: item[0])}
+    name = args['change_output_name'] if args['change_output_name'] else "last"
+    with open(os.path.join(output_folder, f"{name}.txt"), "w") as f:
+        f.write(f"Below is a list of keywords used during the training of {args['change_output_name']}:\n")
+        for k, v in output_list.items():
+            f.write(f"[{v}] {k}\n")
+    print(f"Created a txt file named {name}.txt in the output folder")
+
+
+def get_tags_from_file(file, occurrence_dict):
+    f = open(file)
+    temp = f.read().replace(", ", ",").split(",")
+    f.close()
+    for tag in temp:
+        if tag in occurrence_dict:
+            occurrence_dict[tag] += 1
+        else:
+            occurrence_dict[tag] = 1
+
+
+def ask_file(message, accepted_ext_list, file_path=None):
+    mb.showinfo(message=message)
+    res = ""
+    _initialdir = ""
+    _initialfile = ""
+    if file_path != None:
+        _initialdir = os.path.dirname(file_path) if os.path.exists(file_path) else ""
+        _initialfile = os.path.basename(file_path) if os.path.exists(file_path) else ""
+
+    while res == "":
+        res = fd.askopenfilename(title=message, initialdir=_initialdir, initialfile=_initialfile)
+        if res == "" or type(res) == tuple:
+            ret = mb.askretrycancel(message="Do you want to to cancel training?")
+            if not ret:
+                exit()
+            continue
+        elif not os.path.exists(res):
+            res = ""
+            continue
+        _, name = os.path.split(res)
+        split_name = name.split(".")
+        if split_name[-1] not in accepted_ext_list:
+            res = ""
+    return res
+
+
+def ask_dir(message, dir_path=None):
+    mb.showinfo(message=message)
+    res = ""
+    _initialdir = ""
+    if dir_path != None:
+        _initialdir = dir_path if os.path.exists(dir_path) else ""
+    while res == "":
+        res = fd.askdirectory(title=message, initialdir=_initialdir)
+        if res == "" or type(res) == tuple:
+            ret = mb.askretrycancel(message="Do you want to to cancel training?")
+            if not ret:
+                exit()
+            continue
+        if not os.path.exists(res):
+            res = ""
+    return res
+
+
+def ask_elements_trunc(args: dict):
+    args['base_model'] = ask_file("Select your base model", {"ckpt", "safetensors"}, args['base_model'])
+    args['img_folder'] = ask_dir("Select your image folder", args['img_folder'])
+    args['output_folder'] = ask_dir("Select your output folder", args['output_folder'])
+
+    ret = mb.askyesno(message="Do you want to save a json of your configuration?")
+    if ret:
+        args['save_json_folder'] = ask_dir("Select the folder to save json files to", args['save_json_folder'])
+    else:
+        args['save_json_folder'] = None
+
+    ret = mb.askyesno(message="Are you training on a SD2 based model?")
+    if ret:
+        args['v2'] = True
+
+    ret = mb.askyesno(message="Are you training on an realistic model?")
+    if ret:
+        args['clip_skip'] = 1
+
+    if args['v2']:
+        ret = mb.askyesno(message="Are you training on a model based on the 768x version of SD2?")
+        if ret:
+            args['v_parameterization'] = True
+
+    ret = mb.askyesno(message="Do you want to use regularization images?")
+    if ret:
+        args['reg_img_folder'] = ask_dir("Select your regularization folder", args['reg_img_folder'])
+    else:
+        args['reg_img_folder'] = None
+
+    ret = mb.askyesno(message="Do you want to continue from an earlier version?")
+    if ret:
+        args['lora_model_for_resume'] = ask_file("Select your lora model", {"ckpt", "pt", "safetensors"},
+                                                 args['lora_model_for_resume'])
+    else:
+        args['lora_model_for_resume'] = None
+
+    ret = mb.askyesno(message="Do you want to flip all of your images? It is supposed to reduce biases\n"
+                              "within your dataset but it can also ruin learning an asymmetrical element\n")
+    if ret:
+        args['flip_aug'] = True
+
+    ret = mb.askyesno(message="Do you want to change the name of output checkpoints?")
+    if ret:
+        ret = sd.askstring(title="output_name", prompt="What do you want your output name to be?\n"
+                                                       "Cancel keeps outputs the original")
+        if ret:
+            args['change_output_name'] = ret
+        else:
+            args['change_output_name'] = None
+
+    ret = sd.askstring(title="comment",
+                       prompt="Do you want to set a comment that gets put into the metadata?\nA good use of this would "
+                              "be to include how to use, such as activation keywords.\nCancel will leave empty")
+    if ret is None:
+        args['training_comment'] = ret
+    else:
+        args['training_comment'] = None
+
+    ret = mb.askyesno(message="Do you want to train only one of unet and text encoder?")
+    if ret:
+        button = ButtonBox("Which do you want to train with?", ["unet_only", "text_only"])
+        button.window.mainloop()
+        if button.current_value != "":
+            args[button.current_value] = True
+
+    ret = mb.askyesno(message="Do you want to save a txt file that contains a list\n"
+                              "of all tags that you have used in your training data?\n")
+    if ret:
+        args['tag_occurrence_txt_file'] = True
+        button = ButtonBox("How do you want tags to be ordered?", ["alphabetically", "occurrence-ly"])
+        button.window.mainloop()
+        if button.current_value == "alphabetically":
+            args['sort_tag_occurrence_alphabetically'] = True
+
+    ret = mb.askyesno(message="Do you want to use caption dropout?")
+    if ret:
+        ret = mb.askyesno(message="Do you want full caption files to dropout randomly?")
+        if ret:
+            ret = sd.askinteger(title="Caption_File_Dropout",
+                                prompt="How often do you want caption files to drop out?\n"
+                                       "enter a number from 0 to 100 that is the percentage chance of dropout\n"
+                                       "Cancel sets to 0")
+            if ret and 0 <= ret <= 100:
+                args['caption_dropout_rate'] = ret / 100.0
+
+        ret = mb.askyesno(message="Do you want to have full epochs have no captions?")
+        if ret:
+            ret = sd.askinteger(title="Caption_epoch_dropout", prompt="The number set here is how often you will have an"
+                                                                      "epoch with no captions\nSo if you set 3, then every"
+                                                                      "three epochs will not have captions (3, 6, 9)\n"
+                                                                      "Cancel will set to None")
+            if ret:
+                args['caption_dropout_every_n_epochs'] = ret
+
+        ret = mb.askyesno(message="Do you want to have tags to randomly drop?")
+        if ret:
+            ret = sd.askinteger(title="Caption_tag_dropout", prompt="How often do you want tags to randomly drop out?\n"
+                                                                    "Enter a number between 0 and 100, that is the percentage"
+                                                                    "chance of dropout.\nCancel sets to 0")
+            if ret and 0 <= ret <= 100:
+                args['caption_tag_dropout_rate'] = ret / 100.0
+
+    ret = mb.askyesno(message="Do you want to use noise offset? Noise offset seems to allow for SD to better generate\n"
+                              "darker or lighter images using this than normal.")
+    if ret:
+        ret = sd.askfloat(title="noise_offset", prompt="What value do you want to set? recommended value is 0.1,\n"
+                                                       "but it can go higher. Cancel defaults to 0.1")
+        if ret:
+            args['noise_offset'] = ret
+        else:
+            args['noise_offset'] = 0.1
+    return args
+
+
+def ask_elements(args: dict):
+    # start with file dialog
+    args['base_model'] = ask_file("Select your base model", {"ckpt", "safetensors"}, args['base_model'])
+    args['img_folder'] = ask_dir("Select your image folder", args['img_folder'])
+    args['output_folder'] = ask_dir("Select your output folder", args['output_folder'])
+
+    # optional file dialog
+    ret = mb.askyesno(message="Do you want to save a json of your configuration?")
+    if ret:
+        args['save_json_folder'] = ask_dir("Select the folder to save json files to", args['save_json_folder'])
+    else:
+        args['save_json_folder'] = None
+
+    ret = mb.askyesno(message="Are you training on a SD2 based model?")
+    if ret:
+        args['v2'] = True
+
+    ret = mb.askyesno(message="Are you training on an realistic model?")
+    if ret:
+        args['clip_skip'] = 1
+
+    if args['v2']:
+        ret = mb.askyesno(message="Are you training on a model based on the 768x version of SD2?")
+        if ret:
+            args['v_parameterization'] = True
+
+    ret = mb.askyesno(message="Do you want to use regularization images?")
+    if ret:
+        args['reg_img_folder'] = ask_dir("Select your regularization folder", args['reg_img_folder'])
+    else:
+        args['reg_img_folder'] = None
+
+    ret = mb.askyesno(message="Do you want to continue from an earlier version?")
+    if ret:
+        args['lora_model_for_resume'] = ask_file("Select your lora model", {"ckpt", "pt", "safetensors"},
+                                                 args['lora_model_for_resume'])
+    else:
+        args['lora_model_for_resume'] = None
+
+    ret = mb.askyesno(message="Do you want to flip all of your images? It is supposed to reduce biases\n"
+                              "within your dataset but it can also ruin learning an asymmetrical element\n")
+    if ret:
+        args['flip_aug'] = True
+
+    # text based required elements
+    ret = sd.askinteger(title="batch_size",
+                        prompt="The number of images that get processed at one time, this is directly proportional to "
+                               "your vram and resolution. with 12gb of vram, at 512 reso, you can get a maximum of 6 "
+                               "batch size\nHow large is your batch size going to be?\nCancel will default to 1")
+    if ret is None:
+        args['batch_size'] = 1
+    else:
+        args['batch_size'] = ret
+
+    ret = sd.askinteger(title="num_epochs", prompt="How many epochs do you want?\nCancel will default to 1")
+    if ret is None:
+        args['num_epochs'] = 1
+    else:
+        args['num_epochs'] = ret
+
+    ret = sd.askinteger(title="network_dim", prompt="What is the dim size you want to use?\nCancel will default to 128")
+    if ret is None:
+        args['net_dim'] = 128
+    else:
+        args['net_dim'] = ret
+
+    ret = sd.askfloat(title="alpha", prompt="Alpha is the scalar of the training, generally a good starting point is "
+                                            "0.5x dim size\nWhat Alpha do you want?\nCancel will default to equal to "
+                                            "0.5 x network_dim")
+    if ret is None:
+        args['alpha'] = args['net_dim'] / 2
+    else:
+        args['alpha'] = ret
+
+    ret = sd.askinteger(title="resolution", prompt="How large of a resolution do you want to train at?\n"
+                                                   "Cancel will default to 512")
+    if ret is None:
+        args['train_resolution'] = 512
+    else:
+        args['train_resolution'] = ret
+
+    ret = sd.askfloat(title="learning_rate", prompt="What learning rate do you want to use?\n"
+                                                    "Cancel will default to 1e-4")
+    if ret is None:
+        args['learning_rate'] = 1e-4
+    else:
+        args['learning_rate'] = ret
+
+    ret = sd.askfloat(title="text_encoder_lr", prompt="Do you want to set the text_encoder_lr?\n"
+                                                      "Cancel will default to None")
+    if ret is None:
+        args['text_encoder_lr'] = None
+    else:
+        args['text_encoder_lr'] = ret
+
+    ret = sd.askfloat(title="unet_lr", prompt="Do you want to set the unet_lr?\nCancel will default to None")
+    if ret is None:
+        args['unet_lr'] = None
+    else:
+        args['unet_lr'] = ret
+
+    button = ButtonBox("Which scheduler do you want?", ["cosine_with_restarts", "cosine", "polynomial",
+                                                        "constant", "constant_with_warmup", "linear"])
+    button.window.mainloop()
+    args['scheduler'] = button.current_value if button.current_value != "" else "cosine_with_restarts"
+
+    if args['scheduler'] == "cosine_with_restarts":
+        ret = sd.askinteger(title="Cycle Count",
+                            prompt="How many times do you want cosine to restart?\nThis is the entire amount of times "
+                                   "it will restart for the entire training\nCancel will default to 1")
+        if ret is None:
+            args['cosine_restarts'] = 1
+        else:
+            args['cosine_restarts'] = ret
+
+    if args['scheduler'] == "polynomial":
+        ret = sd.askfloat(title="Poly Strength",
+                          prompt="What power do you want to set your polynomial to?\nhigher power means that the "
+                                 "model reduces the learning more more aggressively from initial training.\n1 = "
+                                 "linear\nCancel sets to 1")
+        if ret is None:
+            args['scheduler_power'] = 1
+        else:
+            args['scheduler_power'] = ret
+
+    ret = mb.askyesno(message="Do you want to save epochs as it trains?")
+    if ret:
+        ret = sd.askinteger(title="save_epoch",
+                            prompt="How often do you want to save epochs?\nCancel will default to 1")
+        if ret is None:
+            args['save_every_n_epochs'] = 1
+        else:
+            args['save_every_n_epochs'] = ret
+
+    ret = mb.askyesno(message="Do you want to shuffle captions?")
+    if ret:
+        args['shuffle_captions'] = True
+    else:
+        args['shuffle_captions'] = False
+
+    ret = mb.askyesno(message="Do you want to keep some tokens at the front of your captions?")
+    if ret:
+        ret = sd.askinteger(title="keep_tokens", prompt="How many do you want to keep at the front?"
+                                                        "\nCancel will default to 1")
+        if ret is None:
+            args['keep_tokens'] = 1
+        else:
+            args['keep_tokens'] = ret
+
+    ret = mb.askyesno(message="Do you want to have a warmup ratio?")
+    if ret:
+        ret = sd.askfloat(title="warmup_ratio", prompt="What is the ratio of steps to use as warmup "
+                                                       "steps?\nCancel will default to None")
+        if ret is None:
+            args['warmup_lr_ratio'] = None
+        else:
+            args['warmup_lr_ratio'] = ret
+
+    ret = mb.askyesno(message="Do you want to change the name of output checkpoints?")
+    if ret:
+        ret = sd.askstring(title="output_name", prompt="What do you want your output name to be?\n"
+                                                       "Cancel keeps outputs the original")
+        if ret:
+            args['change_output_name'] = ret
+        else:
+            args['change_output_name'] = None
+
+    ret = sd.askstring(title="comment",
+                       prompt="Do you want to set a comment that gets put into the metadata?\nA good use of this would "
+                              "be to include how to use, such as activation keywords.\nCancel will leave empty")
+    if ret is None:
+        args['training_comment'] = ret
+    else:
+        args['training_comment'] = None
+
+    ret = mb.askyesno(message="Do you want to train only one of unet and text encoder?")
+    if ret:
+        if ret:
+            button = ButtonBox("Which do you want to train with?", ["unet_only", "text_only"])
+            button.window.mainloop()
+            if button.current_value != "":
+                args[button.current_value] = True
+
+    ret = mb.askyesno(message="Do you want to save a txt file that contains a list\n"
+                              "of all tags that you have used in your training data?\n")
+    if ret:
+        args['tag_occurrence_txt_file'] = True
+        button = ButtonBox("How do you want tags to be ordered?", ["alphabetically", "occurrence-ly"])
+        button.window.mainloop()
+        if button.current_value == "alphabetically":
+            args['sort_tag_occurrence_alphabetically'] = True
+
+    ret = mb.askyesno(message="Do you want to use caption dropout?")
+    if ret:
+        ret = mb.askyesno(message="Do you want full caption files to dropout randomly?")
+        if ret:
+            ret = sd.askinteger(title="Caption_File_Dropout",
+                                prompt="How often do you want caption files to drop out?\n"
+                                       "enter a number from 0 to 100 that is the percentage chance of dropout\n"
+                                       "Cancel sets to 0")
+            if ret and 0 <= ret <= 100:
+                args['caption_dropout_rate'] = ret / 100.0
+
+        ret = mb.askyesno(message="Do you want to have full epochs have no captions?")
+        if ret:
+            ret = sd.askinteger(title="Caption_epoch_dropout", prompt="The number set here is how often you will have an"
+                                                                      "epoch with no captions\nSo if you set 3, then every"
+                                                                      "three epochs will not have captions (3, 6, 9)\n"
+                                                                      "Cancel will set to None")
+            if ret:
+                args['caption_dropout_every_n_epochs'] = ret
+
+        ret = mb.askyesno(message="Do you want to have tags to randomly drop?")
+        if ret:
+            ret = sd.askinteger(title="Caption_tag_dropout", prompt="How often do you want tags to randomly drop out?\n"
+                                                                    "Enter a number between 0 and 100, that is the percentage"
+                                                                    "chance of dropout.\nCancel sets to 0")
+            if ret and 0 <= ret <= 100:
+                args['caption_tag_dropout_rate'] = ret / 100.0
+
+    ret = mb.askyesno(message="Do you want to use noise offset? Noise offset seems to allow for SD to better generate\n"
+                              "darker or lighter images using this than normal.")
+    if ret:
+        ret = sd.askfloat(title="noise_offset", prompt="What value do you want to set? recommended value is 0.1,\n"
+                                                       "but it can go higher. Cancel defaults to 0.1")
+        if ret:
+            args['noise_offset'] = ret
+        else:
+            args['noise_offset'] = 0.1
+    return args
+
+
+def save_json(path, obj: dict) -> None:
+    fp = open(os.path.join(path, f"config-{time.time()}.json"), "w")
+    json.dump(obj, fp=fp, indent=4)
+    fp.close()
+
+
+def load_json(path, obj: dict) -> dict:
+    with open(path) as f:
+        json_obj = json.loads(f.read())
+    print("loaded json, setting variables...")
+    ui_name_scheme = {"pretrained_model_name_or_path": "base_model", "logging_dir": "log_dir",
+                      "train_data_dir": "img_folder", "reg_data_dir": "reg_img_folder",
+                      "output_dir": "output_folder", "max_resolution": "train_resolution",
+                      "lr_scheduler": "scheduler", "lr_warmup": "warmup_lr_ratio",
+                      "train_batch_size": "batch_size", "epoch": "num_epochs",
+                      "save_at_n_epochs": "save_every_n_epochs", "num_cpu_threads_per_process": "num_workers",
+                      "enable_bucket": "buckets", "save_model_as": "save_as", "shuffle_caption": "shuffle_captions",
+                      "resume": "load_previous_save_state", "network_dim": "net_dim",
+                      "gradient_accumulation_steps": "gradient_acc_steps", "output_name": "change_output_name",
+                      "network_alpha": "alpha", "lr_scheduler_num_cycles": "cosine_restarts",
+                      "lr_scheduler_power": "scheduler_power"}
+
+    for key in list(json_obj):
+        if key in ui_name_scheme:
+            json_obj[ui_name_scheme[key]] = json_obj[key]
+            if ui_name_scheme[key] in {"batch_size", "num_epochs"}:
+                try:
+                    json_obj[ui_name_scheme[key]] = int(json_obj[ui_name_scheme[key]])
+                except ValueError:
+                    print(f"attempting to load {key} from json failed as input isn't an integer")
+                    quit(1)
+
+    for key in list(json_obj):
+        if obj["json_load_skip_list"] and key in obj["json_load_skip_list"]:
+            continue
+        if key in obj:
+            if key in {"keep_tokens", "warmup_lr_ratio"}:
+                json_obj[key] = int(json_obj[key]) if json_obj[key] is not None else None
+            if key in {"learning_rate", "unet_lr", "text_encoder_lr"}:
+                json_obj[key] = float(json_obj[key]) if json_obj[key] is not None else None
+            if obj[key] != json_obj[key]:
+                print_change(key, obj[key], json_obj[key])
+                obj[key] = json_obj[key]
+    print("completed changing variables.")
+    return obj
+
+
+def print_change(value, old, new):
+    print(f"{value} changed from {old} to {new}")
+
+
+class ButtonBox:
+    def __init__(self, label: str, button_name_list: list[str]) -> None:
+        self.window = tk.Tk()
+        self.button_list = []
+        self.current_value = ""
+
+        self.window.attributes("-topmost", True)
+        self.window.resizable(False, False)
+        self.window.eval('tk::PlaceWindow . center')
+
+        def del_window():
+            self.window.quit()
+            self.window.destroy()
+
+        self.window.protocol("WM_DELETE_WINDOW", del_window)
+        tk.Label(text=label, master=self.window).pack()
+        for button in button_name_list:
+            self.button_list.append(ttk.Button(text=button, master=self.window,
+                                               command=partial(self.set_current_value, button)))
+            self.button_list[-1].pack()
+
+    def set_current_value(self, value):
+        self.current_value = value
+        self.window.quit()
+        self.window.destroy()
+
+
+root = tk.Tk()
+root.attributes('-topmost', True)
+root.withdraw()
+
+if __name__ == "__main__":
+    main()

lycoris/__init__.py

@@ -0,0 +1,8 @@
+from lycoris import (
+    kohya,
+    kohya_model_utils,
+    kohya_utils,
+    locon,
+    loha,
+    utils,
+)

lycoris/kohya.py

@@ -0,0 +1,276 @@
+# network module for kohya
+# reference:
+# https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
+# https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/lora.py
+# https://github.com/kohya-ss/sd-scripts/blob/main/networks/lora.py
+
+import math
+import os
+from typing import List
+import torch
+
+from .kohya_utils import *
+from .locon import LoConModule
+from .loha import LohaModule
+
+
+def create_network(multiplier, network_dim, network_alpha, vae, text_encoder, unet, **kwargs):
+    if network_dim is None:
+        network_dim = 4                     # default
+    conv_dim = int(kwargs.get('conv_dim', network_dim))
+    conv_alpha = float(kwargs.get('conv_alpha', network_alpha))
+    dropout = float(kwargs.get('dropout', 0.))
+    algo = kwargs.get('algo', 'lora')
+    network_module = {
+        'lora': LoConModule,
+        'loha': LohaModule,
+    }[algo]
+    
+    print(f'Using rank adaptation algo: {algo}')
+    network = LoRANetwork(
+        text_encoder, unet, 
+        multiplier=multiplier, 
+        lora_dim=network_dim, conv_lora_dim=conv_dim, 
+        alpha=network_alpha, conv_alpha=conv_alpha,
+        dropout=dropout,
+        network_module=network_module
+    )
+    
+    return network
+
+
+def create_network_from_weights(multiplier, file, vae, text_encoder, unet, **kwargs):
+    if os.path.splitext(file)[1] == '.safetensors':
+        from safetensors.torch import load_file, safe_open
+        weights_sd = load_file(file)
+    else:
+        weights_sd = torch.load(file, map_location='cpu')
+
+    # get dim (rank)
+    network_alpha = None
+    network_dim = None
+    for key, value in weights_sd.items():
+        if network_alpha is None and 'alpha' in key:
+            network_alpha = value
+        if network_dim is None and 'lora_down' in key and len(value.size()) == 2:
+            network_dim = value.size()[0]
+
+    if network_alpha is None:
+        network_alpha = network_dim
+
+    network = LoRANetwork(text_encoder, unet, multiplier=multiplier, lora_dim=network_dim, alpha=network_alpha)
+    network.weights_sd = weights_sd
+    return network
+
+
+class LoRANetwork(torch.nn.Module):
+    '''
+    LoRA + LoCon
+    '''
+    # Ignore proj_in or proj_out, their channels is only a few.
+    UNET_TARGET_REPLACE_MODULE = [
+        "Transformer2DModel", 
+        "Attention", 
+        "ResnetBlock2D", 
+        "Downsample2D", 
+        "Upsample2D"
+    ]
+    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
+    LORA_PREFIX_UNET = 'lora_unet'
+    LORA_PREFIX_TEXT_ENCODER = 'lora_te'
+
+    def __init__(
+        self, 
+        text_encoder, unet, 
+        multiplier=1.0, 
+        lora_dim=4, conv_lora_dim=4, 
+        alpha=1, conv_alpha=1,
+        dropout = 0, network_module = LoConModule,
+    ) -> None:
+        super().__init__()
+        self.multiplier = multiplier
+        self.lora_dim = lora_dim
+        self.conv_lora_dim = int(conv_lora_dim)
+        if self.conv_lora_dim != self.lora_dim: 
+            print('Apply different lora dim for conv layer')
+            print(f'LoCon Dim: {conv_lora_dim}, LoRA Dim: {lora_dim}')
+            
+        self.alpha = alpha
+        self.conv_alpha = float(conv_alpha)
+        if self.alpha != self.conv_alpha: 
+            print('Apply different alpha value for conv layer')
+            print(f'LoCon alpha: {conv_alpha}, LoRA alpha: {alpha}')
+        
+        if 1 >= dropout >= 0:
+            print(f'Use Dropout value: {dropout}')
+        self.dropout = dropout
+        
+        # create module instances
+        def create_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> List[network_module]:
+            print('Create LoCon Module')
+            loras = []
+            for name, module in root_module.named_modules():
+                if module.__class__.__name__ in target_replace_modules:
+                    for child_name, child_module in module.named_modules():
+                        lora_name = prefix + '.' + name + '.' + child_name
+                        lora_name = lora_name.replace('.', '_')
+                        if child_module.__class__.__name__ == 'Linear' and lora_dim>0:
+                            lora = network_module(
+                                lora_name, child_module, self.multiplier, 
+                                self.lora_dim, self.alpha, self.dropout
+                            )
+                        elif child_module.__class__.__name__ == 'Conv2d':
+                            k_size, *_ = child_module.kernel_size
+                            if k_size==1 and lora_dim>0:
+                                lora = network_module(
+                                    lora_name, child_module, self.multiplier, 
+                                    self.lora_dim, self.alpha, self.dropout
+                                )
+                            elif conv_lora_dim>0:
+                                lora = network_module(
+                                    lora_name, child_module, self.multiplier, 
+                                    self.conv_lora_dim, self.conv_alpha, self.dropout
+                                )
+                            else:
+                                continue
+                        else:
+                            continue
+                        loras.append(lora)
+            return loras
+
+        self.text_encoder_loras = create_modules(
+            LoRANetwork.LORA_PREFIX_TEXT_ENCODER,
+            text_encoder, 
+            LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE
+        )
+        print(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
+
+        self.unet_loras = create_modules(LoRANetwork.LORA_PREFIX_UNET, unet, LoRANetwork.UNET_TARGET_REPLACE_MODULE)
+        print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
+
+        self.weights_sd = None
+
+        # assertion
+        names = set()
+        for lora in self.text_encoder_loras + self.unet_loras:
+            assert lora.lora_name not in names, f"duplicated lora name: {lora.lora_name}"
+            names.add(lora.lora_name)
+
+    def set_multiplier(self, multiplier):
+        self.multiplier = multiplier
+        for lora in self.text_encoder_loras + self.unet_loras:
+            lora.multiplier = self.multiplier
+            
+    def load_weights(self, file):
+        if os.path.splitext(file)[1] == '.safetensors':
+            from safetensors.torch import load_file, safe_open
+            self.weights_sd = load_file(file)
+        else:
+            self.weights_sd = torch.load(file, map_location='cpu')
+
+    def apply_to(self, text_encoder, unet, apply_text_encoder=None, apply_unet=None):
+        if self.weights_sd:
+            weights_has_text_encoder = weights_has_unet = False
+            for key in self.weights_sd.keys():
+                if key.startswith(LoRANetwork.LORA_PREFIX_TEXT_ENCODER):
+                    weights_has_text_encoder = True
+                elif key.startswith(LoRANetwork.LORA_PREFIX_UNET):
+                    weights_has_unet = True
+
+            if apply_text_encoder is None:
+                apply_text_encoder = weights_has_text_encoder
+            else:
+                assert apply_text_encoder == weights_has_text_encoder, f"text encoder weights: {weights_has_text_encoder} but text encoder flag: {apply_text_encoder} / 重みとText Encoderのフラグが矛盾しています"
+
+            if apply_unet is None:
+                apply_unet = weights_has_unet
+            else:
+                assert apply_unet == weights_has_unet, f"u-net weights: {weights_has_unet} but u-net flag: {apply_unet} / 重みとU-Netのフラグが矛盾しています"
+        else:
+            assert apply_text_encoder is not None and apply_unet is not None, f"internal error: flag not set"
+
+        if apply_text_encoder:
+            print("enable LoRA for text encoder")
+        else:
+            self.text_encoder_loras = []
+
+        if apply_unet:
+            print("enable LoRA for U-Net")
+        else:
+            self.unet_loras = []
+
+        for lora in self.text_encoder_loras + self.unet_loras:
+            lora.apply_to()
+            self.add_module(lora.lora_name, lora)
+
+        if self.weights_sd:
+            # if some weights are not in state dict, it is ok because initial LoRA does nothing (lora_up is initialized by zeros)
+            info = self.load_state_dict(self.weights_sd, False)
+            print(f"weights are loaded: {info}")
+
+    def enable_gradient_checkpointing(self):
+        # not supported
+        def make_ckpt(module):
+            if isinstance(module, torch.nn.Module):
+                module.grad_ckpt = True
+        self.apply(make_ckpt)
+        pass
+
+    def prepare_optimizer_params(self, text_encoder_lr, unet_lr):
+        def enumerate_params(loras):
+            params = []
+            for lora in loras:
+                params.extend(lora.parameters())
+            return params
+
+        self.requires_grad_(True)
+        all_params = []
+
+        if self.text_encoder_loras:
+            param_data = {'params': enumerate_params(self.text_encoder_loras)}
+            if text_encoder_lr is not None:
+                param_data['lr'] = text_encoder_lr
+            all_params.append(param_data)
+
+        if self.unet_loras:
+            param_data = {'params': enumerate_params(self.unet_loras)}
+            if unet_lr is not None:
+                param_data['lr'] = unet_lr
+            all_params.append(param_data)
+
+        return all_params
+
+    def prepare_grad_etc(self, text_encoder, unet):
+        self.requires_grad_(True)
+
+    def on_epoch_start(self, text_encoder, unet):
+        self.train()
+
+    def get_trainable_params(self):
+        return self.parameters()
+
+    def save_weights(self, file, dtype, metadata):
+        if metadata is not None and len(metadata) == 0:
+            metadata = None
+
+        state_dict = self.state_dict()
+
+        if dtype is not None:
+            for key in list(state_dict.keys()):
+                v = state_dict[key]
+                v = v.detach().clone().to("cpu").to(dtype)
+                state_dict[key] = v
+
+        if os.path.splitext(file)[1] == '.safetensors':
+            from safetensors.torch import save_file
+
+            # Precalculate model hashes to save time on indexing
+            if metadata is None:
+                metadata = {}
+            model_hash, legacy_hash = precalculate_safetensors_hashes(state_dict, metadata)
+            metadata["sshs_model_hash"] = model_hash
+            metadata["sshs_legacy_hash"] = legacy_hash
+
+            save_file(state_dict, file, metadata)
+        else:
+            torch.save(state_dict, file)

lycoris/kohya_model_utils.py

@@ -0,0 +1,1184 @@
+'''
+https://github.com/kohya-ss/sd-scripts/blob/main/library/model_util.py
+'''
+# v1: split from train_db_fixed.py.
+# v2: support safetensors
+
+import math
+import os
+import torch
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig
+from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNet2DConditionModel
+from safetensors.torch import load_file, save_file
+
+# DiffUsers版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
+
+# V2
+V2_UNET_PARAMS_ATTENTION_HEAD_DIM = [5, 10, 20, 20]
+V2_UNET_PARAMS_CONTEXT_DIM = 1024
+
+# Diffusersの設定を読み込むための参照モデル
+DIFFUSERS_REF_MODEL_ID_V1 = "runwayml/stable-diffusion-v1-5"
+DIFFUSERS_REF_MODEL_ID_V2 = "stabilityai/stable-diffusion-2-1"
+
+
+# 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 linear_transformer_to_conv(checkpoint):
+  keys = list(checkpoint.keys())
+  tf_keys = ["proj_in.weight", "proj_out.weight"]
+  for key in keys:
+    if ".".join(key.split(".")[-2:]) in tf_keys:
+      if checkpoint[key].ndim == 2:
+        checkpoint[key] = checkpoint[key].unsqueeze(2).unsqueeze(2)
+
+
+def convert_ldm_unet_checkpoint(v2, 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"])
+
+      # biasとweightの順番に依存しないようにする：もっといいやり方がありそうだが
+      for l in output_block_list.values():
+        l.sort()
+
+      if ["conv.bias", "conv.weight"] in output_block_list.values():
+        index = list(output_block_list.values()).index(["conv.bias", "conv.weight"])
+        new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[
+            f"output_blocks.{i}.{index}.conv.bias"
+        ]
+        new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[
+            f"output_blocks.{i}.{index}.conv.weight"
+        ]
+
+        # 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]
+
+  # SDのv2では1*1のconv2dがlinearに変わっているので、linear->convに変換する
+  if v2:
+    linear_transformer_to_conv(new_checkpoint)
+
+  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)
+  # if len(vae_state_dict) == 0:
+  #   # 渡されたcheckpointは.ckptから読み込んだcheckpointではなくvaeのstate_dict
+  #   vae_state_dict = checkpoint
+
+  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(v2):
+  """
+  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 if not v2 else V2_UNET_PARAMS_CONTEXT_DIM,
+      attention_head_dim=UNET_PARAMS_NUM_HEADS if not v2 else V2_UNET_PARAMS_ATTENTION_HEAD_DIM,
+  )
+
+  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_v1(checkpoint):
+  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]
+  return text_model_dict
+
+
+def convert_ldm_clip_checkpoint_v2(checkpoint, max_length):
+  # 嫌になるくらい違うぞ！
+  def convert_key(key):
+    if not key.startswith("cond_stage_model"):
+      return None
+
+    # common conversion
+    key = key.replace("cond_stage_model.model.transformer.", "text_model.encoder.")
+    key = key.replace("cond_stage_model.model.", "text_model.")
+
+    if "resblocks" in key:
+      # resblocks conversion
+      key = key.replace(".resblocks.", ".layers.")
+      if ".ln_" in key:
+        key = key.replace(".ln_", ".layer_norm")
+      elif ".mlp." in key:
+        key = key.replace(".c_fc.", ".fc1.")
+        key = key.replace(".c_proj.", ".fc2.")
+      elif '.attn.out_proj' in key:
+        key = key.replace(".attn.out_proj.", ".self_attn.out_proj.")
+      elif '.attn.in_proj' in key:
+        key = None                  # 特殊なので後で処理する
+      else:
+        raise ValueError(f"unexpected key in SD: {key}")
+    elif '.positional_embedding' in key:
+      key = key.replace(".positional_embedding", ".embeddings.position_embedding.weight")
+    elif '.text_projection' in key:
+      key = None    # 使われない???
+    elif '.logit_scale' in key:
+      key = None    # 使われない???
+    elif '.token_embedding' in key:
+      key = key.replace(".token_embedding.weight", ".embeddings.token_embedding.weight")
+    elif '.ln_final' in key:
+      key = key.replace(".ln_final", ".final_layer_norm")
+    return key
+
+  keys = list(checkpoint.keys())
+  new_sd = {}
+  for key in keys:
+    # remove resblocks 23
+    if '.resblocks.23.' in key:
+      continue
+    new_key = convert_key(key)
+    if new_key is None:
+      continue
+    new_sd[new_key] = checkpoint[key]
+
+  # attnの変換
+  for key in keys:
+    if '.resblocks.23.' in key:
+      continue
+    if '.resblocks' in key and '.attn.in_proj_' in key:
+      # 三つに分割
+      values = torch.chunk(checkpoint[key], 3)
+
+      key_suffix = ".weight" if "weight" in key else ".bias"
+      key_pfx = key.replace("cond_stage_model.model.transformer.resblocks.", "text_model.encoder.layers.")
+      key_pfx = key_pfx.replace("_weight", "")
+      key_pfx = key_pfx.replace("_bias", "")
+      key_pfx = key_pfx.replace(".attn.in_proj", ".self_attn.")
+      new_sd[key_pfx + "q_proj" + key_suffix] = values[0]
+      new_sd[key_pfx + "k_proj" + key_suffix] = values[1]
+      new_sd[key_pfx + "v_proj" + key_suffix] = values[2]
+
+  # rename or add position_ids
+  ANOTHER_POSITION_IDS_KEY = "text_model.encoder.text_model.embeddings.position_ids"
+  if ANOTHER_POSITION_IDS_KEY in new_sd:
+    # waifu diffusion v1.4
+    position_ids = new_sd[ANOTHER_POSITION_IDS_KEY]
+    del new_sd[ANOTHER_POSITION_IDS_KEY]
+  else:
+    position_ids = torch.Tensor([list(range(max_length))]).to(torch.int64)
+
+  new_sd["text_model.embeddings.position_ids"] = position_ids
+  return new_sd
+
+# endregion
+
+
+# region Diffusers->StableDiffusion の変換コード
+# convert_diffusers_to_original_stable_diffusion をコピーして修正している（ASL 2.0）
+
+def conv_transformer_to_linear(checkpoint):
+  keys = list(checkpoint.keys())
+  tf_keys = ["proj_in.weight", "proj_out.weight"]
+  for key in keys:
+    if ".".join(key.split(".")[-2:]) in tf_keys:
+      if checkpoint[key].ndim > 2:
+        checkpoint[key] = checkpoint[key][:, :, 0, 0]
+
+
+def convert_unet_state_dict_to_sd(v2, 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()}
+
+  if v2:
+    conv_transformer_to_linear(new_state_dict)
+
+  return new_state_dict
+
+
+# ================#
+# VAE Conversion #
+# ================#
+
+def reshape_weight_for_sd(w):
+    # convert HF linear weights to SD conv2d weights
+  return w.reshape(*w.shape, 1, 1)
+
+
+def convert_vae_state_dict(vae_state_dict):
+  vae_conversion_map = [
+      # (stable-diffusion, HF Diffusers)
+      ("nin_shortcut", "conv_shortcut"),
+      ("norm_out", "conv_norm_out"),
+      ("mid.attn_1.", "mid_block.attentions.0."),
+  ]
+
+  for i in range(4):
+    # down_blocks have two resnets
+    for j in range(2):
+      hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}."
+      sd_down_prefix = f"encoder.down.{i}.block.{j}."
+      vae_conversion_map.append((sd_down_prefix, hf_down_prefix))
+
+    if i < 3:
+      hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0."
+      sd_downsample_prefix = f"down.{i}.downsample."
+      vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))
+
+      hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+      sd_upsample_prefix = f"up.{3-i}.upsample."
+      vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))
+
+    # up_blocks have three resnets
+    # also, up blocks in hf are numbered in reverse from sd
+    for j in range(3):
+      hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}."
+      sd_up_prefix = f"decoder.up.{3-i}.block.{j}."
+      vae_conversion_map.append((sd_up_prefix, hf_up_prefix))
+
+  # this part accounts for mid blocks in both the encoder and the decoder
+  for i in range(2):
+    hf_mid_res_prefix = f"mid_block.resnets.{i}."
+    sd_mid_res_prefix = f"mid.block_{i+1}."
+    vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+  vae_conversion_map_attn = [
+      # (stable-diffusion, HF Diffusers)
+      ("norm.", "group_norm."),
+      ("q.", "query."),
+      ("k.", "key."),
+      ("v.", "value."),
+      ("proj_out.", "proj_attn."),
+  ]
+
+  mapping = {k: k for k in vae_state_dict.keys()}
+  for k, v in mapping.items():
+    for sd_part, hf_part in vae_conversion_map:
+      v = v.replace(hf_part, sd_part)
+    mapping[k] = v
+  for k, v in mapping.items():
+    if "attentions" in k:
+      for sd_part, hf_part in vae_conversion_map_attn:
+        v = v.replace(hf_part, sd_part)
+      mapping[k] = v
+  new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()}
+  weights_to_convert = ["q", "k", "v", "proj_out"]
+  for k, v in new_state_dict.items():
+    for weight_name in weights_to_convert:
+      if f"mid.attn_1.{weight_name}.weight" in k:
+        # print(f"Reshaping {k} for SD format")
+        new_state_dict[k] = reshape_weight_for_sd(v)
+
+  return new_state_dict
+
+
+# endregion
+
+# region 自作のモデル読み書きなど
+
+def is_safetensors(path):
+  return os.path.splitext(path)[1].lower() == '.safetensors'
+
+
+def load_checkpoint_with_text_encoder_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.')
+  ]
+
+  if is_safetensors(ckpt_path):
+    checkpoint = None
+    state_dict = load_file(ckpt_path, "cpu")
+  else:
+    checkpoint = torch.load(ckpt_path, map_location="cpu")
+    if "state_dict" in checkpoint:
+      state_dict = checkpoint["state_dict"]
+    else:
+      state_dict = checkpoint
+      checkpoint = None
+
+  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, state_dict
+
+
+# TODO dtype指定の動作が怪しいので確認する text_encoderを指定形式で作れるか未確認
+def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, dtype=None):
+  _, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path)
+  if dtype is not None:
+    for k, v in state_dict.items():
+      if type(v) is torch.Tensor:
+        state_dict[k] = v.to(dtype)
+
+  # Convert the UNet2DConditionModel model.
+  unet_config = create_unet_diffusers_config(v2)
+  converted_unet_checkpoint = convert_ldm_unet_checkpoint(v2, state_dict, unet_config)
+
+  unet = UNet2DConditionModel(**unet_config)
+  info = unet.load_state_dict(converted_unet_checkpoint)
+  print("loading u-net:", info)
+
+  # 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)
+  info = vae.load_state_dict(converted_vae_checkpoint)
+  print("loading vae:", info)
+
+  # convert text_model
+  if v2:
+    converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v2(state_dict, 77)
+    cfg = CLIPTextConfig(
+        vocab_size=49408,
+        hidden_size=1024,
+        intermediate_size=4096,
+        num_hidden_layers=23,
+        num_attention_heads=16,
+        max_position_embeddings=77,
+        hidden_act="gelu",
+        layer_norm_eps=1e-05,
+        dropout=0.0,
+        attention_dropout=0.0,
+        initializer_range=0.02,
+        initializer_factor=1.0,
+        pad_token_id=1,
+        bos_token_id=0,
+        eos_token_id=2,
+        model_type="clip_text_model",
+        projection_dim=512,
+        torch_dtype="float32",
+        transformers_version="4.25.0.dev0",
+    )
+    text_model = CLIPTextModel._from_config(cfg)
+    info = text_model.load_state_dict(converted_text_encoder_checkpoint)
+  else:
+    converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v1(state_dict)
+    text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+    info = text_model.load_state_dict(converted_text_encoder_checkpoint)
+  print("loading text encoder:", info)
+
+  return text_model, vae, unet
+
+
+def convert_text_encoder_state_dict_to_sd_v2(checkpoint, make_dummy_weights=False):
+  def convert_key(key):
+    # position_idsの除去
+    if ".position_ids" in key:
+      return None
+
+    # common
+    key = key.replace("text_model.encoder.", "transformer.")
+    key = key.replace("text_model.", "")
+    if "layers" in key:
+      # resblocks conversion
+      key = key.replace(".layers.", ".resblocks.")
+      if ".layer_norm" in key:
+        key = key.replace(".layer_norm", ".ln_")
+      elif ".mlp." in key:
+        key = key.replace(".fc1.", ".c_fc.")
+        key = key.replace(".fc2.", ".c_proj.")
+      elif '.self_attn.out_proj' in key:
+        key = key.replace(".self_attn.out_proj.", ".attn.out_proj.")
+      elif '.self_attn.' in key:
+        key = None                  # 特殊なので後で処理する
+      else:
+        raise ValueError(f"unexpected key in DiffUsers model: {key}")
+    elif '.position_embedding' in key:
+      key = key.replace("embeddings.position_embedding.weight", "positional_embedding")
+    elif '.token_embedding' in key:
+      key = key.replace("embeddings.token_embedding.weight", "token_embedding.weight")
+    elif 'final_layer_norm' in key:
+      key = key.replace("final_layer_norm", "ln_final")
+    return key
+
+  keys = list(checkpoint.keys())
+  new_sd = {}
+  for key in keys:
+    new_key = convert_key(key)
+    if new_key is None:
+      continue
+    new_sd[new_key] = checkpoint[key]
+
+  # attnの変換
+  for key in keys:
+    if 'layers' in key and 'q_proj' in key:
+      # 三つを結合
+      key_q = key
+      key_k = key.replace("q_proj", "k_proj")
+      key_v = key.replace("q_proj", "v_proj")
+
+      value_q = checkpoint[key_q]
+      value_k = checkpoint[key_k]
+      value_v = checkpoint[key_v]
+      value = torch.cat([value_q, value_k, value_v])
+
+      new_key = key.replace("text_model.encoder.layers.", "transformer.resblocks.")
+      new_key = new_key.replace(".self_attn.q_proj.", ".attn.in_proj_")
+      new_sd[new_key] = value
+
+  # 最後の層などを捏造するか
+  if make_dummy_weights:
+    print("make dummy weights for resblock.23, text_projection and logit scale.")
+    keys = list(new_sd.keys())
+    for key in keys:
+      if key.startswith("transformer.resblocks.22."):
+        new_sd[key.replace(".22.", ".23.")] = new_sd[key].clone()          # copyしないとsafetensorsの保存で落ちる
+
+    # Diffusersに含まれない重みを作っておく
+    new_sd['text_projection'] = torch.ones((1024, 1024), dtype=new_sd[keys[0]].dtype, device=new_sd[keys[0]].device)
+    new_sd['logit_scale'] = torch.tensor(1)
+
+  return new_sd
+
+
+def save_stable_diffusion_checkpoint(v2, output_file, text_encoder, unet, ckpt_path, epochs, steps, save_dtype=None, vae=None):
+  if ckpt_path is not None:
+    # epoch/stepを参照する。またVAEがメモリ上にないときなど、もう一度VAEを含めて読み込む
+    checkpoint, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path)
+    if checkpoint is None:                # safetensors または state_dictのckpt
+      checkpoint = {}
+      strict = False
+    else:
+      strict = True
+    if "state_dict" in state_dict:
+      del state_dict["state_dict"]
+  else:
+    # 新しく作る
+    assert vae is not None, "VAE is required to save a checkpoint without a given checkpoint"
+    checkpoint = {}
+    state_dict = {}
+    strict = False
+
+  def update_sd(prefix, sd):
+    for k, v in sd.items():
+      key = prefix + k
+      assert not strict or key in state_dict, f"Illegal key in save SD: {key}"
+      if save_dtype is not None:
+        v = v.detach().clone().to("cpu").to(save_dtype)
+      state_dict[key] = v
+
+  # Convert the UNet model
+  unet_state_dict = convert_unet_state_dict_to_sd(v2, unet.state_dict())
+  update_sd("model.diffusion_model.", unet_state_dict)
+
+  # Convert the text encoder model
+  if v2:
+    make_dummy = ckpt_path is None                 # 参照元のcheckpoint���ない場合は最後の層を前の層から複製して作るなどダミーの重みを入れる
+    text_enc_dict = convert_text_encoder_state_dict_to_sd_v2(text_encoder.state_dict(), make_dummy)
+    update_sd("cond_stage_model.model.", text_enc_dict)
+  else:
+    text_enc_dict = text_encoder.state_dict()
+    update_sd("cond_stage_model.transformer.", text_enc_dict)
+
+  # Convert the VAE
+  if vae is not None:
+    vae_dict = convert_vae_state_dict(vae.state_dict())
+    update_sd("first_stage_model.", vae_dict)
+
+  # Put together new checkpoint
+  key_count = len(state_dict.keys())
+  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
+
+  if is_safetensors(output_file):
+    # TODO Tensor以外のdictの値を削除したほうがいいか
+    save_file(state_dict, output_file)
+  else:
+    torch.save(new_ckpt, output_file)
+
+  return key_count
+
+
+def save_diffusers_checkpoint(v2, output_dir, text_encoder, unet, pretrained_model_name_or_path, vae=None, use_safetensors=False):
+  if pretrained_model_name_or_path is None:
+    # load default settings for v1/v2
+    if v2:
+      pretrained_model_name_or_path = DIFFUSERS_REF_MODEL_ID_V2
+    else:
+      pretrained_model_name_or_path = DIFFUSERS_REF_MODEL_ID_V1
+
+  scheduler = DDIMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
+  tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer")
+  if vae is None:
+    vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae")
+
+  pipeline = StableDiffusionPipeline(
+      unet=unet,
+      text_encoder=text_encoder,
+      vae=vae,
+      scheduler=scheduler,
+      tokenizer=tokenizer,
+      safety_checker=None,
+      feature_extractor=None,
+      requires_safety_checker=None,
+  )
+  pipeline.save_pretrained(output_dir, safe_serialization=use_safetensors)
+
+
+VAE_PREFIX = "first_stage_model."
+
+
+def load_vae(vae_id, dtype):
+  print(f"load VAE: {vae_id}")
+  if os.path.isdir(vae_id) or not os.path.isfile(vae_id):
+    # Diffusers local/remote
+    try:
+      vae = AutoencoderKL.from_pretrained(vae_id, subfolder=None, torch_dtype=dtype)
+    except EnvironmentError as e:
+      print(f"exception occurs in loading vae: {e}")
+      print("retry with subfolder='vae'")
+      vae = AutoencoderKL.from_pretrained(vae_id, subfolder="vae", torch_dtype=dtype)
+    return vae
+
+  # local
+  vae_config = create_vae_diffusers_config()
+
+  if vae_id.endswith(".bin"):
+    # SD 1.5 VAE on Huggingface
+    converted_vae_checkpoint = torch.load(vae_id, map_location="cpu")
+  else:
+    # StableDiffusion
+    vae_model = (load_file(vae_id, "cpu") if is_safetensors(vae_id)
+                 else torch.load(vae_id, map_location="cpu"))
+    vae_sd = vae_model['state_dict'] if 'state_dict' in vae_model else vae_model
+
+    # vae only or full model
+    full_model = False
+    for vae_key in vae_sd:
+      if vae_key.startswith(VAE_PREFIX):
+        full_model = True
+        break
+    if not full_model:
+      sd = {}
+      for key, value in vae_sd.items():
+        sd[VAE_PREFIX + key] = value
+      vae_sd = sd
+      del sd
+
+    # Convert the VAE model.
+    converted_vae_checkpoint = convert_ldm_vae_checkpoint(vae_sd, vae_config)
+
+  vae = AutoencoderKL(**vae_config)
+  vae.load_state_dict(converted_vae_checkpoint)
+  return vae
+
+# 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)

lycoris/kohya_utils.py

@@ -0,0 +1,48 @@
+# part of https://github.com/kohya-ss/sd-scripts/blob/main/library/train_util.py
+
+import hashlib
+import safetensors
+from io import BytesIO
+
+
+def addnet_hash_legacy(b):
+    """Old model hash used by sd-webui-additional-networks for .safetensors format files"""
+    m = hashlib.sha256()
+
+    b.seek(0x100000)
+    m.update(b.read(0x10000))
+    return m.hexdigest()[0:8]
+
+
+def addnet_hash_safetensors(b):
+    """New model hash used by sd-webui-additional-networks for .safetensors format files"""
+    hash_sha256 = hashlib.sha256()
+    blksize = 1024 * 1024
+
+    b.seek(0)
+    header = b.read(8)
+    n = int.from_bytes(header, "little")
+
+    offset = n + 8
+    b.seek(offset)
+    for chunk in iter(lambda: b.read(blksize), b""):
+        hash_sha256.update(chunk)
+
+    return hash_sha256.hexdigest()
+
+
+def precalculate_safetensors_hashes(tensors, metadata):
+    """Precalculate the model hashes needed by sd-webui-additional-networks to
+    save time on indexing the model later."""
+
+    # Because writing user metadata to the file can change the result of
+    # sd_models.model_hash(), only retain the training metadata for purposes of
+    # calculating the hash, as they are meant to be immutable
+    metadata = {k: v for k, v in metadata.items() if k.startswith("ss_")}
+
+    bytes = safetensors.torch.save(tensors, metadata)
+    b = BytesIO(bytes)
+
+    model_hash = addnet_hash_safetensors(b)
+    legacy_hash = addnet_hash_legacy(b)
+    return model_hash, legacy_hash

lycoris/locon.py

@@ -0,0 +1,76 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LoConModule(nn.Module):
+    """
+    modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
+    """
+
+    def __init__(self, lora_name, org_module: nn.Module, multiplier=1.0, lora_dim=4, alpha=1, dropout=0.):
+        """ if alpha == 0 or None, alpha is rank (no scaling). """
+        super().__init__()
+        self.lora_name = lora_name
+        self.lora_dim = lora_dim
+
+        if org_module.__class__.__name__ == 'Conv2d':
+            # For general LoCon
+            in_dim = org_module.in_channels
+            k_size = org_module.kernel_size
+            stride = org_module.stride
+            padding = org_module.padding
+            out_dim = org_module.out_channels
+            self.lora_down = nn.Conv2d(in_dim, lora_dim, k_size, stride, padding, bias=False)
+            self.lora_up = nn.Conv2d(lora_dim, out_dim, (1, 1), bias=False)
+            self.op = F.conv2d
+            self.extra_args = {
+                'stride': stride,
+                'padding': padding
+            }
+        else:
+            in_dim = org_module.in_features
+            out_dim = org_module.out_features
+            self.lora_down = nn.Linear(in_dim, lora_dim, bias=False)
+            self.lora_up = nn.Linear(lora_dim, out_dim, bias=False)
+            self.op = F.linear
+            self.extra_args = {}
+        self.shape = org_module.weight.shape
+        
+        if dropout:
+            self.dropout = nn.Dropout(dropout)
+        else:
+            self.dropout = nn.Identity()
+        
+        if type(alpha) == torch.Tensor:
+            alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
+        alpha = lora_dim if alpha is None or alpha == 0 else alpha
+        self.scale = alpha / self.lora_dim
+        self.register_buffer('alpha', torch.tensor(alpha)) # 定数として扱える
+
+        # same as microsoft's
+        torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
+        torch.nn.init.zeros_(self.lora_up.weight)
+
+        self.multiplier = multiplier
+        self.org_module = [org_module]
+
+    def apply_to(self):
+        self.org_module[0].forward = self.forward
+
+    def make_weight(self):
+        wa = self.lora_up.weight
+        wb = self.lora_down.weight
+        return (wa.view(wa.size(0), -1) @ wb.view(wb.size(0), -1)).view(self.shape)
+
+    def forward(self, x):
+        bias = None if self.org_module[0].bias is None else self.org_module[0].bias.data
+        return self.op(
+            x,
+            (self.org_module[0].weight.data 
+             + self.dropout(self.make_weight()) * self.multiplier * self.scale),
+            bias,
+            **self.extra_args,
+        )

lycoris/loha.py

@@ -0,0 +1,116 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class HadaWeight(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, orig_weight, w1a, w1b, w2a, w2b, scale=torch.tensor(1), dropout=nn.Identity()):
+        ctx.save_for_backward(w1a, w1b, w2a, w2b, scale)
+        diff_weight = ((w1a@w1b)*(w2a@w2b)) * scale
+        return orig_weight.reshape(diff_weight.shape) + dropout(diff_weight)
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        (w1a, w1b, w2a, w2b, scale) = ctx.saved_tensors
+        temp = grad_out*(w2a@w2b)*scale
+        grad_w1a = temp @ w1b.T
+        grad_w1b = w1a.T @ temp
+
+        temp = grad_out * (w1a@w1b)*scale
+        grad_w2a = temp @ w2b.T
+        grad_w2b = w2a.T @ temp
+        
+        del temp
+        return grad_out, grad_w1a, grad_w1b, grad_w2a, grad_w2b, None
+
+
+def make_weight(orig_weight, w1a, w1b, w2a, w2b, scale):
+    return HadaWeight.apply(orig_weight, w1a, w1b, w2a, w2b, scale)
+
+
+class LohaModule(nn.Module):
+    """
+    Hadamard product Implementaion for Low Rank Adaptation
+    """
+
+    def __init__(self, lora_name, org_module: nn.Module, multiplier=1.0, lora_dim=4, alpha=1, dropout=0.):
+        """ if alpha == 0 or None, alpha is rank (no scaling). """
+        super().__init__()
+        self.lora_name = lora_name
+        self.lora_dim = lora_dim
+        
+        self.shape = org_module.weight.shape
+        if org_module.__class__.__name__ == 'Conv2d':
+            in_dim = org_module.in_channels
+            k_size = org_module.kernel_size
+            out_dim = org_module.out_channels
+            shape = (out_dim, in_dim*k_size[0]*k_size[1])
+            self.op = F.conv2d
+            self.extra_args = {
+                "stride": org_module.stride,
+                "padding": org_module.padding,
+                "dilation": org_module.dilation,
+                "groups": org_module.groups
+            }
+        else:
+            in_dim = org_module.in_features
+            out_dim = org_module.out_features
+            shape = (out_dim, in_dim)
+            self.op = F.linear
+            self.extra_args = {}
+        
+        self.hada_w1_a = nn.Parameter(torch.empty(shape[0], lora_dim))
+        self.hada_w1_b = nn.Parameter(torch.empty(lora_dim, shape[1]))
+        
+        self.hada_w2_a = nn.Parameter(torch.empty(shape[0], lora_dim))
+        self.hada_w2_b = nn.Parameter(torch.empty(lora_dim, shape[1]))
+        
+        if dropout:
+            self.dropout = nn.Dropout(dropout)
+        else:
+            self.dropout = nn.Identity()
+        
+        if type(alpha) == torch.Tensor:
+            alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
+        alpha = lora_dim if alpha is None or alpha == 0 else alpha
+        self.scale = alpha / self.lora_dim
+        self.register_buffer('alpha', torch.tensor(alpha)) # 定数として扱える
+
+        # Need more experiences on init method
+        torch.nn.init.normal_(self.hada_w1_b, std=1)
+        torch.nn.init.normal_(self.hada_w2_b, std=0.05)
+        torch.nn.init.normal_(self.hada_w1_a, std=1)
+        torch.nn.init.constant_(self.hada_w2_a, 0)
+
+        self.multiplier = multiplier
+        self.org_module = [org_module] # remove in applying
+        self.grad_ckpt = False
+
+    def apply_to(self):
+        self.org_module[0].forward = self.forward
+
+    def get_weight(self):
+        d_weight = self.hada_w1_a @ self.hada_w1_b
+        d_weight *= self.hada_w2_a @ self.hada_w2_b
+        return (d_weight).reshape(self.shape)
+
+    @torch.enable_grad()
+    def forward(self, x):
+        # print(torch.mean(torch.abs(self.orig_w1a.to(x.device) - self.hada_w1_a)), end='\r')
+        weight = make_weight(
+            self.org_module[0].weight.data, 
+            self.hada_w1_a, self.hada_w1_b,
+            self.hada_w2_a, self.hada_w2_b,
+            scale = torch.tensor(self.scale*self.multiplier),
+        )
+        
+        bias = None if self.org_module[0].bias is None else self.org_module[0].bias.data
+        return self.op(
+            x, 
+            weight.view(self.shape),
+            bias,
+            **self.extra_args
+        )

lycoris/utils.py

@@ -0,0 +1,271 @@
+from typing import *
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import torch.linalg as linalg
+
+from tqdm import tqdm
+
+
+def extract_conv(
+    weight: Union[torch.Tensor, nn.Parameter],
+    mode = 'fixed',
+    mode_param = 0,
+    device = 'cpu',
+) -> Tuple[nn.Parameter, nn.Parameter]:
+    out_ch, in_ch, kernel_size, _ = weight.shape
+    
+    U, S, Vh = linalg.svd(weight.reshape(out_ch, -1).to(device))
+    
+    if mode=='fixed':
+        lora_rank = mode_param
+    elif mode=='threshold':
+        assert mode_param>=0
+        lora_rank = torch.sum(S>mode_param)
+    elif mode=='ratio':
+        assert 1>=mode_param>=0
+        min_s = torch.max(S)*mode_param
+        lora_rank = torch.sum(S>min_s)
+    elif mode=='percentile':
+        assert 1>=mode_param>=0
+        s_cum = torch.cumsum(S, dim=0)
+        min_cum_sum = mode_param * torch.sum(S)
+        lora_rank = torch.sum(s_cum<min_cum_sum)
+    lora_rank = max(1, lora_rank)
+    lora_rank = min(out_ch, in_ch, lora_rank)
+    
+    U = U[:, :lora_rank]
+    S = S[:lora_rank]
+    U = U @ torch.diag(S)
+    Vh = Vh[:lora_rank, :]
+    
+    extract_weight_A = Vh.reshape(lora_rank, in_ch, kernel_size, kernel_size).cpu()
+    extract_weight_B = U.reshape(out_ch, lora_rank, 1, 1).cpu()
+    del U, S, Vh, weight
+    return extract_weight_A, extract_weight_B
+
+
+def merge_conv(
+    weight_a: Union[torch.Tensor, nn.Parameter],
+    weight_b: Union[torch.Tensor, nn.Parameter],
+    device = 'cpu'
+):
+    rank, in_ch, kernel_size, k_ = weight_a.shape
+    out_ch, rank_, _, _ = weight_b.shape
+    assert rank == rank_ and kernel_size == k_
+    
+    wa = weight_a.to(device)
+    wb = weight_b.to(device)
+    
+    if device == 'cpu':
+        wa = wa.float()
+        wb = wb.float()
+    
+    merged = wb.reshape(out_ch, -1) @ wa.reshape(rank, -1)
+    weight = merged.reshape(out_ch, in_ch, kernel_size, kernel_size)
+    del wb, wa
+    return weight
+
+
+def extract_linear(
+    weight: Union[torch.Tensor, nn.Parameter],
+    mode = 'fixed',
+    mode_param = 0,
+    device = 'cpu',
+) -> Tuple[nn.Parameter, nn.Parameter]:
+    out_ch, in_ch = weight.shape
+    
+    U, S, Vh = linalg.svd(weight.to(device))
+    
+    if mode=='fixed':
+        lora_rank = mode_param
+    elif mode=='threshold':
+        assert mode_param>=0
+        lora_rank = torch.sum(S>mode_param)
+    elif mode=='ratio':
+        assert 1>=mode_param>=0
+        min_s = torch.max(S)*mode_param
+        lora_rank = torch.sum(S>min_s)
+    elif mode=='percentile':
+        assert 1>=mode_param>=0
+        s_cum = torch.cumsum(S, dim=0)
+        min_cum_sum = mode_param * torch.sum(S)
+        lora_rank = torch.sum(s_cum<min_cum_sum)
+    lora_rank = max(1, lora_rank)
+    lora_rank = min(out_ch, in_ch, lora_rank)
+    
+    U = U[:, :lora_rank]
+    S = S[:lora_rank]
+    U = U @ torch.diag(S)
+    Vh = Vh[:lora_rank, :]
+    
+    extract_weight_A = Vh.reshape(lora_rank, in_ch).cpu()
+    extract_weight_B = U.reshape(out_ch, lora_rank).cpu()
+    del U, S, Vh, weight
+    return extract_weight_A, extract_weight_B
+
+
+def merge_linear(
+    weight_a: Union[torch.Tensor, nn.Parameter],
+    weight_b: Union[torch.Tensor, nn.Parameter],
+    device = 'cpu'
+):
+    rank, in_ch = weight_a.shape
+    out_ch, rank_ = weight_b.shape
+    assert rank == rank_
+    
+    wa = weight_a.to(device)
+    wb = weight_b.to(device)
+    
+    if device == 'cpu':
+        wa = wa.float()
+        wb = wb.float()
+    
+    weight = wb @ wa
+    del wb, wa
+    return weight
+
+
+def extract_diff(
+    base_model,
+    db_model,
+    mode = 'fixed',
+    linear_mode_param = 0,
+    conv_mode_param = 0,
+    extract_device = 'cpu'
+):
+    UNET_TARGET_REPLACE_MODULE = [
+        "Transformer2DModel", 
+        "Attention", 
+        "ResnetBlock2D", 
+        "Downsample2D", 
+        "Upsample2D"
+    ]
+    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
+    LORA_PREFIX_UNET = 'lora_unet'
+    LORA_PREFIX_TEXT_ENCODER = 'lora_te'
+    def make_state_dict(
+        prefix, 
+        root_module: torch.nn.Module,
+        target_module: torch.nn.Module,
+        target_replace_modules
+    ):
+        loras = {}
+        temp = {}
+        
+        for name, module in root_module.named_modules():
+            if module.__class__.__name__ in target_replace_modules:
+                temp[name] = {}
+                for child_name, child_module in module.named_modules():
+                    if child_module.__class__.__name__ not in {'Linear', 'Conv2d'}:
+                        continue
+                    temp[name][child_name] = child_module.weight
+        
+        for name, module in tqdm(list(target_module.named_modules())):
+            if name in temp:
+                weights = temp[name]
+                for child_name, child_module in module.named_modules():
+                    lora_name = prefix + '.' + name + '.' + child_name
+                    lora_name = lora_name.replace('.', '_')
+                    
+                    layer = child_module.__class__.__name__
+                    if layer == 'Linear':
+                        extract_a, extract_b = extract_linear(
+                            (child_module.weight - weights[child_name]),
+                            mode,
+                            linear_mode_param,
+                            device = extract_device,
+                        )
+                    elif layer == 'Conv2d':
+                        is_linear = (child_module.weight.shape[2] == 1
+                                     and child_module.weight.shape[3] == 1)
+                        extract_a, extract_b = extract_conv(
+                            (child_module.weight - weights[child_name]), 
+                            mode,
+                            linear_mode_param if is_linear else conv_mode_param,
+                            device = extract_device,
+                        )
+                    else:
+                        continue
+                    loras[f'{lora_name}.lora_down.weight'] = extract_a.detach().cpu().contiguous().half()
+                    loras[f'{lora_name}.lora_up.weight'] = extract_b.detach().cpu().contiguous().half()
+                    loras[f'{lora_name}.alpha'] = torch.Tensor([extract_a.shape[0]]).half()
+                    del extract_a, extract_b
+        return loras
+    
+    text_encoder_loras = make_state_dict(
+        LORA_PREFIX_TEXT_ENCODER, 
+        base_model[0], db_model[0], 
+        TEXT_ENCODER_TARGET_REPLACE_MODULE
+    )
+    
+    unet_loras = make_state_dict(
+        LORA_PREFIX_UNET,
+        base_model[2], db_model[2], 
+        UNET_TARGET_REPLACE_MODULE
+    )
+    print(len(text_encoder_loras), len(unet_loras))
+    return text_encoder_loras|unet_loras
+
+
+def merge_locon(
+    base_model,
+    locon_state_dict: Dict[str, torch.TensorType],
+    scale: float = 1.0,
+    device = 'cpu'
+):
+    UNET_TARGET_REPLACE_MODULE = [
+        "Transformer2DModel", 
+        "Attention", 
+        "ResnetBlock2D", 
+        "Downsample2D", 
+        "Upsample2D"
+    ]
+    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
+    LORA_PREFIX_UNET = 'lora_unet'
+    LORA_PREFIX_TEXT_ENCODER = 'lora_te'
+    def merge(
+        prefix, 
+        root_module: torch.nn.Module,
+        target_replace_modules
+    ):
+        temp = {}
+        
+        for name, module in tqdm(list(root_module.named_modules())):
+            if module.__class__.__name__ in target_replace_modules:
+                temp[name] = {}
+                for child_name, child_module in module.named_modules():
+                    layer = child_module.__class__.__name__
+                    if layer not in {'Linear', 'Conv2d'}:
+                        continue
+                    lora_name = prefix + '.' + name + '.' + child_name
+                    lora_name = lora_name.replace('.', '_')
+                    
+                    down = locon_state_dict[f'{lora_name}.lora_down.weight'].float()
+                    up = locon_state_dict[f'{lora_name}.lora_up.weight'].float()
+                    alpha = locon_state_dict[f'{lora_name}.alpha'].float()
+                    rank = down.shape[0]
+                    
+                    if layer == 'Conv2d':
+                        delta = merge_conv(down, up, device)
+                        child_module.weight.requires_grad_(False)
+                        child_module.weight += (alpha.to(device)/rank * scale * delta).cpu()
+                        del delta
+                    elif layer == 'Linear':
+                        delta = merge_linear(down, up, device)
+                        child_module.weight.requires_grad_(False)
+                        child_module.weight += (alpha.to(device)/rank * scale * delta).cpu()
+                        del delta
+    
+    merge(
+        LORA_PREFIX_TEXT_ENCODER, 
+        base_model[0], 
+        TEXT_ENCODER_TARGET_REPLACE_MODULE
+    )
+    merge(
+        LORA_PREFIX_UNET,
+        base_model[2], 
+        UNET_TARGET_REPLACE_MODULE
+    )

networks/check_lora_weights.py

@@ -0,0 +1,32 @@
+import argparse
+import os
+import torch
+from safetensors.torch import load_file
+
+
+def main(file):
+  print(f"loading: {file}")
+  if os.path.splitext(file)[1] == '.safetensors':
+    sd = load_file(file)
+  else:
+    sd = torch.load(file, map_location='cpu')
+
+  values = []
+
+  keys = list(sd.keys())
+  for key in keys:
+    if 'lora_up' in key or 'lora_down' in key:
+      values.append((key, sd[key]))
+  print(f"number of LoRA modules: {len(values)}")
+
+  for key, value in values:
+    value = value.to(torch.float32)
+    print(f"{key},{torch.mean(torch.abs(value))},{torch.min(torch.abs(value))}")
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("file", type=str, help="model file to check / 重みを確認するモデルファイル")
+  args = parser.parse_args()
+
+  main(args.file)

networks/extract_lora_from_models.py

@@ -0,0 +1,164 @@
+# extract approximating LoRA by svd from two SD models
+# The code is based on https://github.com/cloneofsimo/lora/blob/develop/lora_diffusion/cli_svd.py
+# Thanks to cloneofsimo!
+
+import argparse
+import os
+import torch
+from safetensors.torch import load_file, save_file
+from tqdm import tqdm
+import library.model_util as model_util
+import lora
+
+
+CLAMP_QUANTILE = 0.99
+MIN_DIFF = 1e-6
+
+
+def save_to_file(file_name, model, state_dict, dtype):
+  if dtype is not None:
+    for key in list(state_dict.keys()):
+      if type(state_dict[key]) == torch.Tensor:
+        state_dict[key] = state_dict[key].to(dtype)
+
+  if os.path.splitext(file_name)[1] == '.safetensors':
+    save_file(model, file_name)
+  else:
+    torch.save(model, file_name)
+
+
+def svd(args):
+  def str_to_dtype(p):
+    if p == 'float':
+      return torch.float
+    if p == 'fp16':
+      return torch.float16
+    if p == 'bf16':
+      return torch.bfloat16
+    return None
+
+  save_dtype = str_to_dtype(args.save_precision)
+
+  print(f"loading SD model : {args.model_org}")
+  text_encoder_o, _, unet_o = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.model_org)
+  print(f"loading SD model : {args.model_tuned}")
+  text_encoder_t, _, unet_t = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.model_tuned)
+
+  # create LoRA network to extract weights: Use dim (rank) as alpha
+  lora_network_o = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_o, unet_o)
+  lora_network_t = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_t, unet_t)
+  assert len(lora_network_o.text_encoder_loras) == len(
+      lora_network_t.text_encoder_loras), f"model version is different (SD1.x vs SD2.x) / それぞれのモデルのバージョンが違います（SD1.xベースとSD2.xベース） "
+
+  # get diffs
+  diffs = {}
+  text_encoder_different = False
+  for i, (lora_o, lora_t) in enumerate(zip(lora_network_o.text_encoder_loras, lora_network_t.text_encoder_loras)):
+    lora_name = lora_o.lora_name
+    module_o = lora_o.org_module
+    module_t = lora_t.org_module
+    diff = module_t.weight - module_o.weight
+
+    # Text Encoder might be same
+    if torch.max(torch.abs(diff)) > MIN_DIFF:
+      text_encoder_different = True
+
+    diff = diff.float()
+    diffs[lora_name] = diff
+
+  if not text_encoder_different:
+    print("Text encoder is same. Extract U-Net only.")
+    lora_network_o.text_encoder_loras = []
+    diffs = {}
+
+  for i, (lora_o, lora_t) in enumerate(zip(lora_network_o.unet_loras, lora_network_t.unet_loras)):
+    lora_name = lora_o.lora_name
+    module_o = lora_o.org_module
+    module_t = lora_t.org_module
+    diff = module_t.weight - module_o.weight
+    diff = diff.float()
+
+    if args.device:
+      diff = diff.to(args.device)
+
+    diffs[lora_name] = diff
+
+  # make LoRA with svd
+  print("calculating by svd")
+  rank = args.dim
+  lora_weights = {}
+  with torch.no_grad():
+    for lora_name, mat in tqdm(list(diffs.items())):
+      conv2d = (len(mat.size()) == 4)
+      if conv2d:
+        mat = mat.squeeze()
+
+      U, S, Vh = torch.linalg.svd(mat)
+
+      U = U[:, :rank]
+      S = S[:rank]
+      U = U @ torch.diag(S)
+
+      Vh = Vh[:rank, :]
+
+      dist = torch.cat([U.flatten(), Vh.flatten()])
+      hi_val = torch.quantile(dist, CLAMP_QUANTILE)
+      low_val = -hi_val
+
+      U = U.clamp(low_val, hi_val)
+      Vh = Vh.clamp(low_val, hi_val)
+
+      lora_weights[lora_name] = (U, Vh)
+
+  # make state dict for LoRA
+  lora_network_o.apply_to(text_encoder_o, unet_o, text_encoder_different, True)   # to make state dict
+  lora_sd = lora_network_o.state_dict()
+  print(f"LoRA has {len(lora_sd)} weights.")
+
+  for key in list(lora_sd.keys()):
+    if "alpha" in key:
+      continue
+
+    lora_name = key.split('.')[0]
+    i = 0 if "lora_up" in key else 1
+
+    weights = lora_weights[lora_name][i]
+    # print(key, i, weights.size(), lora_sd[key].size())
+    if len(lora_sd[key].size()) == 4:
+      weights = weights.unsqueeze(2).unsqueeze(3)
+
+    assert weights.size() == lora_sd[key].size(), f"size unmatch: {key}"
+    lora_sd[key] = weights
+
+  # load state dict to LoRA and save it
+  info = lora_network_o.load_state_dict(lora_sd)
+  print(f"Loading extracted LoRA weights: {info}")
+
+  dir_name = os.path.dirname(args.save_to)
+  if dir_name and not os.path.exists(dir_name):
+    os.makedirs(dir_name, exist_ok=True)
+
+  # minimum metadata
+  metadata = {"ss_network_dim": str(args.dim), "ss_network_alpha": str(args.dim)}
+
+  lora_network_o.save_weights(args.save_to, save_dtype, metadata)
+  print(f"LoRA weights are saved to: {args.save_to}")
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--v2", action='store_true',
+                      help='load Stable Diffusion v2.x model / Stable Diffusion 2.xのモデルを読み込む')
+  parser.add_argument("--save_precision", type=str, default=None,
+                      choices=[None, "float", "fp16", "bf16"], help="precision in saving, same to merging if omitted / 保存時に精度を変更して保存する、省略時はfloat")
+  parser.add_argument("--model_org", type=str, default=None,
+                      help="Stable Diffusion original model: ckpt or safetensors file / 元モデル、ckptまたはsafetensors")
+  parser.add_argument("--model_tuned", type=str, default=None,
+                      help="Stable Diffusion tuned model, LoRA is difference of `original to tuned`: ckpt or safetensors file / 派生モデル（生成されるLoRAは元→派生の差分になります）、ckptまたはsafetensors")
+  parser.add_argument("--save_to", type=str, default=None,
+                      help="destination file name: ckpt or safetensors file / 保存先のファイル名、ckptまたはsafetensors")
+  parser.add_argument("--dim", type=int, default=4, help="dimension (rank) of LoRA (default 4) / LoRAの次元数（rank）（デフォルト4）")
+  parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")
+
+  args = parser.parse_args()
+  svd(args)

networks/lora.py

@@ -0,0 +1,237 @@
+# LoRA network module
+# reference:
+# https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
+# https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/lora.py
+
+import math
+import os
+from typing import List
+import torch
+
+from library import train_util
+
+
+class LoRAModule(torch.nn.Module):
+  """
+  replaces forward method of the original Linear, instead of replacing the original Linear module.
+  """
+
+  def __init__(self, lora_name, org_module: torch.nn.Module, multiplier=1.0, lora_dim=4, alpha=1):
+    """ if alpha == 0 or None, alpha is rank (no scaling). """
+    super().__init__()
+    self.lora_name = lora_name
+    self.lora_dim = lora_dim
+
+    if org_module.__class__.__name__ == 'Conv2d':
+      in_dim = org_module.in_channels
+      out_dim = org_module.out_channels
+      self.lora_down = torch.nn.Conv2d(in_dim, lora_dim, (1, 1), bias=False)
+      self.lora_up = torch.nn.Conv2d(lora_dim, out_dim, (1, 1), bias=False)
+    else:
+      in_dim = org_module.in_features
+      out_dim = org_module.out_features
+      self.lora_down = torch.nn.Linear(in_dim, lora_dim, bias=False)
+      self.lora_up = torch.nn.Linear(lora_dim, out_dim, bias=False)
+
+    if type(alpha) == torch.Tensor:
+      alpha = alpha.detach().float().numpy()                              # without casting, bf16 causes error
+    alpha = lora_dim if alpha is None or alpha == 0 else alpha
+    self.scale = alpha / self.lora_dim
+    self.register_buffer('alpha', torch.tensor(alpha))                    # 定数として扱える
+
+    # same as microsoft's
+    torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
+    torch.nn.init.zeros_(self.lora_up.weight)
+
+    self.multiplier = multiplier
+    self.org_module = org_module                  # remove in applying
+
+  def apply_to(self):
+    self.org_forward = self.org_module.forward
+    self.org_module.forward = self.forward
+    del self.org_module
+
+  def forward(self, x):
+    return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
+
+
+def create_network(multiplier, network_dim, network_alpha, vae, text_encoder, unet, **kwargs):
+  if network_dim is None:
+    network_dim = 4                     # default
+  network = LoRANetwork(text_encoder, unet, multiplier=multiplier, lora_dim=network_dim, alpha=network_alpha)
+  return network
+
+
+def create_network_from_weights(multiplier, file, vae, text_encoder, unet, **kwargs):
+  if os.path.splitext(file)[1] == '.safetensors':
+    from safetensors.torch import load_file, safe_open
+    weights_sd = load_file(file)
+  else:
+    weights_sd = torch.load(file, map_location='cpu')
+
+  # get dim (rank)
+  network_alpha = None
+  network_dim = None
+  for key, value in weights_sd.items():
+    if network_alpha is None and 'alpha' in key:
+      network_alpha = value
+    if network_dim is None and 'lora_down' in key and len(value.size()) == 2:
+      network_dim = value.size()[0]
+
+  if network_alpha is None:
+    network_alpha = network_dim
+
+  network = LoRANetwork(text_encoder, unet, multiplier=multiplier, lora_dim=network_dim, alpha=network_alpha)
+  network.weights_sd = weights_sd
+  return network
+
+
+class LoRANetwork(torch.nn.Module):
+  UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"]
+  TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
+  LORA_PREFIX_UNET = 'lora_unet'
+  LORA_PREFIX_TEXT_ENCODER = 'lora_te'
+
+  def __init__(self, text_encoder, unet, multiplier=1.0, lora_dim=4, alpha=1) -> None:
+    super().__init__()
+    self.multiplier = multiplier
+    self.lora_dim = lora_dim
+    self.alpha = alpha
+
+    # create module instances
+    def create_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> List[LoRAModule]:
+      loras = []
+      for name, module in root_module.named_modules():
+        if module.__class__.__name__ in target_replace_modules:
+          for child_name, child_module in module.named_modules():
+            if child_module.__class__.__name__ == "Linear" or (child_module.__class__.__name__ == "Conv2d" and child_module.kernel_size == (1, 1)):
+              lora_name = prefix + '.' + name + '.' + child_name
+              lora_name = lora_name.replace('.', '_')
+              lora = LoRAModule(lora_name, child_module, self.multiplier, self.lora_dim, self.alpha)
+              loras.append(lora)
+      return loras
+
+    self.text_encoder_loras = create_modules(LoRANetwork.LORA_PREFIX_TEXT_ENCODER,
+                                             text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE)
+    print(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
+
+    self.unet_loras = create_modules(LoRANetwork.LORA_PREFIX_UNET, unet, LoRANetwork.UNET_TARGET_REPLACE_MODULE)
+    print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
+
+    self.weights_sd = None
+
+    # assertion
+    names = set()
+    for lora in self.text_encoder_loras + self.unet_loras:
+      assert lora.lora_name not in names, f"duplicated lora name: {lora.lora_name}"
+      names.add(lora.lora_name)
+
+  def load_weights(self, file):
+    if os.path.splitext(file)[1] == '.safetensors':
+      from safetensors.torch import load_file, safe_open
+      self.weights_sd = load_file(file)
+    else:
+      self.weights_sd = torch.load(file, map_location='cpu')
+
+  def apply_to(self, text_encoder, unet, apply_text_encoder=None, apply_unet=None):
+    if self.weights_sd:
+      weights_has_text_encoder = weights_has_unet = False
+      for key in self.weights_sd.keys():
+        if key.startswith(LoRANetwork.LORA_PREFIX_TEXT_ENCODER):
+          weights_has_text_encoder = True
+        elif key.startswith(LoRANetwork.LORA_PREFIX_UNET):
+          weights_has_unet = True
+
+      if apply_text_encoder is None:
+        apply_text_encoder = weights_has_text_encoder
+      else:
+        assert apply_text_encoder == weights_has_text_encoder, f"text encoder weights: {weights_has_text_encoder} but text encoder flag: {apply_text_encoder} / 重みとText Encoderのフラグが矛盾しています"
+
+      if apply_unet is None:
+        apply_unet = weights_has_unet
+      else:
+        assert apply_unet == weights_has_unet, f"u-net weights: {weights_has_unet} but u-net flag: {apply_unet} / 重みとU-Netのフラグが矛盾しています"
+    else:
+      assert apply_text_encoder is not None and apply_unet is not None, f"internal error: flag not set"
+
+    if apply_text_encoder:
+      print("enable LoRA for text encoder")
+    else:
+      self.text_encoder_loras = []
+
+    if apply_unet:
+      print("enable LoRA for U-Net")
+    else:
+      self.unet_loras = []
+
+    for lora in self.text_encoder_loras + self.unet_loras:
+      lora.apply_to()
+      self.add_module(lora.lora_name, lora)
+
+    if self.weights_sd:
+      # if some weights are not in state dict, it is ok because initial LoRA does nothing (lora_up is initialized by zeros)
+      info = self.load_state_dict(self.weights_sd, False)
+      print(f"weights are loaded: {info}")
+
+  def enable_gradient_checkpointing(self):
+    # not supported
+    pass
+
+  def prepare_optimizer_params(self, text_encoder_lr, unet_lr):
+    def enumerate_params(loras):
+      params = []
+      for lora in loras:
+        params.extend(lora.parameters())
+      return params
+
+    self.requires_grad_(True)
+    all_params = []
+
+    if self.text_encoder_loras:
+      param_data = {'params': enumerate_params(self.text_encoder_loras)}
+      if text_encoder_lr is not None:
+        param_data['lr'] = text_encoder_lr
+      all_params.append(param_data)
+
+    if self.unet_loras:
+      param_data = {'params': enumerate_params(self.unet_loras)}
+      if unet_lr is not None:
+        param_data['lr'] = unet_lr
+      all_params.append(param_data)
+
+    return all_params
+
+  def prepare_grad_etc(self, text_encoder, unet):
+    self.requires_grad_(True)
+
+  def on_epoch_start(self, text_encoder, unet):
+    self.train()
+
+  def get_trainable_params(self):
+    return self.parameters()
+
+  def save_weights(self, file, dtype, metadata):
+    if metadata is not None and len(metadata) == 0:
+      metadata = None
+
+    state_dict = self.state_dict()
+
+    if dtype is not None:
+      for key in list(state_dict.keys()):
+        v = state_dict[key]
+        v = v.detach().clone().to("cpu").to(dtype)
+        state_dict[key] = v
+
+    if os.path.splitext(file)[1] == '.safetensors':
+      from safetensors.torch import save_file
+
+      # Precalculate model hashes to save time on indexing
+      if metadata is None:
+        metadata = {}
+      model_hash, legacy_hash = train_util.precalculate_safetensors_hashes(state_dict, metadata)
+      metadata["sshs_model_hash"] = model_hash
+      metadata["sshs_legacy_hash"] = legacy_hash
+
+      save_file(state_dict, file, metadata)
+    else:
+      torch.save(state_dict, file)

networks/lora_interrogator.py

@@ -0,0 +1,122 @@
+
+
+from tqdm import tqdm
+from library import model_util
+import argparse
+from transformers import CLIPTokenizer
+import torch
+
+import library.model_util as model_util
+import lora
+
+TOKENIZER_PATH = "openai/clip-vit-large-patch14"
+V2_STABLE_DIFFUSION_PATH = "stabilityai/stable-diffusion-2"     # ここからtokenizerだけ使う
+
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def interrogate(args):
+  # いろいろ準備する
+  print(f"loading SD model: {args.sd_model}")
+  text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.sd_model)
+
+  print(f"loading LoRA: {args.model}")
+  network = lora.create_network_from_weights(1.0, args.model, vae, text_encoder, unet)
+
+  # text encoder向けの重みがあるかチェックする：本当はlora側でやるのがいい
+  has_te_weight = False
+  for key in network.weights_sd.keys():
+    if 'lora_te' in key:
+      has_te_weight = True
+      break
+  if not has_te_weight:
+    print("This LoRA does not have modules for Text Encoder, cannot interrogate / このLoRAはText Encoder向けのモジュールがないため調査できません")
+    return
+  del vae
+
+  print("loading tokenizer")
+  if args.v2:
+    tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(V2_STABLE_DIFFUSION_PATH, subfolder="tokenizer")
+  else:
+    tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(TOKENIZER_PATH)  # , model_max_length=max_token_length + 2)
+
+  text_encoder.to(DEVICE)
+  text_encoder.eval()
+  unet.to(DEVICE)
+  unet.eval()               # U-Netは呼び出さないので不要だけど
+
+  # トークンをひとつひとつ当たっていく
+  token_id_start = 0
+  token_id_end = max(tokenizer.all_special_ids)
+  print(f"interrogate tokens are: {token_id_start} to {token_id_end}")
+
+  def get_all_embeddings(text_encoder):
+    embs = []
+    with torch.no_grad():
+      for token_id in tqdm(range(token_id_start, token_id_end + 1, args.batch_size)):
+        batch = []
+        for tid in range(token_id, min(token_id_end + 1, token_id + args.batch_size)):
+          tokens = [tokenizer.bos_token_id, tid, tokenizer.eos_token_id]
+          # tokens = [tid]                                                    # こちらは結果がいまひとつ
+          batch.append(tokens)
+
+        # batch_embs = text_encoder(torch.tensor(batch).to(DEVICE))[0].to("cpu")  # bos/eosも含めたほうが差が出るようだ [:, 1]
+        # clip skip対応
+        batch = torch.tensor(batch).to(DEVICE)
+        if args.clip_skip is None:
+          encoder_hidden_states = text_encoder(batch)[0]
+        else:
+          enc_out = text_encoder(batch, 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.to("cpu")
+
+        embs.extend(encoder_hidden_states)
+    return torch.stack(embs)
+
+  print("get original text encoder embeddings.")
+  orig_embs = get_all_embeddings(text_encoder)
+
+  network.apply_to(text_encoder, unet, True, len(network.unet_loras) > 0)
+  network.to(DEVICE)
+  network.eval()
+
+  print("You can ignore warning messages start with '_IncompatibleKeys' (LoRA model does not have alpha because trained by older script) / '_IncompatibleKeys'の警告は無視して構いません（以前のスクリプトで学習されたLoRAモデルのためalphaの定義がありません）")
+  print("get text encoder embeddings with lora.")
+  lora_embs = get_all_embeddings(text_encoder)
+
+  # 比べる：とりあえず単純に差分の絶対値で
+  print("comparing...")
+  diffs = {}
+  for i, (orig_emb, lora_emb) in enumerate(zip(orig_embs, tqdm(lora_embs))):
+    diff = torch.mean(torch.abs(orig_emb - lora_emb))
+    # diff = torch.mean(torch.cosine_similarity(orig_emb, lora_emb, dim=1))       # うまく検出できない
+    diff = float(diff.detach().to('cpu').numpy())
+    diffs[token_id_start + i] = diff
+
+  diffs_sorted = sorted(diffs.items(), key=lambda x: -x[1])
+
+  # 結果を表示する
+  print("top 100:")
+  for i, (token, diff) in enumerate(diffs_sorted[:100]):
+    # if diff < 1e-6:
+    #   break
+    string = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens([token]))
+    print(f"[{i:3d}]: {token:5d} {string:<20s}: {diff:.5f}")
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--v2", action='store_true',
+                      help='load Stable Diffusion v2.x model / Stable Diffusion 2.xのモデルを読み込む')
+  parser.add_argument("--sd_model", type=str, default=None,
+                      help="Stable Diffusion model to load: ckpt or safetensors file / 読み込むSDのモデル、ckptまたはsafetensors")
+  parser.add_argument("--model", type=str, default=None,
+                      help="LoRA model to interrogate: ckpt or safetensors file / 調査するLoRAモデル、ckptまたはsafetensors")
+  parser.add_argument("--batch_size", type=int, default=16,
+                      help="batch size for processing with Text Encoder / Text Encoderで処理するときのバッチサイズ")
+  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以上）")
+
+  args = parser.parse_args()
+  interrogate(args)

networks/merge_lora.py

@@ -0,0 +1,212 @@
+
+import math
+import argparse
+import os
+import torch
+from safetensors.torch import load_file, save_file
+import library.model_util as model_util
+import lora
+
+
+def load_state_dict(file_name, dtype):
+  if os.path.splitext(file_name)[1] == '.safetensors':
+    sd = load_file(file_name)
+  else:
+    sd = torch.load(file_name, map_location='cpu')
+  for key in list(sd.keys()):
+    if type(sd[key]) == torch.Tensor:
+      sd[key] = sd[key].to(dtype)
+  return sd
+
+
+def save_to_file(file_name, model, state_dict, dtype):
+  if dtype is not None:
+    for key in list(state_dict.keys()):
+      if type(state_dict[key]) == torch.Tensor:
+        state_dict[key] = state_dict[key].to(dtype)
+
+  if os.path.splitext(file_name)[1] == '.safetensors':
+    save_file(model, file_name)
+  else:
+    torch.save(model, file_name)
+
+
+def merge_to_sd_model(text_encoder, unet, models, ratios, merge_dtype):
+  text_encoder.to(merge_dtype)
+  unet.to(merge_dtype)
+
+  # create module map
+  name_to_module = {}
+  for i, root_module in enumerate([text_encoder, unet]):
+    if i == 0:
+      prefix = lora.LoRANetwork.LORA_PREFIX_TEXT_ENCODER
+      target_replace_modules = lora.LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE
+    else:
+      prefix = lora.LoRANetwork.LORA_PREFIX_UNET
+      target_replace_modules = lora.LoRANetwork.UNET_TARGET_REPLACE_MODULE
+
+    for name, module in root_module.named_modules():
+      if module.__class__.__name__ in target_replace_modules:
+        for child_name, child_module in module.named_modules():
+          if child_module.__class__.__name__ == "Linear" or (child_module.__class__.__name__ == "Conv2d" and child_module.kernel_size == (1, 1)):
+            lora_name = prefix + '.' + name + '.' + child_name
+            lora_name = lora_name.replace('.', '_')
+            name_to_module[lora_name] = child_module
+
+  for model, ratio in zip(models, ratios):
+    print(f"loading: {model}")
+    lora_sd = load_state_dict(model, merge_dtype)
+
+    print(f"merging...")
+    for key in lora_sd.keys():
+      if "lora_down" in key:
+        up_key = key.replace("lora_down", "lora_up")
+        alpha_key = key[:key.index("lora_down")] + 'alpha'
+
+        # find original module for this lora
+        module_name = '.'.join(key.split('.')[:-2])               # remove trailing ".lora_down.weight"
+        if module_name not in name_to_module:
+          print(f"no module found for LoRA weight: {key}")
+          continue
+        module = name_to_module[module_name]
+        # print(f"apply {key} to {module}")
+
+        down_weight = lora_sd[key]
+        up_weight = lora_sd[up_key]
+
+        dim = down_weight.size()[0]
+        alpha = lora_sd.get(alpha_key, dim)
+        scale = alpha / dim
+
+        # W <- W + U * D
+        weight = module.weight
+        if len(weight.size()) == 2:
+          # linear
+          weight = weight + ratio * (up_weight @ down_weight) * scale
+        else:
+          # conv2d
+          weight = weight + ratio * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)
+                                     ).unsqueeze(2).unsqueeze(3) * scale
+
+        module.weight = torch.nn.Parameter(weight)
+
+
+def merge_lora_models(models, ratios, merge_dtype):
+  base_alphas = {}                          # alpha for merged model
+  base_dims = {}
+
+  merged_sd = {}
+  for model, ratio in zip(models, ratios):
+    print(f"loading: {model}")
+    lora_sd = load_state_dict(model, merge_dtype)
+
+    # get alpha and dim
+    alphas = {}                             # alpha for current model
+    dims = {}                               # dims for current model
+    for key in lora_sd.keys():
+      if 'alpha' in key:
+        lora_module_name = key[:key.rfind(".alpha")]
+        alpha = float(lora_sd[key].detach().numpy())
+        alphas[lora_module_name] = alpha
+        if lora_module_name not in base_alphas:
+          base_alphas[lora_module_name] = alpha
+      elif "lora_down" in key:
+        lora_module_name = key[:key.rfind(".lora_down")]
+        dim = lora_sd[key].size()[0]
+        dims[lora_module_name] = dim
+        if lora_module_name not in base_dims:
+          base_dims[lora_module_name] = dim
+
+    for lora_module_name in dims.keys():
+      if lora_module_name not in alphas:
+        alpha = dims[lora_module_name]
+        alphas[lora_module_name] = alpha
+        if lora_module_name not in base_alphas:
+          base_alphas[lora_module_name] = alpha
+    
+    print(f"dim: {list(set(dims.values()))}, alpha: {list(set(alphas.values()))}")
+
+    # merge
+    print(f"merging...")
+    for key in lora_sd.keys():
+      if 'alpha' in key:
+        continue
+
+      lora_module_name = key[:key.rfind(".lora_")]
+
+      base_alpha = base_alphas[lora_module_name]
+      alpha = alphas[lora_module_name]
+
+      scale = math.sqrt(alpha / base_alpha) * ratio
+
+      if key in merged_sd:
+        assert merged_sd[key].size() == lora_sd[key].size(
+        ), f"weights shape mismatch merging v1 and v2, different dims? / 重みのサイズ��合いません。v1とv2、または次元数の異なるモデルはマージできません"
+        merged_sd[key] = merged_sd[key] + lora_sd[key] * scale
+      else:
+        merged_sd[key] = lora_sd[key] * scale
+  
+  # set alpha to sd
+  for lora_module_name, alpha in base_alphas.items():
+    key = lora_module_name + ".alpha"
+    merged_sd[key] = torch.tensor(alpha)
+
+  print("merged model")
+  print(f"dim: {list(set(base_dims.values()))}, alpha: {list(set(base_alphas.values()))}")
+
+  return merged_sd
+
+
+def merge(args):
+  assert len(args.models) == len(args.ratios), f"number of models must be equal to number of ratios / モデルの数と重みの数は合わせてください"
+
+  def str_to_dtype(p):
+    if p == 'float':
+      return torch.float
+    if p == 'fp16':
+      return torch.float16
+    if p == 'bf16':
+      return torch.bfloat16
+    return None
+
+  merge_dtype = str_to_dtype(args.precision)
+  save_dtype = str_to_dtype(args.save_precision)
+  if save_dtype is None:
+    save_dtype = merge_dtype
+
+  if args.sd_model is not None:
+    print(f"loading SD model: {args.sd_model}")
+
+    text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.sd_model)
+
+    merge_to_sd_model(text_encoder, unet, args.models, args.ratios, merge_dtype)
+
+    print(f"saving SD model to: {args.save_to}")
+    model_util.save_stable_diffusion_checkpoint(args.v2, args.save_to, text_encoder, unet,
+                                                args.sd_model, 0, 0, save_dtype, vae)
+  else:
+    state_dict = merge_lora_models(args.models, args.ratios, merge_dtype)
+
+    print(f"saving model to: {args.save_to}")
+    save_to_file(args.save_to, state_dict, state_dict, save_dtype)
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--v2", action='store_true',
+                      help='load Stable Diffusion v2.x model / Stable Diffusion 2.xのモデルを読み込む')
+  parser.add_argument("--save_precision", type=str, default=None,
+                      choices=[None, "float", "fp16", "bf16"], help="precision in saving, same to merging if omitted / 保存時に精度を変更して保存する、省略時はマージ時の精度と同じ")
+  parser.add_argument("--precision", type=str, default="float",
+                      choices=["float", "fp16", "bf16"], help="precision in merging (float is recommended) / マージの計算時の精度（floatを推奨）")
+  parser.add_argument("--sd_model", type=str, default=None,
+                      help="Stable Diffusion model to load: ckpt or safetensors file, merge LoRA models if omitted / 読み込むモデル、ckptまたはsafetensors。省略時はLoRAモデル同士をマージする")
+  parser.add_argument("--save_to", type=str, default=None,
+                      help="destination file name: ckpt or safetensors file / 保存先のファイル名、ckptまたはsafetensors")
+  parser.add_argument("--models", type=str, nargs='*',
+                      help="LoRA models to merge: ckpt or safetensors file / マージするLoRAモデル、ckptまたはsafetensors")
+  parser.add_argument("--ratios", type=float, nargs='*',
+                      help="ratios for each model / それぞれのLoRAモデルの比率")
+
+  args = parser.parse_args()
+  merge(args)

networks/merge_lora_old.py

@@ -0,0 +1,179 @@
+
+
+import argparse
+import os
+import torch
+from safetensors.torch import load_file, save_file
+import library.model_util as model_util
+import lora
+
+
+def load_state_dict(file_name, dtype):
+  if os.path.splitext(file_name)[1] == '.safetensors':
+    sd = load_file(file_name)
+  else:
+    sd = torch.load(file_name, map_location='cpu')
+  for key in list(sd.keys()):
+    if type(sd[key]) == torch.Tensor:
+      sd[key] = sd[key].to(dtype)
+  return sd
+
+
+def save_to_file(file_name, model, state_dict, dtype):
+  if dtype is not None:
+    for key in list(state_dict.keys()):
+      if type(state_dict[key]) == torch.Tensor:
+        state_dict[key] = state_dict[key].to(dtype)
+
+  if os.path.splitext(file_name)[1] == '.safetensors':
+    save_file(model, file_name)
+  else:
+    torch.save(model, file_name)
+
+
+def merge_to_sd_model(text_encoder, unet, models, ratios, merge_dtype):
+  text_encoder.to(merge_dtype)
+  unet.to(merge_dtype)
+
+  # create module map
+  name_to_module = {}
+  for i, root_module in enumerate([text_encoder, unet]):
+    if i == 0:
+      prefix = lora.LoRANetwork.LORA_PREFIX_TEXT_ENCODER
+      target_replace_modules = lora.LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE
+    else:
+      prefix = lora.LoRANetwork.LORA_PREFIX_UNET
+      target_replace_modules = lora.LoRANetwork.UNET_TARGET_REPLACE_MODULE
+
+    for name, module in root_module.named_modules():
+      if module.__class__.__name__ in target_replace_modules:
+        for child_name, child_module in module.named_modules():
+          if child_module.__class__.__name__ == "Linear" or (child_module.__class__.__name__ == "Conv2d" and child_module.kernel_size == (1, 1)):
+            lora_name = prefix + '.' + name + '.' + child_name
+            lora_name = lora_name.replace('.', '_')
+            name_to_module[lora_name] = child_module
+
+  for model, ratio in zip(models, ratios):
+    print(f"loading: {model}")
+    lora_sd = load_state_dict(model, merge_dtype)
+
+    print(f"merging...")
+    for key in lora_sd.keys():
+      if "lora_down" in key:
+        up_key = key.replace("lora_down", "lora_up")
+        alpha_key = key[:key.index("lora_down")] + 'alpha'
+
+        # find original module for this lora
+        module_name = '.'.join(key.split('.')[:-2])               # remove trailing ".lora_down.weight"
+        if module_name not in name_to_module:
+          print(f"no module found for LoRA weight: {key}")
+          continue
+        module = name_to_module[module_name]
+        # print(f"apply {key} to {module}")
+
+        down_weight = lora_sd[key]
+        up_weight = lora_sd[up_key]
+
+        dim = down_weight.size()[0]
+        alpha = lora_sd.get(alpha_key, dim)
+        scale = alpha / dim
+
+        # W <- W + U * D
+        weight = module.weight
+        if len(weight.size()) == 2:
+          # linear
+          weight = weight + ratio * (up_weight @ down_weight) * scale
+        else:
+          # conv2d
+          weight = weight + ratio * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3) * scale
+
+        module.weight = torch.nn.Parameter(weight)
+
+
+def merge_lora_models(models, ratios, merge_dtype):
+  merged_sd = {}
+
+  alpha = None
+  dim = None
+  for model, ratio in zip(models, ratios):
+    print(f"loading: {model}")
+    lora_sd = load_state_dict(model, merge_dtype)
+
+    print(f"merging...")
+    for key in lora_sd.keys():
+      if 'alpha' in key:
+        if key in merged_sd:
+          assert merged_sd[key] == lora_sd[key], f"alpha mismatch / alphaが異なる場合、現時点ではマージできません"
+        else:
+          alpha = lora_sd[key].detach().numpy()
+          merged_sd[key] = lora_sd[key]
+      else:
+        if key in merged_sd:
+          assert merged_sd[key].size() == lora_sd[key].size(
+          ), f"weights shape mismatch merging v1 and v2, different dims? / 重みのサイズが合いません。v1とv2、または次元数の異なるモデルはマージできません"
+          merged_sd[key] = merged_sd[key] + lora_sd[key] * ratio
+        else:
+          if "lora_down" in key:
+            dim = lora_sd[key].size()[0]
+          merged_sd[key] = lora_sd[key] * ratio
+
+  print(f"dim (rank): {dim}, alpha: {alpha}")
+  if alpha is None:
+    alpha = dim
+
+  return merged_sd, dim, alpha
+
+
+def merge(args):
+  assert len(args.models) == len(args.ratios), f"number of models must be equal to number of ratios / モデルの数と重みの数は合わせてください"
+
+  def str_to_dtype(p):
+    if p == 'float':
+      return torch.float
+    if p == 'fp16':
+      return torch.float16
+    if p == 'bf16':
+      return torch.bfloat16
+    return None
+
+  merge_dtype = str_to_dtype(args.precision)
+  save_dtype = str_to_dtype(args.save_precision)
+  if save_dtype is None:
+    save_dtype = merge_dtype
+
+  if args.sd_model is not None:
+    print(f"loading SD model: {args.sd_model}")
+
+    text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.sd_model)
+
+    merge_to_sd_model(text_encoder, unet, args.models, args.ratios, merge_dtype)
+
+    print(f"saving SD model to: {args.save_to}")
+    model_util.save_stable_diffusion_checkpoint(args.v2, args.save_to, text_encoder, unet,
+                                                args.sd_model, 0, 0, save_dtype, vae)
+  else:
+    state_dict, _, _ = merge_lora_models(args.models, args.ratios, merge_dtype)
+
+    print(f"saving model to: {args.save_to}")
+    save_to_file(args.save_to, state_dict, state_dict, save_dtype)
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--v2", action='store_true',
+                      help='load Stable Diffusion v2.x model / Stable Diffusion 2.xのモデルを読み込む')
+  parser.add_argument("--save_precision", type=str, default=None,
+                      choices=[None, "float", "fp16", "bf16"], help="precision in saving, same to merging if omitted / 保存時に精度を変更して保存する、省略時はマージ時の精度と同じ")
+  parser.add_argument("--precision", type=str, default="float",
+                      choices=["float", "fp16", "bf16"], help="precision in merging (float is recommended) / マージの計算時の精度（floatを推奨）")
+  parser.add_argument("--sd_model", type=str, default=None,
+                      help="Stable Diffusion model to load: ckpt or safetensors file, merge LoRA models if omitted / 読み込むモデル、ckptまたはsafetensors。省略時はLoRAモデル同士をマージする")
+  parser.add_argument("--save_to", type=str, default=None,
+                      help="destination file name: ckpt or safetensors file / 保存先のファイル名、ckptまたはsafetensors")
+  parser.add_argument("--models", type=str, nargs='*',
+                      help="LoRA models to merge: ckpt or safetensors file / マージするLoRAモデル、ckptまたはsafetensors")
+  parser.add_argument("--ratios", type=float, nargs='*',
+                      help="ratios for each model / それぞれのLoRAモデルの比率")
+
+  args = parser.parse_args()
+  merge(args)

networks/resize_lora.py

@@ -0,0 +1,198 @@
+# Convert LoRA to different rank approximation (should only be used to go to lower rank)
+# This code is based off the extract_lora_from_models.py file which is based on https://github.com/cloneofsimo/lora/blob/develop/lora_diffusion/cli_svd.py
+# Thanks to cloneofsimo and kohya
+
+import argparse
+import os
+import torch
+from safetensors.torch import load_file, save_file, safe_open
+from tqdm import tqdm
+from library import train_util, model_util
+
+
+def load_state_dict(file_name, dtype):
+  if model_util.is_safetensors(file_name):
+    sd = load_file(file_name)
+    with safe_open(file_name, framework="pt") as f:
+      metadata = f.metadata()
+  else:
+    sd = torch.load(file_name, map_location='cpu')
+    metadata = None
+
+  for key in list(sd.keys()):
+    if type(sd[key]) == torch.Tensor:
+      sd[key] = sd[key].to(dtype)
+
+  return sd, metadata
+
+
+def save_to_file(file_name, model, state_dict, dtype, metadata):
+  if dtype is not None:
+    for key in list(state_dict.keys()):
+      if type(state_dict[key]) == torch.Tensor:
+        state_dict[key] = state_dict[key].to(dtype)
+
+  if model_util.is_safetensors(file_name):
+    save_file(model, file_name, metadata)
+  else:
+    torch.save(model, file_name)
+
+
+def resize_lora_model(lora_sd, new_rank, save_dtype, device, verbose):
+  network_alpha = None
+  network_dim = None
+  verbose_str = "\n"
+
+  CLAMP_QUANTILE = 0.99
+
+  # Extract loaded lora dim and alpha
+  for key, value in lora_sd.items():
+    if network_alpha is None and 'alpha' in key:
+      network_alpha = value
+    if network_dim is None and 'lora_down' in key and len(value.size()) == 2:
+      network_dim = value.size()[0]
+    if network_alpha is not None and network_dim is not None:
+      break
+    if network_alpha is None:
+      network_alpha = network_dim
+
+  scale = network_alpha/network_dim
+  new_alpha = float(scale*new_rank)  # calculate new alpha from scale
+
+  print(f"old dimension: {network_dim}, old alpha: {network_alpha}, new alpha: {new_alpha}")
+
+  lora_down_weight = None
+  lora_up_weight = None
+
+  o_lora_sd = lora_sd.copy()
+  block_down_name = None
+  block_up_name = None
+
+  print("resizing lora...")
+  with torch.no_grad():
+    for key, value in tqdm(lora_sd.items()):
+      if 'lora_down' in key:
+        block_down_name = key.split(".")[0]
+        lora_down_weight = value
+      if 'lora_up' in key:
+        block_up_name = key.split(".")[0]
+        lora_up_weight = value
+
+      weights_loaded = (lora_down_weight is not None and lora_up_weight is not None)
+
+      if (block_down_name == block_up_name) and weights_loaded:
+
+        conv2d = (len(lora_down_weight.size()) == 4)
+
+        if conv2d:
+          lora_down_weight = lora_down_weight.squeeze()
+          lora_up_weight = lora_up_weight.squeeze()
+
+        if device:
+          org_device = lora_up_weight.device
+          lora_up_weight = lora_up_weight.to(args.device)
+          lora_down_weight = lora_down_weight.to(args.device)
+
+        full_weight_matrix = torch.matmul(lora_up_weight, lora_down_weight)
+
+        U, S, Vh = torch.linalg.svd(full_weight_matrix)
+
+        if verbose:
+          s_sum = torch.sum(torch.abs(S))
+          s_rank = torch.sum(torch.abs(S[:new_rank]))
+          verbose_str+=f"{block_down_name:76} | "
+          verbose_str+=f"sum(S) retained: {(s_rank)/s_sum:.1%}, max(S) ratio: {S[0]/S[new_rank]:0.1f}\n"
+
+        U = U[:, :new_rank]
+        S = S[:new_rank]
+        U = U @ torch.diag(S)
+
+        Vh = Vh[:new_rank, :]
+
+        dist = torch.cat([U.flatten(), Vh.flatten()])
+        hi_val = torch.quantile(dist, CLAMP_QUANTILE)
+        low_val = -hi_val
+
+        U = U.clamp(low_val, hi_val)
+        Vh = Vh.clamp(low_val, hi_val)
+
+        if conv2d:
+          U = U.unsqueeze(2).unsqueeze(3)
+          Vh = Vh.unsqueeze(2).unsqueeze(3)
+
+        if device:
+          U = U.to(org_device)
+          Vh = Vh.to(org_device)
+
+        o_lora_sd[block_down_name + "." + "lora_down.weight"] = Vh.to(save_dtype).contiguous()
+        o_lora_sd[block_up_name + "." + "lora_up.weight"] = U.to(save_dtype).contiguous()
+        o_lora_sd[block_up_name + "." "alpha"] = torch.tensor(new_alpha).to(save_dtype)
+
+        block_down_name = None
+        block_up_name = None
+        lora_down_weight = None
+        lora_up_weight = None
+        weights_loaded = False
+
+  if verbose:
+    print(verbose_str)
+  print("resizing complete")
+  return o_lora_sd, network_dim, new_alpha
+
+
+def resize(args):
+
+  def str_to_dtype(p):
+    if p == 'float':
+      return torch.float
+    if p == 'fp16':
+      return torch.float16
+    if p == 'bf16':
+      return torch.bfloat16
+    return None
+
+  merge_dtype = str_to_dtype('float')  # matmul method above only seems to work in float32
+  save_dtype = str_to_dtype(args.save_precision)
+  if save_dtype is None:
+    save_dtype = merge_dtype
+
+  print("loading Model...")
+  lora_sd, metadata = load_state_dict(args.model, merge_dtype)
+
+  print("resizing rank...")
+  state_dict, old_dim, new_alpha = resize_lora_model(lora_sd, args.new_rank, save_dtype, args.device, args.verbose)
+
+  # update metadata
+  if metadata is None:
+    metadata = {}
+
+  comment = metadata.get("ss_training_comment", "")
+  metadata["ss_training_comment"] = f"dimension is resized from {old_dim} to {args.new_rank}; {comment}"
+  metadata["ss_network_dim"] = str(args.new_rank)
+  metadata["ss_network_alpha"] = str(new_alpha)
+
+  model_hash, legacy_hash = train_util.precalculate_safetensors_hashes(state_dict, metadata)
+  metadata["sshs_model_hash"] = model_hash
+  metadata["sshs_legacy_hash"] = legacy_hash
+
+  print(f"saving model to: {args.save_to}")
+  save_to_file(args.save_to, state_dict, state_dict, save_dtype, metadata)
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+
+  parser.add_argument("--save_precision", type=str, default=None,
+                      choices=[None, "float", "fp16", "bf16"], help="precision in saving, float if omitted / 保存時の精度、未指定時はfloat")
+  parser.add_argument("--new_rank", type=int, default=4,
+                      help="Specify rank of output LoRA / 出力するLoRAのrank (dim)")
+  parser.add_argument("--save_to", type=str, default=None,
+                      help="destination file name: ckpt or safetensors file / 保存先のファイル名、ckptまたはsafetensors")
+  parser.add_argument("--model", type=str, default=None,
+                      help="LoRA model to resize at to new rank: ckpt or safetensors file / 読み込むLoRAモデル、ckptまたはsafetensors")
+  parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")
+  parser.add_argument("--verbose", action="store_true", 
+                      help="Display verbose resizing information / rank変更時の詳細情報を出力する")
+
+  args = parser.parse_args()
+  resize(args)

networks/svd_merge_lora.py

@@ -0,0 +1,164 @@
+
+import math
+import argparse
+import os
+import torch
+from safetensors.torch import load_file, save_file
+from tqdm import tqdm
+import library.model_util as model_util
+import lora
+
+
+CLAMP_QUANTILE = 0.99
+
+
+def load_state_dict(file_name, dtype):
+  if os.path.splitext(file_name)[1] == '.safetensors':
+    sd = load_file(file_name)
+  else:
+    sd = torch.load(file_name, map_location='cpu')
+  for key in list(sd.keys()):
+    if type(sd[key]) == torch.Tensor:
+      sd[key] = sd[key].to(dtype)
+  return sd
+
+
+def save_to_file(file_name, model, state_dict, dtype):
+  if dtype is not None:
+    for key in list(state_dict.keys()):
+      if type(state_dict[key]) == torch.Tensor:
+        state_dict[key] = state_dict[key].to(dtype)
+
+  if os.path.splitext(file_name)[1] == '.safetensors':
+    save_file(model, file_name)
+  else:
+    torch.save(model, file_name)
+
+
+def merge_lora_models(models, ratios, new_rank, device,  merge_dtype):
+  merged_sd = {}
+  for model, ratio in zip(models, ratios):
+    print(f"loading: {model}")
+    lora_sd = load_state_dict(model, merge_dtype)
+
+    # merge
+    print(f"merging...")
+    for key in tqdm(list(lora_sd.keys())):
+      if 'lora_down' not in key:
+        continue
+
+      lora_module_name = key[:key.rfind(".lora_down")]
+
+      down_weight = lora_sd[key]
+      network_dim = down_weight.size()[0]
+
+      up_weight = lora_sd[lora_module_name + '.lora_up.weight']
+      alpha = lora_sd.get(lora_module_name + '.alpha', network_dim)
+
+      in_dim = down_weight.size()[1]
+      out_dim = up_weight.size()[0]
+      conv2d = len(down_weight.size()) == 4
+      print(lora_module_name, network_dim, alpha, in_dim, out_dim)
+
+      # make original weight if not exist
+      if lora_module_name not in merged_sd:
+        weight = torch.zeros((out_dim, in_dim, 1, 1) if conv2d else (out_dim, in_dim), dtype=merge_dtype)
+        if device:
+          weight = weight.to(device)
+      else:
+        weight = merged_sd[lora_module_name]
+
+      # merge to weight
+      if device:
+        up_weight = up_weight.to(device)
+        down_weight = down_weight.to(device)
+
+      # W <- W + U * D
+      scale = (alpha / network_dim)
+      if not conv2d:        # linear
+        weight = weight + ratio * (up_weight @ down_weight) * scale
+      else:
+        weight = weight + ratio * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)
+                                   ).unsqueeze(2).unsqueeze(3) * scale
+
+      merged_sd[lora_module_name] = weight
+
+  # extract from merged weights
+  print("extract new lora...")
+  merged_lora_sd = {}
+  with torch.no_grad():
+    for lora_module_name, mat in tqdm(list(merged_sd.items())):
+      conv2d = (len(mat.size()) == 4)
+      if conv2d:
+        mat = mat.squeeze()
+
+      U, S, Vh = torch.linalg.svd(mat)
+
+      U = U[:, :new_rank]
+      S = S[:new_rank]
+      U = U @ torch.diag(S)
+
+      Vh = Vh[:new_rank, :]
+
+      dist = torch.cat([U.flatten(), Vh.flatten()])
+      hi_val = torch.quantile(dist, CLAMP_QUANTILE)
+      low_val = -hi_val
+
+      U = U.clamp(low_val, hi_val)
+      Vh = Vh.clamp(low_val, hi_val)
+
+      up_weight = U
+      down_weight = Vh
+
+      if conv2d:
+        up_weight = up_weight.unsqueeze(2).unsqueeze(3)
+        down_weight = down_weight.unsqueeze(2).unsqueeze(3)
+
+      merged_lora_sd[lora_module_name + '.lora_up.weight'] = up_weight.to("cpu").contiguous()
+      merged_lora_sd[lora_module_name + '.lora_down.weight'] = down_weight.to("cpu").contiguous()
+      merged_lora_sd[lora_module_name + '.alpha'] = torch.tensor(new_rank)
+
+  return merged_lora_sd
+
+
+def merge(args):
+  assert len(args.models) == len(args.ratios), f"number of models must be equal to number of ratios / モデルの数と重みの数は合わせてください"
+
+  def str_to_dtype(p):
+    if p == 'float':
+      return torch.float
+    if p == 'fp16':
+      return torch.float16
+    if p == 'bf16':
+      return torch.bfloat16
+    return None
+
+  merge_dtype = str_to_dtype(args.precision)
+  save_dtype = str_to_dtype(args.save_precision)
+  if save_dtype is None:
+    save_dtype = merge_dtype
+
+  state_dict = merge_lora_models(args.models, args.ratios, args.new_rank, args.device, merge_dtype)
+
+  print(f"saving model to: {args.save_to}")
+  save_to_file(args.save_to, state_dict, state_dict, save_dtype)
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--save_precision", type=str, default=None,
+                      choices=[None, "float", "fp16", "bf16"], help="precision in saving, same to merging if omitted / 保存時に精度を変更して保存する、省略時はマージ時の精度と同じ")
+  parser.add_argument("--precision", type=str, default="float",
+                      choices=["float", "fp16", "bf16"], help="precision in merging (float is recommended) / マージの計算時の精度（floatを推奨）")
+  parser.add_argument("--save_to", type=str, default=None,
+                      help="destination file name: ckpt or safetensors file / 保存先のファイル名、ckptまたはsafetensors")
+  parser.add_argument("--models", type=str, nargs='*',
+                      help="LoRA models to merge: ckpt or safetensors file / マージするLoRAモデル、ckptまたはsafetensors")
+  parser.add_argument("--ratios", type=float, nargs='*',
+                      help="ratios for each model / それぞれのLoRAモデルの比率")
+  parser.add_argument("--new_rank", type=int, default=4,
+                      help="Specify rank of output LoRA / 出力するLoRAのrank (dim)")
+  parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")
+
+  args = parser.parse_args()
+  merge(args)

requirements.txt

@@ -0,0 +1,25 @@
+accelerate==0.15.0
+transformers==4.26.0
+ftfy==6.1.1
+albumentations==1.3.0
+opencv-python==4.7.0.68
+einops==0.6.0
+diffusers[torch]==0.10.2
+pytorch-lightning==1.9.0
+bitsandbytes==0.35.0
+tensorboard==2.10.1
+safetensors==0.2.6
+gradio==3.16.2
+altair==4.2.2
+easygui==0.98.3
+# for BLIP captioning
+requests==2.28.2
+timm==0.6.12
+fairscale==0.4.13
+# for WD14 captioning
+# tensorflow<2.11
+tensorflow==2.10.1
+huggingface-hub==0.12.0
+# for kohya_ss library
+#locon.locon_kohya
+.

requirements_startup.txt

@@ -0,0 +1,23 @@
+accelerate==0.15.0
+transformers==4.26.0
+ftfy==6.1.1
+albumentations==1.3.0
+opencv-python==4.7.0.68
+einops==0.6.0
+diffusers[torch]==0.10.2
+pytorch-lightning==1.9.0
+bitsandbytes==0.35.0
+tensorboard==2.10.1
+safetensors==0.2.6
+gradio==3.18.0
+altair==4.2.2
+easygui==0.98.3
+# for BLIP captioning
+requests==2.28.2
+timm==0.4.12
+fairscale==0.4.4
+# for WD14 captioning
+tensorflow==2.10.1
+huggingface-hub==0.12.0
+# for kohya_ss library
+.

setup.py

@@ -0,0 +1,3 @@
+from setuptools import setup, find_packages
+ 
+setup(name = "library", packages = find_packages())

tools/convert_diffusers20_original_sd.py

@@ -0,0 +1,89 @@
+# convert Diffusers v1.x/v2.0 model to original Stable Diffusion
+
+import argparse
+import os
+import torch
+from diffusers import StableDiffusionPipeline
+
+import library.model_util as model_util
+
+
+def convert(args):
+  # 引数を確認する
+  load_dtype = torch.float16 if args.fp16 else None
+
+  save_dtype = None
+  if args.fp16:
+    save_dtype = torch.float16
+  elif args.bf16:
+    save_dtype = torch.bfloat16
+  elif args.float:
+    save_dtype = torch.float
+
+  is_load_ckpt = os.path.isfile(args.model_to_load)
+  is_save_ckpt = len(os.path.splitext(args.model_to_save)[1]) > 0
+
+  assert not is_load_ckpt or args.v1 != args.v2, f"v1 or v2 is required to load checkpoint / checkpointの読み込みにはv1/v2指定が必要です"
+  assert is_save_ckpt or args.reference_model is not None, f"reference model is required to save as Diffusers / Diffusers形式での保存には参照モデルが必要です"
+
+  # モデルを読み込む
+  msg = "checkpoint" if is_load_ckpt else ("Diffusers" + (" as fp16" if args.fp16 else ""))
+  print(f"loading {msg}: {args.model_to_load}")
+
+  if is_load_ckpt:
+    v2_model = args.v2
+    text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(v2_model, args.model_to_load)
+  else:
+    pipe = StableDiffusionPipeline.from_pretrained(args.model_to_load, torch_dtype=load_dtype, tokenizer=None, safety_checker=None)
+    text_encoder = pipe.text_encoder
+    vae = pipe.vae
+    unet = pipe.unet
+
+    if args.v1 == args.v2:
+      # 自動判定する
+      v2_model = unet.config.cross_attention_dim == 1024
+      print("checking model version: model is " + ('v2' if v2_model else 'v1'))
+    else:
+      v2_model = not args.v1
+
+  # 変換して保存する
+  msg = ("checkpoint" + ("" if save_dtype is None else f" in {save_dtype}")) if is_save_ckpt else "Diffusers"
+  print(f"converting and saving as {msg}: {args.model_to_save}")
+
+  if is_save_ckpt:
+    original_model = args.model_to_load if is_load_ckpt else None
+    key_count = model_util.save_stable_diffusion_checkpoint(v2_model, args.model_to_save, text_encoder, unet,
+                                                            original_model, args.epoch, args.global_step, save_dtype, vae)
+    print(f"model saved. total converted state_dict keys: {key_count}")
+  else:
+    print(f"copy scheduler/tokenizer config from: {args.reference_model}")
+    model_util.save_diffusers_checkpoint(v2_model, args.model_to_save, text_encoder, unet, args.reference_model, vae, args.use_safetensors)
+    print(f"model saved.")
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--v1", action='store_true',
+                      help='load v1.x model (v1 or v2 is required to load checkpoint) / 1.xのモデルを読み込む')
+  parser.add_argument("--v2", action='store_true',
+                      help='load v2.0 model (v1 or v2 is required to load checkpoint) / 2.0のモデルを読み込む')
+  parser.add_argument("--fp16", action='store_true',
+                      help='load as fp16 (Diffusers only) and save as fp16 (checkpoint only) / fp16形式で読み込み（Diffusers形式のみ対応）、保存する（checkpointのみ対応）')
+  parser.add_argument("--bf16", action='store_true', help='save as bf16 (checkpoint only) / bf16形式で保存する（checkpointのみ対応）')
+  parser.add_argument("--float", action='store_true',
+                      help='save as float (checkpoint only) / float(float32)形式で保存する（checkpointのみ対応）')
+  parser.add_argument("--epoch", type=int, default=0, help='epoch to write to checkpoint / checkpointに記録するepoch数の値')
+  parser.add_argument("--global_step", type=int, default=0,
+                      help='global_step to write to checkpoint / checkpointに記録するglobal_stepの値')
+  parser.add_argument("--reference_model", type=str, default=None,
+                      help="reference model for schduler/tokenizer, required in saving Diffusers, copy schduler/tokenizer from this / scheduler/tokenizerのコピー元のDiffusersモデル、Diffusers形式で保存するときに必要")
+  parser.add_argument("--use_safetensors", action='store_true',
+                      help="use safetensors format to save Diffusers model (checkpoint depends on the file extension) / Duffusersモデルをsafetensors形式で保存する（checkpointは拡張子で自動判定）")
+
+  parser.add_argument("model_to_load", type=str, default=None,
+                      help="model to load: checkpoint file or Diffusers model's directory / 読み込むモデル、checkpointかDiffusers形式モデルのディレクトリ")
+  parser.add_argument("model_to_save", type=str, default=None,
+                      help="model to save: checkpoint (with extension) or Diffusers model's directory (without extension) / 変換後のモデル、拡張子がある場合はcheckpoint、ない場合はDiffusesモデルとして保存")
+
+  args = parser.parse_args()
+  convert(args)

tools/detect_face_rotate.py

@@ -0,0 +1,239 @@
+# このスクリプトのライセンスは、train_dreambooth.pyと同じくApache License 2.0とします
+# (c) 2022 Kohya S. @kohya_ss
+
+# 横長の画像から顔検出して正立するように回転し、そこを中心に正方形に切り出す
+
+# v2: extract max face if multiple faces are found
+# v3: add crop_ratio option
+# v4: add multiple faces extraction and min/max size
+
+import argparse
+import math
+import cv2
+import glob
+import os
+from anime_face_detector import create_detector
+from tqdm import tqdm
+import numpy as np
+
+KP_REYE = 11
+KP_LEYE = 19
+
+SCORE_THRES = 0.90
+
+
+def detect_faces(detector, image, min_size):
+  preds = detector(image)                     # bgr
+  # print(len(preds))
+
+  faces = []
+  for pred in preds:
+    bb = pred['bbox']
+    score = bb[-1]
+    if score < SCORE_THRES:
+      continue
+
+    left, top, right, bottom = bb[:4]
+    cx = int((left + right) / 2)
+    cy = int((top + bottom) / 2)
+    fw = int(right - left)
+    fh = int(bottom - top)
+
+    lex, ley = pred['keypoints'][KP_LEYE, 0:2]
+    rex, rey = pred['keypoints'][KP_REYE, 0:2]
+    angle = math.atan2(ley - rey, lex - rex)
+    angle = angle / math.pi * 180
+
+    faces.append((cx, cy, fw, fh, angle))
+
+  faces.sort(key=lambda x: max(x[2], x[3]), reverse=True)         # 大きい順
+  return faces
+
+
+def rotate_image(image, angle, cx, cy):
+  h, w = image.shape[0:2]
+  rot_mat = cv2.getRotationMatrix2D((cx, cy), angle, 1.0)
+
+  # # 回転する分、すこし画像サイズを大きくする→とりあえず無効化
+  # nh = max(h, int(w * math.sin(angle)))
+  # nw = max(w, int(h * math.sin(angle)))
+  # if nh > h or nw > w:
+  #   pad_y = nh - h
+  #   pad_t = pad_y // 2
+  #   pad_x = nw - w
+  #   pad_l = pad_x // 2
+  #   m = np.array([[0, 0, pad_l],
+  #                 [0, 0, pad_t]])
+  #   rot_mat = rot_mat + m
+  #   h, w = nh, nw
+  #   cx += pad_l
+  #   cy += pad_t
+
+  result = cv2.warpAffine(image, rot_mat, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)
+  return result, cx, cy
+
+
+def process(args):
+  assert (not args.resize_fit) or args.resize_face_size is None, f"resize_fit and resize_face_size can't be specified both / resize_fitとresize_face_sizeはどちらか片方しか指定できません"
+  assert args.crop_ratio is None or args.resize_face_size is None, f"crop_ratio指定時はresize_face_sizeは指定できません"
+
+  # アニメ顔検出モデルを読み込む
+  print("loading face detector.")
+  detector = create_detector('yolov3')
+
+  # cropの引数を解析する
+  if args.crop_size is None:
+    crop_width = crop_height = None
+  else:
+    tokens = args.crop_size.split(',')
+    assert len(tokens) == 2, f"crop_size must be 'width,height' / crop_sizeは'幅,高さ'で指定してください"
+    crop_width, crop_height = [int(t) for t in tokens]
+
+  if args.crop_ratio is None:
+    crop_h_ratio = crop_v_ratio = None
+  else:
+    tokens = args.crop_ratio.split(',')
+    assert len(tokens) == 2, f"crop_ratio must be 'horizontal,vertical' / crop_ratioは'幅,高さ'の倍率で指定してください"
+    crop_h_ratio, crop_v_ratio = [float(t) for t in tokens]
+
+  # 画像を処理する
+  print("processing.")
+  output_extension = ".png"
+
+  os.makedirs(args.dst_dir, exist_ok=True)
+  paths = glob.glob(os.path.join(args.src_dir, "*.png")) + glob.glob(os.path.join(args.src_dir, "*.jpg")) + \
+      glob.glob(os.path.join(args.src_dir, "*.webp"))
+  for path in tqdm(paths):
+    basename = os.path.splitext(os.path.basename(path))[0]
+
+    # image = cv2.imread(path)        # 日本語ファイル名でエラーになる
+    image = cv2.imdecode(np.fromfile(path, np.uint8), cv2.IMREAD_UNCHANGED)
+    if len(image.shape) == 2:
+      image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+    if image.shape[2] == 4:
+      print(f"image has alpha. ignore / 画像の透明度が設定されているため無視します: {path}")
+      image = image[:, :, :3].copy()                    # copyをしないと内部的に透明度情報が付いたままになるらしい
+
+    h, w = image.shape[:2]
+
+    faces = detect_faces(detector, image, args.multiple_faces)
+    for i, face in enumerate(faces):
+      cx, cy, fw, fh, angle = face
+      face_size = max(fw, fh)
+      if args.min_size is not None and face_size < args.min_size:
+        continue
+      if args.max_size is not None and face_size >= args.max_size:
+        continue
+      face_suffix = f"_{i+1:02d}" if args.multiple_faces else ""
+
+      # オプション指定があれば回転する
+      face_img = image
+      if args.rotate:
+        face_img, cx, cy = rotate_image(face_img, angle, cx, cy)
+
+      # オプション指定があれば顔を中心に切り出す
+      if crop_width is not None or crop_h_ratio is not None:
+        cur_crop_width, cur_crop_height = crop_width, crop_height
+        if crop_h_ratio is not None:
+          cur_crop_width = int(face_size * crop_h_ratio + .5)
+          cur_crop_height = int(face_size * crop_v_ratio + .5)
+
+        # リサイズを必要なら行う
+        scale = 1.0
+        if args.resize_face_size is not None:
+          # 顔サイズを基準にリサイズする
+          scale = args.resize_face_size / face_size
+          if scale < cur_crop_width / w:
+            print(
+                f"image width too small in face size based resizing / 顔を基準にリサイズすると画像の幅がcrop sizeより小さい（顔が相対的に大きすぎる）ので顔サイズが変わります: {path}")
+            scale = cur_crop_width / w
+          if scale < cur_crop_height / h:
+            print(
+                f"image height too small in face size based resizing / 顔を基準にリサイズすると画像の高さがcrop sizeより小さい（顔が相対的に大きすぎる）ので顔サイズが変わります: {path}")
+            scale = cur_crop_height / h
+        elif crop_h_ratio is not None:
+          # 倍率指定の時にはリサイズしない
+          pass
+        else:
+          # 切り出しサイズ指定あり
+          if w < cur_crop_width:
+            print(f"image width too small/ 画像の幅がcrop sizeより小さいので画質が劣化します: {path}")
+            scale = cur_crop_width / w
+          if h < cur_crop_height:
+            print(f"image height too small/ 画像の高さがcrop sizeより小さいので画質が劣化します: {path}")
+            scale = cur_crop_height / h
+          if args.resize_fit:
+            scale = max(cur_crop_width / w, cur_crop_height / h)
+
+        if scale != 1.0:
+          w = int(w * scale + .5)
+          h = int(h * scale + .5)
+          face_img = cv2.resize(face_img, (w, h), interpolation=cv2.INTER_AREA if scale < 1.0 else cv2.INTER_LANCZOS4)
+          cx = int(cx * scale + .5)
+          cy = int(cy * scale + .5)
+          fw = int(fw * scale + .5)
+          fh = int(fh * scale + .5)
+
+        cur_crop_width = min(cur_crop_width, face_img.shape[1])
+        cur_crop_height = min(cur_crop_height, face_img.shape[0])
+
+        x = cx - cur_crop_width // 2
+        cx = cur_crop_width // 2
+        if x < 0:
+          cx = cx + x
+          x = 0
+        elif x + cur_crop_width > w:
+          cx = cx + (x + cur_crop_width - w)
+          x = w - cur_crop_width
+        face_img = face_img[:, x:x+cur_crop_width]
+
+        y = cy - cur_crop_height // 2
+        cy = cur_crop_height // 2
+        if y < 0:
+          cy = cy + y
+          y = 0
+        elif y + cur_crop_height > h:
+          cy = cy + (y + cur_crop_height - h)
+          y = h - cur_crop_height
+        face_img = face_img[y:y + cur_crop_height]
+
+      # # debug
+      # print(path, cx, cy, angle)
+      # crp = cv2.resize(image, (image.shape[1]//8, image.shape[0]//8))
+      # cv2.imshow("image", crp)
+      # if cv2.waitKey() == 27:
+      #   break
+      # cv2.destroyAllWindows()
+
+      # debug
+      if args.debug:
+        cv2.rectangle(face_img, (cx-fw//2, cy-fh//2), (cx+fw//2, cy+fh//2), (255, 0, 255), fw//20)
+
+      _, buf = cv2.imencode(output_extension, face_img)
+      with open(os.path.join(args.dst_dir, f"{basename}{face_suffix}_{cx:04d}_{cy:04d}_{fw:04d}_{fh:04d}{output_extension}"), "wb") as f:
+        buf.tofile(f)
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--src_dir", type=str, help="directory to load images / 画像を読み込むディレクトリ")
+  parser.add_argument("--dst_dir", type=str, help="directory to save images / 画像を保存するディレクトリ")
+  parser.add_argument("--rotate", action="store_true", help="rotate images to align faces / 顔が正立するように画像を回転する")
+  parser.add_argument("--resize_fit", action="store_true",
+                      help="resize to fit smaller side after cropping / 切り出し後の画像の短辺がcrop_sizeにあうようにリサイズする")
+  parser.add_argument("--resize_face_size", type=int, default=None,
+                      help="resize image before cropping by face size / 切り出し前に顔がこのサイズになるようにリサイズする")
+  parser.add_argument("--crop_size", type=str, default=None,
+                      help="crop images with 'width,height' pixels, face centered / 顔を中心として'幅,高さ'のサイズで切り出す")
+  parser.add_argument("--crop_ratio", type=str, default=None,
+                      help="crop images with 'horizontal,vertical' ratio to face, face centered / 顔を中心として顔サイズの'幅倍率,高さ倍率'のサイズで切り出す")
+  parser.add_argument("--min_size", type=int, default=None,
+                      help="minimum face size to output (included) / 処理対象とする顔の最小サイズ（この値以上）")
+  parser.add_argument("--max_size", type=int, default=None,
+                      help="maximum face size to output (excluded) / 処理対象とする顔の最大サイズ（この値未満）")
+  parser.add_argument("--multiple_faces", action="store_true",
+                      help="output each faces / 複数の顔が見つかった場合、それぞれを切り出す")
+  parser.add_argument("--debug", action="store_true", help="render rect for face / 処理後画像の顔位置に矩形を描画します")
+  args = parser.parse_args()
+
+  process(args)

tools/resize_images_to_resolution.py

@@ -0,0 +1,122 @@
+import glob
+import os
+import cv2
+import argparse
+import shutil
+import math
+from PIL import Image
+import numpy as np
+
+
+def resize_images(src_img_folder, dst_img_folder, max_resolution="512x512", divisible_by=2, interpolation=None, save_as_png=False, copy_associated_files=False):
+  # Split the max_resolution string by "," and strip any whitespaces
+  max_resolutions = [res.strip() for res in max_resolution.split(',')]
+
+  # # Calculate max_pixels from max_resolution string
+  # max_pixels = int(max_resolution.split("x")[0]) * int(max_resolution.split("x")[1])
+
+  # Create destination folder if it does not exist
+  if not os.path.exists(dst_img_folder):
+    os.makedirs(dst_img_folder)
+
+  # Select interpolation method
+  if interpolation == 'lanczos4':
+    cv2_interpolation = cv2.INTER_LANCZOS4
+  elif interpolation == 'cubic':
+    cv2_interpolation = cv2.INTER_CUBIC
+  else:
+    cv2_interpolation = cv2.INTER_AREA
+
+  # Iterate through all files in src_img_folder
+  img_exts = (".png", ".jpg", ".jpeg", ".webp", ".bmp")                   # copy from train_util.py
+  for filename in os.listdir(src_img_folder):
+    # Check if the image is png, jpg or webp etc...
+    if not filename.endswith(img_exts):
+      # Copy the file to the destination folder if not png, jpg or webp etc (.txt or .caption or etc.)
+      shutil.copy(os.path.join(src_img_folder, filename), os.path.join(dst_img_folder, filename))
+      continue
+
+    # Load image
+    # img = cv2.imread(os.path.join(src_img_folder, filename))
+    image = Image.open(os.path.join(src_img_folder, filename))
+    if not image.mode == "RGB":
+      image = image.convert("RGB")
+    img = np.array(image, np.uint8)
+
+    base, _ = os.path.splitext(filename)
+    for max_resolution in max_resolutions:
+      # Calculate max_pixels from max_resolution string
+      max_pixels = int(max_resolution.split("x")[0]) * int(max_resolution.split("x")[1])
+
+      # Calculate current number of pixels
+      current_pixels = img.shape[0] * img.shape[1]
+
+      # Check if the image needs resizing
+      if current_pixels > max_pixels:
+        # Calculate scaling factor
+        scale_factor = max_pixels / current_pixels
+
+        # Calculate new dimensions
+        new_height = int(img.shape[0] * math.sqrt(scale_factor))
+        new_width = int(img.shape[1] * math.sqrt(scale_factor))
+
+        # Resize image
+        img = cv2.resize(img, (new_width, new_height), interpolation=cv2_interpolation)
+      else:
+        new_height, new_width = img.shape[0:2]
+
+      # Calculate the new height and width that are divisible by divisible_by (with/without resizing)
+      new_height = new_height if new_height % divisible_by == 0 else new_height - new_height % divisible_by
+      new_width = new_width if new_width % divisible_by == 0 else new_width - new_width % divisible_by
+
+      # Center crop the image to the calculated dimensions
+      y = int((img.shape[0] - new_height) / 2)
+      x = int((img.shape[1] - new_width) / 2)
+      img = img[y:y + new_height, x:x + new_width]
+
+      # Split filename into base and extension
+      new_filename = base + '+' + max_resolution + ('.png' if save_as_png else '.jpg')
+
+      # Save resized image in dst_img_folder
+      # cv2.imwrite(os.path.join(dst_img_folder, new_filename), img, [cv2.IMWRITE_JPEG_QUALITY, 100])
+      image = Image.fromarray(img)
+      image.save(os.path.join(dst_img_folder, new_filename), quality=100)
+
+      proc = "Resized" if current_pixels > max_pixels else "Saved"
+      print(f"{proc} image: {filename} with size {img.shape[0]}x{img.shape[1]} as {new_filename}")
+
+    # If other files with same basename, copy them with resolution suffix
+    if copy_associated_files:
+      asoc_files = glob.glob(os.path.join(src_img_folder, base + ".*"))
+      for asoc_file in asoc_files:
+        ext = os.path.splitext(asoc_file)[1]
+        if ext in img_exts:
+          continue
+        for max_resolution in max_resolutions:
+          new_asoc_file = base + '+' + max_resolution + ext
+          print(f"Copy {asoc_file} as {new_asoc_file}")
+          shutil.copy(os.path.join(src_img_folder, asoc_file), os.path.join(dst_img_folder, new_asoc_file))
+
+
+def main():
+  parser = argparse.ArgumentParser(
+      description='Resize images in a folder to a specified max resolution(s) / 指定されたフォルダ内の画像を指定した最大画像サイズ（面積）以下にアスペクト比を維持したままリサイズします')
+  parser.add_argument('src_img_folder', type=str, help='Source folder containing the images / 元画像のフォルダ')
+  parser.add_argument('dst_img_folder', type=str, help='Destination folder to save the resized images / リサイズ後の画像を保存するフォルダ')
+  parser.add_argument('--max_resolution', type=str,
+                      help='Maximum resolution(s) in the format "512x512,384x384, etc, etc" / 最大画像サイズをカンマ区切りで指定 ("512x512,384x384, etc, etc" など)', default="512x512,384x384,256x256,128x128")
+  parser.add_argument('--divisible_by', type=int,
+                      help='Ensure new dimensions are divisible by this value / リサイズ後の画像のサイズをこの値で割り切れるようにします', default=1)
+  parser.add_argument('--interpolation', type=str, choices=['area', 'cubic', 'lanczos4'],
+                      default='area', help='Interpolation method for resizing / リサイズ時の補完方法')
+  parser.add_argument('--save_as_png', action='store_true', help='Save as png format / png形式で保存')
+  parser.add_argument('--copy_associated_files', action='store_true',
+                      help='Copy files with same base name to images (captions etc) / 画像と同じファイル名（拡張子を除く）のファイルもコピーする')
+
+  args = parser.parse_args()
+  resize_images(args.src_img_folder, args.dst_img_folder, args.max_resolution,
+                args.divisible_by, args.interpolation, args.save_as_png, args.copy_associated_files)
+
+
+if __name__ == '__main__':
+  main()