Merge pull request #830 from kohya-ss/dev2

add extension checking for resize_lora.py

README.md

@@ -18,131 +18,11 @@ This repository contains the scripts for:
 
 __Stable Diffusion web UI now seems to support LoRA trained by ``sd-scripts``.__ Thank you for great work!!! 
 
-## About SDXL training
-
-The feature of SDXL training is now available in sdxl branch as an experimental feature. 
-
-Sep 3, 2023: The feature will be merged into the main branch soon. Following are the changes from the previous version. 
-
-- ControlNet-LLLite is added. See [documentation](./docs/train_lllite_README.md) for details.
-- JPEG XL is supported. [#786](https://github.com/kohya-ss/sd-scripts/pull/786) 
-- Peak memory usage is reduced. [#791](https://github.com/kohya-ss/sd-scripts/pull/791)
-- Input perturbation noise is added. See [#798](https://github.com/kohya-ss/sd-scripts/pull/798) for details.
-- Dataset subset now has `caption_prefix` and `caption_suffix` options. The strings are added to the beginning and the end of the captions before shuffling. You can specify the options in `.toml`.
-- Other minor changes.
-- Thanks for contributions from Isotr0py, vvern999, lansing  and others!
-
-Aug 13, 2023: 
-
-- LoRA-FA is added experimentally. Specify `--network_module networks.lora_fa` option instead of `--network_module networks.lora`. The trained model can be used as a normal LoRA model.
-
-Aug 12, 2023: 
-
-- The default value of noise offset when omitted has been changed to 0 from 0.0357.
-- The different learning rates for each U-Net block are now supported. Specify with `--block_lr` option. Specify 23 values separated by commas like `--block_lr 1e-3,1e-3 ... 1e-3`.
-  - 23 values correspond to `0: time/label embed, 1-9: input blocks 0-8, 10-12: mid blocks 0-2, 13-21: output blocks 0-8, 22: out`.
-
-Aug 6, 2023: 
-
-- [SAI Model Spec](https://github.com/Stability-AI/ModelSpec) metadata is now supported partially. `hash_sha256` is not supported yet.
-  - The main items are set automatically. 
-  - You can set title, author, description, license and tags with `--metadata_xxx` options in each training script.
-  - Merging scripts also support minimum SAI Model Spec metadata. See the help message for the usage.
-  - Metadata editor will be available soon.
-- SDXL LoRA has `sdxl_base_v1-0` now  for `ss_base_model_version` metadata item, instead of `v0-9`.
-
-Aug 4, 2023: 
-
-- `bitsandbytes` is now optional. Please install it if you want to use it. The insructions are in the later section.
-- `albumentations` is not required anymore.
-- An issue for pooled output for Textual Inversion training is fixed.
-- `--v_pred_like_loss ratio` option is added. This option adds the loss like v-prediction loss in SDXL training. `0.1` means that the loss is added 10% of the v-prediction loss. The default value is None (disabled).
-  - In v-prediction, the loss is higher in the early timesteps (near the noise). This option can be used to increase the loss in the early timesteps.
-- Arbitrary options can be used for Diffusers' schedulers. For example `--lr_scheduler_args "lr_end=1e-8"`.
-- `sdxl_gen_imgs.py` supports batch size > 1.
-- Fix ControlNet to work with attention couple and reginal LoRA in `gen_img_diffusers.py`.
-
-Summary of the feature:
-
-- `tools/cache_latents.py` is added. This script can be used to cache the latents to disk in advance. 
-  - The options are almost the same as `sdxl_train.py'. See the help message for the usage.
-  - Please launch the script as follows:
-    `accelerate launch  --num_cpu_threads_per_process 1 tools/cache_latents.py ...`
-  - This script should work with multi-GPU, but it is not tested in my environment.
-
-- `tools/cache_text_encoder_outputs.py` is added. This script can be used to cache the text encoder outputs to disk in advance. 
-  - The options are almost the same as `cache_latents.py' and `sdxl_train.py'. See the help message for the usage.
-
-- `sdxl_train.py` is a script for SDXL fine-tuning. The usage is almost the same as `fine_tune.py`, but it also supports DreamBooth dataset.
-  - `--full_bf16` option is added. Thanks to KohakuBlueleaf!
-    - This option enables the full bfloat16 training (includes gradients). This option is useful to reduce the GPU memory usage. 
-    - However, bitsandbytes==0.35 doesn't seem to support this. Please use a newer version of bitsandbytes or another optimizer.
-    - I cannot find bitsandbytes>0.35.0 that works correctly on Windows.
-    - In addition, the full bfloat16 training might be unstable. Please use it at your own risk.
-- `prepare_buckets_latents.py` now supports SDXL fine-tuning.
-- `sdxl_train_network.py` is a script for LoRA training for SDXL. The usage is almost the same as `train_network.py`.
-- Both scripts has following additional options:
-  - `--cache_text_encoder_outputs` and `--cache_text_encoder_outputs_to_disk`: Cache the outputs of the text encoders. This option is useful to reduce the GPU memory usage. This option cannot be used with options for shuffling or dropping the captions.
-  - `--no_half_vae`: Disable the half-precision (mixed-precision) VAE. VAE for SDXL seems to produce NaNs in some cases. This option is useful to avoid the NaNs.
-- The image generation during training is now available. `--no_half_vae` option also works to avoid black images.
-
-- `--weighted_captions` option is not supported yet for both scripts.
-- `--min_timestep` and `--max_timestep` options are added to each training script. These options can be used to train U-Net with different timesteps. The default values are 0 and 1000.
-
-- `sdxl_train_textual_inversion.py` is a script for Textual Inversion training for SDXL. The usage is almost the same as `train_textual_inversion.py`.
-  - `--cache_text_encoder_outputs` is not supported.
-  - `token_string` must be alphabet only currently, due to the limitation of the open-clip tokenizer.
-  - There are two options for captions:
-    1. Training with captions. All captions must include the token string. The token string is replaced with multiple tokens.
-    2. Use `--use_object_template` or `--use_style_template` option. The captions are generated from the template. The existing captions are ignored.
-  - See below for the format of the embeddings.
-  
-- `sdxl_gen_img.py` is added. This script can be used to generate images with SDXL, including LoRA. See the help message for the usage.
-  - Textual Inversion is supported, but the name for the embeds in the caption becomes alphabet only. For example, `neg_hand_v1.safetensors` can be activated with `neghandv`.
-
-`requirements.txt` is updated to support SDXL training. 
-
-### Tips for SDXL training
-
-- The default resolution of SDXL is 1024x1024.
-- The fine-tuning can be done with 24GB GPU memory with the batch size of 1. For 24GB GPU, the following options are recommended __for the fine-tuning with 24GB GPU memory__:
-  - Train U-Net only.
-  - Use gradient checkpointing.
-  - Use `--cache_text_encoder_outputs` option and caching latents.
-  - Use Adafactor optimizer. RMSprop 8bit or Adagrad 8bit may work. AdamW 8bit doesn't seem to work.
-- The LoRA training can be done with 8GB GPU memory (10GB recommended). For reducing the GPU memory usage, the following options are recommended:
-  - Train U-Net only.
-  - Use gradient checkpointing.
-  - Use `--cache_text_encoder_outputs` option and caching latents.
-  - Use one of 8bit optimizers or Adafactor optimizer.
-  - Use lower dim (-8 for 8GB GPU).
-- `--network_train_unet_only` option is highly recommended for SDXL LoRA. Because SDXL has two text encoders, the result of the training will be unexpected.
-- PyTorch 2 seems to use slightly less GPU memory than PyTorch 1.
-- `--bucket_reso_steps` can be set to 32 instead of the default value 64. Smaller values than 32 will not work for SDXL training.
-
-Example of the optimizer settings for Adafactor with the fixed learning rate:
-```toml
-optimizer_type = "adafactor"
-optimizer_args = [ "scale_parameter=False", "relative_step=False", "warmup_init=False" ]
-lr_scheduler = "constant_with_warmup"
-lr_warmup_steps = 100
-learning_rate = 4e-7 # SDXL original learning rate
-```
-
-### Format of Textual Inversion embeddings
-
-```python
-from safetensors.torch import save_file
-
-state_dict = {"clip_g": embs_for_text_encoder_1280, "clip_l": embs_for_text_encoder_768}
-save_file(state_dict, file)
-```
-
 ## About requirements.txt
 
 These files do not contain requirements for PyTorch. Because the versions of them depend on your environment. Please install PyTorch at first (see installation guide below.) 
 
-The scripts are tested with PyTorch 1.12.1 and 2.0.1, Diffusers 0.18.2.
+The scripts are tested with PyTorch 1.12.1 and 1.13.0, Diffusers 0.10.2.
 
 ## Links to how-to-use documents
 
@@ -188,16 +68,15 @@ pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 --extra-index-url http
 pip install --upgrade -r requirements.txt
 pip install -U -I --no-deps https://github.com/C43H66N12O12S2/stable-diffusion-webui/releases/download/f/xformers-0.0.14.dev0-cp310-cp310-win_amd64.whl
 
-accelerate config
-```
-
-__Note:__ Now bitsandbytes is optional. Please install any version of bitsandbytes as needed. Installation instructions are in the following section.
-
-<!-- 
 cp .\bitsandbytes_windows\*.dll .\venv\Lib\site-packages\bitsandbytes\
 cp .\bitsandbytes_windows\cextension.py .\venv\Lib\site-packages\bitsandbytes\cextension.py
 cp .\bitsandbytes_windows\main.py .\venv\Lib\site-packages\bitsandbytes\cuda_setup\main.py
--->
+
+accelerate config
+```
+
+update: ``python -m venv venv`` is seemed to be safer than ``python -m venv --system-site-packages venv`` (some user have packages in global python).
+
 Answers to accelerate config:
 
 ```txt
@@ -215,68 +94,31 @@ note: Some user reports ``ValueError: fp16 mixed precision requires a GPU`` is o
 
 (Single GPU with id `0` will be used.)
 
-### Experimental: Use PyTorch 2.0
-
-In this case, you need to install PyTorch 2.0 and xformers 0.0.20. Instead of the above, please type the following:
-
-```powershell
-git clone https://github.com/kohya-ss/sd-scripts.git
-cd sd-scripts
-
-python -m venv venv
-.\venv\Scripts\activate
-
-pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 --index-url https://download.pytorch.org/whl/cu118
-pip install --upgrade -r requirements.txt
-pip install xformers==0.0.20
-
-accelerate config
-```
-
-Answers to accelerate config should be the same as above.
-
 ### about PyTorch and xformers
 
 Other versions of PyTorch and xformers seem to have problems with training.
 If there is no other reason, please install the specified version.
 
-### Optional: Use `bitsandbytes` (8bit optimizer)
+### Optional: Use Lion8bit
 
-For 8bit optimizer, you need to install `bitsandbytes`. For Linux, please install `bitsandbytes` as usual (0.41.1 or later is recommended.)
-
-For Windows, there are several versions of `bitsandbytes`:
-
-- `bitsandbytes` 0.35.0: Stable version. AdamW8bit is available. `full_bf16` is not available.
-- `bitsandbytes` 0.41.1: Lion8bit, PagedAdamW8bit and PagedLion8bit are available. `full_bf16` is available.
-
-Note: `bitsandbytes`above 0.35.0 till 0.41.0 seems to have an issue: https://github.com/TimDettmers/bitsandbytes/issues/659
-
-Follow the instructions below to install `bitsandbytes` for Windows.
-
-### bitsandbytes 0.35.0 for Windows
-
-Open a regular Powershell terminal and type the following inside:
+For Lion8bit, you need to upgrade `bitsandbytes` to 0.38.0 or later. Uninstall `bitsandbytes`, and for Windows, install the Windows version whl file from [here](https://github.com/jllllll/bitsandbytes-windows-webui) or other sources, like:
 
 ```powershell
-cd sd-scripts
-.\venv\Scripts\activate
-pip install bitsandbytes==0.35.0
-
-cp .\bitsandbytes_windows\*.dll .\venv\Lib\site-packages\bitsandbytes\
-cp .\bitsandbytes_windows\cextension.py .\venv\Lib\site-packages\bitsandbytes\cextension.py
-cp .\bitsandbytes_windows\main.py .\venv\Lib\site-packages\bitsandbytes\cuda_setup\main.py
+pip install https://github.com/jllllll/bitsandbytes-windows-webui/raw/main/bitsandbytes-0.38.1-py3-none-any.whl
 ```
 
-This will install `bitsandbytes` 0.35.0 and copy the necessary files to the `bitsandbytes` directory.
+For upgrading, upgrade this repo with `pip install .`, and upgrade necessary packages manually.
 
-### bitsandbytes 0.41.1 for Windows
+### Optional: Use PagedAdamW8bit and PagedLion8bit
 
-Install the Windows version whl file from [here](https://github.com/jllllll/bitsandbytes-windows-webui) or other sources, like:
+For PagedAdamW8bit and PagedLion8bit, you need to upgrade `bitsandbytes` to 0.39.0 or later. Uninstall `bitsandbytes`, and for Windows, install the Windows version whl file from [here](https://github.com/jllllll/bitsandbytes-windows-webui) or other sources, like:
 
 ```powershell
-python -m pip install bitsandbytes==0.41.1 --prefer-binary --extra-index-url=https://jllllll.github.io/bitsandbytes-windows-webui
+pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/download/wheels/bitsandbytes-0.39.1-py3-none-win_amd64.whl
 ```
 
+For upgrading, upgrade this repo with `pip install .`, and upgrade necessary packages manually.
+
 ## Upgrade
 
 When a new release comes out you can upgrade your repo with the following command:

XTI_hijack.py

@@ -2,123 +2,132 @@ import torch
 from typing import Union, List, Optional, Dict, Any, Tuple
 from diffusers.models.unet_2d_condition import UNet2DConditionOutput
 
-from library.original_unet import SampleOutput
+def unet_forward_XTI(self,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        class_labels: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[UNet2DConditionOutput, Tuple]:
+        r"""
+        Args:
+            sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor
+            timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps
+            encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.
 
+        Returns:
+            [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
+            [`~models.unet_2d_condition.UNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When
+            returning a tuple, the first element is the sample tensor.
+        """
+        # By default samples have to be AT least a multiple of the overall upsampling factor.
+        # The overall upsampling factor is equal to 2 ** (# num of upsampling layears).
+        # However, the upsampling interpolation output size can be forced to fit any upsampling size
+        # on the fly if necessary.
+        default_overall_up_factor = 2**self.num_upsamplers
 
-def unet_forward_XTI(
-    self,
-    sample: torch.FloatTensor,
-    timestep: Union[torch.Tensor, float, int],
-    encoder_hidden_states: torch.Tensor,
-    class_labels: Optional[torch.Tensor] = None,
-    return_dict: bool = True,
-) -> Union[Dict, Tuple]:
-    r"""
-    Args:
-        sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor
-        timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps
-        encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states
-        return_dict (`bool`, *optional*, defaults to `True`):
-            Whether or not to return a dict instead of a plain tuple.
+        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
+        forward_upsample_size = False
+        upsample_size = None
 
-    Returns:
-        `SampleOutput` or `tuple`:
-        `SampleOutput` if `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is the sample tensor.
-    """
-    # By default samples have to be AT least a multiple of the overall upsampling factor.
-    # The overall upsampling factor is equal to 2 ** (# num of upsampling layears).
-    # However, the upsampling interpolation output size can be forced to fit any upsampling size
-    # on the fly if necessary.
-    # デフォルトではサンプルは「2^アップサンプルの数」、つまり64の倍数である必要がある
-    # ただそれ以外のサイズにも対応できるように、必要ならアップサンプルのサイズを変更する
-    # 多分画質が悪くなるので、64で割り切れるようにしておくのが良い
-    default_overall_up_factor = 2**self.num_upsamplers
+        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
+            logger.info("Forward upsample size to force interpolation output size.")
+            forward_upsample_size = True
 
-    # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
-    # 64で割り切れないときはupsamplerにサイズを伝える
-    forward_upsample_size = False
-    upsample_size = None
+        # 0. center input if necessary
+        if self.config.center_input_sample:
+            sample = 2 * sample - 1.0
 
-    if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
-        # logger.info("Forward upsample size to force interpolation output size.")
-        forward_upsample_size = True
+        # 1. time
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = sample.device.type == "mps"
+            if isinstance(timestep, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
 
-    # 1. time
-    timesteps = timestep
-    timesteps = self.handle_unusual_timesteps(sample, timesteps)  # 変な時だけ処理
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timesteps = timesteps.expand(sample.shape[0])
 
-    t_emb = self.time_proj(timesteps)
+        t_emb = self.time_proj(timesteps)
 
-    # timesteps does not contain any weights and will always return f32 tensors
-    # but time_embedding might actually be running in fp16. so we need to cast here.
-    # there might be better ways to encapsulate this.
-    # timestepsは重みを含まないので常にfloat32のテンソルを返す
-    # しかしtime_embeddingはfp16で動いているかもしれないので、ここでキャストする必要がある
-    # time_projでキャストしておけばいいんじゃね？
-    t_emb = t_emb.to(dtype=self.dtype)
-    emb = self.time_embedding(t_emb)
+        # timesteps does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=self.dtype)
+        emb = self.time_embedding(t_emb)
 
-    # 2. pre-process
-    sample = self.conv_in(sample)
+        if self.config.num_class_embeds is not None:
+            if class_labels is None:
+                raise ValueError("class_labels should be provided when num_class_embeds > 0")
+            class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
+            emb = emb + class_emb
 
-    # 3. down
-    down_block_res_samples = (sample,)
-    down_i = 0
-    for downsample_block in self.down_blocks:
-        # downblockはforwardで必ずencoder_hidden_statesを受け取るようにしても良さそうだけど、
-        # まあこちらのほうがわかりやすいかもしれない
-        if downsample_block.has_cross_attention:
-            sample, res_samples = downsample_block(
-                hidden_states=sample,
-                temb=emb,
-                encoder_hidden_states=encoder_hidden_states[down_i : down_i + 2],
-            )
-            down_i += 2
-        else:
-            sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+        # 2. pre-process
+        sample = self.conv_in(sample)
 
-        down_block_res_samples += res_samples
+        # 3. down
+        down_block_res_samples = (sample,)
+        down_i = 0
+        for downsample_block in self.down_blocks:
+            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states[down_i:down_i+2],
+                )
+                down_i += 2
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
 
-    # 4. mid
-    sample = self.mid_block(sample, emb, encoder_hidden_states=encoder_hidden_states[6])
+            down_block_res_samples += res_samples
 
-    # 5. up
-    up_i = 7
-    for i, upsample_block in enumerate(self.up_blocks):
-        is_final_block = i == len(self.up_blocks) - 1
+        # 4. mid
+        sample = self.mid_block(sample, emb, encoder_hidden_states=encoder_hidden_states[6])
 
-        res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
-        down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]  # skip connection
+        # 5. up
+        up_i = 7
+        for i, upsample_block in enumerate(self.up_blocks):
+            is_final_block = i == len(self.up_blocks) - 1
 
-        # if we have not reached the final block and need to forward the upsample size, we do it here
-        # 前述のように最後のブロック以外ではupsample_sizeを伝える
-        if not is_final_block and forward_upsample_size:
-            upsample_size = down_block_res_samples[-1].shape[2:]
+            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
 
-        if upsample_block.has_cross_attention:
-            sample = upsample_block(
-                hidden_states=sample,
-                temb=emb,
-                res_hidden_states_tuple=res_samples,
-                encoder_hidden_states=encoder_hidden_states[up_i : up_i + 3],
-                upsample_size=upsample_size,
-            )
-            up_i += 3
-        else:
-            sample = upsample_block(
-                hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size
-            )
+            # if we have not reached the final block and need to forward the
+            # upsample size, we do it here
+            if not is_final_block and forward_upsample_size:
+                upsample_size = down_block_res_samples[-1].shape[2:]
 
-    # 6. post-process
-    sample = self.conv_norm_out(sample)
-    sample = self.conv_act(sample)
-    sample = self.conv_out(sample)
+            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    encoder_hidden_states=encoder_hidden_states[up_i:up_i+3],
+                    upsample_size=upsample_size,
+                )
+                up_i += 3
+            else:
+                sample = upsample_block(
+                    hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size
+                )
+        # 6. post-process
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
 
-    if not return_dict:
-        return (sample,)
-
-    return SampleOutput(sample=sample)
+        if not return_dict:
+            return (sample,)
 
+        return UNet2DConditionOutput(sample=sample)
 
 def downblock_forward_XTI(
     self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None, cross_attention_kwargs=None
@@ -157,7 +166,6 @@ def downblock_forward_XTI(
 
     return hidden_states, output_states
 
-
 def upblock_forward_XTI(
     self,
     hidden_states,
@@ -191,11 +199,11 @@ def upblock_forward_XTI(
         else:
             hidden_states = resnet(hidden_states, temb)
             hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states[i]).sample
-
+        
         i += 1
 
     if self.upsamplers is not None:
         for upsampler in self.upsamplers:
             hidden_states = upsampler(hidden_states, upsample_size)
 
-    return hidden_states
+    return hidden_states

bitsandbytes_windows/main.py

@@ -1,166 +1,166 @@
-"""
-extract factors the build is dependent on:
-[X] compute capability
-    [ ] TODO: Q - What if we have multiple GPUs of different makes?
-- CUDA version
-- Software:
-    - CPU-only: only CPU quantization functions (no optimizer, no matrix multiple)
-    - CuBLAS-LT: full-build 8-bit optimizer
-    - no CuBLAS-LT: no 8-bit matrix multiplication (`nomatmul`)
-
-evaluation:
-    - if paths faulty, return meaningful error
-    - else:
-        - determine CUDA version
-        - determine capabilities
-        - based on that set the default path
-"""
-
-import ctypes
-
-from .paths import determine_cuda_runtime_lib_path
-
-
-def check_cuda_result(cuda, result_val):
-    # 3. Check for CUDA errors
-    if result_val != 0:
-        error_str = ctypes.c_char_p()
-        cuda.cuGetErrorString(result_val, ctypes.byref(error_str))
-        print(f"CUDA exception! Error code: {error_str.value.decode()}")
-
-def get_cuda_version(cuda, cudart_path):
-    # https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART____VERSION.html#group__CUDART____VERSION
-    try:
-        cudart = ctypes.CDLL(cudart_path)
-    except OSError:
-        # TODO: shouldn't we error or at least warn here?
-        print(f'ERROR: libcudart.so could not be read from path: {cudart_path}!')
-        return None
-
-    version = ctypes.c_int()
-    check_cuda_result(cuda, cudart.cudaRuntimeGetVersion(ctypes.byref(version)))
-    version = int(version.value)
-    major = version//1000
-    minor = (version-(major*1000))//10
-
-    if major < 11:
-       print('CUDA SETUP: CUDA version lower than 11 are currently not supported for LLM.int8(). You will be only to use 8-bit optimizers and quantization routines!!')
-
-    return f'{major}{minor}'
-
-
-def get_cuda_lib_handle():
-    # 1. find libcuda.so library (GPU driver) (/usr/lib)
-    try:
-        cuda = ctypes.CDLL("libcuda.so")
-    except OSError:
-        # TODO: shouldn't we error or at least warn here?
-        print('CUDA SETUP: WARNING! libcuda.so not found! Do you have a CUDA driver installed? If you are on a cluster, make sure you are on a CUDA machine!')
-        return None
-    check_cuda_result(cuda, cuda.cuInit(0))
-
-    return cuda
-
-
-def get_compute_capabilities(cuda):
-    """
-    1. find libcuda.so library (GPU driver) (/usr/lib)
-       init_device -> init variables -> call function by reference
-    2. call extern C function to determine CC
-       (https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__DEVICE__DEPRECATED.html)
-    3. Check for CUDA errors
-       https://stackoverflow.com/questions/14038589/what-is-the-canonical-way-to-check-for-errors-using-the-cuda-runtime-api
-    # bits taken from https://gist.github.com/f0k/63a664160d016a491b2cbea15913d549
-    """
-
-
-    nGpus = ctypes.c_int()
-    cc_major = ctypes.c_int()
-    cc_minor = ctypes.c_int()
-
-    device = ctypes.c_int()
-
-    check_cuda_result(cuda, cuda.cuDeviceGetCount(ctypes.byref(nGpus)))
-    ccs = []
-    for i in range(nGpus.value):
-        check_cuda_result(cuda, cuda.cuDeviceGet(ctypes.byref(device), i))
-        ref_major = ctypes.byref(cc_major)
-        ref_minor = ctypes.byref(cc_minor)
-        # 2. call extern C function to determine CC
-        check_cuda_result(
-            cuda, cuda.cuDeviceComputeCapability(ref_major, ref_minor, device)
-        )
-        ccs.append(f"{cc_major.value}.{cc_minor.value}")
-
-    return ccs
-
-
-# def get_compute_capability()-> Union[List[str, ...], None]: # FIXME: error
-def get_compute_capability(cuda):
-    """
-    Extracts the highest compute capbility from all available GPUs, as compute
-    capabilities are downwards compatible. If no GPUs are detected, it returns
-    None.
-    """
-    ccs = get_compute_capabilities(cuda)
-    if ccs is not None:
-        # TODO: handle different compute capabilities; for now, take the max
-        return ccs[-1]
-    return None
-
-
-def evaluate_cuda_setup():
-    print('')
-    print('='*35 + 'BUG REPORT' + '='*35)
-    print('Welcome to bitsandbytes. For bug reports, please submit your error trace to: https://github.com/TimDettmers/bitsandbytes/issues')
-    print('For effortless bug reporting copy-paste your error into this form: https://docs.google.com/forms/d/e/1FAIpQLScPB8emS3Thkp66nvqwmjTEgxp8Y9ufuWTzFyr9kJ5AoI47dQ/viewform?usp=sf_link')
-    print('='*80)
-    return "libbitsandbytes_cuda116.dll"            # $$$
-    
-    binary_name = "libbitsandbytes_cpu.so"
-    #if not torch.cuda.is_available():
-        #print('No GPU detected. Loading CPU library...')
-        #return binary_name
-
-    cudart_path = determine_cuda_runtime_lib_path()
-    if cudart_path is None:
-        print(
-            "WARNING: No libcudart.so found! Install CUDA or the cudatoolkit package (anaconda)!"
-        )
-        return binary_name
-
-    print(f"CUDA SETUP: CUDA runtime path found: {cudart_path}")
-    cuda = get_cuda_lib_handle()
-    cc = get_compute_capability(cuda)
-    print(f"CUDA SETUP: Highest compute capability among GPUs detected: {cc}")
-    cuda_version_string = get_cuda_version(cuda, cudart_path)
-
-
-    if cc == '':
-        print(
-            "WARNING: No GPU detected! Check your CUDA paths. Processing to load CPU-only library..."
-        )
-        return binary_name
-
-    # 7.5 is the minimum CC vor cublaslt
-    has_cublaslt = cc in ["7.5", "8.0", "8.6"]
-
-    # TODO:
-    # (1) CUDA missing cases (no CUDA installed by CUDA driver (nvidia-smi accessible)
-    # (2) Multiple CUDA versions installed
-
-    # we use ls -l instead of nvcc to determine the cuda version
-    # since most installations will have the libcudart.so installed, but not the compiler
-    print(f'CUDA SETUP: Detected CUDA version {cuda_version_string}')
-
-    def get_binary_name():
-        "if not has_cublaslt (CC < 7.5), then we have to choose  _nocublaslt.so"
-        bin_base_name = "libbitsandbytes_cuda"
-        if has_cublaslt:
-            return f"{bin_base_name}{cuda_version_string}.so"
-        else:
-            return f"{bin_base_name}{cuda_version_string}_nocublaslt.so"
-
-    binary_name = get_binary_name()
-
-    return binary_name
+"""
+extract factors the build is dependent on:
+[X] compute capability
+    [ ] TODO: Q - What if we have multiple GPUs of different makes?
+- CUDA version
+- Software:
+    - CPU-only: only CPU quantization functions (no optimizer, no matrix multiple)
+    - CuBLAS-LT: full-build 8-bit optimizer
+    - no CuBLAS-LT: no 8-bit matrix multiplication (`nomatmul`)
+
+evaluation:
+    - if paths faulty, return meaningful error
+    - else:
+        - determine CUDA version
+        - determine capabilities
+        - based on that set the default path
+"""
+
+import ctypes
+
+from .paths import determine_cuda_runtime_lib_path
+
+
+def check_cuda_result(cuda, result_val):
+    # 3. Check for CUDA errors
+    if result_val != 0:
+        error_str = ctypes.c_char_p()
+        cuda.cuGetErrorString(result_val, ctypes.byref(error_str))
+        print(f"CUDA exception! Error code: {error_str.value.decode()}")
+
+def get_cuda_version(cuda, cudart_path):
+    # https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART____VERSION.html#group__CUDART____VERSION
+    try:
+        cudart = ctypes.CDLL(cudart_path)
+    except OSError:
+        # TODO: shouldn't we error or at least warn here?
+        print(f'ERROR: libcudart.so could not be read from path: {cudart_path}!')
+        return None
+
+    version = ctypes.c_int()
+    check_cuda_result(cuda, cudart.cudaRuntimeGetVersion(ctypes.byref(version)))
+    version = int(version.value)
+    major = version//1000
+    minor = (version-(major*1000))//10
+
+    if major < 11:
+       print('CUDA SETUP: CUDA version lower than 11 are currently not supported for LLM.int8(). You will be only to use 8-bit optimizers and quantization routines!!')
+
+    return f'{major}{minor}'
+
+
+def get_cuda_lib_handle():
+    # 1. find libcuda.so library (GPU driver) (/usr/lib)
+    try:
+        cuda = ctypes.CDLL("libcuda.so")
+    except OSError:
+        # TODO: shouldn't we error or at least warn here?
+        print('CUDA SETUP: WARNING! libcuda.so not found! Do you have a CUDA driver installed? If you are on a cluster, make sure you are on a CUDA machine!')
+        return None
+    check_cuda_result(cuda, cuda.cuInit(0))
+
+    return cuda
+
+
+def get_compute_capabilities(cuda):
+    """
+    1. find libcuda.so library (GPU driver) (/usr/lib)
+       init_device -> init variables -> call function by reference
+    2. call extern C function to determine CC
+       (https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__DEVICE__DEPRECATED.html)
+    3. Check for CUDA errors
+       https://stackoverflow.com/questions/14038589/what-is-the-canonical-way-to-check-for-errors-using-the-cuda-runtime-api
+    # bits taken from https://gist.github.com/f0k/63a664160d016a491b2cbea15913d549
+    """
+
+
+    nGpus = ctypes.c_int()
+    cc_major = ctypes.c_int()
+    cc_minor = ctypes.c_int()
+
+    device = ctypes.c_int()
+
+    check_cuda_result(cuda, cuda.cuDeviceGetCount(ctypes.byref(nGpus)))
+    ccs = []
+    for i in range(nGpus.value):
+        check_cuda_result(cuda, cuda.cuDeviceGet(ctypes.byref(device), i))
+        ref_major = ctypes.byref(cc_major)
+        ref_minor = ctypes.byref(cc_minor)
+        # 2. call extern C function to determine CC
+        check_cuda_result(
+            cuda, cuda.cuDeviceComputeCapability(ref_major, ref_minor, device)
+        )
+        ccs.append(f"{cc_major.value}.{cc_minor.value}")
+
+    return ccs
+
+
+# def get_compute_capability()-> Union[List[str, ...], None]: # FIXME: error
+def get_compute_capability(cuda):
+    """
+    Extracts the highest compute capbility from all available GPUs, as compute
+    capabilities are downwards compatible. If no GPUs are detected, it returns
+    None.
+    """
+    ccs = get_compute_capabilities(cuda)
+    if ccs is not None:
+        # TODO: handle different compute capabilities; for now, take the max
+        return ccs[-1]
+    return None
+
+
+def evaluate_cuda_setup():
+    print('')
+    print('='*35 + 'BUG REPORT' + '='*35)
+    print('Welcome to bitsandbytes. For bug reports, please submit your error trace to: https://github.com/TimDettmers/bitsandbytes/issues')
+    print('For effortless bug reporting copy-paste your error into this form: https://docs.google.com/forms/d/e/1FAIpQLScPB8emS3Thkp66nvqwmjTEgxp8Y9ufuWTzFyr9kJ5AoI47dQ/viewform?usp=sf_link')
+    print('='*80)
+    return "libbitsandbytes_cuda116.dll"            # $$$
+    
+    binary_name = "libbitsandbytes_cpu.so"
+    #if not torch.cuda.is_available():
+        #print('No GPU detected. Loading CPU library...')
+        #return binary_name
+
+    cudart_path = determine_cuda_runtime_lib_path()
+    if cudart_path is None:
+        print(
+            "WARNING: No libcudart.so found! Install CUDA or the cudatoolkit package (anaconda)!"
+        )
+        return binary_name
+
+    print(f"CUDA SETUP: CUDA runtime path found: {cudart_path}")
+    cuda = get_cuda_lib_handle()
+    cc = get_compute_capability(cuda)
+    print(f"CUDA SETUP: Highest compute capability among GPUs detected: {cc}")
+    cuda_version_string = get_cuda_version(cuda, cudart_path)
+
+
+    if cc == '':
+        print(
+            "WARNING: No GPU detected! Check your CUDA paths. Processing to load CPU-only library..."
+        )
+        return binary_name
+
+    # 7.5 is the minimum CC vor cublaslt
+    has_cublaslt = cc in ["7.5", "8.0", "8.6"]
+
+    # TODO:
+    # (1) CUDA missing cases (no CUDA installed by CUDA driver (nvidia-smi accessible)
+    # (2) Multiple CUDA versions installed
+
+    # we use ls -l instead of nvcc to determine the cuda version
+    # since most installations will have the libcudart.so installed, but not the compiler
+    print(f'CUDA SETUP: Detected CUDA version {cuda_version_string}')
+
+    def get_binary_name():
+        "if not has_cublaslt (CC < 7.5), then we have to choose  _nocublaslt.so"
+        bin_base_name = "libbitsandbytes_cuda"
+        if has_cublaslt:
+            return f"{bin_base_name}{cuda_version_string}.so"
+        else:
+            return f"{bin_base_name}{cuda_version_string}_nocublaslt.so"
+
+    binary_name = get_binary_name()
+
+    return binary_name

docs/config_README-ja.md

@@ -138,13 +138,9 @@ DreamBooth の手法と fine tuning の手法の両方とも利用可能な学
 | `num_repeats` | `10` | o | o | o |
 | `random_crop` | `false` | o | o | o |
 | `shuffle_caption` | `true` | o | o | o |
-| `caption_prefix` | `“masterpiece, best quality, ”` | o | o | o |
-| `caption_suffix` | `“, from side”` | o | o | o |
 
 * `num_repeats`
     * サブセットの画像の繰り返し回数を指定します。fine tuning における `--dataset_repeats` に相当しますが、`num_repeats` はどの学習方法でも指定可能です。
-* `caption_prefix`, `caption_suffix`
-    * キャプションの前、後に付与する文字列を指定します。シャッフルはこれらの文字列を含めた状態で行われます。`keep_tokens` を指定する場合には注意してください。
 
 ### DreamBooth 方式専用のオプション
 

docs/train_lllite_README-ja.md

@@ -1,214 +0,0 @@
-# ControlNet-LLLite について
-
-__きわめて実験的な実装のため、将来的に大きく変更される可能性があります。__
-
-## 概要
-ControlNet-LLLite は、[ControlNet](https://github.com/lllyasviel/ControlNet) の軽量版です。LoRA Like Lite という意味で、LoRAからインスピレーションを得た構造を持つ、軽量なControlNetです。現在はSDXLにのみ対応しています。
-
-## サンプルの重みファイルと推論
-
-こちらにあります: https://huggingface.co/kohya-ss/controlnet-lllite
-
-ComfyUIのカスタムノードを用意しています。: https://github.com/kohya-ss/ControlNet-LLLite-ComfyUI
-
-生成サンプルはこのページの末尾にあります。
-
-## モデル構造
-ひとつのLLLiteモジュールは、制御用画像（以下conditioning image）を潜在空間に写像するconditioning image embeddingと、LoRAにちょっと似た構造を持つ小型のネットワークからなります。LLLiteモジュールを、LoRAと同様にU-NetのLinearやConvに追加します。詳しくはソースコードを参照してください。
-
-推論環境の制限で、現在はCrossAttentionのみ（attn1のq/k/v、attn2のq）に追加されます。
-
-## モデルの学習
-
-### データセットの準備
-通常のdatasetに加え、`conditioning_data_dir` で指定したディレクトリにconditioning imageを格納してください。conditioning imageは学習用画像と同じbasenameを持つ必要があります。また、conditioning imageは学習用画像と同じサイズに自動的にリサイズされます。conditioning imageにはキャプションファイルは不要です。
-
-たとえば DreamBooth 方式でキャプションファイルを用いる場合の設定ファイルは以下のようになります。
-
-```toml
-[[datasets.subsets]]
-image_dir = "path/to/image/dir"
-caption_extension = ".txt"
-conditioning_data_dir = "path/to/conditioning/image/dir"
-```
-
-現時点の制約として、random_cropは使用できません。
-
-学習データとしては、元のモデルで生成した画像を学習用画像として、そこから加工した画像をconditioning imageとした、合成によるデータセットを用いるのがもっとも簡単です（データセットの品質的には問題があるかもしれません）。具体的なデータセットの合成方法については後述します。
-
-なお、元モデルと異なる画風の画像を学習用画像とすると、制御に加えて、その画風についても学ぶ必要が生じます。ControlNet-LLLiteは容量が少ないため、画風学習には不向きです。このような場合には、後述の次元数を多めにしてください。
-
-### 学習
-スクリプトで生成する場合は、`sdxl_train_control_net_lllite.py` を実行してください。`--cond_emb_dim` でconditioning image embeddingの次元数を指定できます。`--network_dim` でLoRA的モジュールのrankを指定できます。その他のオプションは`sdxl_train_network.py`に準じますが、`--network_module`の指定は不要です。
-
-学習時にはメモリを大量に使用しますので、キャッシュやgradient checkpointingなどの省メモリ化のオプションを有効にしてください。また`--full_bf16` オプションで、BFloat16を使用するのも有効です（RTX 30シリーズ以降のGPUが必要です）。24GB VRAMで動作確認しています。
-
-conditioning image embeddingの次元数は、サンプルのCannyでは32を指定しています。LoRA的モジュールのrankは同じく64です。対象とするconditioning imageの特徴に合わせて調整してください。
-
-（サンプルのCannyは恐らくかなり難しいと思われます。depthなどでは半分程度にしてもいいかもしれません。）
-
-以下は .toml の設定例です。
-
-```toml
-pretrained_model_name_or_path = "/path/to/model_trained_on.safetensors"
-max_train_epochs = 12
-max_data_loader_n_workers = 4
-persistent_data_loader_workers = true
-seed = 42
-gradient_checkpointing = true
-mixed_precision = "bf16"
-save_precision = "bf16"
-full_bf16 = true
-optimizer_type = "adamw8bit"
-learning_rate = 2e-4
-xformers = true
-output_dir = "/path/to/output/dir"
-output_name = "output_name"
-save_every_n_epochs = 1
-save_model_as = "safetensors"
-vae_batch_size = 4
-cache_latents = true
-cache_latents_to_disk = true
-cache_text_encoder_outputs = true
-cache_text_encoder_outputs_to_disk = true
-network_dim = 64
-cond_emb_dim = 32
-dataset_config = "/path/to/dataset.toml"
-```
-
-### 推論
-
-スクリプトで生成する場合は、`sdxl_gen_img.py` を実行してください。`--control_net_lllite_models` でLLLiteのモデルファイルを指定できます。次元数はモデルファイルから自動取得します。
-
-`--guide_image_path`で推論に用いるconditioning imageを指定してください。なおpreprocessは行われないため、たとえばCannyならCanny処理を行った画像を指定してください（背景黒に白線）。`--control_net_preps`, `--control_net_weights`, `--control_net_ratios` には未対応です。
-
-## データセットの合成方法
-
-### 学習用画像の生成
-
-学習のベースとなるモデルで画像生成を行います。Web UIやComfyUIなどで生成してください。画像サイズはモデルのデフォルトサイズで良いと思われます（1024x1024など）。bucketingを用いることもできます。その場合は適宜適切な解像度で生成してください。
-
-生成時のキャプション等は、ControlNet-LLLiteの利用時に生成したい画像にあわせるのが良いと思われます。
-
-生成した画像を任意のディレクトリに保存してください。このディレクトリをデータセットの設定ファイルで指定します。
-
-当リポジトリ内の `sdxl_gen_img.py` でも生成できます。例えば以下のように実行します。
-
-```dos
-python sdxl_gen_img.py --ckpt path/to/model.safetensors --n_iter 1 --scale 10 --steps 36 --outdir path/to/output/dir --xformers --W 1024 --H 1024 --original_width 2048 --original_height 2048 --bf16 --sampler ddim --batch_size 4 --vae_batch_size 2 --images_per_prompt 512 --max_embeddings_multiples 1 --prompt "{portrait|digital art|anime screen cap|detailed illustration} of 1girl, {standing|sitting|walking|running|dancing} on {classroom|street|town|beach|indoors|outdoors}, {looking at viewer|looking away|looking at another}, {in|wearing} {shirt and skirt|school uniform|casual wear} { |, dynamic pose}, (solo), teen age, {0-1$$smile,|blush,|kind smile,|expression less,|happy,|sadness,} {0-1$$upper body,|full body,|cowboy shot,|face focus,} trending on pixiv, {0-2$$depth of fields,|8k wallpaper,|highly detailed,|pov,} {0-1$$summer, |winter, |spring, |autumn, } beautiful face { |, from below|, from above|, from side|, from behind|, from back} --n nsfw, bad face, lowres, low quality, worst quality, low effort, watermark, signature, ugly, poorly drawn"
-```
-
-VRAM 24GBの設定です。VRAMサイズにより`--batch_size` `--vae_batch_size`を調整してください。
-
-`--prompt`でワイルドカードを利用してランダムに生成しています。適宜調整してください。
-
-### 画像の加工
-
-外部のプログラムを用いて、生成した画像を加工します。加工した画像を任意のディレクトリに保存してください。これらがconditioning imageになります。
-
-加工にはたとえばCannyなら以下のようなスクリプトが使えます。
-
-```python
-import glob
-import os
-import random
-import cv2
-import numpy as np
-
-IMAGES_DIR = "path/to/generated/images"
-CANNY_DIR = "path/to/canny/images"
-
-os.makedirs(CANNY_DIR, exist_ok=True)
-img_files = glob.glob(IMAGES_DIR + "/*.png")
-for img_file in img_files:
-    can_file = CANNY_DIR + "/" + os.path.basename(img_file)
-    if os.path.exists(can_file):
-        print("Skip: " + img_file)
-        continue
-
-    print(img_file)
-
-    img = cv2.imread(img_file)
-
-    # random threshold
-    # while True:
-    #     threshold1 = random.randint(0, 127)
-    #     threshold2 = random.randint(128, 255)
-    #     if threshold2 - threshold1 > 80:
-    #         break
-
-    # fixed threshold
-    threshold1 = 100
-    threshold2 = 200
-
-    img = cv2.Canny(img, threshold1, threshold2)
-
-    cv2.imwrite(can_file, img)
-```
-
-### キャプションファイルの作成
-
-学習用画像のbasenameと同じ名前で、それぞれの画像に対応したキャプションファイルを作成してください。生成時のプロンプトをそのまま利用すれば良いと思われます。
-
-`sdxl_gen_img.py` で生成した場合は、画像内のメタデータに生成時のプロンプトが記録されていますので、以下のようなスクリプトで学習用画像と同じディレクトリにキャプションファイルを作成できます（拡張子 `.txt`）。
-
-```python
-import glob
-import os
-from PIL import Image
-
-IMAGES_DIR = "path/to/generated/images"
-
-img_files = glob.glob(IMAGES_DIR + "/*.png")
-for img_file in img_files:
-    cap_file = img_file.replace(".png", ".txt")
-    if os.path.exists(cap_file):
-        print(f"Skip: {img_file}")
-        continue
-    print(img_file)
-
-    img = Image.open(img_file)
-    prompt = img.text["prompt"] if "prompt" in img.text else ""
-    if prompt == "":
-        print(f"Prompt not found in {img_file}")
-
-    with open(cap_file, "w") as f:
-        f.write(prompt + "\n")
-```
-
-### データセットの設定ファイルの作成
-
-コマンドラインオプションからの指定も可能ですが、`.toml`ファイルを作成する場合は `conditioning_data_dir` に加工した画像を保存したディレクトリを指定します。
-
-以下は設定ファイルの例です。
-
-```toml
-[general]
-flip_aug = false
-color_aug = false
-resolution = [1024,1024]
-
-[[datasets]]
-batch_size = 8
-enable_bucket = false
-
-    [[datasets.subsets]]
-    image_dir = "path/to/generated/image/dir"
-    caption_extension = ".txt"
-    conditioning_data_dir = "path/to/canny/image/dir"
-```
-
-## 謝辞
-
-ControlNetの作者である lllyasviel 氏、実装上のアドバイスとトラブル解決へのご尽力をいただいた furusu 氏、ControlNetデータセットを実装していただいた ddPn08 氏に感謝いたします。
-
-## サンプル
-Canny
-![kohya_ss_girl_standing_at_classroom_smiling_to_the_viewer_class_78976b3e-0d4d-4ea0-b8e3-053ae493abbc](https://github.com/kohya-ss/sd-scripts/assets/52813779/37e9a736-649b-4c0f-ab26-880a1bf319b5)
-
-![im_20230820104253_000_1](https://github.com/kohya-ss/sd-scripts/assets/52813779/c8896900-ab86-4120-932f-6e2ae17b77c0)
-
-![im_20230820104302_000_1](https://github.com/kohya-ss/sd-scripts/assets/52813779/b12457a0-ee3c-450e-ba9a-b712d0fe86bb)
-
-![im_20230820104310_000_1](https://github.com/kohya-ss/sd-scripts/assets/52813779/8845b8d9-804a-44ac-9618-113a28eac8a1)
-

docs/train_lllite_README.md

@@ -1,217 +0,0 @@
-# About ControlNet-LLLite
-
-__This is an extremely experimental implementation and may change significantly in the future.__
-
-日本語版は[こちら](./train_lllite_README-ja.md)
-
-## Overview
-
-ControlNet-LLLite is a lightweight version of [ControlNet](https://github.com/lllyasviel/ControlNet). It is a "LoRA Like Lite" that is inspired by LoRA and has a lightweight structure. Currently, only SDXL is supported.
-
-## Sample weight file and inference
-
-Sample weight file is available here: https://huggingface.co/kohya-ss/controlnet-lllite
-
-A custom node for ComfyUI is available: https://github.com/kohya-ss/ControlNet-LLLite-ComfyUI
-
-Sample images are at the end of this page.
-
-## Model structure
-
-A single LLLite module consists of a conditioning image embedding that maps a conditioning image to a latent space and a small network with a structure similar to LoRA. The LLLite module is added to U-Net's Linear and Conv in the same way as LoRA. Please refer to the source code for details.
-
-Due to the limitations of the inference environment, only CrossAttention (attn1 q/k/v, attn2 q) is currently added.
-
-## Model training
-
-### Preparing the dataset
-
-In addition to the normal dataset, please store the conditioning image in the directory specified by `conditioning_data_dir`. The conditioning image must have the same basename as the training image. The conditioning image will be automatically resized to the same size as the training image. The conditioning image does not require a caption file.
-
-```toml
-[[datasets.subsets]]
-image_dir = "path/to/image/dir"
-caption_extension = ".txt"
-conditioning_data_dir = "path/to/conditioning/image/dir"
-```
-
-At the moment, random_crop cannot be used.
-
-For training data, it is easiest to use a synthetic dataset with the original model-generated images as training images and processed images as conditioning images (the quality of the dataset may be problematic). See below for specific methods of synthesizing datasets.
-
-Note that if you use an image with a different art style than the original model as a training image, the model will have to learn not only the control but also the art style. ControlNet-LLLite has a small capacity, so it is not suitable for learning art styles. In such cases, increase the number of dimensions as described below.
-
-### Training
-
-Run `sdxl_train_control_net_lllite.py`. You can specify the dimension of the conditioning image embedding with `--cond_emb_dim`. You can specify the rank of the LoRA-like module with `--network_dim`. Other options are the same as `sdxl_train_network.py`, but `--network_module` is not required.
-
-Since a large amount of memory is used during training, please enable memory-saving options such as cache and gradient checkpointing. It is also effective to use BFloat16 with the `--full_bf16` option (requires RTX 30 series or later GPU). It has been confirmed to work with 24GB VRAM.
-
-For the sample Canny, the dimension of the conditioning image embedding is 32. The rank of the LoRA-like module is also 64. Adjust according to the features of the conditioning image you are targeting.
-
-(The sample Canny is probably quite difficult. It may be better to reduce it to about half for depth, etc.)
-
-The following is an example of a .toml configuration.
-
-```toml
-pretrained_model_name_or_path = "/path/to/model_trained_on.safetensors"
-max_train_epochs = 12
-max_data_loader_n_workers = 4
-persistent_data_loader_workers = true
-seed = 42
-gradient_checkpointing = true
-mixed_precision = "bf16"
-save_precision = "bf16"
-full_bf16 = true
-optimizer_type = "adamw8bit"
-learning_rate = 2e-4
-xformers = true
-output_dir = "/path/to/output/dir"
-output_name = "output_name"
-save_every_n_epochs = 1
-save_model_as = "safetensors"
-vae_batch_size = 4
-cache_latents = true
-cache_latents_to_disk = true
-cache_text_encoder_outputs = true
-cache_text_encoder_outputs_to_disk = true
-network_dim = 64
-cond_emb_dim = 32
-dataset_config = "/path/to/dataset.toml"
-```
-
-### Inference
-
-If you want to generate images with a script, run `sdxl_gen_img.py`. You can specify the LLLite model file with `--control_net_lllite_models`. The dimension is automatically obtained from the model file.
-
-Specify the conditioning image to be used for inference with `--guide_image_path`. Since preprocess is not performed, if it is Canny, specify an image processed with Canny (white line on black background). `--control_net_preps`, `--control_net_weights`, and `--control_net_ratios` are not supported.
-
-## How to synthesize a dataset
-
-### Generating training images
-
-Generate images with the base model for training. Please generate them with Web UI or ComfyUI etc. The image size should be the default size of the model (1024x1024, etc.). You can also use bucketing. In that case, please generate it at an arbitrary resolution.
-
-The captions and other settings when generating the images should be the same as when generating the images with the trained ControlNet-LLLite model.
-
-Save the generated images in an arbitrary directory. Specify this directory in the dataset configuration file.
-
-
-You can also generate them with `sdxl_gen_img.py` in this repository. For example, run as follows:
-
-```dos
-python sdxl_gen_img.py --ckpt path/to/model.safetensors --n_iter 1 --scale 10 --steps 36 --outdir path/to/output/dir --xformers --W 1024 --H 1024 --original_width 2048 --original_height 2048 --bf16 --sampler ddim --batch_size 4 --vae_batch_size 2 --images_per_prompt 512 --max_embeddings_multiples 1 --prompt "{portrait|digital art|anime screen cap|detailed illustration} of 1girl, {standing|sitting|walking|running|dancing} on {classroom|street|town|beach|indoors|outdoors}, {looking at viewer|looking away|looking at another}, {in|wearing} {shirt and skirt|school uniform|casual wear} { |, dynamic pose}, (solo), teen age, {0-1$$smile,|blush,|kind smile,|expression less,|happy,|sadness,} {0-1$$upper body,|full body,|cowboy shot,|face focus,} trending on pixiv, {0-2$$depth of fields,|8k wallpaper,|highly detailed,|pov,} {0-1$$summer, |winter, |spring, |autumn, } beautiful face { |, from below|, from above|, from side|, from behind|, from back} --n nsfw, bad face, lowres, low quality, worst quality, low effort, watermark, signature, ugly, poorly drawn"
-```
-
-This is a setting for VRAM 24GB. Adjust `--batch_size` and `--vae_batch_size` according to the VRAM size.
-
-The images are generated randomly using wildcards in `--prompt`. Adjust as necessary.
-
-### Processing images
-
-Use an external program to process the generated images. Save the processed images in an arbitrary directory. These will be the conditioning images.
-
-For example, you can use the following script to process the images with Canny.
-
-```python
-import glob
-import os
-import random
-import cv2
-import numpy as np
-
-IMAGES_DIR = "path/to/generated/images"
-CANNY_DIR = "path/to/canny/images"
-
-os.makedirs(CANNY_DIR, exist_ok=True)
-img_files = glob.glob(IMAGES_DIR + "/*.png")
-for img_file in img_files:
-    can_file = CANNY_DIR + "/" + os.path.basename(img_file)
-    if os.path.exists(can_file):
-        print("Skip: " + img_file)
-        continue
-
-    print(img_file)
-
-    img = cv2.imread(img_file)
-
-    # random threshold
-    # while True:
-    #     threshold1 = random.randint(0, 127)
-    #     threshold2 = random.randint(128, 255)
-    #     if threshold2 - threshold1 > 80:
-    #         break
-
-    # fixed threshold
-    threshold1 = 100
-    threshold2 = 200
-
-    img = cv2.Canny(img, threshold1, threshold2)
-
-    cv2.imwrite(can_file, img)
-```
-
-### Creating caption files
-
-Create a caption file for each image with the same basename as the training image. It is fine to use the same caption as the one used when generating the image. 
-
-If you generated the images with `sdxl_gen_img.py`, you can use the following script to create the caption files (`*.txt`) from the metadata in the generated images.
-
-```python
-import glob
-import os
-from PIL import Image
-
-IMAGES_DIR = "path/to/generated/images"
-
-img_files = glob.glob(IMAGES_DIR + "/*.png")
-for img_file in img_files:
-    cap_file = img_file.replace(".png", ".txt")
-    if os.path.exists(cap_file):
-        print(f"Skip: {img_file}")
-        continue
-    print(img_file)
-
-    img = Image.open(img_file)
-    prompt = img.text["prompt"] if "prompt" in img.text else ""
-    if prompt == "":
-        print(f"Prompt not found in {img_file}")
-
-    with open(cap_file, "w") as f:
-        f.write(prompt + "\n")
-```
-
-### Creating a dataset configuration file
-
-You can use the command line arguments of `sdxl_train_control_net_lllite.py` to specify the conditioning image directory. However, if you want to use a `.toml` file, specify the conditioning image directory in `conditioning_data_dir`.
-
-```toml
-[general]
-flip_aug = false
-color_aug = false
-resolution = [1024,1024]
-
-[[datasets]]
-batch_size = 8
-enable_bucket = false
-
-    [[datasets.subsets]]
-    image_dir = "path/to/generated/image/dir"
-    caption_extension = ".txt"
-    conditioning_data_dir = "path/to/canny/image/dir"
-```
-
-## Credit
-
-I would like to thank lllyasviel, the author of ControlNet, furusu, who provided me with advice on implementation and helped me solve problems, and ddPn08, who implemented the ControlNet dataset.
-
-## Sample
-
-Canny
-![kohya_ss_girl_standing_at_classroom_smiling_to_the_viewer_class_78976b3e-0d4d-4ea0-b8e3-053ae493abbc](https://github.com/kohya-ss/sd-scripts/assets/52813779/37e9a736-649b-4c0f-ab26-880a1bf319b5)
-
-![im_20230820104253_000_1](https://github.com/kohya-ss/sd-scripts/assets/52813779/c8896900-ab86-4120-932f-6e2ae17b77c0)
-
-![im_20230820104302_000_1](https://github.com/kohya-ss/sd-scripts/assets/52813779/b12457a0-ee3c-450e-ba9a-b712d0fe86bb)
-
-![im_20230820104310_000_1](https://github.com/kohya-ss/sd-scripts/assets/52813779/8845b8d9-804a-44ac-9618-113a28eac8a1)

fine_tune.py

@@ -5,12 +5,13 @@ import argparse
 import gc
 import math
 import os
-from multiprocessing import Value
 import toml
+from multiprocessing import Value
 
 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
@@ -24,6 +25,8 @@ from library.custom_train_functions import (
     apply_snr_weight,
     get_weighted_text_embeddings,
     prepare_scheduler_for_custom_training,
+    pyramid_noise_like,
+    apply_noise_offset,
     scale_v_prediction_loss_like_noise_prediction,
 )
 
@@ -41,7 +44,7 @@ def train(args):
 
     # データセットを準備する
     if args.dataset_class is None:
-        blueprint_generator = BlueprintGenerator(ConfigSanitizer(False, True, False, True))
+        blueprint_generator = BlueprintGenerator(ConfigSanitizer(False, True, True))
         if args.dataset_config is not None:
             print(f"Load dataset config from {args.dataset_config}")
             user_config = config_util.load_user_config(args.dataset_config)
@@ -92,7 +95,7 @@ def train(args):
 
     # acceleratorを準備する
     print("prepare accelerator")
-    accelerator = train_util.prepare_accelerator(args)
+    accelerator, unwrap_model = train_util.prepare_accelerator(args)
 
     # mixed precisionに対応した型を用意しておき適宜castする
     weight_dtype, save_dtype = train_util.prepare_dtype(args)
@@ -136,13 +139,13 @@ def train(args):
 
     # モデルに xformers とか memory efficient attention を組み込む
     if args.diffusers_xformers:
-        accelerator.print("Use xformers by Diffusers")
+        print("Use xformers by Diffusers")
         set_diffusers_xformers_flag(unet, True)
     else:
         # Windows版のxformersはfloatで学習できないのでxformersを使わない設定も可能にしておく必要がある
-        accelerator.print("Disable Diffusers' xformers")
+        print("Disable Diffusers' xformers")
         set_diffusers_xformers_flag(unet, False)
-        train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers, args.sdpa)
+        train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
 
     # 学習を準備する
     if cache_latents:
@@ -165,7 +168,7 @@ def train(args):
     training_models.append(unet)
 
     if args.train_text_encoder:
-        accelerator.print("enable text encoder training")
+        print("enable text encoder training")
         if args.gradient_checkpointing:
             text_encoder.gradient_checkpointing_enable()
         training_models.append(text_encoder)
@@ -191,7 +194,7 @@ def train(args):
     params_to_optimize = params
 
     # 学習に必要なクラスを準備する
-    accelerator.print("prepare optimizer, data loader etc.")
+    print("prepare optimizer, data loader etc.")
     _, _, optimizer = train_util.get_optimizer(args, trainable_params=params_to_optimize)
 
     # dataloaderを準備する
@@ -211,7 +214,7 @@ def train(args):
         args.max_train_steps = args.max_train_epochs * math.ceil(
             len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps
         )
-        accelerator.print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
+        print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
 
     # データセット側にも学習ステップを送信
     train_dataset_group.set_max_train_steps(args.max_train_steps)
@@ -224,7 +227,7 @@ def train(args):
         assert (
             args.mixed_precision == "fp16"
         ), "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
-        accelerator.print("enable full fp16 training.")
+        print("enable full fp16 training.")
         unet.to(weight_dtype)
         text_encoder.to(weight_dtype)
 
@@ -254,16 +257,14 @@ def train(args):
 
     # 学習する
     total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
-    accelerator.print("running training / 学習開始")
-    accelerator.print(f"  num examples / サンプル数: {train_dataset_group.num_train_images}")
-    accelerator.print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
-    accelerator.print(f"  num epochs / epoch数: {num_train_epochs}")
-    accelerator.print(f"  batch size per device / バッチサイズ: {args.train_batch_size}")
-    accelerator.print(
-        f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}"
-    )
-    accelerator.print(f"  gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
-    accelerator.print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")
+    print("running training / 学習開始")
+    print(f"  num examples / サンプル数: {train_dataset_group.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 accumulation 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
@@ -272,17 +273,12 @@ def train(args):
         beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False
     )
     prepare_scheduler_for_custom_training(noise_scheduler, accelerator.device)
-    if args.zero_terminal_snr:
-        custom_train_functions.fix_noise_scheduler_betas_for_zero_terminal_snr(noise_scheduler)
 
     if accelerator.is_main_process:
-        init_kwargs = {}
-        if args.log_tracker_config is not None:
-            init_kwargs = toml.load(args.log_tracker_config)
-        accelerator.init_trackers("finetuning" if args.log_tracker_name is None else args.log_tracker_name, init_kwargs=init_kwargs)
+        accelerator.init_trackers("finetuning" if args.log_tracker_name is None else args.log_tracker_name)
 
     for epoch in range(num_train_epochs):
-        accelerator.print(f"\nepoch {epoch+1}/{num_train_epochs}")
+        print(f"\nepoch {epoch+1}/{num_train_epochs}")
         current_epoch.value = epoch + 1
 
         for m in training_models:
@@ -318,9 +314,20 @@ def train(args):
                             args, input_ids, tokenizer, text_encoder, None if not args.full_fp16 else weight_dtype
                         )
 
-                # Sample noise, sample a random timestep for each image, and add noise to the latents,
-                # with noise offset and/or multires noise if specified
-                noise, noisy_latents, timesteps = train_util.get_noise_noisy_latents_and_timesteps(args, noise_scheduler, latents)
+                # Sample noise that we'll add to the latents
+                noise = torch.randn_like(latents, device=latents.device)
+                if args.noise_offset:
+                    noise = apply_noise_offset(latents, noise, args.noise_offset, args.adaptive_noise_scale)
+                elif args.multires_noise_iterations:
+                    noise = pyramid_noise_like(noise, latents.device, args.multires_noise_iterations, args.multires_noise_discount)
+
+                # 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
                 with accelerator.autocast():
@@ -382,17 +389,15 @@ def train(args):
                             epoch,
                             num_train_epochs,
                             global_step,
-                            accelerator.unwrap_model(text_encoder),
-                            accelerator.unwrap_model(unet),
+                            unwrap_model(text_encoder),
+                            unwrap_model(unet),
                             vae,
                         )
 
             current_loss = loss.detach().item()  # 平均なのでbatch sizeは関係ないはず
             if args.logging_dir is not None:
                 logs = {"loss": current_loss, "lr": float(lr_scheduler.get_last_lr()[0])}
-                if (
-                    args.optimizer_type.lower().startswith("DAdapt".lower()) or args.optimizer_type.lower() == "Prodigy".lower()
-                ):  # tracking d*lr value
+                if args.optimizer_type.lower().startswith("DAdapt".lower()) or args.optimizer_type.lower() == "Prodigy".lower():  # tracking d*lr value
                     logs["lr/d*lr"] = (
                         lr_scheduler.optimizers[0].param_groups[0]["d"] * lr_scheduler.optimizers[0].param_groups[0]["lr"]
                     )
@@ -427,8 +432,8 @@ def train(args):
                     epoch,
                     num_train_epochs,
                     global_step,
-                    accelerator.unwrap_model(text_encoder),
-                    accelerator.unwrap_model(unet),
+                    unwrap_model(text_encoder),
+                    unwrap_model(unet),
                     vae,
                 )
 
@@ -436,8 +441,8 @@ def train(args):
 
     is_main_process = accelerator.is_main_process
     if is_main_process:
-        unet = accelerator.unwrap_model(unet)
-        text_encoder = accelerator.unwrap_model(text_encoder)
+        unet = unwrap_model(unet)
+        text_encoder = unwrap_model(text_encoder)
 
     accelerator.end_training()
 

finetune/prepare_buckets_latents.py

@@ -34,7 +34,16 @@ def collate_fn_remove_corrupted(batch):
     return batch
 
 
-def get_npz_filename(data_dir, image_key, is_full_path, recursive):
+def get_latents(vae, images, weight_dtype):
+    img_tensors = [IMAGE_TRANSFORMS(image) for image in images]
+    img_tensors = torch.stack(img_tensors)
+    img_tensors = img_tensors.to(DEVICE, weight_dtype)
+    with torch.no_grad():
+        latents = vae.encode(img_tensors).latent_dist.sample().float().to("cpu").numpy()
+    return latents
+
+
+def get_npz_filename_wo_ext(data_dir, image_key, is_full_path, flip, recursive):
     if is_full_path:
         base_name = os.path.splitext(os.path.basename(image_key))[0]
         relative_path = os.path.relpath(os.path.dirname(image_key), data_dir)
@@ -42,20 +51,19 @@ def get_npz_filename(data_dir, image_key, is_full_path, recursive):
         base_name = image_key
         relative_path = ""
 
+    if flip:
+        base_name += "_flip"
+
     if recursive and relative_path:
-        return os.path.join(data_dir, relative_path, base_name) + ".npz"
+        return os.path.join(data_dir, relative_path, base_name)
     else:
-        return os.path.join(data_dir, base_name) + ".npz"
+        return os.path.join(data_dir, base_name)
 
 
 def main(args):
     # assert args.bucket_reso_steps % 8 == 0, f"bucket_reso_steps must be divisible by 8 / bucket_reso_stepは8で割り切れる必要があります"
     if args.bucket_reso_steps % 8 > 0:
         print(f"resolution of buckets in training time is a multiple of 8 / 学習時の各bucketの解像度は8単位になります")
-    if args.bucket_reso_steps % 32 > 0:
-        print(
-            f"WARNING: bucket_reso_steps is not divisible by 32. It is not working with SDXL / bucket_reso_stepsが32で割り切れません。SDXLでは動作しません"
-        )
 
     train_data_dir_path = Path(args.train_data_dir)
     image_paths: List[str] = [str(p) for p in train_util.glob_images_pathlib(train_data_dir_path, args.recursive)]
@@ -99,7 +107,34 @@ def main(args):
     def process_batch(is_last):
         for bucket in bucket_manager.buckets:
             if (is_last and len(bucket) > 0) or len(bucket) >= args.batch_size:
-                train_util.cache_batch_latents(vae, True, bucket, args.flip_aug, False)
+                latents = get_latents(vae, [img for _, img in bucket], weight_dtype)
+                assert (
+                    latents.shape[2] == bucket[0][1].shape[0] // 8 and latents.shape[3] == bucket[0][1].shape[1] // 8
+                ), f"latent shape {latents.shape}, {bucket[0][1].shape}"
+
+                for (image_key, _), latent in zip(bucket, latents):
+                    npz_file_name = get_npz_filename_wo_ext(args.train_data_dir, image_key, args.full_path, False, args.recursive)
+                    np.savez(npz_file_name, latent)
+
+                # flip
+                if args.flip_aug:
+                    latents = get_latents(vae, [img[:, ::-1].copy() for _, img in bucket], weight_dtype)  # copyがないとTensor変換できない
+
+                    for (image_key, _), latent in zip(bucket, latents):
+                        npz_file_name = get_npz_filename_wo_ext(
+                            args.train_data_dir, image_key, args.full_path, True, args.recursive
+                        )
+                        np.savez(npz_file_name, latent)
+                else:
+                    # remove existing flipped npz
+                    for image_key, _ in bucket:
+                        npz_file_name = (
+                            get_npz_filename_wo_ext(args.train_data_dir, image_key, args.full_path, True, args.recursive) + ".npz"
+                        )
+                        if os.path.isfile(npz_file_name):
+                            print(f"remove existing flipped npz / 既存のflipされたnpzファイルを削除します: {npz_file_name}")
+                            os.remove(npz_file_name)
+
                 bucket.clear()
 
     # 読み込みの高速化のためにDataLoaderを使うオプション
@@ -159,19 +194,50 @@ def main(args):
             resized_size[0] >= reso[0] and resized_size[1] >= reso[1]
         ), f"internal error resized size is small: {resized_size}, {reso}"
 
-        # 既に存在するファイルがあればshape等を確認して同じならskipする
-        npz_file_name = get_npz_filename(args.train_data_dir, image_key, args.full_path, args.recursive)
+        # 既に存在するファイルがあればshapeを確認して同じならskipする
         if args.skip_existing:
-            if train_util.is_disk_cached_latents_is_expected(reso, npz_file_name, args.flip_aug):
+            npz_files = [get_npz_filename_wo_ext(args.train_data_dir, image_key, args.full_path, False, args.recursive) + ".npz"]
+            if args.flip_aug:
+                npz_files.append(
+                    get_npz_filename_wo_ext(args.train_data_dir, image_key, args.full_path, True, args.recursive) + ".npz"
+                )
+
+            found = True
+            for npz_file in npz_files:
+                if not os.path.exists(npz_file):
+                    found = False
+                    break
+
+                dat = np.load(npz_file)["arr_0"]
+                if dat.shape[1] != reso[1] // 8 or dat.shape[2] != reso[0] // 8:  # latentsのshapeを確認
+                    found = False
+                    break
+            if found:
                 continue
 
+        # 画像をリサイズしてトリミングする
+        # PILにinter_areaがないのでcv2で……
+        image = np.array(image)
+        if resized_size[0] != image.shape[1] or resized_size[1] != image.shape[0]:  # リサイズ処理が必要？
+            image = cv2.resize(image, resized_size, interpolation=cv2.INTER_AREA)
+
+        if resized_size[0] > reso[0]:
+            trim_size = resized_size[0] - reso[0]
+            image = image[:, trim_size // 2 : trim_size // 2 + reso[0]]
+
+        if resized_size[1] > reso[1]:
+            trim_size = resized_size[1] - reso[1]
+            image = image[trim_size // 2 : trim_size // 2 + reso[1]]
+
+        assert (
+            image.shape[0] == reso[1] and image.shape[1] == reso[0]
+        ), f"internal error, illegal trimmed size: {image.shape}, {reso}"
+
+        # # debug
+        # cv2.imwrite(f"r:\\test\\img_{len(img_ar_errors)}.jpg", image[:, :, ::-1])
+
         # バッチへ追加
-        image_info = train_util.ImageInfo(image_key, 1, "", False, image_path)
-        image_info.latents_npz = npz_file_name
-        image_info.bucket_reso = reso
-        image_info.resized_size = resized_size
-        image_info.image = image
-        bucket_manager.add_image(reso, image_info)
+        bucket_manager.add_image(reso, (image_key, image))
 
         # バッチを推論するか判定して推論する
         process_batch(False)

gen_img_diffusers.py

@@ -79,10 +79,11 @@ from diffusers import (
     HeunDiscreteScheduler,
     KDPM2DiscreteScheduler,
     KDPM2AncestralDiscreteScheduler,
-    # UNet2DConditionModel,
+    UNet2DConditionModel,
     StableDiffusionPipeline,
 )
 from einops import rearrange
+from torch import einsum
 from tqdm import tqdm
 from torchvision import transforms
 from transformers import CLIPTextModel, CLIPTokenizer, CLIPModel, CLIPTextConfig
@@ -95,11 +96,15 @@ import library.train_util as train_util
 from networks.lora import LoRANetwork
 import tools.original_control_net as original_control_net
 from tools.original_control_net import ControlNetInfo
-from library.original_unet import UNet2DConditionModel
-from library.original_unet import FlashAttentionFunction
 
 from XTI_hijack import unet_forward_XTI, downblock_forward_XTI, upblock_forward_XTI
 
+# Tokenizer: checkpointから読み込むのではなくあらかじめ提供されているものを使う
+TOKENIZER_PATH = "openai/clip-vit-large-patch14"
+V2_STABLE_DIFFUSION_PATH = "stabilityai/stable-diffusion-2"  # ここからtokenizerだけ使う
+
+DEFAULT_TOKEN_LENGTH = 75
+
 # scheduler:
 SCHEDULER_LINEAR_START = 0.00085
 SCHEDULER_LINEAR_END = 0.0120
@@ -131,42 +136,336 @@ USE_CUTOUTS = False
 高速化のためのモジュール入れ替え
 """
 
+# 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
 
-def replace_unet_modules(unet: diffusers.models.unet_2d_condition.UNet2DConditionModel, mem_eff_attn, xformers, sdpa):
-    if mem_eff_attn:
-        print("Enable memory efficient attention for U-Net")
+# constants
 
-        # これはDiffusersのU-Netではなく自前のU-Netなので置き換えなくても良い
-        unet.set_use_memory_efficient_attention(False, True)
-    elif xformers:
-        print("Enable xformers for U-Net")
-        try:
-            import xformers.ops
-        except ImportError:
-            raise ImportError("No xformers / xformersがインストールされていないようです")
+EPSILON = 1e-6
 
-        unet.set_use_memory_efficient_attention(True, False)
-    elif sdpa:
-        print("Enable SDPA for U-Net")
-        unet.set_use_memory_efficient_attention(False, False)
-        unet.set_use_sdpa(True)
+# helper functions
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+# flash attention forwards and backwards
+
+# https://arxiv.org/abs/2205.14135
+
+
+class FlashAttentionFunction(torch.autograd.Function):
+    @staticmethod
+    @torch.no_grad()
+    def forward(ctx, q, k, v, mask, causal, q_bucket_size, k_bucket_size):
+        """Algorithm 2 in the paper"""
+
+        device = q.device
+        dtype = q.dtype
+        max_neg_value = -torch.finfo(q.dtype).max
+        qk_len_diff = max(k.shape[-2] - q.shape[-2], 0)
+
+        o = torch.zeros_like(q)
+        all_row_sums = torch.zeros((*q.shape[:-1], 1), dtype=dtype, device=device)
+        all_row_maxes = torch.full((*q.shape[:-1], 1), max_neg_value, dtype=dtype, device=device)
+
+        scale = q.shape[-1] ** -0.5
+
+        if not exists(mask):
+            mask = (None,) * math.ceil(q.shape[-2] / q_bucket_size)
+        else:
+            mask = rearrange(mask, "b n -> b 1 1 n")
+            mask = mask.split(q_bucket_size, dim=-1)
+
+        row_splits = zip(
+            q.split(q_bucket_size, dim=-2),
+            o.split(q_bucket_size, dim=-2),
+            mask,
+            all_row_sums.split(q_bucket_size, dim=-2),
+            all_row_maxes.split(q_bucket_size, dim=-2),
+        )
+
+        for ind, (qc, oc, row_mask, row_sums, row_maxes) in enumerate(row_splits):
+            q_start_index = ind * q_bucket_size - qk_len_diff
+
+            col_splits = zip(
+                k.split(k_bucket_size, dim=-2),
+                v.split(k_bucket_size, dim=-2),
+            )
+
+            for k_ind, (kc, vc) in enumerate(col_splits):
+                k_start_index = k_ind * k_bucket_size
+
+                attn_weights = einsum("... i d, ... j d -> ... i j", qc, kc) * scale
+
+                if exists(row_mask):
+                    attn_weights.masked_fill_(~row_mask, max_neg_value)
+
+                if causal and q_start_index < (k_start_index + k_bucket_size - 1):
+                    causal_mask = torch.ones((qc.shape[-2], kc.shape[-2]), dtype=torch.bool, device=device).triu(
+                        q_start_index - k_start_index + 1
+                    )
+                    attn_weights.masked_fill_(causal_mask, max_neg_value)
+
+                block_row_maxes = attn_weights.amax(dim=-1, keepdims=True)
+                attn_weights -= block_row_maxes
+                exp_weights = torch.exp(attn_weights)
+
+                if exists(row_mask):
+                    exp_weights.masked_fill_(~row_mask, 0.0)
+
+                block_row_sums = exp_weights.sum(dim=-1, keepdims=True).clamp(min=EPSILON)
+
+                new_row_maxes = torch.maximum(block_row_maxes, row_maxes)
+
+                exp_values = einsum("... i j, ... j d -> ... i d", exp_weights, vc)
+
+                exp_row_max_diff = torch.exp(row_maxes - new_row_maxes)
+                exp_block_row_max_diff = torch.exp(block_row_maxes - new_row_maxes)
+
+                new_row_sums = exp_row_max_diff * row_sums + exp_block_row_max_diff * block_row_sums
+
+                oc.mul_((row_sums / new_row_sums) * exp_row_max_diff).add_((exp_block_row_max_diff / new_row_sums) * exp_values)
+
+                row_maxes.copy_(new_row_maxes)
+                row_sums.copy_(new_row_sums)
+
+        ctx.args = (causal, scale, mask, q_bucket_size, k_bucket_size)
+        ctx.save_for_backward(q, k, v, o, all_row_sums, all_row_maxes)
+
+        return o
+
+    @staticmethod
+    @torch.no_grad()
+    def backward(ctx, do):
+        """Algorithm 4 in the paper"""
+
+        causal, scale, mask, q_bucket_size, k_bucket_size = ctx.args
+        q, k, v, o, l, m = ctx.saved_tensors
+
+        device = q.device
+
+        max_neg_value = -torch.finfo(q.dtype).max
+        qk_len_diff = max(k.shape[-2] - q.shape[-2], 0)
+
+        dq = torch.zeros_like(q)
+        dk = torch.zeros_like(k)
+        dv = torch.zeros_like(v)
+
+        row_splits = zip(
+            q.split(q_bucket_size, dim=-2),
+            o.split(q_bucket_size, dim=-2),
+            do.split(q_bucket_size, dim=-2),
+            mask,
+            l.split(q_bucket_size, dim=-2),
+            m.split(q_bucket_size, dim=-2),
+            dq.split(q_bucket_size, dim=-2),
+        )
+
+        for ind, (qc, oc, doc, row_mask, lc, mc, dqc) in enumerate(row_splits):
+            q_start_index = ind * q_bucket_size - qk_len_diff
+
+            col_splits = zip(
+                k.split(k_bucket_size, dim=-2),
+                v.split(k_bucket_size, dim=-2),
+                dk.split(k_bucket_size, dim=-2),
+                dv.split(k_bucket_size, dim=-2),
+            )
+
+            for k_ind, (kc, vc, dkc, dvc) in enumerate(col_splits):
+                k_start_index = k_ind * k_bucket_size
+
+                attn_weights = einsum("... i d, ... j d -> ... i j", qc, kc) * scale
+
+                if causal and q_start_index < (k_start_index + k_bucket_size - 1):
+                    causal_mask = torch.ones((qc.shape[-2], kc.shape[-2]), dtype=torch.bool, device=device).triu(
+                        q_start_index - k_start_index + 1
+                    )
+                    attn_weights.masked_fill_(causal_mask, max_neg_value)
+
+                exp_attn_weights = torch.exp(attn_weights - mc)
+
+                if exists(row_mask):
+                    exp_attn_weights.masked_fill_(~row_mask, 0.0)
+
+                p = exp_attn_weights / lc
+
+                dv_chunk = einsum("... i j, ... i d -> ... j d", p, doc)
+                dp = einsum("... i d, ... j d -> ... i j", doc, vc)
+
+                D = (doc * oc).sum(dim=-1, keepdims=True)
+                ds = p * scale * (dp - D)
+
+                dq_chunk = einsum("... i j, ... j d -> ... i d", ds, kc)
+                dk_chunk = einsum("... i j, ... i d -> ... j d", ds, qc)
+
+                dqc.add_(dq_chunk)
+                dkc.add_(dk_chunk)
+                dvc.add_(dv_chunk)
+
+        return dq, dk, dv, None, None, None, None
 
 
 # TODO common train_util.py
-def replace_vae_modules(vae: diffusers.models.AutoencoderKL, mem_eff_attn, xformers, sdpa):
+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("CrossAttention.forward has been replaced to FlashAttention (not xformers) and NAI style Hypernetwork")
+    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("CrossAttention.forward has been replaced to enable xformers and NAI style Hypernetwork")
+    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
+
+
+def replace_vae_modules(vae: diffusers.models.AutoencoderKL, mem_eff_attn, xformers):
     if mem_eff_attn:
         replace_vae_attn_to_memory_efficient()
     elif xformers:
-        replace_vae_attn_to_xformers()
-    elif sdpa:
-        replace_vae_attn_to_sdpa()
+        # とりあえずDiffusersのxformersを使う。AttentionがあるのはMidBlockのみ
+        print("Use Diffusers xformers for VAE")
+        vae.set_use_memory_efficient_attention_xformers(True)
+
+    """
+    # VAEがbfloat16でメモリ消費が大きい問題を解決する
+    upsamplers = []
+    for block in vae.decoder.up_blocks:
+        if block.upsamplers is not None:
+            upsamplers.extend(block.upsamplers)
+
+    def forward_upsample(_self, hidden_states, output_size=None):
+        assert hidden_states.shape[1] == _self.channels
+        if _self.use_conv_transpose:
+            return _self.conv(hidden_states)
+
+        dtype = hidden_states.dtype
+        if dtype == torch.bfloat16:
+            assert output_size is None
+            # repeat_interleaveはすごく遅いが、回数はあまり呼ばれないので許容する
+            hidden_states = hidden_states.repeat_interleave(2, dim=-1)
+            hidden_states = hidden_states.repeat_interleave(2, dim=-2)
+        else:
+            if hidden_states.shape[0] >= 64:
+                hidden_states = hidden_states.contiguous()
+
+            # if `output_size` is passed we force the interpolation output
+            # size and do not make use of `scale_factor=2`
+            if output_size is None:
+                hidden_states = torch.nn.functional.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
+            else:
+                hidden_states = torch.nn.functional.interpolate(hidden_states, size=output_size, mode="nearest")
+
+        if _self.use_conv:
+            if _self.name == "conv":
+                hidden_states = _self.conv(hidden_states)
+            else:
+                hidden_states = _self.Conv2d_0(hidden_states)
+        return hidden_states
+
+    # replace upsamplers
+    for upsampler in upsamplers:
+        # make new scope
+        def make_replacer(upsampler):
+            def forward(hidden_states, output_size=None):
+                return forward_upsample(upsampler, hidden_states, output_size)
+
+            return forward
+
+        upsampler.forward = make_replacer(upsampler)
+"""
 
 
 def replace_vae_attn_to_memory_efficient():
-    print("VAE Attention.forward has been replaced to FlashAttention (not xformers)")
+    print("AttentionBlock.forward has been replaced to FlashAttention (not xformers)")
     flash_func = FlashAttentionFunction
 
-    def forward_flash_attn(self, hidden_states, **kwargs):
+    def forward_flash_attn(self, hidden_states):
+        print("forward_flash_attn")
         q_bucket_size = 512
         k_bucket_size = 1024
 
@@ -179,12 +478,12 @@ def replace_vae_attn_to_memory_efficient():
         hidden_states = hidden_states.view(batch, channel, height * width).transpose(1, 2)
 
         # proj to q, k, v
-        query_proj = self.to_q(hidden_states)
-        key_proj = self.to_k(hidden_states)
-        value_proj = self.to_v(hidden_states)
+        query_proj = self.query(hidden_states)
+        key_proj = self.key(hidden_states)
+        value_proj = self.value(hidden_states)
 
         query_proj, key_proj, value_proj = map(
-            lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.heads), (query_proj, key_proj, value_proj)
+            lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.num_heads), (query_proj, key_proj, value_proj)
         )
 
         out = flash_func.apply(query_proj, key_proj, value_proj, None, False, q_bucket_size, k_bucket_size)
@@ -192,140 +491,14 @@ def replace_vae_attn_to_memory_efficient():
         out = rearrange(out, "b h n d -> b n (h d)")
 
         # compute next hidden_states
-        # linear proj
-        hidden_states = self.to_out[0](hidden_states)
-        # dropout
-        hidden_states = self.to_out[1](hidden_states)
-
+        hidden_states = self.proj_attn(hidden_states)
         hidden_states = hidden_states.transpose(-1, -2).reshape(batch, channel, height, width)
 
         # res connect and rescale
         hidden_states = (hidden_states + residual) / self.rescale_output_factor
         return hidden_states
 
-    def forward_flash_attn_0_14(self, hidden_states, **kwargs):
-        if not hasattr(self, "to_q"):
-            self.to_q = self.query
-            self.to_k = self.key
-            self.to_v = self.value
-            self.to_out = [self.proj_attn, torch.nn.Identity()]
-            self.heads = self.num_heads
-        return forward_flash_attn(self, hidden_states, **kwargs)
-
-    if diffusers.__version__ < "0.15.0":
-        diffusers.models.attention.AttentionBlock.forward = forward_flash_attn_0_14
-    else:
-        diffusers.models.attention_processor.Attention.forward = forward_flash_attn
-
-
-def replace_vae_attn_to_xformers():
-    print("VAE: Attention.forward has been replaced to xformers")
-    import xformers.ops
-
-    def forward_xformers(self, hidden_states, **kwargs):
-        residual = hidden_states
-        batch, channel, height, width = hidden_states.shape
-
-        # norm
-        hidden_states = self.group_norm(hidden_states)
-
-        hidden_states = hidden_states.view(batch, channel, height * width).transpose(1, 2)
-
-        # proj to q, k, v
-        query_proj = self.to_q(hidden_states)
-        key_proj = self.to_k(hidden_states)
-        value_proj = self.to_v(hidden_states)
-
-        query_proj, key_proj, value_proj = map(
-            lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.heads), (query_proj, key_proj, value_proj)
-        )
-
-        query_proj = query_proj.contiguous()
-        key_proj = key_proj.contiguous()
-        value_proj = value_proj.contiguous()
-        out = xformers.ops.memory_efficient_attention(query_proj, key_proj, value_proj, attn_bias=None)
-
-        out = rearrange(out, "b h n d -> b n (h d)")
-
-        # compute next hidden_states
-        # linear proj
-        hidden_states = self.to_out[0](hidden_states)
-        # dropout
-        hidden_states = self.to_out[1](hidden_states)
-
-        hidden_states = hidden_states.transpose(-1, -2).reshape(batch, channel, height, width)
-
-        # res connect and rescale
-        hidden_states = (hidden_states + residual) / self.rescale_output_factor
-        return hidden_states
-
-    def forward_xformers_0_14(self, hidden_states, **kwargs):
-        if not hasattr(self, "to_q"):
-            self.to_q = self.query
-            self.to_k = self.key
-            self.to_v = self.value
-            self.to_out = [self.proj_attn, torch.nn.Identity()]
-            self.heads = self.num_heads
-        return forward_xformers(self, hidden_states, **kwargs)
-
-    if diffusers.__version__ < "0.15.0":
-        diffusers.models.attention.AttentionBlock.forward = forward_xformers_0_14
-    else:
-        diffusers.models.attention_processor.Attention.forward = forward_xformers
-
-
-def replace_vae_attn_to_sdpa():
-    print("VAE: Attention.forward has been replaced to sdpa")
-
-    def forward_sdpa(self, hidden_states, **kwargs):
-        residual = hidden_states
-        batch, channel, height, width = hidden_states.shape
-
-        # norm
-        hidden_states = self.group_norm(hidden_states)
-
-        hidden_states = hidden_states.view(batch, channel, height * width).transpose(1, 2)
-
-        # proj to q, k, v
-        query_proj = self.to_q(hidden_states)
-        key_proj = self.to_k(hidden_states)
-        value_proj = self.to_v(hidden_states)
-
-        query_proj, key_proj, value_proj = map(
-            lambda t: rearrange(t, "b n (h d) -> b n h d", h=self.heads), (query_proj, key_proj, value_proj)
-        )
-
-        out = torch.nn.functional.scaled_dot_product_attention(
-            query_proj, key_proj, value_proj, attn_mask=None, dropout_p=0.0, is_causal=False
-        )
-
-        out = rearrange(out, "b n h d -> b n (h d)")
-
-        # compute next hidden_states
-        # linear proj
-        hidden_states = self.to_out[0](hidden_states)
-        # dropout
-        hidden_states = self.to_out[1](hidden_states)
-
-        hidden_states = hidden_states.transpose(-1, -2).reshape(batch, channel, height, width)
-
-        # res connect and rescale
-        hidden_states = (hidden_states + residual) / self.rescale_output_factor
-        return hidden_states
-
-    def forward_sdpa_0_14(self, hidden_states, **kwargs):
-        if not hasattr(self, "to_q"):
-            self.to_q = self.query
-            self.to_k = self.key
-            self.to_v = self.value
-            self.to_out = [self.proj_attn, torch.nn.Identity()]
-            self.heads = self.num_heads
-        return forward_sdpa(self, hidden_states, **kwargs)
-
-    if diffusers.__version__ < "0.15.0":
-        diffusers.models.attention.AttentionBlock.forward = forward_sdpa_0_14
-    else:
-        diffusers.models.attention_processor.Attention.forward = forward_sdpa
+    diffusers.models.attention.AttentionBlock.forward = forward_flash_attn
 
 
 # endregion
@@ -937,17 +1110,6 @@ class PipelineLike:
         if self.control_nets:
             guided_hints = original_control_net.get_guided_hints(self.control_nets, num_latent_input, batch_size, clip_guide_images)
 
-        if reginonal_network:
-            num_sub_and_neg_prompts = len(text_embeddings) // batch_size
-            # last subprompt and negative prompt
-            text_emb_last = []
-            for j in range(batch_size):
-                text_emb_last.append(text_embeddings[(j + 1) * num_sub_and_neg_prompts - 2])
-                text_emb_last.append(text_embeddings[(j + 1) * num_sub_and_neg_prompts - 1])
-            text_emb_last = torch.stack(text_emb_last)
-        else:
-            text_emb_last = text_embeddings
-            
         for i, t in enumerate(tqdm(timesteps)):
             # expand the latents if we are doing classifier free guidance
             latent_model_input = latents.repeat((num_latent_input, 1, 1, 1))
@@ -955,6 +1117,11 @@ class PipelineLike:
 
             # predict the noise residual
             if self.control_nets and self.control_net_enabled:
+                if reginonal_network:
+                    num_sub_and_neg_prompts = len(text_embeddings) // batch_size
+                    text_emb_last = text_embeddings[num_sub_and_neg_prompts - 2 :: num_sub_and_neg_prompts]  # last subprompt
+                else:
+                    text_emb_last = text_embeddings
                 noise_pred = original_control_net.call_unet_and_control_net(
                     i,
                     num_latent_input,
@@ -964,7 +1131,6 @@ class PipelineLike:
                     i / len(timesteps),
                     latent_model_input,
                     t,
-                    text_embeddings,
                     text_emb_last,
                 ).sample
             else:
@@ -2176,17 +2342,6 @@ def main(args):
         tokenizer = loading_pipe.tokenizer
         del loading_pipe
 
-        # Diffusers U-Net to original U-Net
-        original_unet = UNet2DConditionModel(
-            unet.config.sample_size,
-            unet.config.attention_head_dim,
-            unet.config.cross_attention_dim,
-            unet.config.use_linear_projection,
-            unet.config.upcast_attention,
-        )
-        original_unet.load_state_dict(unet.state_dict())
-        unet = original_unet
-
     # VAEを読み込む
     if args.vae is not None:
         vae = model_util.load_vae(args.vae, dtype)
@@ -2211,9 +2366,8 @@ def main(args):
 
     # xformers、Hypernetwork対応
     if not args.diffusers_xformers:
-        mem_eff = not (args.xformers or args.sdpa)
-        replace_unet_modules(unet, mem_eff, args.xformers, args.sdpa)
-        replace_vae_modules(vae, mem_eff, args.xformers, args.sdpa)
+        replace_unet_modules(unet, not args.xformers, args.xformers)
+        replace_vae_modules(vae, not args.xformers, args.xformers)
 
     # tokenizerを読み込む
     print("loading tokenizer")
@@ -2753,10 +2907,6 @@ def main(args):
         print(f"iteration {gen_iter+1}/{args.n_iter}")
         iter_seed = random.randint(0, 0x7FFFFFFF)
 
-        # shuffle prompt list
-        if args.shuffle_prompts:
-            random.shuffle(prompt_list)
-
         # バッチ処理の関数
         def process_batch(batch: List[BatchData], highres_fix, highres_1st=False):
             batch_size = len(batch)
@@ -2974,8 +3124,6 @@ def main(args):
             for i, (image, prompt, negative_prompts, seed, clip_prompt) in enumerate(
                 zip(images, prompts, negative_prompts, seeds, clip_prompts)
             ):
-                if highres_fix:
-                    seed -= 1  # record original seed
                 metadata = PngInfo()
                 metadata.add_text("prompt", prompt)
                 metadata.add_text("seed", str(seed))
@@ -3332,15 +3480,9 @@ def setup_parser() -> argparse.ArgumentParser:
         action="store_true",
         help="use same seed for all prompts in iteration if no seed specified / 乱数seedの指定がないとき繰り返し内はすべて同じseedを使う（プロンプト間の差異の比較用）",
     )
-    parser.add_argument(
-        "--shuffle_prompts",
-        action="store_true",
-        help="shuffle prompts in iteration / 繰り返し内のプロンプトをシャッフルする",
-    )
     parser.add_argument("--fp16", action="store_true", help="use fp16 / fp16を指定し省メモリ化する")
     parser.add_argument("--bf16", action="store_true", help="use bfloat16 / bfloat16を指定し省メモリ化する")
     parser.add_argument("--xformers", action="store_true", help="use xformers / xformersを使用し高速化する")
-    parser.add_argument("--sdpa", action="store_true", help="use sdpa in PyTorch 2 / sdpa")
     parser.add_argument(
         "--diffusers_xformers",
         action="store_true",

library/attention_processors.py

@@ -1,227 +0,0 @@
-import math
-from typing import Any
-from einops import rearrange
-import torch
-from diffusers.models.attention_processor import Attention
-
-
-# flash attention forwards and backwards
-
-# https://arxiv.org/abs/2205.14135
-
-EPSILON = 1e-6
-
-
-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 mask is None:
-            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 = (
-                    torch.einsum("... i d, ... j d -> ... i j", qc, kc) * scale
-                )
-
-                if row_mask is not None:
-                    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 row_mask is not None:
-                    exp_weights.masked_fill_(~row_mask, 0.0)
-
-                block_row_sums = exp_weights.sum(dim=-1, keepdims=True).clamp(
-                    min=EPSILON
-                )
-
-                new_row_maxes = torch.maximum(block_row_maxes, row_maxes)
-
-                exp_values = torch.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 = (
-                    torch.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 row_mask is not None:
-                    exp_attn_weights.masked_fill_(~row_mask, 0.0)
-
-                p = exp_attn_weights / lc
-
-                dv_chunk = torch.einsum("... i j, ... i d -> ... j d", p, doc)
-                dp = torch.einsum("... i d, ... j d -> ... i j", doc, vc)
-
-                D = (doc * oc).sum(dim=-1, keepdims=True)
-                ds = p * scale * (dp - D)
-
-                dq_chunk = torch.einsum("... i j, ... j d -> ... i d", ds, kc)
-                dk_chunk = torch.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
-
-
-class FlashAttnProcessor:
-    def __call__(
-        self,
-        attn: Attention,
-        hidden_states,
-        encoder_hidden_states=None,
-        attention_mask=None,
-    ) -> Any:
-        q_bucket_size = 512
-        k_bucket_size = 1024
-
-        h = attn.heads
-        q = attn.to_q(hidden_states)
-
-        encoder_hidden_states = (
-            encoder_hidden_states
-            if encoder_hidden_states is not None
-            else hidden_states
-        )
-        encoder_hidden_states = encoder_hidden_states.to(hidden_states.dtype)
-
-        if hasattr(attn, "hypernetwork") and attn.hypernetwork is not None:
-            context_k, context_v = attn.hypernetwork.forward(
-                hidden_states, encoder_hidden_states
-            )
-            context_k = context_k.to(hidden_states.dtype)
-            context_v = context_v.to(hidden_states.dtype)
-        else:
-            context_k = encoder_hidden_states
-            context_v = encoder_hidden_states
-
-        k = attn.to_k(context_k)
-        v = attn.to_v(context_v)
-        del encoder_hidden_states, hidden_states
-
-        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v))
-
-        out = FlashAttentionFunction.apply(
-            q, k, v, attention_mask, False, q_bucket_size, k_bucket_size
-        )
-
-        out = rearrange(out, "b h n d -> b n (h d)")
-
-        out = attn.to_out[0](out)
-        out = attn.to_out[1](out)
-        return out

library/config_util.py

@@ -33,10 +33,8 @@ from . import train_util
 from .train_util import (
   DreamBoothSubset,
   FineTuningSubset,
-  ControlNetSubset,
   DreamBoothDataset,
   FineTuningDataset,
-  ControlNetDataset,
   DatasetGroup,
 )
 
@@ -56,8 +54,6 @@ class BaseSubsetParams:
   flip_aug: bool = False
   face_crop_aug_range: Optional[Tuple[float, float]] = None
   random_crop: bool = False
-  caption_prefix: Optional[str] = None
-  caption_suffix: Optional[str] = None
   caption_dropout_rate: float = 0.0
   caption_dropout_every_n_epochs: int = 0
   caption_tag_dropout_rate: float = 0.0
@@ -74,14 +70,9 @@ class DreamBoothSubsetParams(BaseSubsetParams):
 class FineTuningSubsetParams(BaseSubsetParams):
   metadata_file: Optional[str] = None
 
-@dataclass
-class ControlNetSubsetParams(BaseSubsetParams):
-  conditioning_data_dir: str = None
-  caption_extension: str = ".caption"
-
 @dataclass
 class BaseDatasetParams:
-  tokenizer: Union[CLIPTokenizer, List[CLIPTokenizer]] = None
+  tokenizer: CLIPTokenizer = None
   max_token_length: int = None
   resolution: Optional[Tuple[int, int]] = None
   debug_dataset: bool = False
@@ -105,15 +96,6 @@ class FineTuningDatasetParams(BaseDatasetParams):
   bucket_reso_steps: int = 64
   bucket_no_upscale: bool = False
 
-@dataclass
-class ControlNetDatasetParams(BaseDatasetParams):
-  batch_size: int = 1
-  enable_bucket: bool = False
-  min_bucket_reso: int = 256
-  max_bucket_reso: int = 1024
-  bucket_reso_steps: int = 64
-  bucket_no_upscale: bool = False
-
 @dataclass
 class SubsetBlueprint:
   params: Union[DreamBoothSubsetParams, FineTuningSubsetParams]
@@ -121,7 +103,6 @@ class SubsetBlueprint:
 @dataclass
 class DatasetBlueprint:
   is_dreambooth: bool
-  is_controlnet: bool
   params: Union[DreamBoothDatasetParams, FineTuningDatasetParams]
   subsets: Sequence[SubsetBlueprint]
 
@@ -161,8 +142,6 @@ class ConfigSanitizer:
     "keep_tokens": int,
     "token_warmup_min": int,
     "token_warmup_step": Any(float,int),
-    "caption_prefix": str,
-    "caption_suffix": str,
   }
   # DO means DropOut
   DO_SUBSET_ASCENDABLE_SCHEMA = {
@@ -184,13 +163,6 @@ class ConfigSanitizer:
     Required("metadata_file"): str,
     "image_dir": str,
   }
-  CN_SUBSET_ASCENDABLE_SCHEMA = {
-    "caption_extension": str,
-  }
-  CN_SUBSET_DISTINCT_SCHEMA = {
-    Required("image_dir"): str,
-    Required("conditioning_data_dir"): str,
-  }
 
   # datasets schema
   DATASET_ASCENDABLE_SCHEMA = {
@@ -220,8 +192,8 @@ class ConfigSanitizer:
     "dataset_repeats": "num_repeats",
   }
 
-  def __init__(self, support_dreambooth: bool, support_finetuning: bool, support_controlnet: bool, support_dropout: bool) -> None:
-    assert support_dreambooth or support_finetuning or support_controlnet, "Neither DreamBooth mode nor fine tuning mode specified. Please specify one mode or more. / DreamBooth モードか fine tuning モードのどちらも指定されていません。1つ以上指定してください。"
+  def __init__(self, support_dreambooth: bool, support_finetuning: bool, support_dropout: bool) -> None:
+    assert support_dreambooth or support_finetuning, "Neither DreamBooth mode nor fine tuning mode specified. Please specify one mode or more. / DreamBooth モードか fine tuning モードのどちらも指定されていません。1つ以上指定してください。"
 
     self.db_subset_schema = self.__merge_dict(
       self.SUBSET_ASCENDABLE_SCHEMA,
@@ -236,13 +208,6 @@ class ConfigSanitizer:
       self.DO_SUBSET_ASCENDABLE_SCHEMA if support_dropout else {},
     )
 
-    self.cn_subset_schema = self.__merge_dict(
-      self.SUBSET_ASCENDABLE_SCHEMA,
-      self.CN_SUBSET_DISTINCT_SCHEMA,
-      self.CN_SUBSET_ASCENDABLE_SCHEMA,
-      self.DO_SUBSET_ASCENDABLE_SCHEMA if support_dropout else {},
-    )
-
     self.db_dataset_schema = self.__merge_dict(
       self.DATASET_ASCENDABLE_SCHEMA,
       self.SUBSET_ASCENDABLE_SCHEMA,
@@ -258,23 +223,13 @@ class ConfigSanitizer:
       {"subsets": [self.ft_subset_schema]},
     )
 
-    self.cn_dataset_schema = self.__merge_dict(
-      self.DATASET_ASCENDABLE_SCHEMA,
-      self.SUBSET_ASCENDABLE_SCHEMA,
-      self.CN_SUBSET_ASCENDABLE_SCHEMA,
-      self.DO_SUBSET_ASCENDABLE_SCHEMA if support_dropout else {},
-      {"subsets": [self.cn_subset_schema]},
-    )
-
     if support_dreambooth and support_finetuning:
       def validate_flex_dataset(dataset_config: dict):
         subsets_config = dataset_config.get("subsets", [])
 
-        if support_controlnet and all(["conditioning_data_dir" in subset for subset in subsets_config]):
-          return Schema(self.cn_dataset_schema)(dataset_config)
         # check dataset meets FT style
         # NOTE: all FT subsets should have "metadata_file"
-        elif all(["metadata_file" in subset for subset in subsets_config]):
+        if all(["metadata_file" in subset for subset in subsets_config]):
           return Schema(self.ft_dataset_schema)(dataset_config)
         # check dataset meets DB style
         # NOTE: all DB subsets should have no "metadata_file"
@@ -286,16 +241,13 @@ class ConfigSanitizer:
       self.dataset_schema = validate_flex_dataset
     elif support_dreambooth:
       self.dataset_schema = self.db_dataset_schema
-    elif support_finetuning:
+    else:
       self.dataset_schema = self.ft_dataset_schema
-    elif support_controlnet:
-      self.dataset_schema = self.cn_dataset_schema
 
     self.general_schema = self.__merge_dict(
       self.DATASET_ASCENDABLE_SCHEMA,
       self.SUBSET_ASCENDABLE_SCHEMA,
       self.DB_SUBSET_ASCENDABLE_SCHEMA if support_dreambooth else {},
-      self.CN_SUBSET_ASCENDABLE_SCHEMA if support_controlnet else {},
       self.DO_SUBSET_ASCENDABLE_SCHEMA if support_dropout else {},
     )
 
@@ -366,11 +318,7 @@ class BlueprintGenerator:
       # NOTE: if subsets have no "metadata_file", these are DreamBooth datasets/subsets
       subsets = dataset_config.get("subsets", [])
       is_dreambooth = all(["metadata_file" not in subset for subset in subsets])
-      is_controlnet = all(["conditioning_data_dir" in subset for subset in subsets])
-      if is_controlnet:
-        subset_params_klass = ControlNetSubsetParams
-        dataset_params_klass = ControlNetDatasetParams
-      elif is_dreambooth:
+      if is_dreambooth:
         subset_params_klass = DreamBoothSubsetParams
         dataset_params_klass = DreamBoothDatasetParams
       else:
@@ -385,7 +333,7 @@ class BlueprintGenerator:
 
       params = self.generate_params_by_fallbacks(dataset_params_klass,
                                                  [dataset_config, general_config, argparse_config, runtime_params])
-      dataset_blueprints.append(DatasetBlueprint(is_dreambooth, is_controlnet, params, subset_blueprints))
+      dataset_blueprints.append(DatasetBlueprint(is_dreambooth, params, subset_blueprints))
 
     dataset_group_blueprint = DatasetGroupBlueprint(dataset_blueprints)
 
@@ -413,13 +361,10 @@ class BlueprintGenerator:
 
 
 def generate_dataset_group_by_blueprint(dataset_group_blueprint: DatasetGroupBlueprint):
-  datasets: List[Union[DreamBoothDataset, FineTuningDataset, ControlNetDataset]] = []
+  datasets: List[Union[DreamBoothDataset, FineTuningDataset]] = []
 
   for dataset_blueprint in dataset_group_blueprint.datasets:
-    if dataset_blueprint.is_controlnet:
-      subset_klass = ControlNetSubset
-      dataset_klass = ControlNetDataset
-    elif dataset_blueprint.is_dreambooth:
+    if dataset_blueprint.is_dreambooth:
       subset_klass = DreamBoothSubset
       dataset_klass = DreamBoothDataset
     else:
@@ -434,7 +379,6 @@ def generate_dataset_group_by_blueprint(dataset_group_blueprint: DatasetGroupBlu
   info = ""
   for i, dataset in enumerate(datasets):
     is_dreambooth = isinstance(dataset, DreamBoothDataset)
-    is_controlnet = isinstance(dataset, ControlNetDataset)
     info += dedent(f"""\
       [Dataset {i}]
         batch_size: {dataset.batch_size}
@@ -463,8 +407,6 @@ def generate_dataset_group_by_blueprint(dataset_group_blueprint: DatasetGroupBlu
           caption_dropout_rate: {subset.caption_dropout_rate}
           caption_dropout_every_n_epoches: {subset.caption_dropout_every_n_epochs}
           caption_tag_dropout_rate: {subset.caption_tag_dropout_rate}
-          caption_prefix: {subset.caption_prefix}
-          caption_suffix: {subset.caption_suffix}
           color_aug: {subset.color_aug}
           flip_aug: {subset.flip_aug}
           face_crop_aug_range: {subset.face_crop_aug_range}
@@ -479,7 +421,7 @@ def generate_dataset_group_by_blueprint(dataset_group_blueprint: DatasetGroupBlu
           class_tokens: {subset.class_tokens}
           caption_extension: {subset.caption_extension}
         \n"""), "    ")
-      elif not is_controlnet:
+      else:
         info += indent(dedent(f"""\
           metadata_file: {subset.metadata_file}
         \n"""), "    ")
@@ -537,31 +479,6 @@ def generate_dreambooth_subsets_config_by_subdirs(train_data_dir: Optional[str]
   return subsets_config
 
 
-def generate_controlnet_subsets_config_by_subdirs(train_data_dir: Optional[str] = None, conditioning_data_dir: Optional[str] = None, caption_extension: str = ".txt"):
-  def generate(base_dir: Optional[str]):
-    if base_dir is None:
-      return []
-
-    base_dir: Path = Path(base_dir)
-    if not base_dir.is_dir():
-      return []
-
-    subsets_config = []
-    for subdir in base_dir.iterdir():
-      if not subdir.is_dir():
-        continue
-
-      subset_config = {"image_dir": str(subdir), "conditioning_data_dir": conditioning_data_dir, "caption_extension": caption_extension, "num_repeats": 1}
-      subsets_config.append(subset_config)
-
-    return subsets_config
-
-  subsets_config = []
-  subsets_config += generate(train_data_dir, False)
-
-  return subsets_config
-
-
 def load_user_config(file: str) -> dict:
   file: Path = Path(file)
   if not file.is_file():
@@ -590,7 +507,6 @@ if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   parser.add_argument("--support_dreambooth", action="store_true")
   parser.add_argument("--support_finetuning", action="store_true")
-  parser.add_argument("--support_controlnet", action="store_true")
   parser.add_argument("--support_dropout", action="store_true")
   parser.add_argument("dataset_config")
   config_args, remain = parser.parse_known_args()
@@ -609,7 +525,7 @@ if __name__ == "__main__":
   print("\n[user_config]")
   print(user_config)
 
-  sanitizer = ConfigSanitizer(config_args.support_dreambooth, config_args.support_finetuning, config_args.support_controlnet, config_args.support_dropout)
+  sanitizer = ConfigSanitizer(config_args.support_dreambooth, config_args.support_finetuning, config_args.support_dropout)
   sanitized_user_config = sanitizer.sanitize_user_config(user_config)
 
   print("\n[sanitized_user_config]")

library/custom_train_functions.py

@@ -19,71 +19,20 @@ def prepare_scheduler_for_custom_training(noise_scheduler, device):
     noise_scheduler.all_snr = all_snr.to(device)
 
 
-def fix_noise_scheduler_betas_for_zero_terminal_snr(noise_scheduler):
-    # fix beta: zero terminal SNR
-    print(f"fix noise scheduler betas: https://arxiv.org/abs/2305.08891")
-
-    def enforce_zero_terminal_snr(betas):
-        # Convert betas to alphas_bar_sqrt
-        alphas = 1 - betas
-        alphas_bar = alphas.cumprod(0)
-        alphas_bar_sqrt = alphas_bar.sqrt()
-
-        # Store old values.
-        alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
-        alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
-        # Shift so last timestep is zero.
-        alphas_bar_sqrt -= alphas_bar_sqrt_T
-        # Scale so first timestep is back to old value.
-        alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
-
-        # Convert alphas_bar_sqrt to betas
-        alphas_bar = alphas_bar_sqrt**2
-        alphas = alphas_bar[1:] / alphas_bar[:-1]
-        alphas = torch.cat([alphas_bar[0:1], alphas])
-        betas = 1 - alphas
-        return betas
-
-    betas = noise_scheduler.betas
-    betas = enforce_zero_terminal_snr(betas)
-    alphas = 1.0 - betas
-    alphas_cumprod = torch.cumprod(alphas, dim=0)
-
-    # print("original:", noise_scheduler.betas)
-    # print("fixed:", betas)
-
-    noise_scheduler.betas = betas
-    noise_scheduler.alphas = alphas
-    noise_scheduler.alphas_cumprod = alphas_cumprod
-
-
 def apply_snr_weight(loss, timesteps, noise_scheduler, gamma):
     snr = torch.stack([noise_scheduler.all_snr[t] for t in timesteps])
     gamma_over_snr = torch.div(torch.ones_like(snr) * gamma, snr)
-    snr_weight = torch.minimum(gamma_over_snr, torch.ones_like(gamma_over_snr)).float().to(loss.device)  # from paper
+    snr_weight = torch.minimum(gamma_over_snr, torch.ones_like(gamma_over_snr)).float()  # from paper
     loss = loss * snr_weight
     return loss
 
 
 def scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler):
-    scale = get_snr_scale(timesteps, noise_scheduler)
-    loss = loss * scale
-    return loss
-
-
-def get_snr_scale(timesteps, noise_scheduler):
     snr_t = torch.stack([noise_scheduler.all_snr[t] for t in timesteps])  # batch_size
     snr_t = torch.minimum(snr_t, torch.ones_like(snr_t) * 1000)  # if timestep is 0, snr_t is inf, so limit it to 1000
     scale = snr_t / (snr_t + 1)
-    # # show debug info
-    # print(f"timesteps: {timesteps}, snr_t: {snr_t}, scale: {scale}")
-    return scale
 
-
-def add_v_prediction_like_loss(loss, timesteps, noise_scheduler, v_pred_like_loss):
-    scale = get_snr_scale(timesteps, noise_scheduler)
-    # print(f"add v-prediction like loss: {v_pred_like_loss}, scale: {scale}, loss: {loss}, time: {timesteps}")
-    loss = loss + loss / scale * v_pred_like_loss
+    loss = loss * scale
     return loss
 
 
@@ -102,12 +51,6 @@ def add_custom_train_arguments(parser: argparse.ArgumentParser, support_weighted
         action="store_true",
         help="scale v-prediction loss like noise prediction loss / v-prediction lossをnoise prediction lossと同じようにスケーリングする",
     )
-    parser.add_argument(
-        "--v_pred_like_loss",
-        type=float,
-        default=None,
-        help="add v-prediction like loss multiplied by this value / v-prediction lossをこの値をかけたものをlossに加算する",
-    )
     if support_weighted_captions:
         parser.add_argument(
             "--weighted_captions",

library/huggingface_util.py

@@ -26,7 +26,7 @@ def upload(
     repo_id = args.huggingface_repo_id
     repo_type = args.huggingface_repo_type
     token = args.huggingface_token
-    path_in_repo = args.huggingface_path_in_repo + dest_suffix if args.huggingface_path_in_repo is not None else None
+    path_in_repo = args.huggingface_path_in_repo + dest_suffix
     private = args.huggingface_repo_visibility is None or args.huggingface_repo_visibility != "public"
     api = HfApi(token=token)
     if not exists_repo(repo_id=repo_id, repo_type=repo_type, token=token):

library/hypernetwork.py

@@ -1,223 +0,0 @@
-import torch
-import torch.nn.functional as F
-from diffusers.models.attention_processor import (
-    Attention,
-    AttnProcessor2_0,
-    SlicedAttnProcessor,
-    XFormersAttnProcessor
-)
-
-try:
-    import xformers.ops
-except:
-    xformers = None
-
-
-loaded_networks = []
-
-
-def apply_single_hypernetwork(
-    hypernetwork, hidden_states, encoder_hidden_states
-):
-    context_k, context_v = hypernetwork.forward(hidden_states, encoder_hidden_states)
-    return context_k, context_v
-
-
-def apply_hypernetworks(context_k, context_v, layer=None):
-    if len(loaded_networks) == 0:
-        return context_v, context_v
-    for hypernetwork in loaded_networks:
-        context_k, context_v = hypernetwork.forward(context_k, context_v)
-
-    context_k = context_k.to(dtype=context_k.dtype)
-    context_v = context_v.to(dtype=context_k.dtype)
-
-    return context_k, context_v
-
-
-
-def xformers_forward(
-    self: XFormersAttnProcessor,
-    attn: Attention,
-    hidden_states: torch.Tensor,
-    encoder_hidden_states: torch.Tensor = None,
-    attention_mask: torch.Tensor = None,
-):
-    batch_size, sequence_length, _ = (
-        hidden_states.shape
-        if encoder_hidden_states is None
-        else encoder_hidden_states.shape
-    )
-
-    attention_mask = attn.prepare_attention_mask(
-        attention_mask, sequence_length, batch_size
-    )
-
-    query = attn.to_q(hidden_states)
-
-    if encoder_hidden_states is None:
-        encoder_hidden_states = hidden_states
-    elif attn.norm_cross:
-        encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
-
-    context_k, context_v = apply_hypernetworks(hidden_states, encoder_hidden_states)
-
-    key = attn.to_k(context_k)
-    value = attn.to_v(context_v)
-
-    query = attn.head_to_batch_dim(query).contiguous()
-    key = attn.head_to_batch_dim(key).contiguous()
-    value = attn.head_to_batch_dim(value).contiguous()
-
-    hidden_states = xformers.ops.memory_efficient_attention(
-        query,
-        key,
-        value,
-        attn_bias=attention_mask,
-        op=self.attention_op,
-        scale=attn.scale,
-    )
-    hidden_states = hidden_states.to(query.dtype)
-    hidden_states = attn.batch_to_head_dim(hidden_states)
-
-    # linear proj
-    hidden_states = attn.to_out[0](hidden_states)
-    # dropout
-    hidden_states = attn.to_out[1](hidden_states)
-    return hidden_states
-
-
-def sliced_attn_forward(
-    self: SlicedAttnProcessor,
-    attn: Attention,
-    hidden_states: torch.Tensor,
-    encoder_hidden_states: torch.Tensor = None,
-    attention_mask: torch.Tensor = None,
-):
-    batch_size, sequence_length, _ = (
-        hidden_states.shape
-        if encoder_hidden_states is None
-        else encoder_hidden_states.shape
-    )
-    attention_mask = attn.prepare_attention_mask(
-        attention_mask, sequence_length, batch_size
-    )
-
-    query = attn.to_q(hidden_states)
-    dim = query.shape[-1]
-    query = attn.head_to_batch_dim(query)
-
-    if encoder_hidden_states is None:
-        encoder_hidden_states = hidden_states
-    elif attn.norm_cross:
-        encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
-
-    context_k, context_v = apply_hypernetworks(hidden_states, encoder_hidden_states)
-
-    key = attn.to_k(context_k)
-    value = attn.to_v(context_v)
-    key = attn.head_to_batch_dim(key)
-    value = attn.head_to_batch_dim(value)
-
-    batch_size_attention, query_tokens, _ = query.shape
-    hidden_states = torch.zeros(
-        (batch_size_attention, query_tokens, dim // attn.heads),
-        device=query.device,
-        dtype=query.dtype,
-    )
-
-    for i in range(batch_size_attention // self.slice_size):
-        start_idx = i * self.slice_size
-        end_idx = (i + 1) * self.slice_size
-
-        query_slice = query[start_idx:end_idx]
-        key_slice = key[start_idx:end_idx]
-        attn_mask_slice = (
-            attention_mask[start_idx:end_idx] if attention_mask is not None else None
-        )
-
-        attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice)
-
-        attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx])
-
-        hidden_states[start_idx:end_idx] = attn_slice
-
-    hidden_states = attn.batch_to_head_dim(hidden_states)
-
-    # linear proj
-    hidden_states = attn.to_out[0](hidden_states)
-    # dropout
-    hidden_states = attn.to_out[1](hidden_states)
-
-    return hidden_states
-
-
-def v2_0_forward(
-    self: AttnProcessor2_0,
-    attn: Attention,
-    hidden_states,
-    encoder_hidden_states=None,
-    attention_mask=None,
-):
-    batch_size, sequence_length, _ = (
-        hidden_states.shape
-        if encoder_hidden_states is None
-        else encoder_hidden_states.shape
-    )
-    inner_dim = hidden_states.shape[-1]
-
-    if attention_mask is not None:
-        attention_mask = attn.prepare_attention_mask(
-            attention_mask, sequence_length, batch_size
-        )
-        # scaled_dot_product_attention expects attention_mask shape to be
-        # (batch, heads, source_length, target_length)
-        attention_mask = attention_mask.view(
-            batch_size, attn.heads, -1, attention_mask.shape[-1]
-        )
-
-    query = attn.to_q(hidden_states)
-
-    if encoder_hidden_states is None:
-        encoder_hidden_states = hidden_states
-    elif attn.norm_cross:
-        encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
-
-    context_k, context_v = apply_hypernetworks(hidden_states, encoder_hidden_states)
-
-    key = attn.to_k(context_k)
-    value = attn.to_v(context_v)
-
-    head_dim = inner_dim // attn.heads
-    query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-    key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-    value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-
-    # the output of sdp = (batch, num_heads, seq_len, head_dim)
-    # TODO: add support for attn.scale when we move to Torch 2.1
-    hidden_states = F.scaled_dot_product_attention(
-        query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
-    )
-
-    hidden_states = hidden_states.transpose(1, 2).reshape(
-        batch_size, -1, attn.heads * head_dim
-    )
-    hidden_states = hidden_states.to(query.dtype)
-
-    # linear proj
-    hidden_states = attn.to_out[0](hidden_states)
-    # dropout
-    hidden_states = attn.to_out[1](hidden_states)
-    return hidden_states
-
-
-def replace_attentions_for_hypernetwork():
-    import diffusers.models.attention_processor
-
-    diffusers.models.attention_processor.XFormersAttnProcessor.__call__ = (
-        xformers_forward
-    )
-    diffusers.models.attention_processor.SlicedAttnProcessor.__call__ = (
-        sliced_attn_forward
-    )
-    diffusers.models.attention_processor.AttnProcessor2_0.__call__ = v2_0_forward

library/lpw_stable_diffusion.py

@@ -6,7 +6,7 @@ import re
 from typing import Callable, List, Optional, Union
 
 import numpy as np
-import PIL.Image
+import PIL
 import torch
 from packaging import version
 from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
@@ -426,58 +426,6 @@ def preprocess_mask(mask, scale_factor=8):
     return mask
 
 
-def prepare_controlnet_image(
-    image: PIL.Image.Image,
-    width: int,
-    height: int,
-    batch_size: int,
-    num_images_per_prompt: int,
-    device: torch.device,
-    dtype: torch.dtype,
-    do_classifier_free_guidance: bool = False,
-    guess_mode: bool = False,
-):
-    if not isinstance(image, torch.Tensor):
-        if isinstance(image, PIL.Image.Image):
-            image = [image]
-
-        if isinstance(image[0], PIL.Image.Image):
-            images = []
-
-            for image_ in image:
-                image_ = image_.convert("RGB")
-                image_ = image_.resize((width, height), resample=PIL_INTERPOLATION["lanczos"])
-                image_ = np.array(image_)
-                image_ = image_[None, :]
-                images.append(image_)
-
-            image = images
-
-            image = np.concatenate(image, axis=0)
-            image = np.array(image).astype(np.float32) / 255.0
-            image = image.transpose(0, 3, 1, 2)
-            image = torch.from_numpy(image)
-        elif isinstance(image[0], torch.Tensor):
-            image = torch.cat(image, dim=0)
-
-    image_batch_size = image.shape[0]
-
-    if image_batch_size == 1:
-        repeat_by = batch_size
-    else:
-        # image batch size is the same as prompt batch size
-        repeat_by = num_images_per_prompt
-
-    image = image.repeat_interleave(repeat_by, dim=0)
-
-    image = image.to(device=device, dtype=dtype)
-
-    if do_classifier_free_guidance and not guess_mode:
-        image = torch.cat([image] * 2)
-
-    return image
-
-
 class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
     r"""
     Pipeline for text-to-image generation using Stable Diffusion without tokens length limit, and support parsing
@@ -516,11 +464,10 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
         tokenizer: CLIPTokenizer,
         unet: UNet2DConditionModel,
         scheduler: SchedulerMixin,
-        # clip_skip: int,
+        clip_skip: int,
         safety_checker: StableDiffusionSafetyChecker,
         feature_extractor: CLIPFeatureExtractor,
         requires_safety_checker: bool = True,
-        clip_skip: int = 1,
     ):
         super().__init__(
             vae=vae,
@@ -760,8 +707,6 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
         max_embeddings_multiples: Optional[int] = 3,
         output_type: Optional[str] = "pil",
         return_dict: bool = True,
-        controlnet=None,
-        controlnet_image=None,
         callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
         is_cancelled_callback: Optional[Callable[[], bool]] = None,
         callback_steps: int = 1,
@@ -822,11 +767,6 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
                 plain tuple.
-            controlnet (`diffusers.ControlNetModel`, *optional*):
-                A controlnet model to be used for the inference. If not provided, controlnet will be disabled.
-            controlnet_image (`torch.FloatTensor` or `PIL.Image.Image`, *optional*):
-                `Image`, or tensor representing an image batch, to be used as the starting point for the controlnet
-                inference.
             callback (`Callable`, *optional*):
                 A function that will be called every `callback_steps` steps during inference. The function will be
                 called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
@@ -845,9 +785,6 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
             list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
             (nsfw) content, according to the `safety_checker`.
         """
-        if controlnet is not None and controlnet_image is None:
-            raise ValueError("controlnet_image must be provided if controlnet is not None.")
-
         # 0. Default height and width to unet
         height = height or self.unet.config.sample_size * self.vae_scale_factor
         width = width or self.unet.config.sample_size * self.vae_scale_factor
@@ -887,11 +824,6 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
         else:
             mask = None
 
-        if controlnet_image is not None:
-            controlnet_image = prepare_controlnet_image(
-                controlnet_image, width, height, batch_size, 1, self.device, controlnet.dtype, do_classifier_free_guidance, False
-            )
-
         # 5. set timesteps
         self.scheduler.set_timesteps(num_inference_steps, device=device)
         timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device, image is None)
@@ -919,22 +851,8 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
             latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
             latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
 
-            unet_additional_args = {}
-            if controlnet is not None:
-                down_block_res_samples, mid_block_res_sample = controlnet(
-                    latent_model_input,
-                    t,
-                    encoder_hidden_states=text_embeddings,
-                    controlnet_cond=controlnet_image,
-                    conditioning_scale=1.0,
-                    guess_mode=False,
-                    return_dict=False,
-                )
-                unet_additional_args["down_block_additional_residuals"] = down_block_res_samples
-                unet_additional_args["mid_block_additional_residual"] = mid_block_res_sample
-
             # predict the noise residual
-            noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings, **unet_additional_args).sample
+            noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
 
             # perform guidance
             if do_classifier_free_guidance:
@@ -956,13 +874,20 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
                 if is_cancelled_callback is not None and is_cancelled_callback():
                     return None
 
-        return latents
-
-    def latents_to_image(self, latents):
         # 9. Post-processing
-        image = self.decode_latents(latents.to(self.vae.dtype))
-        image = self.numpy_to_pil(image)
-        return image
+        image = self.decode_latents(latents)
+
+        # 10. Run safety checker
+        image, has_nsfw_concept = self.run_safety_checker(image, device, text_embeddings.dtype)
+
+        # 11. Convert to PIL
+        if output_type == "pil":
+            image = self.numpy_to_pil(image)
+
+        if not return_dict:
+            return image, has_nsfw_concept
+
+        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
 
     def text2img(
         self,

library/model_util.py

@@ -4,11 +4,9 @@
 import math
 import os
 import torch
-import diffusers
 from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig, logging
-from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline  # , UNet2DConditionModel
+from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNet2DConditionModel
 from safetensors.torch import load_file, save_file
-from library.original_unet import UNet2DConditionModel
 
 # DiffUsers版StableDiffusionのモデルパラメータ
 NUM_TRAIN_TIMESTEPS = 1000
@@ -128,30 +126,17 @@ def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0):
         new_item = new_item.replace("norm.weight", "group_norm.weight")
         new_item = new_item.replace("norm.bias", "group_norm.bias")
 
-        if diffusers.__version__ < "0.17.0":
-            new_item = new_item.replace("q.weight", "query.weight")
-            new_item = new_item.replace("q.bias", "query.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("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("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")
-        else:
-            new_item = new_item.replace("q.weight", "to_q.weight")
-            new_item = new_item.replace("q.bias", "to_q.bias")
-
-            new_item = new_item.replace("k.weight", "to_k.weight")
-            new_item = new_item.replace("k.bias", "to_k.bias")
-
-            new_item = new_item.replace("v.weight", "to_v.weight")
-            new_item = new_item.replace("v.bias", "to_v.bias")
-
-            new_item = new_item.replace("proj_out.weight", "to_out.0.weight")
-            new_item = new_item.replace("proj_out.bias", "to_out.0.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)
 
@@ -206,16 +191,8 @@ def assign_to_checkpoint(
                 new_path = new_path.replace(replacement["old"], replacement["new"])
 
         # proj_attn.weight has to be converted from conv 1D to linear
-        reshaping = False
-        if diffusers.__version__ < "0.17.0":
-            if "proj_attn.weight" in new_path:
-                reshaping = True
-        else:
-            if ".attentions." in new_path and ".0.to_" in new_path and old_checkpoint[path["old"]].ndim > 2:
-                reshaping = True
-
-        if reshaping:
-            checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0, 0]
+        if "proj_attn.weight" in new_path:
+            checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0]
         else:
             checkpoint[new_path] = old_checkpoint[path["old"]]
 
@@ -384,7 +361,7 @@ def convert_ldm_unet_checkpoint(v2, checkpoint, config):
 
     # SDのv2では1*1のconv2dがlinearに変わっている
     # 誤って Diffusers 側を conv2d のままにしてしまったので、変換必要
-    if v2 and not config.get("use_linear_projection", False):
+    if v2 and not config.get('use_linear_projection', False):
         linear_transformer_to_conv(new_checkpoint)
 
     return new_checkpoint
@@ -563,10 +540,10 @@ def convert_ldm_clip_checkpoint_v1(checkpoint):
     for key in keys:
         if key.startswith("cond_stage_model.transformer"):
             text_model_dict[key[len("cond_stage_model.transformer.") :]] = checkpoint[key]
-
+    
     # support checkpoint without position_ids (invalid checkpoint)
     if "text_model.embeddings.position_ids" not in text_model_dict:
-        text_model_dict["text_model.embeddings.position_ids"] = torch.arange(77).unsqueeze(0)  # 77 is the max length of the text
+        text_model_dict["text_model.embeddings.position_ids"] = torch.arange(77).unsqueeze(0) # 77 is the max length of the text
 
     return text_model_dict
 
@@ -760,105 +737,6 @@ def convert_unet_state_dict_to_sd(v2, unet_state_dict):
     return new_state_dict
 
 
-def controlnet_conversion_map():
-    unet_conversion_map = [
-        ("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"),
-        ("middle_block_out.0.weight", "controlnet_mid_block.weight"),
-        ("middle_block_out.0.bias", "controlnet_mid_block.bias"),
-    ]
-
-    unet_conversion_map_resnet = [
-        ("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):
-        for j in range(2):
-            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:
-                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))
-
-        if i < 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))
-
-    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))
-
-    controlnet_cond_embedding_names = ["conv_in"] + [f"blocks.{i}" for i in range(6)] + ["conv_out"]
-    for i, hf_prefix in enumerate(controlnet_cond_embedding_names):
-        hf_prefix = f"controlnet_cond_embedding.{hf_prefix}."
-        sd_prefix = f"input_hint_block.{i*2}."
-        unet_conversion_map_layer.append((sd_prefix, hf_prefix))
-
-    for i in range(12):
-        hf_prefix = f"controlnet_down_blocks.{i}."
-        sd_prefix = f"zero_convs.{i}.0."
-        unet_conversion_map_layer.append((sd_prefix, hf_prefix))
-
-    return unet_conversion_map, unet_conversion_map_resnet, unet_conversion_map_layer
-
-
-def convert_controlnet_state_dict_to_sd(controlnet_state_dict):
-    unet_conversion_map, unet_conversion_map_resnet, unet_conversion_map_layer = controlnet_conversion_map()
-
-    mapping = {k: k for k in controlnet_state_dict.keys()}
-    for sd_name, diffusers_name in unet_conversion_map:
-        mapping[diffusers_name] = sd_name
-    for k, v in mapping.items():
-        if "resnets" in k:
-            for sd_part, diffusers_part in unet_conversion_map_resnet:
-                v = v.replace(diffusers_part, sd_part)
-            mapping[k] = v
-    for k, v in mapping.items():
-        for sd_part, diffusers_part in unet_conversion_map_layer:
-            v = v.replace(diffusers_part, sd_part)
-        mapping[k] = v
-    new_state_dict = {v: controlnet_state_dict[k] for k, v in mapping.items()}
-    return new_state_dict
-
-
-def convert_controlnet_state_dict_to_diffusers(controlnet_state_dict):
-    unet_conversion_map, unet_conversion_map_resnet, unet_conversion_map_layer = controlnet_conversion_map()
-
-    mapping = {k: k for k in controlnet_state_dict.keys()}
-    for sd_name, diffusers_name in unet_conversion_map:
-        mapping[sd_name] = diffusers_name
-    for k, v in mapping.items():
-        for sd_part, diffusers_part in unet_conversion_map_layer:
-            v = v.replace(sd_part, diffusers_part)
-        mapping[k] = v
-    for k, v in mapping.items():
-        if "resnets" in v:
-            for sd_part, diffusers_part in unet_conversion_map_resnet:
-                v = v.replace(sd_part, diffusers_part)
-            mapping[k] = v
-    new_state_dict = {v: controlnet_state_dict[k] for k, v in mapping.items()}
-    return new_state_dict
-
-
 # ================#
 # VAE Conversion #
 # ================#
@@ -906,24 +784,14 @@ def convert_vae_state_dict(vae_state_dict):
         sd_mid_res_prefix = f"mid.block_{i+1}."
         vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))
 
-    if diffusers.__version__ < "0.17.0":
-        vae_conversion_map_attn = [
-            # (stable-diffusion, HF Diffusers)
-            ("norm.", "group_norm."),
-            ("q.", "query."),
-            ("k.", "key."),
-            ("v.", "value."),
-            ("proj_out.", "proj_attn."),
-        ]
-    else:
-        vae_conversion_map_attn = [
-            # (stable-diffusion, HF Diffusers)
-            ("norm.", "group_norm."),
-            ("q.", "to_q."),
-            ("k.", "to_k."),
-            ("v.", "to_v."),
-            ("proj_out.", "to_out.0."),
-        ]
+    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():
@@ -940,7 +808,7 @@ def convert_vae_state_dict(vae_state_dict):
     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: shape {v.shape} -> {v.shape} x 1 x 1")
+                # print(f"Reshaping {k} for SD format")
                 new_state_dict[k] = reshape_weight_for_sd(v)
 
     return new_state_dict
@@ -989,7 +857,7 @@ def load_checkpoint_with_text_encoder_conversion(ckpt_path, device="cpu"):
 
 
 # TODO dtype指定の動作が怪しいので確認する text_encoderを指定形式で作れるか未確認
-def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, device="cpu", dtype=None, unet_use_linear_projection_in_v2=True):
+def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, device="cpu", dtype=None, unet_use_linear_projection_in_v2=False):
     _, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path, device)
 
     # Convert the UNet2DConditionModel model.
@@ -1037,49 +905,16 @@ def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, device="cpu", dt
     else:
         converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v1(state_dict)
 
-        # logging.set_verbosity_error()  # don't show annoying warning
-        # text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to(device)
-        # logging.set_verbosity_warning()
-        # print(f"config: {text_model.config}")
-        cfg = CLIPTextConfig(
-            vocab_size=49408,
-            hidden_size=768,
-            intermediate_size=3072,
-            num_hidden_layers=12,
-            num_attention_heads=12,
-            max_position_embeddings=77,
-            hidden_act="quick_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=768,
-            torch_dtype="float32",
-        )
-        text_model = CLIPTextModel._from_config(cfg)
+        logging.set_verbosity_error()  # don't show annoying warning
+        text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to(device)
+        logging.set_verbosity_warning()
+
         info = text_model.load_state_dict(converted_text_encoder_checkpoint)
     print("loading text encoder:", info)
 
     return text_model, vae, unet
 
 
-def get_model_version_str_for_sd1_sd2(v2, v_parameterization):
-    # only for reference
-    version_str = "sd"
-    if v2:
-        version_str += "_v2"
-    else:
-        version_str += "_v1"
-    if v_parameterization:
-        version_str += "_v"
-    return version_str
-
-
 def convert_text_encoder_state_dict_to_sd_v2(checkpoint, make_dummy_weights=False):
     def convert_key(key):
         # position_idsの除去
@@ -1151,9 +986,7 @@ def convert_text_encoder_state_dict_to_sd_v2(checkpoint, make_dummy_weights=Fals
     return new_sd
 
 
-def save_stable_diffusion_checkpoint(
-    v2, output_file, text_encoder, unet, ckpt_path, epochs, steps, metadata, save_dtype=None, vae=None
-):
+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)
@@ -1215,7 +1048,7 @@ def save_stable_diffusion_checkpoint(
 
     if is_safetensors(output_file):
         # TODO Tensor以外のdictの値を削除したほうがいいか
-        save_file(state_dict, output_file, metadata)
+        save_file(state_dict, output_file)
     else:
         torch.save(new_ckpt, output_file)
 

library/sai_model_spec.py

@@ -1,305 +0,0 @@
-# based on https://github.com/Stability-AI/ModelSpec
-import datetime
-import hashlib
-from io import BytesIO
-import os
-from typing import List, Optional, Tuple, Union
-import safetensors
-
-r"""
-# Metadata Example
-metadata = {
-    # === Must ===
-    "modelspec.sai_model_spec": "1.0.0", # Required version ID for the spec
-    "modelspec.architecture": "stable-diffusion-xl-v1-base", # Architecture, reference the ID of the original model of the arch to match the ID
-    "modelspec.implementation": "sgm",
-    "modelspec.title": "Example Model Version 1.0", # Clean, human-readable title. May use your own phrasing/language/etc
-    # === Should ===
-    "modelspec.author": "Example Corp", # Your name or company name
-    "modelspec.description": "This is my example model to show you how to do it!", # Describe the model in your own words/language/etc. Focus on what users need to know
-    "modelspec.date": "2023-07-20", # ISO-8601 compliant date of when the model was created
-    # === Can ===
-    "modelspec.license": "ExampleLicense-1.0", # eg CreativeML Open RAIL, etc.
-    "modelspec.usage_hint": "Use keyword 'example'" # In your own language, very short hints about how the user should use the model
-}
-"""
-
-BASE_METADATA = {
-    # === Must ===
-    "modelspec.sai_model_spec": "1.0.0",  # Required version ID for the spec
-    "modelspec.architecture": None,
-    "modelspec.implementation": None,
-    "modelspec.title": None,
-    "modelspec.resolution": None,
-    # === Should ===
-    "modelspec.description": None,
-    "modelspec.author": None,
-    "modelspec.date": None,
-    # === Can ===
-    "modelspec.license": None,
-    "modelspec.tags": None,
-    "modelspec.merged_from": None,
-    "modelspec.prediction_type": None,
-    "modelspec.timestep_range": None,
-    "modelspec.encoder_layer": None,
-}
-
-# 別に使うやつだけ定義
-MODELSPEC_TITLE = "modelspec.title"
-
-ARCH_SD_V1 = "stable-diffusion-v1"
-ARCH_SD_V2_512 = "stable-diffusion-v2-512"
-ARCH_SD_V2_768_V = "stable-diffusion-v2-768-v"
-ARCH_SD_XL_V1_BASE = "stable-diffusion-xl-v1-base"
-
-ADAPTER_LORA = "lora"
-ADAPTER_TEXTUAL_INVERSION = "textual-inversion"
-
-IMPL_STABILITY_AI = "https://github.com/Stability-AI/generative-models"
-IMPL_DIFFUSERS = "diffusers"
-
-PRED_TYPE_EPSILON = "epsilon"
-PRED_TYPE_V = "v"
-
-
-def load_bytes_in_safetensors(tensors):
-    bytes = safetensors.torch.save(tensors)
-    b = BytesIO(bytes)
-
-    b.seek(0)
-    header = b.read(8)
-    n = int.from_bytes(header, "little")
-
-    offset = n + 8
-    b.seek(offset)
-
-    return b.read()
-
-
-def precalculate_safetensors_hashes(state_dict):
-    # calculate each tensor one by one to reduce memory usage
-    hash_sha256 = hashlib.sha256()
-    for tensor in state_dict.values():
-        single_tensor_sd = {"tensor": tensor}
-        bytes_for_tensor = load_bytes_in_safetensors(single_tensor_sd)
-        hash_sha256.update(bytes_for_tensor)
-
-    return f"0x{hash_sha256.hexdigest()}"
-
-
-def update_hash_sha256(metadata: dict, state_dict: dict):
-    raise NotImplementedError
-
-
-def build_metadata(
-    state_dict: Optional[dict],
-    v2: bool,
-    v_parameterization: bool,
-    sdxl: bool,
-    lora: bool,
-    textual_inversion: bool,
-    timestamp: float,
-    title: Optional[str] = None,
-    reso: Optional[Union[int, Tuple[int, int]]] = None,
-    is_stable_diffusion_ckpt: Optional[bool] = None,
-    author: Optional[str] = None,
-    description: Optional[str] = None,
-    license: Optional[str] = None,
-    tags: Optional[str] = None,
-    merged_from: Optional[str] = None,
-    timesteps: Optional[Tuple[int, int]] = None,
-    clip_skip: Optional[int] = None,
-):
-    # if state_dict is None, hash is not calculated
-
-    metadata = {}
-    metadata.update(BASE_METADATA)
-
-    # TODO メモリを消費せずかつ正しいハッシュ計算の方法がわかったら実装する
-    # if state_dict is not None:
-    # hash = precalculate_safetensors_hashes(state_dict)
-    # metadata["modelspec.hash_sha256"] = hash
-
-    if sdxl:
-        arch = ARCH_SD_XL_V1_BASE
-    elif v2:
-        if v_parameterization:
-            arch = ARCH_SD_V2_768_V
-        else:
-            arch = ARCH_SD_V2_512
-    else:
-        arch = ARCH_SD_V1
-
-    if lora:
-        arch += f"/{ADAPTER_LORA}"
-    elif textual_inversion:
-        arch += f"/{ADAPTER_TEXTUAL_INVERSION}"
-
-    metadata["modelspec.architecture"] = arch
-
-    if not lora and not textual_inversion and is_stable_diffusion_ckpt is None:
-        is_stable_diffusion_ckpt = True # default is stable diffusion ckpt if not lora and not textual_inversion
-
-    if (lora and sdxl) or textual_inversion or is_stable_diffusion_ckpt:
-        # Stable Diffusion ckpt, TI, SDXL LoRA
-        impl = IMPL_STABILITY_AI
-    else:
-        # v1/v2 LoRA or Diffusers
-        impl = IMPL_DIFFUSERS
-    metadata["modelspec.implementation"] = impl
-
-    if title is None:
-        if lora:
-            title = "LoRA"
-        elif textual_inversion:
-            title = "TextualInversion"
-        else:
-            title = "Checkpoint"
-        title += f"@{timestamp}"
-    metadata[MODELSPEC_TITLE] = title
-
-    if author is not None:
-        metadata["modelspec.author"] = author
-    else:
-        del metadata["modelspec.author"]
-
-    if description is not None:
-        metadata["modelspec.description"] = description
-    else:
-        del metadata["modelspec.description"]
-
-    if merged_from is not None:
-        metadata["modelspec.merged_from"] = merged_from
-    else:
-        del metadata["modelspec.merged_from"]
-
-    if license is not None:
-        metadata["modelspec.license"] = license
-    else:
-        del metadata["modelspec.license"]
-
-    if tags is not None:
-        metadata["modelspec.tags"] = tags
-    else:
-        del metadata["modelspec.tags"]
-
-    # remove microsecond from time
-    int_ts = int(timestamp)
-
-    # time to iso-8601 compliant date
-    date = datetime.datetime.fromtimestamp(int_ts).isoformat()
-    metadata["modelspec.date"] = date
-
-    if reso is not None:
-        # comma separated to tuple
-        if isinstance(reso, str):
-            reso = tuple(map(int, reso.split(",")))
-        if len(reso) == 1:
-            reso = (reso[0], reso[0])
-    else:
-        # resolution is defined in dataset, so use default
-        if sdxl:
-            reso = 1024
-        elif v2 and v_parameterization:
-            reso = 768
-        else:
-            reso = 512
-    if isinstance(reso, int):
-        reso = (reso, reso)
-
-    metadata["modelspec.resolution"] = f"{reso[0]}x{reso[1]}"
-
-    if v_parameterization:
-        metadata["modelspec.prediction_type"] = PRED_TYPE_V
-    else:
-        metadata["modelspec.prediction_type"] = PRED_TYPE_EPSILON
-
-    if timesteps is not None:
-        if isinstance(timesteps, str) or isinstance(timesteps, int):
-            timesteps = (timesteps, timesteps)
-        if len(timesteps) == 1:
-            timesteps = (timesteps[0], timesteps[0])
-        metadata["modelspec.timestep_range"] = f"{timesteps[0]},{timesteps[1]}"
-    else:
-        del metadata["modelspec.timestep_range"]
-
-    if clip_skip is not None:
-        metadata["modelspec.encoder_layer"] = f"{clip_skip}"
-    else:
-        del metadata["modelspec.encoder_layer"]
-
-    # # assert all values are filled
-    # assert all([v is not None for v in metadata.values()]), metadata
-    if not all([v is not None for v in metadata.values()]):
-        print(f"Internal error: some metadata values are None: {metadata}")
-    
-    return metadata
-
-
-# region utils
-
-
-def get_title(metadata: dict) -> Optional[str]:
-    return metadata.get(MODELSPEC_TITLE, None)
-
-
-def load_metadata_from_safetensors(model: str) -> dict:
-    if not model.endswith(".safetensors"):
-        return {}
-    
-    with safetensors.safe_open(model, framework="pt") as f:
-        metadata = f.metadata()
-    if metadata is None:
-        metadata = {}
-    return metadata
-
-
-def build_merged_from(models: List[str]) -> str:
-    def get_title(model: str):
-        metadata = load_metadata_from_safetensors(model)
-        title = metadata.get(MODELSPEC_TITLE, None)
-        if title is None:
-            title = os.path.splitext(os.path.basename(model))[0]  # use filename
-        return title
-
-    titles = [get_title(model) for model in models]
-    return ", ".join(titles)
-
-
-# endregion
-
-
-r"""
-if __name__ == "__main__":
-    import argparse
-    import torch
-    from safetensors.torch import load_file
-    from library import train_util
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--ckpt", type=str, required=True)
-    args = parser.parse_args()
-
-    print(f"Loading {args.ckpt}")
-    state_dict = load_file(args.ckpt)
-
-    print(f"Calculating metadata")
-    metadata = get(state_dict, False, False, False, False, "sgm", False, False, "title", "date", 256, 1000, 0)
-    print(metadata)
-    del state_dict
-
-    # by reference implementation
-    with open(args.ckpt, mode="rb") as file_data:
-        file_hash = hashlib.sha256()
-        head_len = struct.unpack("Q", file_data.read(8))  # int64 header length prefix
-        header = json.loads(file_data.read(head_len[0]))  # header itself, json string
-        content = (
-            file_data.read()
-        )  # All other content is tightly packed tensors. Copy to RAM for simplicity, but you can avoid this read with a more careful FS-dependent impl.
-        file_hash.update(content)
-        # ===== Update the hash for modelspec =====
-        by_ref = f"0x{file_hash.hexdigest()}"
-    print(by_ref)
-    print("is same?", by_ref == metadata["modelspec.hash_sha256"])
-
-"""

library/sdxl_model_util.py

@@ -1,572 +0,0 @@
-import torch
-from accelerate import init_empty_weights
-from accelerate.utils.modeling import set_module_tensor_to_device
-from safetensors.torch import load_file, save_file
-from transformers import CLIPTextModel, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer
-from typing import List
-from diffusers import AutoencoderKL, EulerDiscreteScheduler, UNet2DConditionModel
-from library import model_util
-from library import sdxl_original_unet
-
-
-VAE_SCALE_FACTOR = 0.13025
-MODEL_VERSION_SDXL_BASE_V1_0 = "sdxl_base_v1-0"
-
-# Diffusersの設定を読み込むための参照モデル
-DIFFUSERS_REF_MODEL_ID_SDXL = "stabilityai/stable-diffusion-xl-base-1.0"
-
-DIFFUSERS_SDXL_UNET_CONFIG = {
-    "act_fn": "silu",
-    "addition_embed_type": "text_time",
-    "addition_embed_type_num_heads": 64,
-    "addition_time_embed_dim": 256,
-    "attention_head_dim": [5, 10, 20],
-    "block_out_channels": [320, 640, 1280],
-    "center_input_sample": False,
-    "class_embed_type": None,
-    "class_embeddings_concat": False,
-    "conv_in_kernel": 3,
-    "conv_out_kernel": 3,
-    "cross_attention_dim": 2048,
-    "cross_attention_norm": None,
-    "down_block_types": ["DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D"],
-    "downsample_padding": 1,
-    "dual_cross_attention": False,
-    "encoder_hid_dim": None,
-    "encoder_hid_dim_type": None,
-    "flip_sin_to_cos": True,
-    "freq_shift": 0,
-    "in_channels": 4,
-    "layers_per_block": 2,
-    "mid_block_only_cross_attention": None,
-    "mid_block_scale_factor": 1,
-    "mid_block_type": "UNetMidBlock2DCrossAttn",
-    "norm_eps": 1e-05,
-    "norm_num_groups": 32,
-    "num_attention_heads": None,
-    "num_class_embeds": None,
-    "only_cross_attention": False,
-    "out_channels": 4,
-    "projection_class_embeddings_input_dim": 2816,
-    "resnet_out_scale_factor": 1.0,
-    "resnet_skip_time_act": False,
-    "resnet_time_scale_shift": "default",
-    "sample_size": 128,
-    "time_cond_proj_dim": None,
-    "time_embedding_act_fn": None,
-    "time_embedding_dim": None,
-    "time_embedding_type": "positional",
-    "timestep_post_act": None,
-    "transformer_layers_per_block": [1, 2, 10],
-    "up_block_types": ["CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D"],
-    "upcast_attention": False,
-    "use_linear_projection": True,
-}
-
-
-def convert_sdxl_text_encoder_2_checkpoint(checkpoint, max_length):
-    SDXL_KEY_PREFIX = "conditioner.embedders.1.model."
-
-    # SD2のと、基本的には同じ。logit_scaleを後で使うので、それを追加で返す
-    # logit_scaleはcheckpointの保存時に使用する
-    def convert_key(key):
-        # common conversion
-        key = key.replace(SDXL_KEY_PREFIX + "transformer.", "text_model.encoder.")
-        key = key.replace(SDXL_KEY_PREFIX, "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 = key.replace("text_model.text_projection", "text_projection.weight")
-        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")
-        # ckpt from comfy has this key: text_model.encoder.text_model.embeddings.position_ids
-        elif ".embeddings.position_ids" in key:
-            key = None  # remove this key: make position_ids by ourselves
-        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 ".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(SDXL_KEY_PREFIX + "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]
-
-    # original SD にはないので、position_idsを追加
-    position_ids = torch.Tensor([list(range(max_length))]).to(torch.int64)
-    new_sd["text_model.embeddings.position_ids"] = position_ids
-
-    # logit_scale はDiffusersには含まれないが、保存時に戻したいので別途返す
-    logit_scale = checkpoint.get(SDXL_KEY_PREFIX + "logit_scale", None)
-
-    return new_sd, logit_scale
-
-
-# load state_dict without allocating new tensors
-def _load_state_dict_on_device(model, state_dict, device, dtype=None):
-    # dtype will use fp32 as default
-    missing_keys = list(model.state_dict().keys() - state_dict.keys())
-    unexpected_keys = list(state_dict.keys() - model.state_dict().keys())
-
-    # similar to model.load_state_dict()
-    if not missing_keys and not unexpected_keys:
-        for k in list(state_dict.keys()):
-            set_module_tensor_to_device(model, k, device, value=state_dict.pop(k), dtype=dtype)
-        return "<All keys matched successfully>"
-
-    # error_msgs
-    error_msgs: List[str] = []
-    if missing_keys:
-        error_msgs.insert(0, "Missing key(s) in state_dict: {}. ".format(", ".join('"{}"'.format(k) for k in missing_keys)))
-    if unexpected_keys:
-        error_msgs.insert(0, "Unexpected key(s) in state_dict: {}. ".format(", ".join('"{}"'.format(k) for k in unexpected_keys)))
-
-    raise RuntimeError("Error(s) in loading state_dict for {}:\n\t{}".format(model.__class__.__name__, "\n\t".join(error_msgs)))
-
-
-def load_models_from_sdxl_checkpoint(model_version, ckpt_path, map_location, dtype=None):
-    # model_version is reserved for future use
-    # dtype is used for full_fp16/bf16 integration. Text Encoder will remain fp32, because it runs on CPU when caching
-
-    # Load the state dict
-    if model_util.is_safetensors(ckpt_path):
-        checkpoint = None
-        try:
-            state_dict = load_file(ckpt_path, device=map_location)
-        except:
-            state_dict = load_file(ckpt_path)  # prevent device invalid Error
-        epoch = None
-        global_step = None
-    else:
-        checkpoint = torch.load(ckpt_path, map_location=map_location)
-        if "state_dict" in checkpoint:
-            state_dict = checkpoint["state_dict"]
-            epoch = checkpoint.get("epoch", 0)
-            global_step = checkpoint.get("global_step", 0)
-        else:
-            state_dict = checkpoint
-            epoch = 0
-            global_step = 0
-        checkpoint = None
-
-    # U-Net
-    print("building U-Net")
-    with init_empty_weights():
-        unet = sdxl_original_unet.SdxlUNet2DConditionModel()
-
-    print("loading U-Net from checkpoint")
-    unet_sd = {}
-    for k in list(state_dict.keys()):
-        if k.startswith("model.diffusion_model."):
-            unet_sd[k.replace("model.diffusion_model.", "")] = state_dict.pop(k)
-    info = _load_state_dict_on_device(unet, unet_sd, device=map_location, dtype=dtype)
-    print("U-Net: ", info)
-
-    # Text Encoders
-    print("building text encoders")
-
-    # Text Encoder 1 is same to Stability AI's SDXL
-    text_model1_cfg = CLIPTextConfig(
-        vocab_size=49408,
-        hidden_size=768,
-        intermediate_size=3072,
-        num_hidden_layers=12,
-        num_attention_heads=12,
-        max_position_embeddings=77,
-        hidden_act="quick_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=768,
-        # torch_dtype="float32",
-        # transformers_version="4.25.0.dev0",
-    )
-    with init_empty_weights():
-        text_model1 = CLIPTextModel._from_config(text_model1_cfg)
-
-    # Text Encoder 2 is different from Stability AI's SDXL. SDXL uses open clip, but we use the model from HuggingFace.
-    # Note: Tokenizer from HuggingFace is different from SDXL. We must use open clip's tokenizer.
-    text_model2_cfg = CLIPTextConfig(
-        vocab_size=49408,
-        hidden_size=1280,
-        intermediate_size=5120,
-        num_hidden_layers=32,
-        num_attention_heads=20,
-        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=1280,
-        # torch_dtype="float32",
-        # transformers_version="4.25.0.dev0",
-    )
-    with init_empty_weights():
-        text_model2 = CLIPTextModelWithProjection(text_model2_cfg)
-
-    print("loading text encoders from checkpoint")
-    te1_sd = {}
-    te2_sd = {}
-    for k in list(state_dict.keys()):
-        if k.startswith("conditioner.embedders.0.transformer."):
-            te1_sd[k.replace("conditioner.embedders.0.transformer.", "")] = state_dict.pop(k)
-        elif k.startswith("conditioner.embedders.1.model."):
-            te2_sd[k] = state_dict.pop(k)
-    
-    # 一部のposition_idsがないモデルへの対応 / add position_ids for some models
-    if "text_model.embeddings.position_ids" not in te1_sd:
-        te1_sd["text_model.embeddings.position_ids"] = torch.arange(77).unsqueeze(0)
-
-    info1 = _load_state_dict_on_device(text_model1, te1_sd, device=map_location)  # remain fp32
-    print("text encoder 1:", info1)
-
-    converted_sd, logit_scale = convert_sdxl_text_encoder_2_checkpoint(te2_sd, max_length=77)
-    info2 = _load_state_dict_on_device(text_model2, converted_sd, device=map_location)  # remain fp32
-    print("text encoder 2:", info2)
-
-    # prepare vae
-    print("building VAE")
-    vae_config = model_util.create_vae_diffusers_config()
-    with init_empty_weights():
-        vae = AutoencoderKL(**vae_config)
-
-    print("loading VAE from checkpoint")
-    converted_vae_checkpoint = model_util.convert_ldm_vae_checkpoint(state_dict, vae_config)
-    info = _load_state_dict_on_device(vae, converted_vae_checkpoint, device=map_location, dtype=dtype)
-    print("VAE:", info)
-
-    ckpt_info = (epoch, global_step) if epoch is not None else None
-    return text_model1, text_model2, vae, unet, logit_scale, ckpt_info
-
-
-def make_unet_conversion_map():
-    unet_conversion_map_layer = []
-
-    for i in range(3):  # num_blocks is 3 in sdxl
-        # 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: commentout for sdxl
-            # 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}.{2}."  # change for sdxl
-            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))
-
-    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 = []
-    for sd, hf in unet_conversion_map_layer:
-        if "resnets" in hf:
-            for sd_res, hf_res in unet_conversion_map_resnet:
-                unet_conversion_map.append((sd + sd_res, hf + hf_res))
-        else:
-            unet_conversion_map.append((sd, hf))
-
-    for j in range(2):
-        hf_time_embed_prefix = f"time_embedding.linear_{j+1}."
-        sd_time_embed_prefix = f"time_embed.{j*2}."
-        unet_conversion_map.append((sd_time_embed_prefix, hf_time_embed_prefix))
-
-    for j in range(2):
-        hf_label_embed_prefix = f"add_embedding.linear_{j+1}."
-        sd_label_embed_prefix = f"label_emb.0.{j*2}."
-        unet_conversion_map.append((sd_label_embed_prefix, hf_label_embed_prefix))
-
-    unet_conversion_map.append(("input_blocks.0.0.", "conv_in."))
-    unet_conversion_map.append(("out.0.", "conv_norm_out."))
-    unet_conversion_map.append(("out.2.", "conv_out."))
-
-    return unet_conversion_map
-
-
-def convert_diffusers_unet_state_dict_to_sdxl(du_sd):
-    unet_conversion_map = make_unet_conversion_map()
-
-    conversion_map = {hf: sd for sd, hf in unet_conversion_map}
-    return convert_unet_state_dict(du_sd, conversion_map)
-
-
-def convert_unet_state_dict(src_sd, conversion_map):
-    converted_sd = {}
-    for src_key, value in src_sd.items():
-        # さすがに全部回すのは時間がかかるので右から要素を削りつつprefixを探す
-        src_key_fragments = src_key.split(".")[:-1]  # remove weight/bias
-        while len(src_key_fragments) > 0:
-            src_key_prefix = ".".join(src_key_fragments) + "."
-            if src_key_prefix in conversion_map:
-                converted_prefix = conversion_map[src_key_prefix]
-                converted_key = converted_prefix + src_key[len(src_key_prefix) :]
-                converted_sd[converted_key] = value
-                break
-            src_key_fragments.pop(-1)
-        assert len(src_key_fragments) > 0, f"key {src_key} not found in conversion map"
-
-    return converted_sd
-
-
-def convert_sdxl_unet_state_dict_to_diffusers(sd):
-    unet_conversion_map = make_unet_conversion_map()
-
-    conversion_dict = {sd: hf for sd, hf in unet_conversion_map}
-    return convert_unet_state_dict(sd, conversion_dict)
-
-
-def convert_text_encoder_2_state_dict_to_sdxl(checkpoint, logit_scale):
-    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 "text_projection" in key:  # no dot in key
-            key = key.replace("text_projection.weight", "text_projection")
-        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 logit_scale is not None:
-        new_sd["logit_scale"] = logit_scale
-
-    return new_sd
-
-
-def save_stable_diffusion_checkpoint(
-    output_file,
-    text_encoder1,
-    text_encoder2,
-    unet,
-    epochs,
-    steps,
-    ckpt_info,
-    vae,
-    logit_scale,
-    metadata,
-    save_dtype=None,
-):
-    state_dict = {}
-
-    def update_sd(prefix, sd):
-        for k, v in sd.items():
-            key = prefix + k
-            if save_dtype is not None:
-                v = v.detach().clone().to("cpu").to(save_dtype)
-            state_dict[key] = v
-
-    # Convert the UNet model
-    update_sd("model.diffusion_model.", unet.state_dict())
-
-    # Convert the text encoders
-    update_sd("conditioner.embedders.0.transformer.", text_encoder1.state_dict())
-
-    text_enc2_dict = convert_text_encoder_2_state_dict_to_sdxl(text_encoder2.state_dict(), logit_scale)
-    update_sd("conditioner.embedders.1.model.", text_enc2_dict)
-
-    # Convert the VAE
-    vae_dict = model_util.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}
-
-    # epoch and global_step are sometimes not int
-    if ckpt_info is not None:
-        epochs += ckpt_info[0]
-        steps += ckpt_info[1]
-
-    new_ckpt["epoch"] = epochs
-    new_ckpt["global_step"] = steps
-
-    if model_util.is_safetensors(output_file):
-        save_file(state_dict, output_file, metadata)
-    else:
-        torch.save(new_ckpt, output_file)
-
-    return key_count
-
-
-def save_diffusers_checkpoint(
-    output_dir, text_encoder1, text_encoder2, unet, pretrained_model_name_or_path, vae=None, use_safetensors=False, save_dtype=None
-):
-    from diffusers import StableDiffusionXLPipeline
-
-    # convert U-Net
-    unet_sd = unet.state_dict()
-    du_unet_sd = convert_sdxl_unet_state_dict_to_diffusers(unet_sd)
-
-    diffusers_unet = UNet2DConditionModel(**DIFFUSERS_SDXL_UNET_CONFIG)
-    if save_dtype is not None:
-        diffusers_unet.to(save_dtype)
-    diffusers_unet.load_state_dict(du_unet_sd)
-
-    # create pipeline to save
-    if pretrained_model_name_or_path is None:
-        pretrained_model_name_or_path = DIFFUSERS_REF_MODEL_ID_SDXL
-
-    scheduler = EulerDiscreteScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
-    tokenizer1 = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer")
-    tokenizer2 = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer_2")
-    if vae is None:
-        vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae")
-
-    # prevent local path from being saved
-    def remove_name_or_path(model):
-        if hasattr(model, "config"):
-            model.config._name_or_path = None
-            model.config._name_or_path = None
-
-    remove_name_or_path(diffusers_unet)
-    remove_name_or_path(text_encoder1)
-    remove_name_or_path(text_encoder2)
-    remove_name_or_path(scheduler)
-    remove_name_or_path(tokenizer1)
-    remove_name_or_path(tokenizer2)
-    remove_name_or_path(vae)
-
-    pipeline = StableDiffusionXLPipeline(
-        unet=diffusers_unet,
-        text_encoder=text_encoder1,
-        text_encoder_2=text_encoder2,
-        vae=vae,
-        scheduler=scheduler,
-        tokenizer=tokenizer1,
-        tokenizer_2=tokenizer2,
-    )
-    if save_dtype is not None:
-        pipeline.to(None, save_dtype)
-    pipeline.save_pretrained(output_dir, safe_serialization=use_safetensors)

library/sdxl_train_util.py

@@ -1,369 +0,0 @@
-import argparse
-import gc
-import math
-import os
-from typing import Optional
-import torch
-from accelerate import init_empty_weights
-from tqdm import tqdm
-from transformers import CLIPTokenizer
-from library import model_util, sdxl_model_util, train_util, sdxl_original_unet
-from library.sdxl_lpw_stable_diffusion import SdxlStableDiffusionLongPromptWeightingPipeline
-
-TOKENIZER1_PATH = "openai/clip-vit-large-patch14"
-TOKENIZER2_PATH = "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"
-
-# DEFAULT_NOISE_OFFSET = 0.0357
-
-
-def load_target_model(args, accelerator, model_version: str, weight_dtype):
-    # load models for each process
-    model_dtype = match_mixed_precision(args, weight_dtype)  # prepare fp16/bf16
-    for pi in range(accelerator.state.num_processes):
-        if pi == accelerator.state.local_process_index:
-            print(f"loading model for process {accelerator.state.local_process_index}/{accelerator.state.num_processes}")
-
-            (
-                load_stable_diffusion_format,
-                text_encoder1,
-                text_encoder2,
-                vae,
-                unet,
-                logit_scale,
-                ckpt_info,
-            ) = _load_target_model(
-                args.pretrained_model_name_or_path,
-                args.vae,
-                model_version,
-                weight_dtype,
-                accelerator.device if args.lowram else "cpu",
-                model_dtype,
-            )
-
-            # work on low-ram device
-            if args.lowram:
-                text_encoder1.to(accelerator.device)
-                text_encoder2.to(accelerator.device)
-                unet.to(accelerator.device)
-                vae.to(accelerator.device)
-
-            gc.collect()
-            torch.cuda.empty_cache()
-        accelerator.wait_for_everyone()
-
-    text_encoder1, text_encoder2, unet = train_util.transform_models_if_DDP([text_encoder1, text_encoder2, unet])
-
-    return load_stable_diffusion_format, text_encoder1, text_encoder2, vae, unet, logit_scale, ckpt_info
-
-
-def _load_target_model(
-    name_or_path: str, vae_path: Optional[str], model_version: str, weight_dtype, device="cpu", model_dtype=None
-):
-    # model_dtype only work with full fp16/bf16
-    name_or_path = os.readlink(name_or_path) if os.path.islink(name_or_path) else name_or_path
-    load_stable_diffusion_format = os.path.isfile(name_or_path)  # determine SD or Diffusers
-
-    if load_stable_diffusion_format:
-        print(f"load StableDiffusion checkpoint: {name_or_path}")
-        (
-            text_encoder1,
-            text_encoder2,
-            vae,
-            unet,
-            logit_scale,
-            ckpt_info,
-        ) = sdxl_model_util.load_models_from_sdxl_checkpoint(model_version, name_or_path, device, model_dtype)
-    else:
-        # Diffusers model is loaded to CPU
-        from diffusers import StableDiffusionXLPipeline
-
-        variant = "fp16" if weight_dtype == torch.float16 else None
-        print(f"load Diffusers pretrained models: {name_or_path}, variant={variant}")
-        try:
-            try:
-                pipe = StableDiffusionXLPipeline.from_pretrained(
-                    name_or_path, torch_dtype=model_dtype, variant=variant, tokenizer=None
-                )
-            except EnvironmentError as ex:
-                if variant is not None:
-                    print("try to load fp32 model")
-                    pipe = StableDiffusionXLPipeline.from_pretrained(name_or_path, variant=None, tokenizer=None)
-                else:
-                    raise ex
-        except EnvironmentError as ex:
-            print(
-                f"model is not found as a file or in Hugging Face, perhaps file name is wrong? / 指定したモデル名のファイル、またはHugging Faceのモデルが見つかりません。ファイル名が誤っているかもしれません: {name_or_path}"
-            )
-            raise ex
-
-        text_encoder1 = pipe.text_encoder
-        text_encoder2 = pipe.text_encoder_2
-
-        # convert to fp32 for cache text_encoders outputs
-        if text_encoder1.dtype != torch.float32:
-            text_encoder1 = text_encoder1.to(dtype=torch.float32)
-        if text_encoder2.dtype != torch.float32:
-            text_encoder2 = text_encoder2.to(dtype=torch.float32)
-
-        vae = pipe.vae
-        unet = pipe.unet
-        del pipe
-
-        # Diffusers U-Net to original U-Net
-        state_dict = sdxl_model_util.convert_diffusers_unet_state_dict_to_sdxl(unet.state_dict())
-        with init_empty_weights():
-            unet = sdxl_original_unet.SdxlUNet2DConditionModel()  # overwrite unet
-        sdxl_model_util._load_state_dict_on_device(unet, state_dict, device=device, dtype=model_dtype)
-        print("U-Net converted to original U-Net")
-
-        logit_scale = None
-        ckpt_info = None
-
-    # VAEを読み込む
-    if vae_path is not None:
-        vae = model_util.load_vae(vae_path, weight_dtype)
-        print("additional VAE loaded")
-
-    return load_stable_diffusion_format, text_encoder1, text_encoder2, vae, unet, logit_scale, ckpt_info
-
-
-def load_tokenizers(args: argparse.Namespace):
-    print("prepare tokenizers")
-
-    original_paths = [TOKENIZER1_PATH, TOKENIZER2_PATH]
-    tokeniers = []
-    for i, original_path in enumerate(original_paths):
-        tokenizer: CLIPTokenizer = None
-        if args.tokenizer_cache_dir:
-            local_tokenizer_path = os.path.join(args.tokenizer_cache_dir, original_path.replace("/", "_"))
-            if os.path.exists(local_tokenizer_path):
-                print(f"load tokenizer from cache: {local_tokenizer_path}")
-                tokenizer = CLIPTokenizer.from_pretrained(local_tokenizer_path)
-
-        if tokenizer is None:
-            tokenizer = CLIPTokenizer.from_pretrained(original_path)
-
-        if args.tokenizer_cache_dir and not os.path.exists(local_tokenizer_path):
-            print(f"save Tokenizer to cache: {local_tokenizer_path}")
-            tokenizer.save_pretrained(local_tokenizer_path)
-
-        if i == 1:
-            tokenizer.pad_token_id = 0  # fix pad token id to make same as open clip tokenizer
-
-        tokeniers.append(tokenizer)
-
-    if hasattr(args, "max_token_length") and args.max_token_length is not None:
-        print(f"update token length: {args.max_token_length}")
-
-    return tokeniers
-
-
-def match_mixed_precision(args, weight_dtype):
-    if args.full_fp16:
-        assert (
-            weight_dtype == torch.float16
-        ), "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
-        return weight_dtype
-    elif args.full_bf16:
-        assert (
-            weight_dtype == torch.bfloat16
-        ), "full_bf16 requires mixed precision='bf16' / full_bf16を使う場合はmixed_precision='bf16'を指定してください。"
-        return weight_dtype
-    else:
-        return None
-
-
-def timestep_embedding(timesteps, dim, max_period=10000):
-    """
-    Create sinusoidal timestep embeddings.
-    :param timesteps: a 1-D Tensor of N indices, one per batch element.
-                      These may be fractional.
-    :param dim: the dimension of the output.
-    :param max_period: controls the minimum frequency of the embeddings.
-    :return: an [N x dim] Tensor of positional embeddings.
-    """
-    half = dim // 2
-    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
-        device=timesteps.device
-    )
-    args = timesteps[:, None].float() * freqs[None]
-    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-    if dim % 2:
-        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-    return embedding
-
-
-def get_timestep_embedding(x, outdim):
-    assert len(x.shape) == 2
-    b, dims = x.shape[0], x.shape[1]
-    x = torch.flatten(x)
-    emb = timestep_embedding(x, outdim)
-    emb = torch.reshape(emb, (b, dims * outdim))
-    return emb
-
-
-def get_size_embeddings(orig_size, crop_size, target_size, device):
-    emb1 = get_timestep_embedding(orig_size, 256)
-    emb2 = get_timestep_embedding(crop_size, 256)
-    emb3 = get_timestep_embedding(target_size, 256)
-    vector = torch.cat([emb1, emb2, emb3], dim=1).to(device)
-    return vector
-
-
-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_encoder1,
-    text_encoder2,
-    unet,
-    vae,
-    logit_scale,
-    ckpt_info,
-):
-    def sd_saver(ckpt_file, epoch_no, global_step):
-        sai_metadata = train_util.get_sai_model_spec(None, args, True, False, False, is_stable_diffusion_ckpt=True)
-        sdxl_model_util.save_stable_diffusion_checkpoint(
-            ckpt_file,
-            text_encoder1,
-            text_encoder2,
-            unet,
-            epoch_no,
-            global_step,
-            ckpt_info,
-            vae,
-            logit_scale,
-            sai_metadata,
-            save_dtype,
-        )
-
-    def diffusers_saver(out_dir):
-        sdxl_model_util.save_diffusers_checkpoint(
-            out_dir,
-            text_encoder1,
-            text_encoder2,
-            unet,
-            src_path,
-            vae,
-            use_safetensors=use_safetensors,
-            save_dtype=save_dtype,
-        )
-
-    train_util.save_sd_model_on_train_end_common(
-        args, save_stable_diffusion_format, use_safetensors, epoch, global_step, sd_saver, diffusers_saver
-    )
-
-
-# epochとstepの保存、メタデータにepoch/stepが含まれ引数が同じになるため、統合している
-# on_epoch_end: Trueならepoch終了時、Falseならstep経過時
-def save_sd_model_on_epoch_end_or_stepwise(
-    args: argparse.Namespace,
-    on_epoch_end: bool,
-    accelerator,
-    src_path,
-    save_stable_diffusion_format: bool,
-    use_safetensors: bool,
-    save_dtype: torch.dtype,
-    epoch: int,
-    num_train_epochs: int,
-    global_step: int,
-    text_encoder1,
-    text_encoder2,
-    unet,
-    vae,
-    logit_scale,
-    ckpt_info,
-):
-    def sd_saver(ckpt_file, epoch_no, global_step):
-        sai_metadata = train_util.get_sai_model_spec(None, args, True, False, False, is_stable_diffusion_ckpt=True)
-        sdxl_model_util.save_stable_diffusion_checkpoint(
-            ckpt_file,
-            text_encoder1,
-            text_encoder2,
-            unet,
-            epoch_no,
-            global_step,
-            ckpt_info,
-            vae,
-            logit_scale,
-            sai_metadata,
-            save_dtype,
-        )
-
-    def diffusers_saver(out_dir):
-        sdxl_model_util.save_diffusers_checkpoint(
-            out_dir,
-            text_encoder1,
-            text_encoder2,
-            unet,
-            src_path,
-            vae,
-            use_safetensors=use_safetensors,
-            save_dtype=save_dtype,
-        )
-
-    train_util.save_sd_model_on_epoch_end_or_stepwise_common(
-        args,
-        on_epoch_end,
-        accelerator,
-        save_stable_diffusion_format,
-        use_safetensors,
-        epoch,
-        num_train_epochs,
-        global_step,
-        sd_saver,
-        diffusers_saver,
-    )
-
-
-def add_sdxl_training_arguments(parser: argparse.ArgumentParser):
-    parser.add_argument(
-        "--cache_text_encoder_outputs", action="store_true", help="cache text encoder outputs / text encoderの出力をキャッシュする"
-    )
-    parser.add_argument(
-        "--cache_text_encoder_outputs_to_disk",
-        action="store_true",
-        help="cache text encoder outputs to disk / text encoderの出力をディスクにキャッシュする",
-    )
-
-
-def verify_sdxl_training_args(args: argparse.Namespace, supportTextEncoderCaching: bool = True):
-    assert not args.v2, "v2 cannot be enabled in SDXL training / SDXL学習ではv2を有効にすることはできません"
-    if args.v_parameterization:
-        print("v_parameterization will be unexpected / SDXL学習ではv_parameterizationは想定外の動作になります")
-
-    if args.clip_skip is not None:
-        print("clip_skip will be unexpected / SDXL学習ではclip_skipは動作しません")
-
-    # if args.multires_noise_iterations:
-    #     print(
-    #         f"Warning: SDXL has been trained with noise_offset={DEFAULT_NOISE_OFFSET}, but noise_offset is disabled due to multires_noise_iterations / SDXLはnoise_offset={DEFAULT_NOISE_OFFSET}で学習されていますが、multires_noise_iterationsが有効になっているためnoise_offsetは無効になります"
-    #     )
-    # else:
-    #     if args.noise_offset is None:
-    #         args.noise_offset = DEFAULT_NOISE_OFFSET
-    #     elif args.noise_offset != DEFAULT_NOISE_OFFSET:
-    #         print(
-    #             f"Warning: SDXL has been trained with noise_offset={DEFAULT_NOISE_OFFSET} / SDXLはnoise_offset={DEFAULT_NOISE_OFFSET}で学習されています"
-    #         )
-    #     print(f"noise_offset is set to {args.noise_offset} / noise_offsetが{args.noise_offset}に設定されました")
-
-    assert (
-        not hasattr(args, "weighted_captions") or not args.weighted_captions
-    ), "weighted_captions cannot be enabled in SDXL training currently / SDXL学習では今のところweighted_captionsを有効にすることはできません"
-
-    if supportTextEncoderCaching:
-        if args.cache_text_encoder_outputs_to_disk and not args.cache_text_encoder_outputs:
-            args.cache_text_encoder_outputs = True
-            print(
-                "cache_text_encoder_outputs is enabled because cache_text_encoder_outputs_to_disk is enabled / "
-                + "cache_text_encoder_outputs_to_diskが有効になっているためcache_text_encoder_outputsが有効になりました"
-            )
-
-
-def sample_images(*args, **kwargs):
-    return train_util.sample_images_common(SdxlStableDiffusionLongPromptWeightingPipeline, *args, **kwargs)

library/slicing_vae.py

@@ -22,10 +22,10 @@ import torch.nn as nn
 
 
 from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.models.modeling_utils import ModelMixin
-from diffusers.models.unet_2d_blocks import UNetMidBlock2D, get_down_block, get_up_block
-from diffusers.models.vae import DecoderOutput, DiagonalGaussianDistribution
-from diffusers.models.autoencoder_kl import AutoencoderKLOutput
+from diffusers.modeling_utils import ModelMixin
+from diffusers.utils import BaseOutput
+from diffusers.models.unet_2d_blocks import UNetMidBlock2D, get_down_block, get_up_block, ResnetBlock2D
+from diffusers.models.vae import DecoderOutput, Encoder, AutoencoderKLOutput, DiagonalGaussianDistribution
 
 
 def slice_h(x, num_slices):
@@ -209,7 +209,7 @@ class SlicingEncoder(nn.Module):
                 downsample_padding=0,
                 resnet_act_fn=act_fn,
                 resnet_groups=norm_num_groups,
-                attention_head_dim=output_channel,
+                attn_num_head_channels=None,
                 temb_channels=None,
             )
             self.down_blocks.append(down_block)
@@ -221,7 +221,7 @@ class SlicingEncoder(nn.Module):
             resnet_act_fn=act_fn,
             output_scale_factor=1,
             resnet_time_scale_shift="default",
-            attention_head_dim=block_out_channels[-1],
+            attn_num_head_channels=None,
             resnet_groups=norm_num_groups,
             temb_channels=None,
         )
@@ -381,7 +381,7 @@ class SlicingDecoder(nn.Module):
             resnet_act_fn=act_fn,
             output_scale_factor=1,
             resnet_time_scale_shift="default",
-            attention_head_dim=block_out_channels[-1],
+            attn_num_head_channels=None,
             resnet_groups=norm_num_groups,
             temb_channels=None,
         )
@@ -406,7 +406,7 @@ class SlicingDecoder(nn.Module):
                 resnet_eps=1e-6,
                 resnet_act_fn=act_fn,
                 resnet_groups=norm_num_groups,
-                attention_head_dim=output_channel,
+                attn_num_head_channels=None,
                 temb_channels=None,
             )
             self.up_blocks.append(up_block)

networks/check_lora_weights.py

@@ -5,41 +5,35 @@ 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")
+  print(f"loading: {file}")
+  if os.path.splitext(file)[1] == '.safetensors':
+    sd = load_file(file)
+  else:
+    sd = torch.load(file, map_location='cpu')
 
-    values = []
+  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)}")
+  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)}")
 
-    if args.show_all_keys:
-        for key in [k for k in keys if k not in values]:
-            values.append((key, sd[key]))
-        print(f"number of all modules: {len(values)}")
-
-    for key, value in values:
-        value = value.to(torch.float32)
-        print(f"{key},{str(tuple(value.size())).replace(', ', '-')},{torch.mean(torch.abs(value))},{torch.min(torch.abs(value))}")
+  for key, value in values:
+    value = value.to(torch.float32)
+    print(f"{key},{str(tuple(value.size())).replace(', ', '-')},{torch.mean(torch.abs(value))},{torch.min(torch.abs(value))}")
 
 
 def setup_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("file", type=str, help="model file to check / 重みを確認するモデルファイル")
-    parser.add_argument("-s", "--show_all_keys", action="store_true", help="show all keys / 全てのキーを表示する")
+  parser = argparse.ArgumentParser()
+  parser.add_argument("file", type=str, help="model file to check / 重みを確認するモデルファイル")
 
-    return parser
+  return parser
 
 
-if __name__ == "__main__":
-    parser = setup_parser()
+if __name__ == '__main__':
+  parser = setup_parser()
 
-    args = parser.parse_args()
+  args = parser.parse_args()
 
-    main(args.file)
+  main(args.file)

networks/control_net_lllite.py

@@ -1,446 +0,0 @@
-import os
-from typing import Optional, List, Type
-import torch
-from library import sdxl_original_unet
-
-
-# input_blocksに適用するかどうか / if True, input_blocks are not applied
-SKIP_INPUT_BLOCKS = False
-
-# output_blocksに適用するかどうか / if True, output_blocks are not applied
-SKIP_OUTPUT_BLOCKS = True
-
-# conv2dに適用するかどうか / if True, conv2d are not applied
-SKIP_CONV2D = False
-
-# transformer_blocksのみに適用するかどうか。Trueの場合、ResBlockには適用されない
-# if True, only transformer_blocks are applied, and ResBlocks are not applied
-TRANSFORMER_ONLY = True  # if True, SKIP_CONV2D is ignored because conv2d is not used in transformer_blocks
-
-# Trueならattn1とattn2にのみ適用し、ffなどには適用しない / if True, apply only to attn1 and attn2, not to ff etc.
-ATTN1_2_ONLY = True
-
-# Trueならattn1のQKV、attn2のQにのみ適用する、ATTN1_2_ONLY指定時のみ有効 / if True, apply only to attn1 QKV and attn2 Q, only valid when ATTN1_2_ONLY is specified
-ATTN_QKV_ONLY = True
-
-# Trueならattn1やffなどにのみ適用し、attn2などには適用しない / if True, apply only to attn1 and ff, not to attn2
-# ATTN1_2_ONLYと同時にTrueにできない / cannot be True at the same time as ATTN1_2_ONLY
-ATTN1_ETC_ONLY = False  # True
-
-# transformer_blocksの最大インデックス。Noneなら全てのtransformer_blocksに適用
-# max index of transformer_blocks. if None, apply to all transformer_blocks
-TRANSFORMER_MAX_BLOCK_INDEX = None
-
-
-class LLLiteModule(torch.nn.Module):
-    def __init__(self, depth, cond_emb_dim, name, org_module, mlp_dim, dropout=None, multiplier=1.0):
-        super().__init__()
-
-        self.is_conv2d = org_module.__class__.__name__ == "Conv2d"
-        self.lllite_name = name
-        self.cond_emb_dim = cond_emb_dim
-        self.org_module = [org_module]
-        self.dropout = dropout
-        self.multiplier = multiplier
-
-        if self.is_conv2d:
-            in_dim = org_module.in_channels
-        else:
-            in_dim = org_module.in_features
-
-        # conditioning1はconditioning imageを embedding する。timestepごとに呼ばれない
-        # conditioning1 embeds conditioning image. it is not called for each timestep
-        modules = []
-        modules.append(torch.nn.Conv2d(3, cond_emb_dim // 2, kernel_size=4, stride=4, padding=0))  # to latent (from VAE) size
-        if depth == 1:
-            modules.append(torch.nn.ReLU(inplace=True))
-            modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=2, stride=2, padding=0))
-        elif depth == 2:
-            modules.append(torch.nn.ReLU(inplace=True))
-            modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=4, stride=4, padding=0))
-        elif depth == 3:
-            # kernel size 8は大きすぎるので、4にする / kernel size 8 is too large, so set it to 4
-            modules.append(torch.nn.ReLU(inplace=True))
-            modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim // 2, kernel_size=4, stride=4, padding=0))
-            modules.append(torch.nn.ReLU(inplace=True))
-            modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=2, stride=2, padding=0))
-
-        self.conditioning1 = torch.nn.Sequential(*modules)
-
-        # downで入力の次元数を削減する。LoRAにヒントを得ていることにする
-        # midでconditioning image embeddingと入力を結合する
-        # upで元の次元数に戻す
-        # これらはtimestepごとに呼ばれる
-        # reduce the number of input dimensions with down. inspired by LoRA
-        # combine conditioning image embedding and input with mid
-        # restore to the original dimension with up
-        # these are called for each timestep
-
-        if self.is_conv2d:
-            self.down = torch.nn.Sequential(
-                torch.nn.Conv2d(in_dim, mlp_dim, kernel_size=1, stride=1, padding=0),
-                torch.nn.ReLU(inplace=True),
-            )
-            self.mid = torch.nn.Sequential(
-                torch.nn.Conv2d(mlp_dim + cond_emb_dim, mlp_dim, kernel_size=1, stride=1, padding=0),
-                torch.nn.ReLU(inplace=True),
-            )
-            self.up = torch.nn.Sequential(
-                torch.nn.Conv2d(mlp_dim, in_dim, kernel_size=1, stride=1, padding=0),
-            )
-        else:
-            # midの前にconditioningをreshapeすること / reshape conditioning before mid
-            self.down = torch.nn.Sequential(
-                torch.nn.Linear(in_dim, mlp_dim),
-                torch.nn.ReLU(inplace=True),
-            )
-            self.mid = torch.nn.Sequential(
-                torch.nn.Linear(mlp_dim + cond_emb_dim, mlp_dim),
-                torch.nn.ReLU(inplace=True),
-            )
-            self.up = torch.nn.Sequential(
-                torch.nn.Linear(mlp_dim, in_dim),
-            )
-
-        # Zero-Convにする / set to Zero-Conv
-        torch.nn.init.zeros_(self.up[0].weight)  # zero conv
-
-        self.depth = depth  # 1~3
-        self.cond_emb = None
-        self.batch_cond_only = False  # Trueなら推論時のcondにのみ適用する / if True, apply only to cond at inference
-        self.use_zeros_for_batch_uncond = False  # Trueならuncondのconditioningを0にする / if True, set uncond conditioning to 0
-
-        # batch_cond_onlyとuse_zeros_for_batch_uncondはどちらも適用すると生成画像の色味がおかしくなるので実際には使えそうにない
-        # Controlの種類によっては使えるかも
-        # both batch_cond_only and use_zeros_for_batch_uncond make the color of the generated image strange, so it doesn't seem to be usable in practice
-        # it may be available depending on the type of Control
-
-    def set_cond_image(self, cond_image):
-        r"""
-        中でモデルを呼び出すので必要ならwith torch.no_grad()で囲む
-        / call the model inside, so if necessary, surround it with torch.no_grad()
-        """
-        if cond_image is None:
-            self.cond_emb = None
-            return
-
-        # timestepごとに呼ばれないので、あらかじめ計算しておく / it is not called for each timestep, so calculate it in advance
-        # print(f"C {self.lllite_name}, cond_image.shape={cond_image.shape}")
-        cx = self.conditioning1(cond_image)
-        if not self.is_conv2d:
-            # reshape / b,c,h,w -> b,h*w,c
-            n, c, h, w = cx.shape
-            cx = cx.view(n, c, h * w).permute(0, 2, 1)
-        self.cond_emb = cx
-
-    def set_batch_cond_only(self, cond_only, zeros):
-        self.batch_cond_only = cond_only
-        self.use_zeros_for_batch_uncond = zeros
-
-    def apply_to(self):
-        self.org_forward = self.org_module[0].forward
-        self.org_module[0].forward = self.forward
-
-    def forward(self, x):
-        r"""
-        学習用の便利forward。元のモジュールのforwardを呼び出す
-        / convenient forward for training. call the forward of the original module
-        """
-        if self.multiplier == 0.0 or self.cond_emb is None:
-            return self.org_forward(x)
-
-        cx = self.cond_emb
-
-        if not self.batch_cond_only and x.shape[0] // 2 == cx.shape[0]:  # inference only
-            cx = cx.repeat(2, 1, 1, 1) if self.is_conv2d else cx.repeat(2, 1, 1)
-            if self.use_zeros_for_batch_uncond:
-                cx[0::2] = 0.0  # uncond is zero
-        # print(f"C {self.lllite_name}, x.shape={x.shape}, cx.shape={cx.shape}")
-
-        # downで入力の次元数を削減し、conditioning image embeddingと結合する
-        # 加算ではなくchannel方向に結合することで、うまいこと混ぜてくれることを期待している
-        # down reduces the number of input dimensions and combines it with conditioning image embedding
-        # we expect that it will mix well by combining in the channel direction instead of adding
-
-        cx = torch.cat([cx, self.down(x if not self.batch_cond_only else x[1::2])], dim=1 if self.is_conv2d else 2)
-        cx = self.mid(cx)
-
-        if self.dropout is not None and self.training:
-            cx = torch.nn.functional.dropout(cx, p=self.dropout)
-
-        cx = self.up(cx) * self.multiplier
-
-        # residual (x) を加算して元のforwardを呼び出す / add residual (x) and call the original forward
-        if self.batch_cond_only:
-            zx = torch.zeros_like(x)
-            zx[1::2] += cx
-            cx = zx
-
-        x = self.org_forward(x + cx)  # ここで元のモジュールを呼び出す / call the original module here
-        return x
-
-
-class ControlNetLLLite(torch.nn.Module):
-    UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel"]
-    UNET_TARGET_REPLACE_MODULE_CONV2D_3X3 = ["ResnetBlock2D", "Downsample2D", "Upsample2D"]
-
-    def __init__(
-        self,
-        unet: sdxl_original_unet.SdxlUNet2DConditionModel,
-        cond_emb_dim: int = 16,
-        mlp_dim: int = 16,
-        dropout: Optional[float] = None,
-        varbose: Optional[bool] = False,
-        multiplier: Optional[float] = 1.0,
-    ) -> None:
-        super().__init__()
-        # self.unets = [unet]
-
-        def create_modules(
-            root_module: torch.nn.Module,
-            target_replace_modules: List[torch.nn.Module],
-            module_class: Type[object],
-        ) -> List[torch.nn.Module]:
-            prefix = "lllite_unet"
-
-            modules = []
-            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():
-                        is_linear = child_module.__class__.__name__ == "Linear"
-                        is_conv2d = child_module.__class__.__name__ == "Conv2d"
-
-                        if is_linear or (is_conv2d and not SKIP_CONV2D):
-                            # block indexからdepthを計算: depthはconditioningのサイズやチャネルを計算するのに使う
-                            # block index to depth: depth is using to calculate conditioning size and channels
-                            block_name, index1, index2 = (name + "." + child_name).split(".")[:3]
-                            index1 = int(index1)
-                            if block_name == "input_blocks":
-                                if SKIP_INPUT_BLOCKS:
-                                    continue
-                                depth = 1 if index1 <= 2 else (2 if index1 <= 5 else 3)
-                            elif block_name == "middle_block":
-                                depth = 3
-                            elif block_name == "output_blocks":
-                                if SKIP_OUTPUT_BLOCKS:
-                                    continue
-                                depth = 3 if index1 <= 2 else (2 if index1 <= 5 else 1)
-                                if int(index2) >= 2:
-                                    depth -= 1
-                            else:
-                                raise NotImplementedError()
-
-                            lllite_name = prefix + "." + name + "." + child_name
-                            lllite_name = lllite_name.replace(".", "_")
-
-                            if TRANSFORMER_MAX_BLOCK_INDEX is not None:
-                                p = lllite_name.find("transformer_blocks")
-                                if p >= 0:
-                                    tf_index = int(lllite_name[p:].split("_")[2])
-                                    if tf_index > TRANSFORMER_MAX_BLOCK_INDEX:
-                                        continue
-
-                            #  time embは適用外とする
-                            # attn2のconditioning (CLIPからの入力) はshapeが違うので適用できない
-                            # time emb is not applied
-                            # attn2 conditioning (input from CLIP) cannot be applied because the shape is different
-                            if "emb_layers" in lllite_name or (
-                                "attn2" in lllite_name and ("to_k" in lllite_name or "to_v" in lllite_name)
-                            ):
-                                continue
-
-                            if ATTN1_2_ONLY:
-                                if not ("attn1" in lllite_name or "attn2" in lllite_name):
-                                    continue
-                                if ATTN_QKV_ONLY:
-                                    if "to_out" in lllite_name:
-                                        continue
-
-                            if ATTN1_ETC_ONLY:
-                                if "proj_out" in lllite_name:
-                                    pass
-                                elif "attn1" in lllite_name and (
-                                    "to_k" in lllite_name or "to_v" in lllite_name or "to_out" in lllite_name
-                                ):
-                                    pass
-                                elif "ff_net_2" in lllite_name:
-                                    pass
-                                else:
-                                    continue
-
-                            module = module_class(
-                                depth,
-                                cond_emb_dim,
-                                lllite_name,
-                                child_module,
-                                mlp_dim,
-                                dropout=dropout,
-                                multiplier=multiplier,
-                            )
-                            modules.append(module)
-            return modules
-
-        target_modules = ControlNetLLLite.UNET_TARGET_REPLACE_MODULE
-        if not TRANSFORMER_ONLY:
-            target_modules = target_modules + ControlNetLLLite.UNET_TARGET_REPLACE_MODULE_CONV2D_3X3
-
-        # create module instances
-        self.unet_modules: List[LLLiteModule] = create_modules(unet, target_modules, LLLiteModule)
-        print(f"create ControlNet LLLite for U-Net: {len(self.unet_modules)} modules.")
-
-    def forward(self, x):
-        return x  # dummy
-
-    def set_cond_image(self, cond_image):
-        r"""
-        中でモデルを呼び出すので必要ならwith torch.no_grad()で囲む
-        / call the model inside, so if necessary, surround it with torch.no_grad()
-        """
-        for module in self.unet_modules:
-            module.set_cond_image(cond_image)
-
-    def set_batch_cond_only(self, cond_only, zeros):
-        for module in self.unet_modules:
-            module.set_batch_cond_only(cond_only, zeros)
-
-    def set_multiplier(self, multiplier):
-        for module in self.unet_modules:
-            module.multiplier = multiplier
-
-    def load_weights(self, file):
-        if os.path.splitext(file)[1] == ".safetensors":
-            from safetensors.torch import load_file
-
-            weights_sd = load_file(file)
-        else:
-            weights_sd = torch.load(file, map_location="cpu")
-
-        info = self.load_state_dict(weights_sd, False)
-        return info
-
-    def apply_to(self):
-        print("applying LLLite for U-Net...")
-        for module in self.unet_modules:
-            module.apply_to()
-            self.add_module(module.lllite_name, module)
-
-    # マージできるかどうかを返す
-    def is_mergeable(self):
-        return False
-
-    def merge_to(self, text_encoder, unet, weights_sd, dtype, device):
-        raise NotImplementedError()
-
-    def enable_gradient_checkpointing(self):
-        # not supported
-        pass
-
-    def prepare_optimizer_params(self):
-        self.requires_grad_(True)
-        return self.parameters()
-
-    def prepare_grad_etc(self):
-        self.requires_grad_(True)
-
-    def on_epoch_start(self):
-        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
-
-            save_file(state_dict, file, metadata)
-        else:
-            torch.save(state_dict, file)
-
-
-if __name__ == "__main__":
-    # デバッグ用 / for debug
-
-    # sdxl_original_unet.USE_REENTRANT = False
-
-    # test shape etc
-    print("create unet")
-    unet = sdxl_original_unet.SdxlUNet2DConditionModel()
-    unet.to("cuda").to(torch.float16)
-
-    print("create ControlNet-LLLite")
-    control_net = ControlNetLLLite(unet, 32, 64)
-    control_net.apply_to()
-    control_net.to("cuda")
-
-    print(control_net)
-
-    # print number of parameters
-    print("number of parameters", sum(p.numel() for p in control_net.parameters() if p.requires_grad))
-
-    input()
-
-    unet.set_use_memory_efficient_attention(True, False)
-    unet.set_gradient_checkpointing(True)
-    unet.train()  # for gradient checkpointing
-
-    control_net.train()
-
-    # # visualize
-    # import torchviz
-    # print("run visualize")
-    # controlnet.set_control(conditioning_image)
-    # output = unet(x, t, ctx, y)
-    # print("make_dot")
-    # image = torchviz.make_dot(output, params=dict(controlnet.named_parameters()))
-    # print("render")
-    # image.format = "svg" # "png"
-    # image.render("NeuralNet") # すごく時間がかかるので注意 / be careful because it takes a long time
-    # input()
-
-    import bitsandbytes
-
-    optimizer = bitsandbytes.adam.Adam8bit(control_net.prepare_optimizer_params(), 1e-3)
-
-    scaler = torch.cuda.amp.GradScaler(enabled=True)
-
-    print("start training")
-    steps = 10
-
-    sample_param = [p for p in control_net.named_parameters() if "up" in p[0]][0]
-    for step in range(steps):
-        print(f"step {step}")
-
-        batch_size = 1
-        conditioning_image = torch.rand(batch_size, 3, 1024, 1024).cuda() * 2.0 - 1.0
-        x = torch.randn(batch_size, 4, 128, 128).cuda()
-        t = torch.randint(low=0, high=10, size=(batch_size,)).cuda()
-        ctx = torch.randn(batch_size, 77, 2048).cuda()
-        y = torch.randn(batch_size, sdxl_original_unet.ADM_IN_CHANNELS).cuda()
-
-        with torch.cuda.amp.autocast(enabled=True):
-            control_net.set_cond_image(conditioning_image)
-
-            output = unet(x, t, ctx, y)
-            target = torch.randn_like(output)
-            loss = torch.nn.functional.mse_loss(output, target)
-
-        scaler.scale(loss).backward()
-        scaler.step(optimizer)
-        scaler.update()
-        optimizer.zero_grad(set_to_none=True)
-        print(sample_param)
-
-    # from safetensors.torch import save_file
-
-    # save_file(control_net.state_dict(), "logs/control_net.safetensors")

networks/control_net_lllite_for_train.py

@@ -1,502 +0,0 @@
-# cond_imageをU-Netのforwardで渡すバージョンのControlNet-LLLite検証用実装
-# ControlNet-LLLite implementation for verification with cond_image passed in U-Net's forward
-
-import os
-import re
-from typing import Optional, List, Type
-import torch
-from library import sdxl_original_unet
-
-
-# input_blocksに適用するかどうか / if True, input_blocks are not applied
-SKIP_INPUT_BLOCKS = False
-
-# output_blocksに適用するかどうか / if True, output_blocks are not applied
-SKIP_OUTPUT_BLOCKS = True
-
-# conv2dに適用するかどうか / if True, conv2d are not applied
-SKIP_CONV2D = False
-
-# transformer_blocksのみに適用するかどうか。Trueの場合、ResBlockには適用されない
-# if True, only transformer_blocks are applied, and ResBlocks are not applied
-TRANSFORMER_ONLY = True  # if True, SKIP_CONV2D is ignored because conv2d is not used in transformer_blocks
-
-# Trueならattn1とattn2にのみ適用し、ffなどには適用しない / if True, apply only to attn1 and attn2, not to ff etc.
-ATTN1_2_ONLY = True
-
-# Trueならattn1のQKV、attn2のQにのみ適用する、ATTN1_2_ONLY指定時のみ有効 / if True, apply only to attn1 QKV and attn2 Q, only valid when ATTN1_2_ONLY is specified
-ATTN_QKV_ONLY = True
-
-# Trueならattn1やffなどにのみ適用し、attn2などには適用しない / if True, apply only to attn1 and ff, not to attn2
-# ATTN1_2_ONLYと同時にTrueにできない / cannot be True at the same time as ATTN1_2_ONLY
-ATTN1_ETC_ONLY = False  # True
-
-# transformer_blocksの最大インデックス。Noneなら全てのtransformer_blocksに適用
-# max index of transformer_blocks. if None, apply to all transformer_blocks
-TRANSFORMER_MAX_BLOCK_INDEX = None
-
-ORIGINAL_LINEAR = torch.nn.Linear
-ORIGINAL_CONV2D = torch.nn.Conv2d
-
-
-def add_lllite_modules(module: torch.nn.Module, in_dim: int, depth, cond_emb_dim, mlp_dim) -> None:
-    # conditioning1はconditioning imageを embedding する。timestepごとに呼ばれない
-    # conditioning1 embeds conditioning image. it is not called for each timestep
-    modules = []
-    modules.append(ORIGINAL_CONV2D(3, cond_emb_dim // 2, kernel_size=4, stride=4, padding=0))  # to latent (from VAE) size
-    if depth == 1:
-        modules.append(torch.nn.ReLU(inplace=True))
-        modules.append(ORIGINAL_CONV2D(cond_emb_dim // 2, cond_emb_dim, kernel_size=2, stride=2, padding=0))
-    elif depth == 2:
-        modules.append(torch.nn.ReLU(inplace=True))
-        modules.append(ORIGINAL_CONV2D(cond_emb_dim // 2, cond_emb_dim, kernel_size=4, stride=4, padding=0))
-    elif depth == 3:
-        # kernel size 8は大きすぎるので、4にする / kernel size 8 is too large, so set it to 4
-        modules.append(torch.nn.ReLU(inplace=True))
-        modules.append(ORIGINAL_CONV2D(cond_emb_dim // 2, cond_emb_dim // 2, kernel_size=4, stride=4, padding=0))
-        modules.append(torch.nn.ReLU(inplace=True))
-        modules.append(ORIGINAL_CONV2D(cond_emb_dim // 2, cond_emb_dim, kernel_size=2, stride=2, padding=0))
-
-    module.lllite_conditioning1 = torch.nn.Sequential(*modules)
-
-    # downで入力の次元数を削減する。LoRAにヒントを得ていることにする
-    # midでconditioning image embeddingと入力を結合する
-    # upで元の次元数に戻す
-    # これらはtimestepごとに呼ばれる
-    # reduce the number of input dimensions with down. inspired by LoRA
-    # combine conditioning image embedding and input with mid
-    # restore to the original dimension with up
-    # these are called for each timestep
-
-    module.lllite_down = torch.nn.Sequential(
-        ORIGINAL_LINEAR(in_dim, mlp_dim),
-        torch.nn.ReLU(inplace=True),
-    )
-    module.lllite_mid = torch.nn.Sequential(
-        ORIGINAL_LINEAR(mlp_dim + cond_emb_dim, mlp_dim),
-        torch.nn.ReLU(inplace=True),
-    )
-    module.lllite_up = torch.nn.Sequential(
-        ORIGINAL_LINEAR(mlp_dim, in_dim),
-    )
-
-    # Zero-Convにする / set to Zero-Conv
-    torch.nn.init.zeros_(module.lllite_up[0].weight)  # zero conv
-
-
-class LLLiteLinear(ORIGINAL_LINEAR):
-    def __init__(self, in_features: int, out_features: int, **kwargs):
-        super().__init__(in_features, out_features, **kwargs)
-        self.enabled = False
-
-    def set_lllite(self, depth, cond_emb_dim, name, mlp_dim, dropout=None, multiplier=1.0):
-        self.enabled = True
-        self.lllite_name = name
-        self.cond_emb_dim = cond_emb_dim
-        self.dropout = dropout
-        self.multiplier = multiplier  # ignored
-
-        in_dim = self.in_features
-        add_lllite_modules(self, in_dim, depth, cond_emb_dim, mlp_dim)
-
-        self.cond_image = None
-        self.cond_emb = None
-
-    def set_cond_image(self, cond_image):
-        self.cond_image = cond_image
-        self.cond_emb = None
-
-    def forward(self, x):
-        if not self.enabled:
-            return super().forward(x)
-
-        if self.cond_emb is None:
-            self.cond_emb = self.lllite_conditioning1(self.cond_image)
-        cx = self.cond_emb
-
-        # reshape / b,c,h,w -> b,h*w,c
-        n, c, h, w = cx.shape
-        cx = cx.view(n, c, h * w).permute(0, 2, 1)
-
-        cx = torch.cat([cx, self.lllite_down(x)], dim=2)
-        cx = self.lllite_mid(cx)
-
-        if self.dropout is not None and self.training:
-            cx = torch.nn.functional.dropout(cx, p=self.dropout)
-
-        cx = self.lllite_up(cx) * self.multiplier
-
-        x = super().forward(x + cx)  # ここで元のモジュールを呼び出す / call the original module here
-        return x
-
-
-class LLLiteConv2d(ORIGINAL_CONV2D):
-    def __init__(self, in_channels: int, out_channels: int, kernel_size, **kwargs):
-        super().__init__(in_channels, out_channels, kernel_size, **kwargs)
-        self.enabled = False
-
-    def set_lllite(self, depth, cond_emb_dim, name, mlp_dim, dropout=None, multiplier=1.0):
-        self.enabled = True
-        self.lllite_name = name
-        self.cond_emb_dim = cond_emb_dim
-        self.dropout = dropout
-        self.multiplier = multiplier  # ignored
-
-        in_dim = self.in_channels
-        add_lllite_modules(self, in_dim, depth, cond_emb_dim, mlp_dim)
-
-        self.cond_image = None
-        self.cond_emb = None
-
-    def set_cond_image(self, cond_image):
-        self.cond_image = cond_image
-        self.cond_emb = None
-
-    def forward(self, x):  # , cond_image=None):
-        if not self.enabled:
-            return super().forward(x)
-
-        if self.cond_emb is None:
-            self.cond_emb = self.lllite_conditioning1(self.cond_image)
-        cx = self.cond_emb
-
-        cx = torch.cat([cx, self.down(x)], dim=1)
-        cx = self.mid(cx)
-
-        if self.dropout is not None and self.training:
-            cx = torch.nn.functional.dropout(cx, p=self.dropout)
-
-        cx = self.up(cx) * self.multiplier
-
-        x = super().forward(x + cx)  # ここで元のモジュールを呼び出す / call the original module here
-        return x
-
-
-class SdxlUNet2DConditionModelControlNetLLLite(sdxl_original_unet.SdxlUNet2DConditionModel):
-    UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel"]
-    UNET_TARGET_REPLACE_MODULE_CONV2D_3X3 = ["ResnetBlock2D", "Downsample2D", "Upsample2D"]
-    LLLITE_PREFIX = "lllite_unet"
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def apply_lllite(
-        self,
-        cond_emb_dim: int = 16,
-        mlp_dim: int = 16,
-        dropout: Optional[float] = None,
-        varbose: Optional[bool] = False,
-        multiplier: Optional[float] = 1.0,
-    ) -> None:
-        def apply_to_modules(
-            root_module: torch.nn.Module,
-            target_replace_modules: List[torch.nn.Module],
-        ) -> List[torch.nn.Module]:
-            prefix = "lllite_unet"
-
-            modules = []
-            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():
-                        is_linear = child_module.__class__.__name__ == "LLLiteLinear"
-                        is_conv2d = child_module.__class__.__name__ == "LLLiteConv2d"
-
-                        if is_linear or (is_conv2d and not SKIP_CONV2D):
-                            # block indexからdepthを計算: depthはconditioningのサイズやチャネルを計算するのに使う
-                            # block index to depth: depth is using to calculate conditioning size and channels
-                            block_name, index1, index2 = (name + "." + child_name).split(".")[:3]
-                            index1 = int(index1)
-                            if block_name == "input_blocks":
-                                if SKIP_INPUT_BLOCKS:
-                                    continue
-                                depth = 1 if index1 <= 2 else (2 if index1 <= 5 else 3)
-                            elif block_name == "middle_block":
-                                depth = 3
-                            elif block_name == "output_blocks":
-                                if SKIP_OUTPUT_BLOCKS:
-                                    continue
-                                depth = 3 if index1 <= 2 else (2 if index1 <= 5 else 1)
-                                if int(index2) >= 2:
-                                    depth -= 1
-                            else:
-                                raise NotImplementedError()
-
-                            lllite_name = prefix + "." + name + "." + child_name
-                            lllite_name = lllite_name.replace(".", "_")
-
-                            if TRANSFORMER_MAX_BLOCK_INDEX is not None:
-                                p = lllite_name.find("transformer_blocks")
-                                if p >= 0:
-                                    tf_index = int(lllite_name[p:].split("_")[2])
-                                    if tf_index > TRANSFORMER_MAX_BLOCK_INDEX:
-                                        continue
-
-                            #  time embは適用外とする
-                            # attn2のconditioning (CLIPからの入力) はshapeが違うので適用できない
-                            # time emb is not applied
-                            # attn2 conditioning (input from CLIP) cannot be applied because the shape is different
-                            if "emb_layers" in lllite_name or (
-                                "attn2" in lllite_name and ("to_k" in lllite_name or "to_v" in lllite_name)
-                            ):
-                                continue
-
-                            if ATTN1_2_ONLY:
-                                if not ("attn1" in lllite_name or "attn2" in lllite_name):
-                                    continue
-                                if ATTN_QKV_ONLY:
-                                    if "to_out" in lllite_name:
-                                        continue
-
-                            if ATTN1_ETC_ONLY:
-                                if "proj_out" in lllite_name:
-                                    pass
-                                elif "attn1" in lllite_name and (
-                                    "to_k" in lllite_name or "to_v" in lllite_name or "to_out" in lllite_name
-                                ):
-                                    pass
-                                elif "ff_net_2" in lllite_name:
-                                    pass
-                                else:
-                                    continue
-
-                            child_module.set_lllite(depth, cond_emb_dim, lllite_name, mlp_dim, dropout, multiplier)
-                            modules.append(child_module)
-
-            return modules
-
-        target_modules = SdxlUNet2DConditionModelControlNetLLLite.UNET_TARGET_REPLACE_MODULE
-        if not TRANSFORMER_ONLY:
-            target_modules = target_modules + SdxlUNet2DConditionModelControlNetLLLite.UNET_TARGET_REPLACE_MODULE_CONV2D_3X3
-
-        # create module instances
-        self.lllite_modules = apply_to_modules(self, target_modules)
-        print(f"enable ControlNet LLLite for U-Net: {len(self.lllite_modules)} modules.")
-
-    # def prepare_optimizer_params(self):
-    def prepare_params(self):
-        train_params = []
-        non_train_params = []
-        for name, p in self.named_parameters():
-            if "lllite" in name:
-                train_params.append(p)
-            else:
-                non_train_params.append(p)
-        print(f"count of trainable parameters: {len(train_params)}")
-        print(f"count of non-trainable parameters: {len(non_train_params)}")
-
-        for p in non_train_params:
-            p.requires_grad_(False)
-
-        # without this, an error occurs in the optimizer
-        #       RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
-        non_train_params[0].requires_grad_(True)
-
-        for p in train_params:
-            p.requires_grad_(True)
-
-        return train_params
-
-    # def prepare_grad_etc(self):
-    #     self.requires_grad_(True)
-
-    # def on_epoch_start(self):
-    #     self.train()
-
-    def get_trainable_params(self):
-        return [p[1] for p in self.named_parameters() if "lllite" in p[0]]
-
-    def save_lllite_weights(self, file, dtype, metadata):
-        if metadata is not None and len(metadata) == 0:
-            metadata = None
-
-        org_state_dict = self.state_dict()
-
-        # copy LLLite keys from org_state_dict to state_dict with key conversion
-        state_dict = {}
-        for key in org_state_dict.keys():
-            # split with ".lllite"
-            pos = key.find(".lllite")
-            if pos < 0:
-                continue
-            lllite_key = SdxlUNet2DConditionModelControlNetLLLite.LLLITE_PREFIX + "." + key[:pos]
-            lllite_key = lllite_key.replace(".", "_") + key[pos:]
-            lllite_key = lllite_key.replace(".lllite_", ".")
-            state_dict[lllite_key] = org_state_dict[key]
-
-        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
-
-            save_file(state_dict, file, metadata)
-        else:
-            torch.save(state_dict, file)
-
-    def load_lllite_weights(self, file, non_lllite_unet_sd=None):
-        r"""
-        LLLiteの重みを読み込まない（initされた値を使う）場合はfileにNoneを指定する。
-        この場合、non_lllite_unet_sdにはU-Netのstate_dictを指定する。
-
-        If you do not want to load LLLite weights (use initialized values), specify None for file.
-        In this case, specify the state_dict of U-Net for non_lllite_unet_sd.
-        """
-        if not file:
-            state_dict = self.state_dict()
-            for key in non_lllite_unet_sd:
-                if key in state_dict:
-                    state_dict[key] = non_lllite_unet_sd[key]
-            info = self.load_state_dict(state_dict, False)
-            return info
-
-        if os.path.splitext(file)[1] == ".safetensors":
-            from safetensors.torch import load_file
-
-            weights_sd = load_file(file)
-        else:
-            weights_sd = torch.load(file, map_location="cpu")
-
-        # module_name = module_name.replace("_block", "@blocks")
-        # module_name = module_name.replace("_layer", "@layer")
-        # module_name = module_name.replace("to_", "to@")
-        # module_name = module_name.replace("time_embed", "time@embed")
-        # module_name = module_name.replace("label_emb", "label@emb")
-        # module_name = module_name.replace("skip_connection", "skip@connection")
-        # module_name = module_name.replace("proj_in", "proj@in")
-        # module_name = module_name.replace("proj_out", "proj@out")
-        pattern = re.compile(r"(_block|_layer|to_|time_embed|label_emb|skip_connection|proj_in|proj_out)")
-
-        # convert to lllite with U-Net state dict
-        state_dict = non_lllite_unet_sd.copy() if non_lllite_unet_sd is not None else {}
-        for key in weights_sd.keys():
-            # split with "."
-            pos = key.find(".")
-            if pos < 0:
-                continue
-
-            module_name = key[:pos]
-            weight_name = key[pos + 1 :]  # exclude "."
-            module_name = module_name.replace(SdxlUNet2DConditionModelControlNetLLLite.LLLITE_PREFIX + "_", "")
-
-            # これはうまくいかない。逆変換を考えなかった設計が悪い / this does not work well. bad design because I didn't think about inverse conversion
-            # module_name = module_name.replace("_", ".")
-
-            # ださいけどSDXLのU-Netの "_" を "@" に変換する / ugly but convert "_" of SDXL U-Net to "@"
-            matches = pattern.findall(module_name)
-            if matches is not None:
-                for m in matches:
-                    print(module_name, m)
-                    module_name = module_name.replace(m, m.replace("_", "@"))
-            module_name = module_name.replace("_", ".")
-            module_name = module_name.replace("@", "_")
-
-            lllite_key = module_name + ".lllite_" + weight_name
-
-            state_dict[lllite_key] = weights_sd[key]
-
-        info = self.load_state_dict(state_dict, False)
-        return info
-
-    def forward(self, x, timesteps=None, context=None, y=None, cond_image=None, **kwargs):
-        for m in self.lllite_modules:
-            m.set_cond_image(cond_image)
-        return super().forward(x, timesteps, context, y, **kwargs)
-
-
-def replace_unet_linear_and_conv2d():
-    print("replace torch.nn.Linear and torch.nn.Conv2d to LLLiteLinear and LLLiteConv2d in U-Net")
-    sdxl_original_unet.torch.nn.Linear = LLLiteLinear
-    sdxl_original_unet.torch.nn.Conv2d = LLLiteConv2d
-
-
-if __name__ == "__main__":
-    # デバッグ用 / for debug
-
-    # sdxl_original_unet.USE_REENTRANT = False
-    replace_unet_linear_and_conv2d()
-
-    # test shape etc
-    print("create unet")
-    unet = SdxlUNet2DConditionModelControlNetLLLite()
-
-    print("enable ControlNet-LLLite")
-    unet.apply_lllite(32, 64, None, False, 1.0)
-    unet.to("cuda")  # .to(torch.float16)
-
-    # from safetensors.torch import load_file
-
-    # model_sd = load_file(r"E:\Work\SD\Models\sdxl\sd_xl_base_1.0_0.9vae.safetensors")
-    # unet_sd = {}
-
-    # # copy U-Net keys from unet_state_dict to state_dict
-    # prefix = "model.diffusion_model."
-    # for key in model_sd.keys():
-    #     if key.startswith(prefix):
-    #         converted_key = key[len(prefix) :]
-    #         unet_sd[converted_key] = model_sd[key]
-
-    # info = unet.load_lllite_weights("r:/lllite_from_unet.safetensors", unet_sd)
-    # print(info)
-
-    # print(unet)
-
-    # print number of parameters
-    params = unet.prepare_params()
-    print("number of parameters", sum(p.numel() for p in params))
-    # print("type any key to continue")
-    # input()
-
-    unet.set_use_memory_efficient_attention(True, False)
-    unet.set_gradient_checkpointing(True)
-    unet.train()  # for gradient checkpointing
-
-    # # visualize
-    # import torchviz
-    # print("run visualize")
-    # controlnet.set_control(conditioning_image)
-    # output = unet(x, t, ctx, y)
-    # print("make_dot")
-    # image = torchviz.make_dot(output, params=dict(controlnet.named_parameters()))
-    # print("render")
-    # image.format = "svg" # "png"
-    # image.render("NeuralNet") # すごく時間がかかるので注意 / be careful because it takes a long time
-    # input()
-
-    import bitsandbytes
-
-    optimizer = bitsandbytes.adam.Adam8bit(params, 1e-3)
-
-    scaler = torch.cuda.amp.GradScaler(enabled=True)
-
-    print("start training")
-    steps = 10
-    batch_size = 1
-
-    sample_param = [p for p in unet.named_parameters() if ".lllite_up." in p[0]][0]
-    for step in range(steps):
-        print(f"step {step}")
-
-        conditioning_image = torch.rand(batch_size, 3, 1024, 1024).cuda() * 2.0 - 1.0
-        x = torch.randn(batch_size, 4, 128, 128).cuda()
-        t = torch.randint(low=0, high=10, size=(batch_size,)).cuda()
-        ctx = torch.randn(batch_size, 77, 2048).cuda()
-        y = torch.randn(batch_size, sdxl_original_unet.ADM_IN_CHANNELS).cuda()
-
-        with torch.cuda.amp.autocast(enabled=True, dtype=torch.bfloat16):
-            output = unet(x, t, ctx, y, conditioning_image)
-            target = torch.randn_like(output)
-            loss = torch.nn.functional.mse_loss(output, target)
-
-        scaler.scale(loss).backward()
-        scaler.step(optimizer)
-        scaler.update()
-        optimizer.zero_grad(set_to_none=True)
-        print(sample_param)
-
-    # from safetensors.torch import save_file
-
-    # print("save weights")
-    # unet.save_lllite_weights("r:/lllite_from_unet.safetensors", torch.float16, None)

networks/dylora.py

@@ -239,7 +239,7 @@ def create_network_from_weights(multiplier, file, vae, text_encoder, unet, weigh
 
 
 class DyLoRANetwork(torch.nn.Module):
-    UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel"]
+    UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"]
     UNET_TARGET_REPLACE_MODULE_CONV2D_3X3 = ["ResnetBlock2D", "Downsample2D", "Upsample2D"]
     TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
     LORA_PREFIX_UNET = "lora_unet"

networks/extract_lora_from_models.py

@@ -3,265 +3,187 @@
 # Thanks to cloneofsimo!
 
 import argparse
-import json
 import os
-import time
 import torch
 from safetensors.torch import load_file, save_file
 from tqdm import tqdm
-from library import sai_model_spec, model_util, sdxl_model_util
+import library.model_util as model_util
 import lora
 
 
 CLAMP_QUANTILE = 0.99
-MIN_DIFF = 1e-4
+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 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)
+  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
+  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
 
-    assert args.v2 != args.sdxl or (
-        not args.v2 and not args.sdxl
-    ), "v2 and sdxl cannot be specified at the same time / v2とsdxlは同時に指定できません"
-    if args.v_parameterization is None:
-        args.v_parameterization = args.v2
+  save_dtype = str_to_dtype(args.save_precision)
 
-    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)
 
-    # load models
-    if not args.sdxl:
-        print(f"loading original SD model : {args.model_org}")
-        text_encoder_o, _, unet_o = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.model_org)
-        text_encoders_o = [text_encoder_o]
-        print(f"loading tuned SD model : {args.model_tuned}")
-        text_encoder_t, _, unet_t = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.model_tuned)
-        text_encoders_t = [text_encoder_t]
-        model_version = model_util.get_model_version_str_for_sd1_sd2(args.v2, args.v_parameterization)
-    else:
-        print(f"loading original SDXL model : {args.model_org}")
-        text_encoder_o1, text_encoder_o2, _, unet_o, _, _ = sdxl_model_util.load_models_from_sdxl_checkpoint(
-            sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0, args.model_org, "cpu"
-        )
-        text_encoders_o = [text_encoder_o1, text_encoder_o2]
-        print(f"loading original SDXL model : {args.model_tuned}")
-        text_encoder_t1, text_encoder_t2, _, unet_t, _, _ = sdxl_model_util.load_models_from_sdxl_checkpoint(
-            sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0, args.model_tuned, "cpu"
-        )
-        text_encoders_t = [text_encoder_t1, text_encoder_t2]
-        model_version = sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0
+  # create LoRA network to extract weights: Use dim (rank) as alpha
+  if args.conv_dim is None:
+    kwargs = {}
+  else:
+    kwargs = {"conv_dim": args.conv_dim, "conv_alpha": args.conv_dim}
 
-    # create LoRA network to extract weights: Use dim (rank) as alpha
-    if args.conv_dim is None:
-        kwargs = {}
-    else:
-        kwargs = {"conv_dim": args.conv_dim, "conv_alpha": args.conv_dim}
+  lora_network_o = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_o, unet_o, **kwargs)
+  lora_network_t = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_t, unet_t, **kwargs)
+  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ベース） "
 
-    lora_network_o = lora.create_network(1.0, args.dim, args.dim, None, text_encoders_o, unet_o, **kwargs)
-    lora_network_t = lora.create_network(1.0, args.dim, args.dim, None, text_encoders_t, unet_t, **kwargs)
-    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
 
-    # get diffs
+    # 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 = {}
-    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 not text_encoder_different and torch.max(torch.abs(diff)) > MIN_DIFF:
-            text_encoder_different = True
-            print(f"Text encoder is different. {torch.max(torch.abs(diff))} > {MIN_DIFF}")
+  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()
 
-        diff = diff.float()
-        diffs[lora_name] = diff
+    if args.device:
+      diff = diff.to(args.device)
 
-    if not text_encoder_different:
-        print("Text encoder is same. Extract U-Net only.")
-        lora_network_o.text_encoder_loras = []
-        diffs = {}
+    diffs[lora_name] = diff
 
-    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()
+  # make LoRA with svd
+  print("calculating by svd")
+  lora_weights = {}
+  with torch.no_grad():
+    for lora_name, mat in tqdm(list(diffs.items())):
+      # if args.conv_dim is None, diffs do not include LoRAs for conv2d-3x3
+      conv2d = (len(mat.size()) == 4)
+      kernel_size = None if not conv2d else mat.size()[2:4]
+      conv2d_3x3 = conv2d and kernel_size != (1, 1)
 
-        if args.device:
-            diff = diff.to(args.device)
+      rank = args.dim if not conv2d_3x3 or args.conv_dim is None else args.conv_dim
+      out_dim, in_dim = mat.size()[0:2]
 
-        diffs[lora_name] = diff
+      if args.device:
+        mat = mat.to(args.device)
 
-    # make LoRA with svd
-    print("calculating by svd")
-    lora_weights = {}
-    with torch.no_grad():
-        for lora_name, mat in tqdm(list(diffs.items())):
-            # if args.conv_dim is None, diffs do not include LoRAs for conv2d-3x3
-            conv2d = len(mat.size()) == 4
-            kernel_size = None if not conv2d else mat.size()[2:4]
-            conv2d_3x3 = conv2d and kernel_size != (1, 1)
+      # print(lora_name, mat.size(), mat.device, rank, in_dim, out_dim)
+      rank = min(rank, in_dim, out_dim)                           # LoRA rank cannot exceed the original dim
 
-            rank = args.dim if not conv2d_3x3 or args.conv_dim is None else args.conv_dim
-            out_dim, in_dim = mat.size()[0:2]
+      if conv2d:
+        if conv2d_3x3:
+          mat = mat.flatten(start_dim=1)
+        else:
+          mat = mat.squeeze()
 
-            if args.device:
-                mat = mat.to(args.device)
+      U, S, Vh = torch.linalg.svd(mat)
 
-            # print(lora_name, mat.size(), mat.device, rank, in_dim, out_dim)
-            rank = min(rank, in_dim, out_dim)  # LoRA rank cannot exceed the original dim
+      U = U[:, :rank]
+      S = S[:rank]
+      U = U @ torch.diag(S)
 
-            if conv2d:
-                if conv2d_3x3:
-                    mat = mat.flatten(start_dim=1)
-                else:
-                    mat = mat.squeeze()
+      Vh = Vh[:rank, :]
 
-            U, S, Vh = torch.linalg.svd(mat)
+      dist = torch.cat([U.flatten(), Vh.flatten()])
+      hi_val = torch.quantile(dist, CLAMP_QUANTILE)
+      low_val = -hi_val
 
-            U = U[:, :rank]
-            S = S[:rank]
-            U = U @ torch.diag(S)
+      U = U.clamp(low_val, hi_val)
+      Vh = Vh.clamp(low_val, hi_val)
 
-            Vh = Vh[:rank, :]
+      if conv2d:
+        U = U.reshape(out_dim, rank, 1, 1)
+        Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1])
 
-            dist = torch.cat([U.flatten(), Vh.flatten()])
-            hi_val = torch.quantile(dist, CLAMP_QUANTILE)
-            low_val = -hi_val
+      U = U.to("cpu").contiguous()
+      Vh = Vh.to("cpu").contiguous()
 
-            U = U.clamp(low_val, hi_val)
-            Vh = Vh.clamp(low_val, hi_val)
+      lora_weights[lora_name] = (U, Vh)
 
-            if conv2d:
-                U = U.reshape(out_dim, rank, 1, 1)
-                Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1])
+  # make state dict for LoRA
+  lora_sd = {}
+  for lora_name, (up_weight, down_weight) in lora_weights.items():
+    lora_sd[lora_name + '.lora_up.weight'] = up_weight
+    lora_sd[lora_name + '.lora_down.weight'] = down_weight
+    lora_sd[lora_name + '.alpha'] = torch.tensor(down_weight.size()[0])
 
-            U = U.to("cpu").contiguous()
-            Vh = Vh.to("cpu").contiguous()
+  # load state dict to LoRA and save it
+  lora_network_save, lora_sd = lora.create_network_from_weights(1.0, None, None, text_encoder_o, unet_o, weights_sd=lora_sd)
+  lora_network_save.apply_to(text_encoder_o, unet_o)  # create internal module references for state_dict  
 
-            lora_weights[lora_name] = (U, Vh)
+  info = lora_network_save.load_state_dict(lora_sd)
+  print(f"Loading extracted LoRA weights: {info}")
 
-    # make state dict for LoRA
-    lora_sd = {}
-    for lora_name, (up_weight, down_weight) in lora_weights.items():
-        lora_sd[lora_name + ".lora_up.weight"] = up_weight
-        lora_sd[lora_name + ".lora_down.weight"] = down_weight
-        lora_sd[lora_name + ".alpha"] = torch.tensor(down_weight.size()[0])
+  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)
 
-    # load state dict to LoRA and save it
-    lora_network_save, lora_sd = lora.create_network_from_weights(1.0, None, None, text_encoders_o, unet_o, weights_sd=lora_sd)
-    lora_network_save.apply_to(text_encoders_o, unet_o)  # create internal module references for state_dict
+  # minimum metadata
+  metadata = {"ss_network_module": "networks.lora", "ss_network_dim": str(args.dim), "ss_network_alpha": str(args.dim)}
 
-    info = lora_network_save.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
-    net_kwargs = {}
-    if args.conv_dim is not None:
-        net_kwargs["conv_dim"] = args.conv_dim
-        net_kwargs["conv_alpha"] = args.conv_dim
-
-    metadata = {
-        "ss_v2": str(args.v2),
-        "ss_base_model_version": model_version,
-        "ss_network_module": "networks.lora",
-        "ss_network_dim": str(args.dim),
-        "ss_network_alpha": str(args.dim),
-        "ss_network_args": json.dumps(net_kwargs),
-    }
-
-    if not args.no_metadata:
-        title = os.path.splitext(os.path.basename(args.save_to))[0]
-        sai_metadata = sai_model_spec.build_metadata(
-            None, args.v2, args.v_parameterization, False, True, False, time.time(), title=title
-        )
-        metadata.update(sai_metadata)
-
-    lora_network_save.save_weights(args.save_to, save_dtype, metadata)
-    print(f"LoRA weights are saved to: {args.save_to}")
+  lora_network_save.save_weights(args.save_to, save_dtype, metadata)
+  print(f"LoRA weights are saved to: {args.save_to}")
 
 
 def setup_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--v2", action="store_true", help="load Stable Diffusion v2.x model / Stable Diffusion 2.xのモデルを読み込む")
-    parser.add_argument(
-        "--v_parameterization",
-        type=bool,
-        default=None,
-        help="make LoRA metadata for v-parameterization (default is same to v2) / 作成するLoRAのメタデータにv-parameterization用と設定する（省略時はv2と同じ）",
-    )
-    parser.add_argument(
-        "--sdxl", action="store_true", help="load Stable Diffusion SDXL base model / Stable Diffusion SDXL baseのモデルを読み込む"
-    )
-    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(
-        "--conv_dim",
-        type=int,
-        default=None,
-        help="dimension (rank) of LoRA for Conv2d-3x3 (default None, disabled) / LoRAのConv2d-3x3の次元数（rank）（デフォルトNone、適用なし）",
-    )
-    parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")
-    parser.add_argument(
-        "--no_metadata",
-        action="store_true",
-        help="do not save sai modelspec metadata (minimum ss_metadata for LoRA is saved) / "
-        + "sai modelspecのメタデータを保存しない（LoRAの最低限のss_metadataは保存される）",
-    )
+  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("--conv_dim", type=int, default=None,
+                      help="dimension (rank) of LoRA for Conv2d-3x3 (default None, disabled) / LoRAのConv2d-3x3の次元数（rank）（デフォルトNone、適用なし）")
+  parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")
 
-    return parser
+  return parser
 
 
-if __name__ == "__main__":
-    parser = setup_parser()
+if __name__ == '__main__':
+  parser = setup_parser()
 
-    args = parser.parse_args()
-    svd(args)
+  args = parser.parse_args()
+  svd(args)

networks/lora.py

@@ -5,9 +5,7 @@
 
 import math
 import os
-from typing import Dict, List, Optional, Tuple, Type, Union
-from diffusers import AutoencoderKL
-from transformers import CLIPTextModel
+from typing import List, Tuple, Union
 import numpy as np
 import torch
 import re
@@ -241,13 +239,9 @@ class LoRAInfModule(LoRAModule):
         else:
             area = x.size()[1]
 
-        mask = self.network.mask_dic.get(area, None)
+        mask = self.network.mask_dic[area]
         if mask is None:
-            # raise ValueError(f"mask is None for resolution {area}")
-            # emb_layers in SDXL doesn't have mask
-            # print(f"mask is None for resolution {area}, {x.size()}")
-            mask_size = (1, x.size()[1]) if len(x.size()) == 2 else (1, *x.size()[1:-1], 1)
-            return torch.ones(mask_size, dtype=x.dtype, device=x.device) / self.network.num_sub_prompts
+            raise ValueError(f"mask is None for resolution {area}")
         if len(x.size()) != 4:
             mask = torch.reshape(mask, (1, -1, 1))
         return mask
@@ -352,10 +346,9 @@ class LoRAInfModule(LoRAModule):
             out[-self.network.batch_size :] = x[-self.network.batch_size :]  # real_uncond
 
         # print("to_out_forward", self.lora_name, self.network.sub_prompt_index, self.network.num_sub_prompts)
-        # if num_sub_prompts > num of LoRAs, fill with zero
-        for i in range(len(masks)):
-            if masks[i] is None:
-                masks[i] = torch.zeros_like(masks[0])
+        # for i in range(len(masks)):
+        #     if masks[i] is None:
+        #         masks[i] = torch.zeros_like(masks[-1])
 
         mask = torch.cat(masks)
         mask_sum = torch.sum(mask, dim=0) + 1e-4
@@ -407,16 +400,7 @@ def parse_block_lr_kwargs(nw_kwargs):
     return down_lr_weight, mid_lr_weight, up_lr_weight
 
 
-def create_network(
-    multiplier: float,
-    network_dim: Optional[int],
-    network_alpha: Optional[float],
-    vae: AutoencoderKL,
-    text_encoder: Union[CLIPTextModel, List[CLIPTextModel]],
-    unet,
-    neuron_dropout: Optional[float] = None,
-    **kwargs,
-):
+def create_network(multiplier, network_dim, network_alpha, vae, text_encoder, unet, neuron_dropout=None, **kwargs):
     if network_dim is None:
         network_dim = 4  # default
     if network_alpha is None:
@@ -735,36 +719,33 @@ def create_network_from_weights(multiplier, file, vae, text_encoder, unet, weigh
 class LoRANetwork(torch.nn.Module):
     NUM_OF_BLOCKS = 12  # フルモデル相当でのup,downの層の数
 
-    UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel"]
+    # is it possible to apply conv_in and conv_out? -> yes, newer LoCon supports it (^^;)
+    UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"]
     UNET_TARGET_REPLACE_MODULE_CONV2D_3X3 = ["ResnetBlock2D", "Downsample2D", "Upsample2D"]
     TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
     LORA_PREFIX_UNET = "lora_unet"
     LORA_PREFIX_TEXT_ENCODER = "lora_te"
 
-    # SDXL: must starts with LORA_PREFIX_TEXT_ENCODER
-    LORA_PREFIX_TEXT_ENCODER1 = "lora_te1"
-    LORA_PREFIX_TEXT_ENCODER2 = "lora_te2"
-
     def __init__(
         self,
-        text_encoder: Union[List[CLIPTextModel], CLIPTextModel],
+        text_encoder,
         unet,
-        multiplier: float = 1.0,
-        lora_dim: int = 4,
-        alpha: float = 1,
-        dropout: Optional[float] = None,
-        rank_dropout: Optional[float] = None,
-        module_dropout: Optional[float] = None,
-        conv_lora_dim: Optional[int] = None,
-        conv_alpha: Optional[float] = None,
-        block_dims: Optional[List[int]] = None,
-        block_alphas: Optional[List[float]] = None,
-        conv_block_dims: Optional[List[int]] = None,
-        conv_block_alphas: Optional[List[float]] = None,
-        modules_dim: Optional[Dict[str, int]] = None,
-        modules_alpha: Optional[Dict[str, int]] = None,
-        module_class: Type[object] = LoRAModule,
-        varbose: Optional[bool] = False,
+        multiplier=1.0,
+        lora_dim=4,
+        alpha=1,
+        dropout=None,
+        rank_dropout=None,
+        module_dropout=None,
+        conv_lora_dim=None,
+        conv_alpha=None,
+        block_dims=None,
+        block_alphas=None,
+        conv_block_dims=None,
+        conv_block_alphas=None,
+        modules_dim=None,
+        modules_alpha=None,
+        module_class=LoRAModule,
+        varbose=False,
     ) -> None:
         """
         LoRA network: すごく引数が多いが、パターンは以下の通り
@@ -802,21 +783,8 @@ class LoRANetwork(torch.nn.Module):
                 print(f"apply LoRA to Conv2d with kernel size (3,3). dim (rank): {self.conv_lora_dim}, alpha: {self.conv_alpha}")
 
         # create module instances
-        def create_modules(
-            is_unet: bool,
-            text_encoder_idx: Optional[int],  # None, 1, 2
-            root_module: torch.nn.Module,
-            target_replace_modules: List[torch.nn.Module],
-        ) -> List[LoRAModule]:
-            prefix = (
-                self.LORA_PREFIX_UNET
-                if is_unet
-                else (
-                    self.LORA_PREFIX_TEXT_ENCODER
-                    if text_encoder_idx is None
-                    else (self.LORA_PREFIX_TEXT_ENCODER1 if text_encoder_idx == 1 else self.LORA_PREFIX_TEXT_ENCODER2)
-                )
-            )
+        def create_modules(is_unet, root_module: torch.nn.Module, target_replace_modules) -> List[LoRAModule]:
+            prefix = LoRANetwork.LORA_PREFIX_UNET if is_unet else LoRANetwork.LORA_PREFIX_TEXT_ENCODER
             loras = []
             skipped = []
             for name, module in root_module.named_modules():
@@ -832,14 +800,11 @@ class LoRANetwork(torch.nn.Module):
 
                             dim = None
                             alpha = None
-
                             if modules_dim is not None:
-                                # モジュール指定あり
                                 if lora_name in modules_dim:
                                     dim = modules_dim[lora_name]
                                     alpha = modules_alpha[lora_name]
                             elif is_unet and block_dims is not None:
-                                # U-Netでblock_dims指定あり
                                 block_idx = get_block_index(lora_name)
                                 if is_linear or is_conv2d_1x1:
                                     dim = block_dims[block_idx]
@@ -848,7 +813,6 @@ class LoRANetwork(torch.nn.Module):
                                     dim = conv_block_dims[block_idx]
                                     alpha = conv_block_alphas[block_idx]
                             else:
-                                # 通常、すべて対象とする
                                 if is_linear or is_conv2d_1x1:
                                     dim = self.lora_dim
                                     alpha = self.alpha
@@ -857,7 +821,6 @@ class LoRANetwork(torch.nn.Module):
                                     alpha = self.conv_alpha
 
                             if dim is None or dim == 0:
-                                # skipした情報を出力
                                 if is_linear or is_conv2d_1x1 or (self.conv_lora_dim is not None or conv_block_dims is not None):
                                     skipped.append(lora_name)
                                 continue
@@ -875,23 +838,7 @@ class LoRANetwork(torch.nn.Module):
                             loras.append(lora)
             return loras, skipped
 
-        text_encoders = text_encoder if type(text_encoder) == list else [text_encoder]
-
-        # create LoRA for text encoder
-        # 毎回すべてのモジュールを作るのは無駄なので要検討
-        self.text_encoder_loras = []
-        skipped_te = []
-        for i, text_encoder in enumerate(text_encoders):
-            if len(text_encoders) > 1:
-                index = i + 1
-                print(f"create LoRA for Text Encoder {index}:")
-            else:
-                index = None
-                print(f"create LoRA for Text Encoder:")
-
-            text_encoder_loras, skipped = create_modules(False, index, text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE)
-            self.text_encoder_loras.extend(text_encoder_loras)
-            skipped_te += skipped
+        self.text_encoder_loras, skipped_te = create_modules(False, text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE)
         print(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
 
         # extend U-Net target modules if conv2d 3x3 is enabled, or load from weights
@@ -899,7 +846,7 @@ class LoRANetwork(torch.nn.Module):
         if modules_dim is not None or self.conv_lora_dim is not None or conv_block_dims is not None:
             target_modules += LoRANetwork.UNET_TARGET_REPLACE_MODULE_CONV2D_3X3
 
-        self.unet_loras, skipped_un = create_modules(True, None, unet, target_modules)
+        self.unet_loras, skipped_un = create_modules(True, unet, target_modules)
         print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
 
         skipped = skipped_te + skipped_un
@@ -1014,7 +961,6 @@ class LoRANetwork(torch.nn.Module):
 
         return lr_weight
 
-    # 二つのText Encoderに別々の学習率を設定できるようにするといいかも
     def prepare_optimizer_params(self, text_encoder_lr, unet_lr, default_lr):
         self.requires_grad_(True)
         all_params = []

networks/lora_diffusers.py

@@ -1,604 +0,0 @@
-# Diffusersで動くLoRA。このファイル単独で完結する。
-# LoRA module for Diffusers. This file works independently.
-
-import bisect
-import math
-import random
-from typing import Any, Dict, List, Mapping, Optional, Union
-from diffusers import UNet2DConditionModel
-import numpy as np
-from tqdm import tqdm
-from transformers import CLIPTextModel
-import torch
-
-
-def make_unet_conversion_map() -> Dict[str, str]:
-    unet_conversion_map_layer = []
-
-    for i in range(3):  # num_blocks is 3 in sdxl
-        # 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: commentout for sdxl
-            # 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}.{2}."  # change for sdxl
-            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))
-
-    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 = []
-    for sd, hf in unet_conversion_map_layer:
-        if "resnets" in hf:
-            for sd_res, hf_res in unet_conversion_map_resnet:
-                unet_conversion_map.append((sd + sd_res, hf + hf_res))
-        else:
-            unet_conversion_map.append((sd, hf))
-
-    for j in range(2):
-        hf_time_embed_prefix = f"time_embedding.linear_{j+1}."
-        sd_time_embed_prefix = f"time_embed.{j*2}."
-        unet_conversion_map.append((sd_time_embed_prefix, hf_time_embed_prefix))
-
-    for j in range(2):
-        hf_label_embed_prefix = f"add_embedding.linear_{j+1}."
-        sd_label_embed_prefix = f"label_emb.0.{j*2}."
-        unet_conversion_map.append((sd_label_embed_prefix, hf_label_embed_prefix))
-
-    unet_conversion_map.append(("input_blocks.0.0.", "conv_in."))
-    unet_conversion_map.append(("out.0.", "conv_norm_out."))
-    unet_conversion_map.append(("out.2.", "conv_out."))
-
-    sd_hf_conversion_map = {sd.replace(".", "_")[:-1]: hf.replace(".", "_")[:-1] for sd, hf in unet_conversion_map}
-    return sd_hf_conversion_map
-
-
-UNET_CONVERSION_MAP = make_unet_conversion_map()
-
-
-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
-
-        if org_module.__class__.__name__ == "Conv2d":
-            in_dim = org_module.in_channels
-            out_dim = org_module.out_channels
-        else:
-            in_dim = org_module.in_features
-            out_dim = org_module.out_features
-
-        self.lora_dim = lora_dim
-
-        if org_module.__class__.__name__ == "Conv2d":
-            kernel_size = org_module.kernel_size
-            stride = org_module.stride
-            padding = org_module.padding
-            self.lora_down = torch.nn.Conv2d(in_dim, self.lora_dim, kernel_size, stride, padding, bias=False)
-            self.lora_up = torch.nn.Conv2d(self.lora_dim, out_dim, (1, 1), (1, 1), bias=False)
-        else:
-            self.lora_down = torch.nn.Linear(in_dim, self.lora_dim, bias=False)
-            self.lora_up = torch.nn.Linear(self.lora_dim, out_dim, bias=False)
-
-        if type(alpha) == torch.Tensor:
-            alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
-        alpha = self.lora_dim if alpha is None or alpha == 0 else alpha
-        self.scale = alpha / self.lora_dim
-        self.register_buffer("alpha", torch.tensor(alpha))  # 勾配計算に含めない / not included in gradient calculation
-
-        # 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]
-        self.enabled = True
-        self.network: LoRANetwork = None
-        self.org_forward = None
-
-    # override org_module's forward method
-    def apply_to(self, multiplier=None):
-        if multiplier is not None:
-            self.multiplier = multiplier
-        if self.org_forward is None:
-            self.org_forward = self.org_module[0].forward
-            self.org_module[0].forward = self.forward
-
-    # restore org_module's forward method
-    def unapply_to(self):
-        if self.org_forward is not None:
-            self.org_module[0].forward = self.org_forward
-
-    # forward with lora
-    def forward(self, x):
-        if not self.enabled:
-            return self.org_forward(x)
-        return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
-
-    def set_network(self, network):
-        self.network = network
-
-    # merge lora weight to org weight
-    def merge_to(self, multiplier=1.0):
-        # get lora weight
-        lora_weight = self.get_weight(multiplier)
-
-        # get org weight
-        org_sd = self.org_module[0].state_dict()
-        org_weight = org_sd["weight"]
-        weight = org_weight + lora_weight.to(org_weight.device, dtype=org_weight.dtype)
-
-        # set weight to org_module
-        org_sd["weight"] = weight
-        self.org_module[0].load_state_dict(org_sd)
-
-    # restore org weight from lora weight
-    def restore_from(self, multiplier=1.0):
-        # get lora weight
-        lora_weight = self.get_weight(multiplier)
-
-        # get org weight
-        org_sd = self.org_module[0].state_dict()
-        org_weight = org_sd["weight"]
-        weight = org_weight - lora_weight.to(org_weight.device, dtype=org_weight.dtype)
-
-        # set weight to org_module
-        org_sd["weight"] = weight
-        self.org_module[0].load_state_dict(org_sd)
-
-    # return lora weight
-    def get_weight(self, multiplier=None):
-        if multiplier is None:
-            multiplier = self.multiplier
-
-        # get up/down weight from module
-        up_weight = self.lora_up.weight.to(torch.float)
-        down_weight = self.lora_down.weight.to(torch.float)
-
-        # pre-calculated weight
-        if len(down_weight.size()) == 2:
-            # linear
-            weight = self.multiplier * (up_weight @ down_weight) * self.scale
-        elif down_weight.size()[2:4] == (1, 1):
-            # conv2d 1x1
-            weight = (
-                self.multiplier
-                * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3)
-                * self.scale
-            )
-        else:
-            # conv2d 3x3
-            conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
-            weight = self.multiplier * conved * self.scale
-
-        return weight
-
-
-# Create network from weights for inference, weights are not loaded here
-def create_network_from_weights(
-    text_encoder: Union[CLIPTextModel, List[CLIPTextModel]], unet: UNet2DConditionModel, weights_sd: Dict, multiplier: float = 1.0
-):
-    # get dim/alpha mapping
-    modules_dim = {}
-    modules_alpha = {}
-    for key, value in weights_sd.items():
-        if "." not in key:
-            continue
-
-        lora_name = key.split(".")[0]
-        if "alpha" in key:
-            modules_alpha[lora_name] = value
-        elif "lora_down" in key:
-            dim = value.size()[0]
-            modules_dim[lora_name] = dim
-            # print(lora_name, value.size(), dim)
-
-    # support old LoRA without alpha
-    for key in modules_dim.keys():
-        if key not in modules_alpha:
-            modules_alpha[key] = modules_dim[key]
-
-    return LoRANetwork(text_encoder, unet, multiplier=multiplier, modules_dim=modules_dim, modules_alpha=modules_alpha)
-
-
-def merge_lora_weights(pipe, weights_sd: Dict, multiplier: float = 1.0):
-    text_encoders = [pipe.text_encoder, pipe.text_encoder_2] if hasattr(pipe, "text_encoder_2") else [pipe.text_encoder]
-    unet = pipe.unet
-
-    lora_network = create_network_from_weights(text_encoders, unet, weights_sd, multiplier=multiplier)
-    lora_network.load_state_dict(weights_sd)
-    lora_network.merge_to(multiplier=multiplier)
-
-
-# block weightや学習に対応しない簡易版 / simple version without block weight and training
-class LoRANetwork(torch.nn.Module):
-    UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel"]
-    UNET_TARGET_REPLACE_MODULE_CONV2D_3X3 = ["ResnetBlock2D", "Downsample2D", "Upsample2D"]
-    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
-    LORA_PREFIX_UNET = "lora_unet"
-    LORA_PREFIX_TEXT_ENCODER = "lora_te"
-
-    # SDXL: must starts with LORA_PREFIX_TEXT_ENCODER
-    LORA_PREFIX_TEXT_ENCODER1 = "lora_te1"
-    LORA_PREFIX_TEXT_ENCODER2 = "lora_te2"
-
-    def __init__(
-        self,
-        text_encoder: Union[List[CLIPTextModel], CLIPTextModel],
-        unet: UNet2DConditionModel,
-        multiplier: float = 1.0,
-        modules_dim: Optional[Dict[str, int]] = None,
-        modules_alpha: Optional[Dict[str, int]] = None,
-        varbose: Optional[bool] = False,
-    ) -> None:
-        super().__init__()
-        self.multiplier = multiplier
-
-        print(f"create LoRA network from weights")
-
-        # convert SDXL Stability AI's U-Net modules to Diffusers
-        converted = self.convert_unet_modules(modules_dim, modules_alpha)
-        if converted:
-            print(f"converted {converted} Stability AI's U-Net LoRA modules to Diffusers (SDXL)")
-
-        # create module instances
-        def create_modules(
-            is_unet: bool,
-            text_encoder_idx: Optional[int],  # None, 1, 2
-            root_module: torch.nn.Module,
-            target_replace_modules: List[torch.nn.Module],
-        ) -> List[LoRAModule]:
-            prefix = (
-                self.LORA_PREFIX_UNET
-                if is_unet
-                else (
-                    self.LORA_PREFIX_TEXT_ENCODER
-                    if text_encoder_idx is None
-                    else (self.LORA_PREFIX_TEXT_ENCODER1 if text_encoder_idx == 1 else self.LORA_PREFIX_TEXT_ENCODER2)
-                )
-            )
-            loras = []
-            skipped = []
-            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():
-                        is_linear = child_module.__class__.__name__ == "Linear"
-                        is_conv2d = child_module.__class__.__name__ == "Conv2d"
-
-                        if is_linear or is_conv2d:
-                            lora_name = prefix + "." + name + "." + child_name
-                            lora_name = lora_name.replace(".", "_")
-
-                            if lora_name not in modules_dim:
-                                # print(f"skipped {lora_name} (not found in modules_dim)")
-                                skipped.append(lora_name)
-                                continue
-
-                            dim = modules_dim[lora_name]
-                            alpha = modules_alpha[lora_name]
-                            lora = LoRAModule(
-                                lora_name,
-                                child_module,
-                                self.multiplier,
-                                dim,
-                                alpha,
-                            )
-                            loras.append(lora)
-            return loras, skipped
-
-        text_encoders = text_encoder if type(text_encoder) == list else [text_encoder]
-
-        # create LoRA for text encoder
-        # 毎回すべてのモジュールを作るのは無駄なので要検討 / it is wasteful to create all modules every time, need to consider
-        self.text_encoder_loras: List[LoRAModule] = []
-        skipped_te = []
-        for i, text_encoder in enumerate(text_encoders):
-            if len(text_encoders) > 1:
-                index = i + 1
-            else:
-                index = None
-
-            text_encoder_loras, skipped = create_modules(False, index, text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE)
-            self.text_encoder_loras.extend(text_encoder_loras)
-            skipped_te += skipped
-        print(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
-        if len(skipped_te) > 0:
-            print(f"skipped {len(skipped_te)} modules because of missing weight.")
-
-        # extend U-Net target modules to include Conv2d 3x3
-        target_modules = LoRANetwork.UNET_TARGET_REPLACE_MODULE + LoRANetwork.UNET_TARGET_REPLACE_MODULE_CONV2D_3X3
-
-        self.unet_loras: List[LoRAModule]
-        self.unet_loras, skipped_un = create_modules(True, None, unet, target_modules)
-        print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
-        if len(skipped_un) > 0:
-            print(f"skipped {len(skipped_un)} modules because of missing weight.")
-
-        # assertion
-        names = set()
-        for lora in self.text_encoder_loras + self.unet_loras:
-            names.add(lora.lora_name)
-        for lora_name in modules_dim.keys():
-            assert lora_name in names, f"{lora_name} is not found in created LoRA modules."
-
-        # make to work load_state_dict
-        for lora in self.text_encoder_loras + self.unet_loras:
-            self.add_module(lora.lora_name, lora)
-
-    # SDXL: convert SDXL Stability AI's U-Net modules to Diffusers
-    def convert_unet_modules(self, modules_dim, modules_alpha):
-        converted_count = 0
-        not_converted_count = 0
-
-        map_keys = list(UNET_CONVERSION_MAP.keys())
-        map_keys.sort()
-
-        for key in list(modules_dim.keys()):
-            if key.startswith(LoRANetwork.LORA_PREFIX_UNET + "_"):
-                search_key = key.replace(LoRANetwork.LORA_PREFIX_UNET + "_", "")
-                position = bisect.bisect_right(map_keys, search_key)
-                map_key = map_keys[position - 1]
-                if search_key.startswith(map_key):
-                    new_key = key.replace(map_key, UNET_CONVERSION_MAP[map_key])
-                    modules_dim[new_key] = modules_dim[key]
-                    modules_alpha[new_key] = modules_alpha[key]
-                    del modules_dim[key]
-                    del modules_alpha[key]
-                    converted_count += 1
-                else:
-                    not_converted_count += 1
-        assert (
-            converted_count == 0 or not_converted_count == 0
-        ), f"some modules are not converted: {converted_count} converted, {not_converted_count} not converted"
-        return converted_count
-
-    def set_multiplier(self, multiplier):
-        self.multiplier = multiplier
-        for lora in self.text_encoder_loras + self.unet_loras:
-            lora.multiplier = self.multiplier
-
-    def apply_to(self, multiplier=1.0, apply_text_encoder=True, apply_unet=True):
-        if apply_text_encoder:
-            print("enable LoRA for text encoder")
-            for lora in self.text_encoder_loras:
-                lora.apply_to(multiplier)
-        if apply_unet:
-            print("enable LoRA for U-Net")
-            for lora in self.unet_loras:
-                lora.apply_to(multiplier)
-
-    def unapply_to(self):
-        for lora in self.text_encoder_loras + self.unet_loras:
-            lora.unapply_to()
-
-    def merge_to(self, multiplier=1.0):
-        print("merge LoRA weights to original weights")
-        for lora in tqdm(self.text_encoder_loras + self.unet_loras):
-            lora.merge_to(multiplier)
-        print(f"weights are merged")
-
-    def restore_from(self, multiplier=1.0):
-        print("restore LoRA weights from original weights")
-        for lora in tqdm(self.text_encoder_loras + self.unet_loras):
-            lora.restore_from(multiplier)
-        print(f"weights are restored")
-
-    def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True):
-        # convert SDXL Stability AI's state dict to Diffusers' based state dict
-        map_keys = list(UNET_CONVERSION_MAP.keys())  # prefix of U-Net modules
-        map_keys.sort()
-        for key in list(state_dict.keys()):
-            if key.startswith(LoRANetwork.LORA_PREFIX_UNET + "_"):
-                search_key = key.replace(LoRANetwork.LORA_PREFIX_UNET + "_", "")
-                position = bisect.bisect_right(map_keys, search_key)
-                map_key = map_keys[position - 1]
-                if search_key.startswith(map_key):
-                    new_key = key.replace(map_key, UNET_CONVERSION_MAP[map_key])
-                    state_dict[new_key] = state_dict[key]
-                    del state_dict[key]
-
-        # in case of V2, some weights have different shape, so we need to convert them
-        # because V2 LoRA is based on U-Net created by use_linear_projection=False
-        my_state_dict = self.state_dict()
-        for key in state_dict.keys():
-            if state_dict[key].size() != my_state_dict[key].size():
-                # print(f"convert {key} from {state_dict[key].size()} to {my_state_dict[key].size()}")
-                state_dict[key] = state_dict[key].view(my_state_dict[key].size())
-
-        return super().load_state_dict(state_dict, strict)
-
-
-if __name__ == "__main__":
-    # sample code to use LoRANetwork
-    import os
-    import argparse
-    from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
-    import torch
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model_id", type=str, default=None, help="model id for huggingface")
-    parser.add_argument("--lora_weights", type=str, default=None, help="path to LoRA weights")
-    parser.add_argument("--sdxl", action="store_true", help="use SDXL model")
-    parser.add_argument("--prompt", type=str, default="A photo of cat", help="prompt text")
-    parser.add_argument("--negative_prompt", type=str, default="", help="negative prompt text")
-    parser.add_argument("--seed", type=int, default=0, help="random seed")
-    args = parser.parse_args()
-
-    image_prefix = args.model_id.replace("/", "_") + "_"
-
-    # load Diffusers model
-    print(f"load model from {args.model_id}")
-    pipe: Union[StableDiffusionPipeline, StableDiffusionXLPipeline]
-    if args.sdxl:
-        # use_safetensors=True does not work with 0.18.2
-        pipe = StableDiffusionXLPipeline.from_pretrained(args.model_id, variant="fp16", torch_dtype=torch.float16)
-    else:
-        pipe = StableDiffusionPipeline.from_pretrained(args.model_id, variant="fp16", torch_dtype=torch.float16)
-    pipe.to(device)
-    pipe.set_use_memory_efficient_attention_xformers(True)
-
-    text_encoders = [pipe.text_encoder, pipe.text_encoder_2] if args.sdxl else [pipe.text_encoder]
-
-    # load LoRA weights
-    print(f"load LoRA weights from {args.lora_weights}")
-    if os.path.splitext(args.lora_weights)[1] == ".safetensors":
-        from safetensors.torch import load_file
-
-        lora_sd = load_file(args.lora_weights)
-    else:
-        lora_sd = torch.load(args.lora_weights)
-
-    # create by LoRA weights and load weights
-    print(f"create LoRA network")
-    lora_network: LoRANetwork = create_network_from_weights(text_encoders, pipe.unet, lora_sd, multiplier=1.0)
-
-    print(f"load LoRA network weights")
-    lora_network.load_state_dict(lora_sd)
-
-    lora_network.to(device, dtype=pipe.unet.dtype)  # required to apply_to. merge_to works without this
-
-    # 必要があれば、元のモデルの重みをバックアップしておく
-    # back-up unet/text encoder weights if necessary
-    def detach_and_move_to_cpu(state_dict):
-        for k, v in state_dict.items():
-            state_dict[k] = v.detach().cpu()
-        return state_dict
-
-    org_unet_sd = pipe.unet.state_dict()
-    detach_and_move_to_cpu(org_unet_sd)
-
-    org_text_encoder_sd = pipe.text_encoder.state_dict()
-    detach_and_move_to_cpu(org_text_encoder_sd)
-
-    if args.sdxl:
-        org_text_encoder_2_sd = pipe.text_encoder_2.state_dict()
-        detach_and_move_to_cpu(org_text_encoder_2_sd)
-
-    def seed_everything(seed):
-        torch.manual_seed(seed)
-        torch.cuda.manual_seed_all(seed)
-        np.random.seed(seed)
-        random.seed(seed)
-
-    # create image with original weights
-    print(f"create image with original weights")
-    seed_everything(args.seed)
-    image = pipe(args.prompt, negative_prompt=args.negative_prompt).images[0]
-    image.save(image_prefix + "original.png")
-
-    # apply LoRA network to the model: slower than merge_to, but can be reverted easily
-    print(f"apply LoRA network to the model")
-    lora_network.apply_to(multiplier=1.0)
-
-    print(f"create image with applied LoRA")
-    seed_everything(args.seed)
-    image = pipe(args.prompt, negative_prompt=args.negative_prompt).images[0]
-    image.save(image_prefix + "applied_lora.png")
-
-    # unapply LoRA network to the model
-    print(f"unapply LoRA network to the model")
-    lora_network.unapply_to()
-
-    print(f"create image with unapplied LoRA")
-    seed_everything(args.seed)
-    image = pipe(args.prompt, negative_prompt=args.negative_prompt).images[0]
-    image.save(image_prefix + "unapplied_lora.png")
-
-    # merge LoRA network to the model: faster than apply_to, but requires back-up of original weights (or unmerge_to)
-    print(f"merge LoRA network to the model")
-    lora_network.merge_to(multiplier=1.0)
-
-    print(f"create image with LoRA")
-    seed_everything(args.seed)
-    image = pipe(args.prompt, negative_prompt=args.negative_prompt).images[0]
-    image.save(image_prefix + "merged_lora.png")
-
-    # restore (unmerge) LoRA weights: numerically unstable
-    # マージされた重みを元に戻す。計算誤差のため、元の重みと完全に一致しないことがあるかもしれない
-    # 保存したstate_dictから元の重みを復元するのが確実
-    print(f"restore (unmerge) LoRA weights")
-    lora_network.restore_from(multiplier=1.0)
-
-    print(f"create image without LoRA")
-    seed_everything(args.seed)
-    image = pipe(args.prompt, negative_prompt=args.negative_prompt).images[0]
-    image.save(image_prefix + "unmerged_lora.png")
-
-    # restore original weights
-    print(f"restore original weights")
-    pipe.unet.load_state_dict(org_unet_sd)
-    pipe.text_encoder.load_state_dict(org_text_encoder_sd)
-    if args.sdxl:
-        pipe.text_encoder_2.load_state_dict(org_text_encoder_2_sd)
-
-    print(f"create image with restored original weights")
-    seed_everything(args.seed)
-    image = pipe(args.prompt, negative_prompt=args.negative_prompt).images[0]
-    image.save(image_prefix + "restore_original.png")
-
-    # use convenience function to merge LoRA weights
-    print(f"merge LoRA weights with convenience function")
-    merge_lora_weights(pipe, lora_sd, multiplier=1.0)
-
-    print(f"create image with merged LoRA weights")
-    seed_everything(args.seed)
-    image = pipe(args.prompt, negative_prompt=args.negative_prompt).images[0]
-    image.save(image_prefix + "convenience_merged_lora.png")

networks/merge_lora.py

@@ -1,10 +1,8 @@
 import math
 import argparse
 import os
-import time
 import torch
 from safetensors.torch import load_file, save_file
-from library import sai_model_spec, train_util
 import library.model_util as model_util
 import lora
 
@@ -12,26 +10,22 @@ import lora
 def load_state_dict(file_name, dtype):
     if os.path.splitext(file_name)[1] == ".safetensors":
         sd = load_file(file_name)
-        metadata = train_util.load_metadata_from_safetensors(file_name)
     else:
         sd = torch.load(file_name, map_location="cpu")
-        metadata = {}
-
     for key in list(sd.keys()):
         if type(sd[key]) == torch.Tensor:
             sd[key] = sd[key].to(dtype)
-
-    return sd, metadata
+    return sd
 
 
-def save_to_file(file_name, model, state_dict, dtype, metadata):
+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, metadata=metadata)
+        save_file(model, file_name)
     else:
         torch.save(model, file_name)
 
@@ -62,7 +56,7 @@ def merge_to_sd_model(text_encoder, unet, models, ratios, merge_dtype):
 
     for model, ratio in zip(models, ratios):
         print(f"loading: {model}")
-        lora_sd, _ = load_state_dict(model, merge_dtype)
+        lora_sd = load_state_dict(model, merge_dtype)
 
         print(f"merging...")
         for key in lora_sd.keys():
@@ -87,11 +81,9 @@ def merge_to_sd_model(text_encoder, unet, models, ratios, merge_dtype):
 
                 # W <- W + U * D
                 weight = module.weight
+                # print(module_name, down_weight.size(), up_weight.size())
                 if len(weight.size()) == 2:
                     # linear
-                    if len(up_weight.size()) == 4:  # use linear projection mismatch
-                        up_weight = up_weight.squeeze(3).squeeze(2)
-                        down_weight = down_weight.squeeze(3).squeeze(2)
                     weight = weight + ratio * (up_weight @ down_weight) * scale
                 elif down_weight.size()[2:4] == (1, 1):
                     # conv2d 1x1
@@ -115,17 +107,9 @@ def merge_lora_models(models, ratios, merge_dtype):
     base_dims = {}
 
     merged_sd = {}
-    v2 = None
-    base_model = None
     for model, ratio in zip(models, ratios):
         print(f"loading: {model}")
-        lora_sd, lora_metadata = load_state_dict(model, merge_dtype)
-
-        if lora_metadata is not None:
-            if v2 is None:
-                v2 = lora_metadata.get(train_util.SS_METADATA_KEY_V2, None)  # return string
-            if base_model is None:
-                base_model = lora_metadata.get(train_util.SS_METADATA_KEY_BASE_MODEL_VERSION, None)
+        lora_sd = load_state_dict(model, merge_dtype)
 
         # get alpha and dim
         alphas = {}  # alpha for current model
@@ -182,26 +166,7 @@ def merge_lora_models(models, ratios, merge_dtype):
     print("merged model")
     print(f"dim: {list(set(base_dims.values()))}, alpha: {list(set(base_alphas.values()))}")
 
-    # check all dims are same
-    dims_list = list(set(base_dims.values()))
-    alphas_list = list(set(base_alphas.values()))
-    all_same_dims = True
-    all_same_alphas = True
-    for dims in dims_list:
-        if dims != dims_list[0]:
-            all_same_dims = False
-            break
-    for alphas in alphas_list:
-        if alphas != alphas_list[0]:
-            all_same_alphas = False
-            break
-
-    # build minimum metadata
-    dims = f"{dims_list[0]}" if all_same_dims else "Dynamic"
-    alphas = f"{alphas_list[0]}" if all_same_alphas else "Dynamic"
-    metadata = train_util.build_minimum_network_metadata(v2, base_model, "networks.lora", dims, alphas, None)
-
-    return merged_sd, metadata, v2 == "True"
+    return merged_sd
 
 
 def merge(args):
@@ -228,57 +193,13 @@ def merge(args):
 
         merge_to_sd_model(text_encoder, unet, args.models, args.ratios, merge_dtype)
 
-        if args.no_metadata:
-            sai_metadata = None
-        else:
-            merged_from = sai_model_spec.build_merged_from([args.sd_model] + args.models)
-            title = os.path.splitext(os.path.basename(args.save_to))[0]
-            sai_metadata = sai_model_spec.build_metadata(
-                None,
-                args.v2,
-                args.v2,
-                False,
-                False,
-                False,
-                time.time(),
-                title=title,
-                merged_from=merged_from,
-                is_stable_diffusion_ckpt=True,
-            )
-            if args.v2:
-                # TODO read sai modelspec
-                print(
-                    "Cannot determine if model is for v-prediction, so save metadata as v-prediction / modelがv-prediction用か否か不明なため、仮にv-prediction用としてmetadataを保存します"
-                )
-
         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, sai_metadata, save_dtype, vae
-        )
+        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, metadata, v2 = merge_lora_models(args.models, args.ratios, merge_dtype)
-
-        print(f"calculating hashes and creating 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
-
-        if not args.no_metadata:
-            merged_from = sai_model_spec.build_merged_from(args.models)
-            title = os.path.splitext(os.path.basename(args.save_to))[0]
-            sai_metadata = sai_model_spec.build_metadata(
-                state_dict, v2, v2, False, True, False, time.time(), title=title, merged_from=merged_from
-            )
-            if v2:
-                # TODO read sai modelspec
-                print(
-                    "Cannot determine if LoRA is for v-prediction, so save metadata as v-prediction / LoRAがv-prediction用か否か不明なため、仮にv-prediction用としてmetadataを保存します"
-                )
-            metadata.update(sai_metadata)
+        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, metadata)
+        save_to_file(args.save_to, state_dict, state_dict, save_dtype)
 
 
 def setup_parser() -> argparse.ArgumentParser:
@@ -311,12 +232,6 @@ def setup_parser() -> argparse.ArgumentParser:
         "--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(
-        "--no_metadata",
-        action="store_true",
-        help="do not save sai modelspec metadata (minimum ss_metadata for LoRA is saved) / "
-        + "sai modelspecのメタデータを保存しない（LoRAの最低限のss_metadataは保存される）",
-    )
 
     return parser
 

networks/resize_lora.py

@@ -219,8 +219,8 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dyn
     for key, value in tqdm(lora_sd.items()):
       weight_name = None
       if 'lora_down' in key:
-        block_down_name = key.split(".")[0]
-        weight_name = key.split(".")[-1]
+        block_down_name = key.rsplit('.lora_down', 1)[0]
+        weight_name = key.rsplit(".", 1)[-1]
         lora_down_weight = value
       else:
         continue
@@ -283,7 +283,10 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dyn
 
 
 def resize(args):
+  if args.save_to is None or not (args.save_to.endswith('.ckpt') or args.save_to.endswith('.pt') or args.save_to.endswith('.pth') or args.save_to.endswith('.safetensors')):
+    raise Exception("The --save_to argument must be specified and must be a .ckpt , .pt, .pth or .safetensors file.")
 
+    
   def str_to_dtype(p):
     if p == 'float':
       return torch.float

networks/sdxl_merge_lora.py

@@ -1,318 +0,0 @@
-import math
-import argparse
-import os
-import time
-import torch
-from safetensors.torch import load_file, save_file
-from tqdm import tqdm
-from library import sai_model_spec, sdxl_model_util, train_util
-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)
-        metadata = train_util.load_metadata_from_safetensors(file_name)
-    else:
-        sd = torch.load(file_name, map_location="cpu")
-        metadata = {}
-
-    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 os.path.splitext(file_name)[1] == ".safetensors":
-        save_file(model, file_name, metadata=metadata)
-    else:
-        torch.save(model, file_name)
-
-
-def merge_to_sd_model(text_encoder1, text_encoder2, unet, models, ratios, merge_dtype):
-    text_encoder1.to(merge_dtype)
-    text_encoder1.to(merge_dtype)
-    unet.to(merge_dtype)
-
-    # create module map
-    name_to_module = {}
-    for i, root_module in enumerate([text_encoder1, text_encoder2, unet]):
-        if i <= 1:
-            if i == 0:
-                prefix = lora.LoRANetwork.LORA_PREFIX_TEXT_ENCODER1
-            else:
-                prefix = lora.LoRANetwork.LORA_PREFIX_TEXT_ENCODER2
-            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 + lora.LoRANetwork.UNET_TARGET_REPLACE_MODULE_CONV2D_3X3
-            )
-
-        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":
-                        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 tqdm(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
-                # print(module_name, down_weight.size(), up_weight.size())
-                if len(weight.size()) == 2:
-                    # linear
-                    weight = weight + ratio * (up_weight @ down_weight) * scale
-                elif down_weight.size()[2:4] == (1, 1):
-                    # conv2d 1x1
-                    weight = (
-                        weight
-                        + ratio
-                        * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3)
-                        * scale
-                    )
-                else:
-                    # conv2d 3x3
-                    conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
-                    # print(conved.size(), weight.size(), module.stride, module.padding)
-                    weight = weight + ratio * conved * 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 = {}
-    v2 = None
-    base_model = None
-    for model, ratio in zip(models, ratios):
-        print(f"loading: {model}")
-        lora_sd, lora_metadata = load_state_dict(model, merge_dtype)
-
-        if lora_metadata is not None:
-            if v2 is None:
-                v2 = lora_metadata.get(train_util.SS_METADATA_KEY_V2, None)  # returns string, SDXLはv2がないのでFalseのはず
-            if base_model is None:
-                base_model = lora_metadata.get(train_util.SS_METADATA_KEY_BASE_MODEL_VERSION, None)
-
-        # 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 tqdm(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()))}")
-
-    # check all dims are same
-    dims_list = list(set(base_dims.values()))
-    alphas_list = list(set(base_alphas.values()))
-    all_same_dims = True
-    all_same_alphas = True
-    for dims in dims_list:
-        if dims != dims_list[0]:
-            all_same_dims = False
-            break
-    for alphas in alphas_list:
-        if alphas != alphas_list[0]:
-            all_same_alphas = False
-            break
-
-    # build minimum metadata
-    dims = f"{dims_list[0]}" if all_same_dims else "Dynamic"
-    alphas = f"{alphas_list[0]}" if all_same_alphas else "Dynamic"
-    metadata = train_util.build_minimum_network_metadata(v2, base_model, "networks.lora", dims, alphas, None)
-
-    return merged_sd, metadata
-
-
-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_model1,
-            text_model2,
-            vae,
-            unet,
-            logit_scale,
-            ckpt_info,
-        ) = sdxl_model_util.load_models_from_sdxl_checkpoint(sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0, args.sd_model, "cpu")
-
-        merge_to_sd_model(text_model1, text_model2, unet, args.models, args.ratios, merge_dtype)
-
-        if args.no_metadata:
-            sai_metadata = None
-        else:
-            merged_from = sai_model_spec.build_merged_from([args.sd_model] + args.models)
-            title = os.path.splitext(os.path.basename(args.save_to))[0]
-            sai_metadata = sai_model_spec.build_metadata(
-                None, False, False, True, False, False, time.time(), title=title, merged_from=merged_from
-            )
-
-        print(f"saving SD model to: {args.save_to}")
-        sdxl_model_util.save_stable_diffusion_checkpoint(
-            args.save_to, text_model1, text_model2, unet, 0, 0, ckpt_info, vae, logit_scale, sai_metadata, save_dtype
-        )
-    else:
-        state_dict, metadata = merge_lora_models(args.models, args.ratios, merge_dtype)
-
-        print(f"calculating hashes and creating 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
-
-        if not args.no_metadata:
-            merged_from = sai_model_spec.build_merged_from(args.models)
-            title = os.path.splitext(os.path.basename(args.save_to))[0]
-            sai_metadata = sai_model_spec.build_metadata(
-                state_dict, False, False, True, True, False, time.time(), title=title, merged_from=merged_from
-            )
-            metadata.update(sai_metadata)
-
-        print(f"saving model to: {args.save_to}")
-        save_to_file(args.save_to, state_dict, state_dict, save_dtype, metadata)
-
-
-def setup_parser() -> argparse.ArgumentParser:
-    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(
-        "--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モデルの比率")
-    parser.add_argument(
-        "--no_metadata",
-        action="store_true",
-        help="do not save sai modelspec metadata (minimum ss_metadata for LoRA is saved) / "
-        + "sai modelspecのメタデータを保存しない（LoRAの最低限のss_metadataは保存される）",
-    )
-
-    return parser
-
-
-if __name__ == "__main__":
-    parser = setup_parser()
-
-    args = parser.parse_args()
-    merge(args)

networks/svd_merge_lora.py

@@ -1,11 +1,10 @@
+
 import math
 import argparse
 import os
-import time
 import torch
 from safetensors.torch import load_file, save_file
 from tqdm import tqdm
-from library import sai_model_spec, train_util
 import library.model_util as model_util
 import lora
 
@@ -14,247 +13,180 @@ CLAMP_QUANTILE = 0.99
 
 
 def load_state_dict(file_name, dtype):
-    if os.path.splitext(file_name)[1] == ".safetensors":
-        sd = load_file(file_name)
-        metadata = train_util.load_metadata_from_safetensors(file_name)
-    else:
-        sd = torch.load(file_name, map_location="cpu")
-        metadata = {}
-
-    for key in list(sd.keys()):
-        if type(sd[key]) == torch.Tensor:
-            sd[key] = sd[key].to(dtype)
-
-    return sd, metadata
+  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, 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)
+def save_to_file(file_name, 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(state_dict, file_name, metadata=metadata)
-    else:
-        torch.save(state_dict, file_name)
+  if os.path.splitext(file_name)[1] == '.safetensors':
+    save_file(state_dict, file_name)
+  else:
+    torch.save(state_dict, file_name)
 
 
 def merge_lora_models(models, ratios, new_rank, new_conv_rank, device, merge_dtype):
-    print(f"new rank: {new_rank}, new conv rank: {new_conv_rank}")
-    merged_sd = {}
-    v2 = None
-    base_model = None
-    for model, ratio in zip(models, ratios):
-        print(f"loading: {model}")
-        lora_sd, lora_metadata = load_state_dict(model, merge_dtype)
+  print(f"new rank: {new_rank}, new conv rank: {new_conv_rank}")
+  merged_sd = {}
+  for model, ratio in zip(models, ratios):
+    print(f"loading: {model}")
+    lora_sd = load_state_dict(model, merge_dtype)
 
-        if lora_metadata is not None:
-            if v2 is None:
-                v2 = lora_metadata.get(train_util.SS_METADATA_KEY_V2, None)  # return string
-            if base_model is None:
-                base_model = lora_metadata.get(train_util.SS_METADATA_KEY_BASE_MODEL_VERSION, None)
+    # merge
+    print(f"merging...")
+    for key in tqdm(list(lora_sd.keys())):
+      if 'lora_down' not in key:
+        continue
 
-        # 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")]
 
-            lora_module_name = key[: key.rfind(".lora_down")]
+      down_weight = lora_sd[key]
+      network_dim = down_weight.size()[0]
 
-            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)
 
-            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
+      kernel_size = None if not conv2d else down_weight.size()[2:4]
+      # print(lora_module_name, network_dim, alpha, in_dim, out_dim, kernel_size)
 
-            in_dim = down_weight.size()[1]
-            out_dim = up_weight.size()[0]
-            conv2d = len(down_weight.size()) == 4
-            kernel_size = None if not conv2d else down_weight.size()[2:4]
-            # print(lora_module_name, network_dim, alpha, in_dim, out_dim, kernel_size)
+      # make original weight if not exist
+      if lora_module_name not in merged_sd:
+        weight = torch.zeros((out_dim, in_dim, *kernel_size) if conv2d else (out_dim, in_dim), dtype=merge_dtype)
+        if device:
+          weight = weight.to(device)
+      else:
+        weight = merged_sd[lora_module_name]
 
-            # make original weight if not exist
-            if lora_module_name not in merged_sd:
-                weight = torch.zeros((out_dim, in_dim, *kernel_size) 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)
 
-            # 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)
 
-            # W <- W + U * D
-            scale = alpha / network_dim
+      if device:                      # and isinstance(scale, torch.Tensor):
+        scale = scale.to(device)
 
-            if device:  # and isinstance(scale, torch.Tensor):
-                scale = scale.to(device)
+      if not conv2d:        # linear
+        weight = weight + ratio * (up_weight @ down_weight) * scale
+      elif kernel_size == (1, 1):
+        weight = weight + ratio * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)
+                                   ).unsqueeze(2).unsqueeze(3) * scale
+      else:
+        conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
+        weight = weight + ratio * conved * scale
 
-            if not conv2d:  # linear
-                weight = weight + ratio * (up_weight @ down_weight) * scale
-            elif kernel_size == (1, 1):
-                weight = (
-                    weight
-                    + ratio
-                    * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3)
-                    * scale
-                )
-            else:
-                conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
-                weight = weight + ratio * conved * scale
+      merged_sd[lora_module_name] = weight
 
-            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)
+      kernel_size = None if not conv2d else mat.size()[2:4]
+      conv2d_3x3 = conv2d and kernel_size != (1, 1)
+      out_dim, in_dim = mat.size()[0:2]
 
-    # 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
-            kernel_size = None if not conv2d else mat.size()[2:4]
-            conv2d_3x3 = conv2d and kernel_size != (1, 1)
-            out_dim, in_dim = mat.size()[0:2]
+      if conv2d:
+        if conv2d_3x3:
+          mat = mat.flatten(start_dim=1)
+        else:
+          mat = mat.squeeze()
 
-            if conv2d:
-                if conv2d_3x3:
-                    mat = mat.flatten(start_dim=1)
-                else:
-                    mat = mat.squeeze()
+      module_new_rank = new_conv_rank if conv2d_3x3 else new_rank
+      module_new_rank = min(module_new_rank, in_dim, out_dim)                           # LoRA rank cannot exceed the original dim
 
-            module_new_rank = new_conv_rank if conv2d_3x3 else new_rank
-            module_new_rank = min(module_new_rank, in_dim, out_dim)  # LoRA rank cannot exceed the original dim
+      U, S, Vh = torch.linalg.svd(mat)
 
-            U, S, Vh = torch.linalg.svd(mat)
+      U = U[:, :module_new_rank]
+      S = S[:module_new_rank]
+      U = U @ torch.diag(S)
 
-            U = U[:, :module_new_rank]
-            S = S[:module_new_rank]
-            U = U @ torch.diag(S)
+      Vh = Vh[:module_new_rank, :]
 
-            Vh = Vh[:module_new_rank, :]
+      dist = torch.cat([U.flatten(), Vh.flatten()])
+      hi_val = torch.quantile(dist, CLAMP_QUANTILE)
+      low_val = -hi_val
 
-            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)
 
-            U = U.clamp(low_val, hi_val)
-            Vh = Vh.clamp(low_val, hi_val)
+      if conv2d:
+        U = U.reshape(out_dim, module_new_rank, 1, 1)
+        Vh = Vh.reshape(module_new_rank, in_dim, kernel_size[0], kernel_size[1])
 
-            if conv2d:
-                U = U.reshape(out_dim, module_new_rank, 1, 1)
-                Vh = Vh.reshape(module_new_rank, in_dim, kernel_size[0], kernel_size[1])
+      up_weight = U
+      down_weight = Vh
 
-            up_weight = U
-            down_weight = Vh
+      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(module_new_rank)
 
-            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(module_new_rank)
-
-    # build minimum metadata
-    dims = f"{new_rank}"
-    alphas = f"{new_rank}"
-    if new_conv_rank is not None:
-        network_args = {"conv_dim": new_conv_rank, "conv_alpha": new_conv_rank}
-    else:
-        network_args = None
-    metadata = train_util.build_minimum_network_metadata(v2, base_model, "networks.lora", dims, alphas, network_args)
-
-    return merged_lora_sd, metadata, v2 == "True", base_model
+  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 / モデルの数と重みの数は合わせてください"
+  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
+  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
+  merge_dtype = str_to_dtype(args.precision)
+  save_dtype = str_to_dtype(args.save_precision)
+  if save_dtype is None:
+    save_dtype = merge_dtype
 
-    new_conv_rank = args.new_conv_rank if args.new_conv_rank is not None else args.new_rank
-    state_dict, metadata, v2, base_model = merge_lora_models(
-        args.models, args.ratios, args.new_rank, new_conv_rank, args.device, merge_dtype
-    )
+  new_conv_rank = args.new_conv_rank if args.new_conv_rank is not None else args.new_rank
+  state_dict = merge_lora_models(args.models, args.ratios, args.new_rank, new_conv_rank, args.device, merge_dtype)
 
-    print(f"calculating hashes and creating 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
-
-    if not args.no_metadata:
-        is_sdxl = base_model is not None and base_model.lower().startswith("sdxl")
-        merged_from = sai_model_spec.build_merged_from(args.models)
-        title = os.path.splitext(os.path.basename(args.save_to))[0]
-        sai_metadata = sai_model_spec.build_metadata(
-            state_dict, v2, v2, is_sdxl, True, False, time.time(), title=title, merged_from=merged_from
-        )
-        if v2:
-            # TODO read sai modelspec
-            print(
-                "Cannot determine if LoRA is for v-prediction, so save metadata as v-prediction / LoRAがv-prediction用か否か不明なため、仮にv-prediction用としてmetadataを保存します"
-            )
-        metadata.update(sai_metadata)
-
-    print(f"saving model to: {args.save_to}")
-    save_to_file(args.save_to, state_dict, save_dtype, metadata)
+  print(f"saving model to: {args.save_to}")
+  save_to_file(args.save_to, state_dict, save_dtype)
 
 
 def setup_parser() -> argparse.ArgumentParser:
-    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(
-        "--new_conv_rank",
-        type=int,
-        default=None,
-        help="Specify rank of output LoRA for Conv2d 3x3, None for same as new_rank / 出力するConv2D 3x3 LoRAのrank (dim)、Noneでnew_rankと同じ",
-    )
-    parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")
-    parser.add_argument(
-        "--no_metadata",
-        action="store_true",
-        help="do not save sai modelspec metadata (minimum ss_metadata for LoRA is saved) / "
-        + "sai modelspecのメタデータを保存しない（LoRAの最低限のss_metadataは保存される）",
-    )
+  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("--new_conv_rank", type=int, default=None,
+                      help="Specify rank of output LoRA for Conv2d 3x3, None for same as new_rank / 出力するConv2D 3x3 LoRAのrank (dim)、Noneでnew_rankと同じ")
+  parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")
 
-    return parser
+  return parser
 
 
-if __name__ == "__main__":
-    parser = setup_parser()
+if __name__ == '__main__':
+  parser = setup_parser()
 
-    args = parser.parse_args()
-    merge(args)
+  args = parser.parse_args()
+  merge(args)

requirements.txt

@@ -1,29 +1,26 @@
-accelerate==0.19.0
-transformers==4.30.2
-diffusers[torch]==0.18.2
+accelerate==0.15.0
+transformers==4.26.0
 ftfy==6.1.1
-# albumentations==1.3.0
+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.39.1
+bitsandbytes==0.35.0
 tensorboard==2.10.1
-safetensors==0.3.1
+safetensors==0.2.6
 # gradio==3.16.2
 altair==4.2.2
 easygui==0.98.3
 toml==0.10.2
 voluptuous==0.13.1
-huggingface-hub==0.15.1
-# for loading Diffusers' SDXL
-invisible-watermark==0.2.0 
 # for BLIP captioning
-# requests==2.28.2
-# timm==0.6.12
-# fairscale==0.4.13
+requests==2.28.2
+timm==0.6.12
+fairscale==0.4.13
 # for WD14 captioning
-# tensorflow==2.10.1
-# open clip for SDXL
-open-clip-torch==2.20.0
+# tensorflow<2.11
+tensorflow==2.10.1
+huggingface-hub==0.15.1
 # for kohya_ss library
--e .
+.

sdxl_minimal_inference.py

@@ -1,321 +0,0 @@
-# 手元で推論を行うための最低限のコード。HuggingFace／DiffusersのCLIP、schedulerとVAEを使う
-# Minimal code for performing inference at local. Use HuggingFace/Diffusers CLIP, scheduler and VAE
-
-import argparse
-import datetime
-import math
-import os
-import random
-from einops import repeat
-import numpy as np
-import torch
-from tqdm import tqdm
-from transformers import CLIPTokenizer
-from diffusers import EulerDiscreteScheduler
-from PIL import Image
-import open_clip
-from safetensors.torch import load_file
-
-from library import model_util, sdxl_model_util
-import networks.lora as lora
-
-# scheduler: このあたりの設定はSD1/2と同じでいいらしい
-# scheduler: The settings around here seem to be the same as SD1/2
-SCHEDULER_LINEAR_START = 0.00085
-SCHEDULER_LINEAR_END = 0.0120
-SCHEDULER_TIMESTEPS = 1000
-SCHEDLER_SCHEDULE = "scaled_linear"
-
-
-# Time EmbeddingはDiffusersからのコピー
-# Time Embedding is copied from Diffusers
-
-
-def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
-    """
-    Create sinusoidal timestep embeddings.
-    :param timesteps: a 1-D Tensor of N indices, one per batch element.
-                      These may be fractional.
-    :param dim: the dimension of the output.
-    :param max_period: controls the minimum frequency of the embeddings.
-    :return: an [N x dim] Tensor of positional embeddings.
-    """
-    if not repeat_only:
-        half = dim // 2
-        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
-            device=timesteps.device
-        )
-        args = timesteps[:, None].float() * freqs[None]
-        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-        if dim % 2:
-            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-    else:
-        embedding = repeat(timesteps, "b -> b d", d=dim)
-    return embedding
-
-
-def get_timestep_embedding(x, outdim):
-    assert len(x.shape) == 2
-    b, dims = x.shape[0], x.shape[1]
-    # x = rearrange(x, "b d -> (b d)")
-    x = torch.flatten(x)
-    emb = timestep_embedding(x, outdim)
-    # emb = rearrange(emb, "(b d) d2 -> b (d d2)", b=b, d=dims, d2=outdim)
-    emb = torch.reshape(emb, (b, dims * outdim))
-    return emb
-
-
-if __name__ == "__main__":
-    # 画像生成条件を変更する場合はここを変更 / change here to change image generation conditions
-
-    # SDXLの追加のvector embeddingへ渡す値 / Values to pass to additional vector embedding of SDXL
-    target_height = 1024
-    target_width = 1024
-    original_height = target_height
-    original_width = target_width
-    crop_top = 0
-    crop_left = 0
-
-    steps = 50
-    guidance_scale = 7
-    seed = None  # 1
-
-    DEVICE = "cuda"
-    DTYPE = torch.float16  # bfloat16 may work
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--ckpt_path", type=str, required=True)
-    parser.add_argument("--prompt", type=str, default="A photo of a cat")
-    parser.add_argument("--prompt2", type=str, default=None)
-    parser.add_argument("--negative_prompt", type=str, default="")
-    parser.add_argument("--output_dir", type=str, default=".")
-    parser.add_argument(
-        "--lora_weights",
-        type=str,
-        nargs="*",
-        default=[],
-        help="LoRA weights, only supports networks.lora, each arguement is a `path;multiplier` (semi-colon separated)",
-    )
-    parser.add_argument("--interactive", action="store_true")
-    args = parser.parse_args()
-
-    if args.prompt2 is None:
-        args.prompt2 = args.prompt
-
-    # HuggingFaceのmodel id
-    text_encoder_1_name = "openai/clip-vit-large-patch14"
-    text_encoder_2_name = "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"
-
-    # checkpointを読み込む。モデル変換についてはそちらの関数を参照
-    # Load checkpoint. For model conversion, see this function
-
-    # 本体RAMが少ない場合はGPUにロードするといいかも
-    # If the main RAM is small, it may be better to load it on the GPU
-    text_model1, text_model2, vae, unet, _, _ = sdxl_model_util.load_models_from_sdxl_checkpoint(
-        sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0, args.ckpt_path, "cpu"
-    )
-
-    # Text Encoder 1はSDXL本体でもHuggingFaceのものを使っている
-    # In SDXL, Text Encoder 1 is also using HuggingFace's
-
-    # Text Encoder 2はSDXL本体ではopen_clipを使っている
-    # それを使ってもいいが、SD2のDiffusers版に合わせる形で、HuggingFaceのものを使う
-    # 重みの変換コードはSD2とほぼ同じ
-    # In SDXL, Text Encoder 2 is using open_clip
-    # It's okay to use it, but to match the Diffusers version of SD2, use HuggingFace's
-    # The weight conversion code is almost the same as SD2
-
-    # VAEの構造はSDXLもSD1/2と同じだが、重みは異なるようだ。何より謎のscale値が違う
-    # fp16でNaNが出やすいようだ
-    # The structure of VAE is the same as SD1/2, but the weights seem to be different. Above all, the mysterious scale value is different.
-    # NaN seems to be more likely to occur in fp16
-
-    unet.to(DEVICE, dtype=DTYPE)
-    unet.eval()
-
-    vae_dtype = DTYPE
-    if DTYPE == torch.float16:
-        print("use float32 for vae")
-        vae_dtype = torch.float32
-    vae.to(DEVICE, dtype=vae_dtype)
-    vae.eval()
-
-    text_model1.to(DEVICE, dtype=DTYPE)
-    text_model1.eval()
-    text_model2.to(DEVICE, dtype=DTYPE)
-    text_model2.eval()
-
-    unet.set_use_memory_efficient_attention(True, False)
-    if torch.__version__ >= "2.0.0": # PyTorch 2.0.0 以上対応のxformersなら以下が使える
-        vae.set_use_memory_efficient_attention_xformers(True)
-
-    # Tokenizers
-    tokenizer1 = CLIPTokenizer.from_pretrained(text_encoder_1_name)
-    tokenizer2 = lambda x: open_clip.tokenize(x, context_length=77)
-
-    # LoRA
-    for weights_file in args.lora_weights:
-        if ";" in weights_file:
-            weights_file, multiplier = weights_file.split(";")
-            multiplier = float(multiplier)
-        else:
-            multiplier = 1.0
-
-        lora_model, weights_sd = lora.create_network_from_weights(
-            multiplier, weights_file, vae, [text_model1, text_model2], unet, None, True
-        )
-        lora_model.merge_to([text_model1, text_model2], unet, weights_sd, DTYPE, DEVICE)
-
-    # scheduler
-    scheduler = EulerDiscreteScheduler(
-        num_train_timesteps=SCHEDULER_TIMESTEPS,
-        beta_start=SCHEDULER_LINEAR_START,
-        beta_end=SCHEDULER_LINEAR_END,
-        beta_schedule=SCHEDLER_SCHEDULE,
-    )
-
-    def generate_image(prompt, prompt2, negative_prompt, seed=None):
-        # 将来的にサイズ情報も変えられるようにする / Make it possible to change the size information in the future
-        # prepare embedding
-        with torch.no_grad():
-            # vector
-            emb1 = get_timestep_embedding(torch.FloatTensor([original_height, original_width]).unsqueeze(0), 256)
-            emb2 = get_timestep_embedding(torch.FloatTensor([crop_top, crop_left]).unsqueeze(0), 256)
-            emb3 = get_timestep_embedding(torch.FloatTensor([target_height, target_width]).unsqueeze(0), 256)
-            # print("emb1", emb1.shape)
-            c_vector = torch.cat([emb1, emb2, emb3], dim=1).to(DEVICE, dtype=DTYPE)
-            uc_vector = c_vector.clone().to(DEVICE, dtype=DTYPE)  # ちょっとここ正しいかどうかわからない I'm not sure if this is right
-
-            # crossattn
-
-        # Text Encoderを二つ呼ぶ関数  Function to call two Text Encoders
-        def call_text_encoder(text, text2):
-            # text encoder 1
-            batch_encoding = tokenizer1(
-                text,
-                truncation=True,
-                return_length=True,
-                return_overflowing_tokens=False,
-                padding="max_length",
-                return_tensors="pt",
-            )
-            tokens = batch_encoding["input_ids"].to(DEVICE)
-
-            with torch.no_grad():
-                enc_out = text_model1(tokens, output_hidden_states=True, return_dict=True)
-                text_embedding1 = enc_out["hidden_states"][11]
-                # text_embedding = pipe.text_encoder.text_model.final_layer_norm(text_embedding)    # layer normは通さないらしい
-
-            # text encoder 2
-            with torch.no_grad():
-                tokens = tokenizer2(text2).to(DEVICE)
-
-                enc_out = text_model2(tokens, output_hidden_states=True, return_dict=True)
-                text_embedding2_penu = enc_out["hidden_states"][-2]
-                # print("hidden_states2", text_embedding2_penu.shape)
-                text_embedding2_pool = enc_out["text_embeds"]   # do not support Textual Inversion
-
-            # 連結して終了 concat and finish
-            text_embedding = torch.cat([text_embedding1, text_embedding2_penu], dim=2)
-            return text_embedding, text_embedding2_pool
-
-        # cond
-        c_ctx, c_ctx_pool = call_text_encoder(prompt, prompt2)
-        # print(c_ctx.shape, c_ctx_p.shape, c_vector.shape)
-        c_vector = torch.cat([c_ctx_pool, c_vector], dim=1)
-
-        # uncond
-        uc_ctx, uc_ctx_pool = call_text_encoder(negative_prompt, negative_prompt)
-        uc_vector = torch.cat([uc_ctx_pool, uc_vector], dim=1)
-
-        text_embeddings = torch.cat([uc_ctx, c_ctx])
-        vector_embeddings = torch.cat([uc_vector, c_vector])
-
-        # メモリ使用量を減らすにはここでText Encoderを削除するかCPUへ移動する
-
-        if seed is not None:
-            random.seed(seed)
-            np.random.seed(seed)
-            torch.manual_seed(seed)
-            torch.cuda.manual_seed_all(seed)
-
-            # # random generator for initial noise
-            # generator = torch.Generator(device="cuda").manual_seed(seed)
-            generator = None
-        else:
-            generator = None
-
-        # get the initial random noise unless the user supplied it
-        # SDXLはCPUでlatentsを作成しているので一応合わせておく、Diffusersはtarget deviceでlatentsを作成している
-        # SDXL creates latents in CPU, Diffusers creates latents in target device
-        latents_shape = (1, 4, target_height // 8, target_width // 8)
-        latents = torch.randn(
-            latents_shape,
-            generator=generator,
-            device="cpu",
-            dtype=torch.float32,
-        ).to(DEVICE, dtype=DTYPE)
-
-        # scale the initial noise by the standard deviation required by the scheduler
-        latents = latents * scheduler.init_noise_sigma
-
-        # set timesteps
-        scheduler.set_timesteps(steps, DEVICE)
-
-        # このへんはDiffusersからのコピペ
-        # Copy from Diffusers
-        timesteps = scheduler.timesteps.to(DEVICE)  # .to(DTYPE)
-        num_latent_input = 2
-        with torch.no_grad():
-            for i, t in enumerate(tqdm(timesteps)):
-                # expand the latents if we are doing classifier free guidance
-                latent_model_input = latents.repeat((num_latent_input, 1, 1, 1))
-                latent_model_input = scheduler.scale_model_input(latent_model_input, t)
-
-                noise_pred = unet(latent_model_input, t, text_embeddings, vector_embeddings)
-
-                noise_pred_uncond, noise_pred_text = noise_pred.chunk(num_latent_input)  # uncond by negative prompt
-                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-
-                # compute the previous noisy sample x_t -> x_t-1
-                # latents = scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
-                latents = scheduler.step(noise_pred, t, latents).prev_sample
-
-            # latents = 1 / 0.18215 * latents
-            latents = 1 / sdxl_model_util.VAE_SCALE_FACTOR * latents
-            latents = latents.to(vae_dtype)
-            image = vae.decode(latents).sample
-            image = (image / 2 + 0.5).clamp(0, 1)
-
-        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
-        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
-
-        # image = self.numpy_to_pil(image)
-        image = (image * 255).round().astype("uint8")
-        image = [Image.fromarray(im) for im in image]
-
-        # 保存して終了 save and finish
-        timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
-        for i, img in enumerate(image):
-            img.save(os.path.join(args.output_dir, f"image_{timestamp}_{i:03d}.png"))
-
-    if not args.interactive:
-        generate_image(args.prompt, args.prompt2, args.negative_prompt, seed)
-    else:
-        # loop for interactive
-        while True:
-            prompt = input("prompt: ")
-            if prompt == "":
-                break
-            prompt2 = input("prompt2: ")
-            if prompt2 == "":
-                prompt2 = prompt
-            negative_prompt = input("negative prompt: ")
-            seed = input("seed: ")
-            if seed == "":
-                seed = None
-            else:
-                seed = int(seed)
-            generate_image(prompt, prompt2, negative_prompt, seed)
-
-    print("Done!")

sdxl_train.py

@@ -1,746 +0,0 @@
-# training with captions
-
-import argparse
-import gc
-import math
-import os
-from multiprocessing import Value
-from typing import List
-import toml
-
-from tqdm import tqdm
-import torch
-from accelerate.utils import set_seed
-from diffusers import DDPMScheduler
-from library import sdxl_model_util
-
-import library.train_util as train_util
-import library.config_util as config_util
-import library.sdxl_train_util as sdxl_train_util
-from library.config_util import (
-    ConfigSanitizer,
-    BlueprintGenerator,
-)
-import library.custom_train_functions as custom_train_functions
-from library.custom_train_functions import (
-    apply_snr_weight,
-    prepare_scheduler_for_custom_training,
-    scale_v_prediction_loss_like_noise_prediction,
-    add_v_prediction_like_loss,
-)
-from library.sdxl_original_unet import SdxlUNet2DConditionModel
-
-
-UNET_NUM_BLOCKS_FOR_BLOCK_LR = 23
-
-
-def get_block_params_to_optimize(unet: SdxlUNet2DConditionModel, block_lrs: List[float]) -> List[dict]:
-    block_params = [[] for _ in range(len(block_lrs))]
-
-    for i, (name, param) in enumerate(unet.named_parameters()):
-        if name.startswith("time_embed.") or name.startswith("label_emb."):
-            block_index = 0  # 0
-        elif name.startswith("input_blocks."):  # 1-9
-            block_index = 1 + int(name.split(".")[1])
-        elif name.startswith("middle_block."):  # 10-12
-            block_index = 10 + int(name.split(".")[1])
-        elif name.startswith("output_blocks."):  # 13-21
-            block_index = 13 + int(name.split(".")[1])
-        elif name.startswith("out."):  # 22
-            block_index = 22
-        else:
-            raise ValueError(f"unexpected parameter name: {name}")
-
-        block_params[block_index].append(param)
-
-    params_to_optimize = []
-    for i, params in enumerate(block_params):
-        if block_lrs[i] == 0:  # 0のときは学習しない do not optimize when lr is 0
-            continue
-        params_to_optimize.append({"params": params, "lr": block_lrs[i]})
-
-    return params_to_optimize
-
-
-def append_block_lr_to_logs(block_lrs, logs, lr_scheduler, optimizer_type):
-    lrs = lr_scheduler.get_last_lr()
-
-    lr_index = 0
-    block_index = 0
-    while lr_index < len(lrs):
-        if block_index < UNET_NUM_BLOCKS_FOR_BLOCK_LR:
-            name = f"block{block_index}"
-            if block_lrs[block_index] == 0:
-                block_index += 1
-                continue
-        elif block_index == UNET_NUM_BLOCKS_FOR_BLOCK_LR:
-            name = "text_encoder1"
-        elif block_index == UNET_NUM_BLOCKS_FOR_BLOCK_LR + 1:
-            name = "text_encoder2"
-        else:
-            raise ValueError(f"unexpected block_index: {block_index}")
-
-        block_index += 1
-
-        logs["lr/" + name] = float(lrs[lr_index])
-
-        if optimizer_type.lower().startswith("DAdapt".lower()) or optimizer_type.lower() == "Prodigy".lower():
-            logs["lr/d*lr/" + name] = (
-                lr_scheduler.optimizers[-1].param_groups[lr_index]["d"] * lr_scheduler.optimizers[-1].param_groups[lr_index]["lr"]
-            )
-
-        lr_index += 1
-
-
-def train(args):
-    train_util.verify_training_args(args)
-    train_util.prepare_dataset_args(args, True)
-    sdxl_train_util.verify_sdxl_training_args(args)
-
-    assert not args.weighted_captions, "weighted_captions is not supported currently / weighted_captionsは現在サポートされていません"
-    assert (
-        not args.train_text_encoder or not args.cache_text_encoder_outputs
-    ), "cache_text_encoder_outputs is not supported when training text encoder / text encoderを学習するときはcache_text_encoder_outputsはサポートされていません"
-
-    if args.block_lr:
-        block_lrs = [float(lr) for lr in args.block_lr.split(",")]
-        assert (
-            len(block_lrs) == UNET_NUM_BLOCKS_FOR_BLOCK_LR
-        ), f"block_lr must have {UNET_NUM_BLOCKS_FOR_BLOCK_LR} values / block_lrは{UNET_NUM_BLOCKS_FOR_BLOCK_LR}個の値を指定してください"
-    else:
-        block_lrs = None
-
-    cache_latents = args.cache_latents
-    use_dreambooth_method = args.in_json is None
-
-    if args.seed is not None:
-        set_seed(args.seed)  # 乱数系列を初期化する
-
-    tokenizer1, tokenizer2 = sdxl_train_util.load_tokenizers(args)
-
-    # データセットを準備する
-    if args.dataset_class is None:
-        blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, False, True))
-        if args.dataset_config is not None:
-            print(f"Load dataset config from {args.dataset_config}")
-            user_config = config_util.load_user_config(args.dataset_config)
-            ignored = ["train_data_dir", "in_json"]
-            if any(getattr(args, attr) is not None for attr in ignored):
-                print(
-                    "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
-                        ", ".join(ignored)
-                    )
-                )
-        else:
-            if use_dreambooth_method:
-                print("Using DreamBooth method.")
-                user_config = {
-                    "datasets": [
-                        {
-                            "subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(
-                                args.train_data_dir, args.reg_data_dir
-                            )
-                        }
-                    ]
-                }
-            else:
-                print("Training with captions.")
-                user_config = {
-                    "datasets": [
-                        {
-                            "subsets": [
-                                {
-                                    "image_dir": args.train_data_dir,
-                                    "metadata_file": args.in_json,
-                                }
-                            ]
-                        }
-                    ]
-                }
-
-        blueprint = blueprint_generator.generate(user_config, args, tokenizer=[tokenizer1, tokenizer2])
-        train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
-    else:
-        train_dataset_group = train_util.load_arbitrary_dataset(args, [tokenizer1, tokenizer2])
-
-    current_epoch = Value("i", 0)
-    current_step = Value("i", 0)
-    ds_for_collater = train_dataset_group if args.max_data_loader_n_workers == 0 else None
-    collater = train_util.collater_class(current_epoch, current_step, ds_for_collater)
-
-    train_dataset_group.verify_bucket_reso_steps(32)
-
-    if args.debug_dataset:
-        train_util.debug_dataset(train_dataset_group, True)
-        return
-    if len(train_dataset_group) == 0:
-        print(
-            "No data found. Please verify the metadata file and train_data_dir option. / 画像がありません。メタデータおよびtrain_data_dirオプションを確認してください。"
-        )
-        return
-
-    if cache_latents:
-        assert (
-            train_dataset_group.is_latent_cacheable()
-        ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
-
-    if args.cache_text_encoder_outputs:
-        assert (
-            train_dataset_group.is_text_encoder_output_cacheable()
-        ), "when caching text encoder output, either caption_dropout_rate, shuffle_caption, token_warmup_step or caption_tag_dropout_rate cannot be used / text encoderの出力をキャッシュするときはcaption_dropout_rate, shuffle_caption, token_warmup_step, caption_tag_dropout_rateは使えません"
-
-    # acceleratorを準備する
-    print("prepare accelerator")
-    accelerator = train_util.prepare_accelerator(args)
-
-    # mixed precisionに対応した型を用意しておき適宜castする
-    weight_dtype, save_dtype = train_util.prepare_dtype(args)
-    vae_dtype = torch.float32 if args.no_half_vae else weight_dtype
-
-    # モデルを読み込む
-    (
-        load_stable_diffusion_format,
-        text_encoder1,
-        text_encoder2,
-        vae,
-        unet,
-        logit_scale,
-        ckpt_info,
-    ) = sdxl_train_util.load_target_model(args, accelerator, "sdxl", weight_dtype)
-    # logit_scale = logit_scale.to(accelerator.device, dtype=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())
-        # assert save_stable_diffusion_format, "save_model_as must be ckpt or safetensors / save_model_asはckptかsafetensorsである必要があります"
-
-    # Diffusers版のxformers使用フラグを設定する関数
-    def set_diffusers_xformers_flag(model, valid):
-        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:
-        # もうU-Netを独自にしたので動かないけどVAEのxformersは動くはず
-        accelerator.print("Use xformers by Diffusers")
-        # set_diffusers_xformers_flag(unet, True)
-        set_diffusers_xformers_flag(vae, True)
-    else:
-        # Windows版のxformersはfloatで学習できなかったりするのでxformersを使わない設定も可能にしておく必要がある
-        accelerator.print("Disable Diffusers' xformers")
-        train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers, args.sdpa)
-        if torch.__version__ >= "2.0.0":  # PyTorch 2.0.0 以上対応のxformersなら以下が使える
-            vae.set_use_memory_efficient_attention_xformers(args.xformers)
-
-    # 学習を準備する
-    if cache_latents:
-        vae.to(accelerator.device, dtype=vae_dtype)
-        vae.requires_grad_(False)
-        vae.eval()
-        with torch.no_grad():
-            train_dataset_group.cache_latents(vae, args.vae_batch_size, args.cache_latents_to_disk, accelerator.is_main_process)
-        vae.to("cpu")
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        gc.collect()
-
-        accelerator.wait_for_everyone()
-
-    # 学習を準備する：モデルを適切な状態にする
-    training_models = []
-    if args.gradient_checkpointing:
-        unet.enable_gradient_checkpointing()
-    training_models.append(unet)
-
-    if args.train_text_encoder:
-        # TODO each option for two text encoders?
-        accelerator.print("enable text encoder training")
-        if args.gradient_checkpointing:
-            text_encoder1.gradient_checkpointing_enable()
-            text_encoder2.gradient_checkpointing_enable()
-        training_models.append(text_encoder1)
-        training_models.append(text_encoder2)
-        # set require_grad=True later
-    else:
-        text_encoder1.requires_grad_(False)
-        text_encoder2.requires_grad_(False)
-        text_encoder1.eval()
-        text_encoder2.eval()
-
-        # TextEncoderの出力をキャッシュする
-        if args.cache_text_encoder_outputs:
-            # Text Encodes are eval and no grad
-            with torch.no_grad():
-                train_dataset_group.cache_text_encoder_outputs(
-                    (tokenizer1, tokenizer2),
-                    (text_encoder1, text_encoder2),
-                    accelerator.device,
-                    None,
-                    args.cache_text_encoder_outputs_to_disk,
-                    accelerator.is_main_process,
-                )
-            accelerator.wait_for_everyone()
-
-    if not cache_latents:
-        vae.requires_grad_(False)
-        vae.eval()
-        vae.to(accelerator.device, dtype=vae_dtype)
-
-    for m in training_models:
-        m.requires_grad_(True)
-
-    if block_lrs is None:
-        params = []
-        for m in training_models:
-            params.extend(m.parameters())
-        params_to_optimize = params
-
-        # calculate number of trainable parameters
-        n_params = 0
-        for p in params:
-            n_params += p.numel()
-    else:
-        params_to_optimize = get_block_params_to_optimize(training_models[0], block_lrs)  # U-Net
-        for m in training_models[1:]:  # Text Encoders if exists
-            params_to_optimize.append({"params": m.parameters(), "lr": args.learning_rate})
-
-        # calculate number of trainable parameters
-        n_params = 0
-        for params in params_to_optimize:
-            for p in params["params"]:
-                n_params += p.numel()
-
-    accelerator.print(f"number of models: {len(training_models)}")
-    accelerator.print(f"number of trainable parameters: {n_params}")
-
-    # 学習に必要なクラスを準備する
-    accelerator.print("prepare optimizer, data loader etc.")
-    _, _, optimizer = train_util.get_optimizer(args, trainable_params=params_to_optimize)
-
-    # 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_group,
-        batch_size=1,
-        shuffle=True,
-        collate_fn=collater,
-        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 * math.ceil(
-            len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps
-        )
-        accelerator.print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
-
-    # データセット側にも学習ステップを送信
-    train_dataset_group.set_max_train_steps(args.max_train_steps)
-
-    # lr schedulerを用意する
-    lr_scheduler = train_util.get_scheduler_fix(args, optimizer, accelerator.num_processes)
-
-    # 実験的機能：勾配も含めたfp16/bf16学習を行う　モデル全体をfp16/bf16にする
-    if args.full_fp16:
-        assert (
-            args.mixed_precision == "fp16"
-        ), "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
-        accelerator.print("enable full fp16 training.")
-        unet.to(weight_dtype)
-        text_encoder1.to(weight_dtype)
-        text_encoder2.to(weight_dtype)
-    elif args.full_bf16:
-        assert (
-            args.mixed_precision == "bf16"
-        ), "full_bf16 requires mixed precision='bf16' / full_bf16を使う場合はmixed_precision='bf16'を指定してください。"
-        accelerator.print("enable full bf16 training.")
-        unet.to(weight_dtype)
-        text_encoder1.to(weight_dtype)
-        text_encoder2.to(weight_dtype)
-
-    # acceleratorがなんかよろしくやってくれるらしい
-    if args.train_text_encoder:
-        unet, text_encoder1, text_encoder2, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-            unet, text_encoder1, text_encoder2, optimizer, train_dataloader, lr_scheduler
-        )
-
-        # transform DDP after prepare
-        text_encoder1, text_encoder2, unet = train_util.transform_models_if_DDP([text_encoder1, text_encoder2, unet])
-    else:
-        unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(unet, optimizer, train_dataloader, lr_scheduler)
-        (unet,) = train_util.transform_models_if_DDP([unet])
-        text_encoder1.to(weight_dtype)
-        text_encoder2.to(weight_dtype)
-
-    # TextEncoderの出力をキャッシュするときにはCPUへ移動する
-    if args.cache_text_encoder_outputs:
-        # move Text Encoders for sampling images. Text Encoder doesn't work on CPU with fp16
-        text_encoder1.to("cpu", dtype=torch.float32)
-        text_encoder2.to("cpu", dtype=torch.float32)
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-    else:
-        # make sure Text Encoders are on GPU
-        text_encoder1.to(accelerator.device)
-        text_encoder2.to(accelerator.device)
-
-    # 実験的機能：勾配も含めたfp16学習を行う　PyTorchにパッチを当ててfp16でのgrad scaleを有効にする
-    if args.full_fp16:
-        train_util.patch_accelerator_for_fp16_training(accelerator)
-
-    # resumeする
-    train_util.resume_from_local_or_hf_if_specified(accelerator, args)
-
-    # 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
-    accelerator.print("running training / 学習開始")
-    accelerator.print(f"  num examples / サンプル数: {train_dataset_group.num_train_images}")
-    accelerator.print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
-    accelerator.print(f"  num epochs / epoch数: {num_train_epochs}")
-    accelerator.print(f"  batch size per device / バッチサイズ: {', '.join([str(d.batch_size) for d in train_dataset_group.datasets])}")
-    # accelerator.print(
-    #     f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}"
-    # )
-    accelerator.print(f"  gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
-    accelerator.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
-    )
-    prepare_scheduler_for_custom_training(noise_scheduler, accelerator.device)
-    if args.zero_terminal_snr:
-        custom_train_functions.fix_noise_scheduler_betas_for_zero_terminal_snr(noise_scheduler)
-
-    if accelerator.is_main_process:
-        init_kwargs = {}
-        if args.log_tracker_config is not None:
-            init_kwargs = toml.load(args.log_tracker_config)
-        accelerator.init_trackers("finetuning" if args.log_tracker_name is None else args.log_tracker_name, init_kwargs=init_kwargs)
-
-    for epoch in range(num_train_epochs):
-        accelerator.print(f"\nepoch {epoch+1}/{num_train_epochs}")
-        current_epoch.value = epoch + 1
-
-        for m in training_models:
-            m.train()
-
-        loss_total = 0
-        for step, batch in enumerate(train_dataloader):
-            current_step.value = global_step
-            with accelerator.accumulate(training_models[0]):  # 複数モデルに対応していない模様だがとりあえずこうしておく
-                if "latents" in batch and batch["latents"] is not None:
-                    latents = batch["latents"].to(accelerator.device).to(dtype=weight_dtype)
-                else:
-                    with torch.no_grad():
-                        # latentに変換
-                        latents = vae.encode(batch["images"].to(vae_dtype)).latent_dist.sample().to(weight_dtype)
-
-                        # NaNが含まれていれば警告を表示し0に置き換える
-                        if torch.any(torch.isnan(latents)):
-                            accelerator.print("NaN found in latents, replacing with zeros")
-                            latents = torch.where(torch.isnan(latents), torch.zeros_like(latents), latents)
-                latents = latents * sdxl_model_util.VAE_SCALE_FACTOR
-
-                if "text_encoder_outputs1_list" not in batch or batch["text_encoder_outputs1_list"] is None:
-                    input_ids1 = batch["input_ids"]
-                    input_ids2 = batch["input_ids2"]
-                    with torch.set_grad_enabled(args.train_text_encoder):
-                        # Get the text embedding for conditioning
-                        # TODO support weighted captions
-                        # if args.weighted_captions:
-                        #     encoder_hidden_states = get_weighted_text_embeddings(
-                        #         tokenizer,
-                        #         text_encoder,
-                        #         batch["captions"],
-                        #         accelerator.device,
-                        #         args.max_token_length // 75 if args.max_token_length else 1,
-                        #         clip_skip=args.clip_skip,
-                        #     )
-                        # else:
-                        input_ids1 = input_ids1.to(accelerator.device)
-                        input_ids2 = input_ids2.to(accelerator.device)
-                        encoder_hidden_states1, encoder_hidden_states2, pool2 = train_util.get_hidden_states_sdxl(
-                            args.max_token_length,
-                            input_ids1,
-                            input_ids2,
-                            tokenizer1,
-                            tokenizer2,
-                            text_encoder1,
-                            text_encoder2,
-                            None if not args.full_fp16 else weight_dtype,
-                        )
-                else:
-                    encoder_hidden_states1 = batch["text_encoder_outputs1_list"].to(accelerator.device).to(weight_dtype)
-                    encoder_hidden_states2 = batch["text_encoder_outputs2_list"].to(accelerator.device).to(weight_dtype)
-                    pool2 = batch["text_encoder_pool2_list"].to(accelerator.device).to(weight_dtype)
-
-                    # # verify that the text encoder outputs are correct
-                    # ehs1, ehs2, p2 = train_util.get_hidden_states_sdxl(
-                    #     args.max_token_length,
-                    #     batch["input_ids"].to(text_encoder1.device),
-                    #     batch["input_ids2"].to(text_encoder1.device),
-                    #     tokenizer1,
-                    #     tokenizer2,
-                    #     text_encoder1,
-                    #     text_encoder2,
-                    #     None if not args.full_fp16 else weight_dtype,
-                    # )
-                    # b_size = encoder_hidden_states1.shape[0]
-                    # assert ((encoder_hidden_states1.to("cpu") - ehs1.to(dtype=weight_dtype)).abs().max() > 1e-2).sum() <= b_size * 2
-                    # assert ((encoder_hidden_states2.to("cpu") - ehs2.to(dtype=weight_dtype)).abs().max() > 1e-2).sum() <= b_size * 2
-                    # assert ((pool2.to("cpu") - p2.to(dtype=weight_dtype)).abs().max() > 1e-2).sum() <= b_size * 2
-                    # print("text encoder outputs verified")
-
-                # get size embeddings
-                orig_size = batch["original_sizes_hw"]
-                crop_size = batch["crop_top_lefts"]
-                target_size = batch["target_sizes_hw"]
-                embs = sdxl_train_util.get_size_embeddings(orig_size, crop_size, target_size, accelerator.device).to(weight_dtype)
-
-                # concat embeddings
-                vector_embedding = torch.cat([pool2, embs], dim=1).to(weight_dtype)
-                text_embedding = torch.cat([encoder_hidden_states1, encoder_hidden_states2], dim=2).to(weight_dtype)
-
-                # Sample noise, sample a random timestep for each image, and add noise to the latents,
-                # with noise offset and/or multires noise if specified
-                noise, noisy_latents, timesteps = train_util.get_noise_noisy_latents_and_timesteps(args, noise_scheduler, latents)
-
-                noisy_latents = noisy_latents.to(weight_dtype)  # TODO check why noisy_latents is not weight_dtype
-
-                # Predict the noise residual
-                with accelerator.autocast():
-                    noise_pred = unet(noisy_latents, timesteps, text_embedding, vector_embedding)
-
-                target = noise
-
-                if args.min_snr_gamma or args.scale_v_pred_loss_like_noise_pred or args.v_pred_like_loss:
-                    # do not mean over batch dimension for snr weight or scale v-pred loss
-                    loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="none")
-                    loss = loss.mean([1, 2, 3])
-
-                    if args.min_snr_gamma:
-                        loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma)
-                    if args.scale_v_pred_loss_like_noise_pred:
-                        loss = scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler)
-                    if args.v_pred_like_loss:
-                        loss = add_v_prediction_like_loss(loss, timesteps, noise_scheduler, args.v_pred_like_loss)
-
-                    loss = loss.mean()  # mean over batch dimension
-                else:
-                    loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="mean")
-
-                accelerator.backward(loss)
-                if accelerator.sync_gradients and args.max_grad_norm != 0.0:
-                    params_to_clip = []
-                    for m in training_models:
-                        params_to_clip.extend(m.parameters())
-                    accelerator.clip_grad_norm_(params_to_clip, 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
-
-                sdxl_train_util.sample_images(
-                    accelerator,
-                    args,
-                    None,
-                    global_step,
-                    accelerator.device,
-                    vae,
-                    [tokenizer1, tokenizer2],
-                    [text_encoder1, text_encoder2],
-                    unet,
-                )
-
-                # 指定ステップごとにモデルを保存
-                if args.save_every_n_steps is not None and global_step % args.save_every_n_steps == 0:
-                    accelerator.wait_for_everyone()
-                    if accelerator.is_main_process:
-                        src_path = src_stable_diffusion_ckpt if save_stable_diffusion_format else src_diffusers_model_path
-                        sdxl_train_util.save_sd_model_on_epoch_end_or_stepwise(
-                            args,
-                            False,
-                            accelerator,
-                            src_path,
-                            save_stable_diffusion_format,
-                            use_safetensors,
-                            save_dtype,
-                            epoch,
-                            num_train_epochs,
-                            global_step,
-                            accelerator.unwrap_model(text_encoder1),
-                            accelerator.unwrap_model(text_encoder2),
-                            accelerator.unwrap_model(unet),
-                            vae,
-                            logit_scale,
-                            ckpt_info,
-                        )
-
-            current_loss = loss.detach().item()  # 平均なのでbatch sizeは関係ないはず
-            if args.logging_dir is not None:
-                logs = {"loss": current_loss}
-                if block_lrs is None:
-                    logs["lr"] = float(lr_scheduler.get_last_lr()[0])
-                    if (
-                        args.optimizer_type.lower().startswith("DAdapt".lower()) or args.optimizer_type.lower() == "Prodigy".lower()
-                    ):  # tracking d*lr value
-                        logs["lr/d*lr"] = (
-                            lr_scheduler.optimizers[0].param_groups[0]["d"] * lr_scheduler.optimizers[0].param_groups[0]["lr"]
-                        )
-                else:
-                    append_block_lr_to_logs(block_lrs, logs, lr_scheduler, args.optimizer_type)
-
-                accelerator.log(logs, step=global_step)
-
-            # TODO moving averageにする
-            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 = {"loss/epoch": loss_total / len(train_dataloader)}
-            accelerator.log(logs, step=epoch + 1)
-
-        accelerator.wait_for_everyone()
-
-        if args.save_every_n_epochs is not None:
-            if accelerator.is_main_process:
-                src_path = src_stable_diffusion_ckpt if save_stable_diffusion_format else src_diffusers_model_path
-                sdxl_train_util.save_sd_model_on_epoch_end_or_stepwise(
-                    args,
-                    True,
-                    accelerator,
-                    src_path,
-                    save_stable_diffusion_format,
-                    use_safetensors,
-                    save_dtype,
-                    epoch,
-                    num_train_epochs,
-                    global_step,
-                    accelerator.unwrap_model(text_encoder1),
-                    accelerator.unwrap_model(text_encoder2),
-                    accelerator.unwrap_model(unet),
-                    vae,
-                    logit_scale,
-                    ckpt_info,
-                )
-
-        sdxl_train_util.sample_images(
-            accelerator,
-            args,
-            epoch + 1,
-            global_step,
-            accelerator.device,
-            vae,
-            [tokenizer1, tokenizer2],
-            [text_encoder1, text_encoder2],
-            unet,
-        )
-
-    is_main_process = accelerator.is_main_process
-    # if is_main_process:
-    unet = accelerator.unwrap_model(unet)
-    text_encoder1 = accelerator.unwrap_model(text_encoder1)
-    text_encoder2 = accelerator.unwrap_model(text_encoder2)
-
-    accelerator.end_training()
-
-    if args.save_state:  # and is_main_process:
-        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
-        sdxl_train_util.save_sd_model_on_train_end(
-            args,
-            src_path,
-            save_stable_diffusion_format,
-            use_safetensors,
-            save_dtype,
-            epoch,
-            global_step,
-            text_encoder1,
-            text_encoder2,
-            unet,
-            vae,
-            logit_scale,
-            ckpt_info,
-        )
-        print("model saved.")
-
-
-def setup_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser()
-
-    train_util.add_sd_models_arguments(parser)
-    train_util.add_dataset_arguments(parser, True, True, True)
-    train_util.add_training_arguments(parser, False)
-    train_util.add_sd_saving_arguments(parser)
-    train_util.add_optimizer_arguments(parser)
-    config_util.add_config_arguments(parser)
-    custom_train_functions.add_custom_train_arguments(parser)
-    sdxl_train_util.add_sdxl_training_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も学習する")
-    parser.add_argument(
-        "--no_half_vae",
-        action="store_true",
-        help="do not use fp16/bf16 VAE in mixed precision (use float VAE) / mixed precisionでも fp16/bf16 VAEを使わずfloat VAEを使う",
-    )
-    parser.add_argument(
-        "--block_lr",
-        type=str,
-        default=None,
-        help=f"learning rates for each block of U-Net, comma-separated, {UNET_NUM_BLOCKS_FOR_BLOCK_LR} values / "
-        + f"U-Netの各ブロックの学習率、カンマ区切り、{UNET_NUM_BLOCKS_FOR_BLOCK_LR}個の値",
-    )
-
-    return parser
-
-
-if __name__ == "__main__":
-    parser = setup_parser()
-
-    args = parser.parse_args()
-    args = train_util.read_config_from_file(args, parser)
-
-    train(args)

sdxl_train_control_net_lllite.py

@@ -1,572 +0,0 @@
-import argparse
-import gc
-import json
-import math
-import os
-import random
-import time
-from multiprocessing import Value
-from types import SimpleNamespace
-import toml
-
-from tqdm import tqdm
-import torch
-from torch.nn.parallel import DistributedDataParallel as DDP
-from accelerate.utils import set_seed
-from diffusers import DDPMScheduler, ControlNetModel
-from safetensors.torch import load_file
-from library import sai_model_spec, sdxl_model_util, sdxl_original_unet, sdxl_train_util
-
-import library.model_util as model_util
-import library.train_util as train_util
-import library.config_util as config_util
-from library.config_util import (
-    ConfigSanitizer,
-    BlueprintGenerator,
-)
-import library.huggingface_util as huggingface_util
-import library.custom_train_functions as custom_train_functions
-from library.custom_train_functions import (
-    add_v_prediction_like_loss,
-    apply_snr_weight,
-    prepare_scheduler_for_custom_training,
-    pyramid_noise_like,
-    apply_noise_offset,
-    scale_v_prediction_loss_like_noise_prediction,
-)
-import networks.control_net_lllite as control_net_lllite
-
-
-# TODO 他のスクリプトと共通化する
-def generate_step_logs(args: argparse.Namespace, current_loss, avr_loss, lr_scheduler):
-    logs = {
-        "loss/current": current_loss,
-        "loss/average": avr_loss,
-        "lr": lr_scheduler.get_last_lr()[0],
-    }
-
-    if args.optimizer_type.lower().startswith("DAdapt".lower()):
-        logs["lr/d*lr"] = lr_scheduler.optimizers[-1].param_groups[0]["d"] * lr_scheduler.optimizers[-1].param_groups[0]["lr"]
-
-    return logs
-
-
-def train(args):
-    train_util.verify_training_args(args)
-    train_util.prepare_dataset_args(args, True)
-    sdxl_train_util.verify_sdxl_training_args(args)
-
-    cache_latents = args.cache_latents
-    use_user_config = args.dataset_config is not None
-
-    if args.seed is None:
-        args.seed = random.randint(0, 2**32)
-    set_seed(args.seed)
-
-    tokenizer1, tokenizer2 = sdxl_train_util.load_tokenizers(args)
-
-    # データセットを準備する
-    blueprint_generator = BlueprintGenerator(ConfigSanitizer(False, False, True, True))
-    if use_user_config:
-        print(f"Load dataset config from {args.dataset_config}")
-        user_config = config_util.load_user_config(args.dataset_config)
-        ignored = ["train_data_dir", "conditioning_data_dir"]
-        if any(getattr(args, attr) is not None for attr in ignored):
-            print(
-                "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
-                    ", ".join(ignored)
-                )
-            )
-    else:
-        user_config = {
-            "datasets": [
-                {
-                    "subsets": config_util.generate_controlnet_subsets_config_by_subdirs(
-                        args.train_data_dir,
-                        args.conditioning_data_dir,
-                        args.caption_extension,
-                    )
-                }
-            ]
-        }
-
-    blueprint = blueprint_generator.generate(user_config, args, tokenizer=[tokenizer1, tokenizer2])
-    train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
-
-    current_epoch = Value("i", 0)
-    current_step = Value("i", 0)
-    ds_for_collater = train_dataset_group if args.max_data_loader_n_workers == 0 else None
-    collater = train_util.collater_class(current_epoch, current_step, ds_for_collater)
-
-    train_dataset_group.verify_bucket_reso_steps(32)
-
-    if args.debug_dataset:
-        train_util.debug_dataset(train_dataset_group)
-        return
-    if len(train_dataset_group) == 0:
-        print(
-            "No data found. Please verify arguments (train_data_dir must be the parent of folders with images) / 画像がありません。引数指定を確認してください（train_data_dirには画像があるフォルダではなく、画像があるフォルダの親フォルダを指定する必要があります）"
-        )
-        return
-
-    if cache_latents:
-        assert (
-            train_dataset_group.is_latent_cacheable()
-        ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
-    else:
-        print("WARNING: random_crop is not supported yet for ControlNet training / ControlNetの学習ではrandom_cropはまだサポートされていません")
-
-    if args.cache_text_encoder_outputs:
-        assert (
-            train_dataset_group.is_text_encoder_output_cacheable()
-        ), "when caching Text Encoder output, either caption_dropout_rate, shuffle_caption, token_warmup_step or caption_tag_dropout_rate cannot be used / Text Encoderの出力をキャッシュするときはcaption_dropout_rate, shuffle_caption, token_warmup_step, caption_tag_dropout_rateは使えません"
-
-    # acceleratorを準備する
-    print("prepare accelerator")
-    accelerator = train_util.prepare_accelerator(args)
-    is_main_process = accelerator.is_main_process
-
-    # mixed precisionに対応した型を用意しておき適宜castする
-    weight_dtype, save_dtype = train_util.prepare_dtype(args)
-    vae_dtype = torch.float32 if args.no_half_vae else weight_dtype
-
-    # モデルを読み込む
-    (
-        load_stable_diffusion_format,
-        text_encoder1,
-        text_encoder2,
-        vae,
-        unet,
-        logit_scale,
-        ckpt_info,
-    ) = sdxl_train_util.load_target_model(args, accelerator, sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0, weight_dtype)
-
-    # モデルに xformers とか memory efficient attention を組み込む
-    train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers, args.sdpa)
-
-    # 学習を準備する
-    if cache_latents:
-        vae.to(accelerator.device, dtype=vae_dtype)
-        vae.requires_grad_(False)
-        vae.eval()
-        with torch.no_grad():
-            train_dataset_group.cache_latents(
-                vae,
-                args.vae_batch_size,
-                args.cache_latents_to_disk,
-                accelerator.is_main_process,
-            )
-        vae.to("cpu")
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        gc.collect()
-
-        accelerator.wait_for_everyone()
-
-    # TextEncoderの出力をキャッシュする
-    if args.cache_text_encoder_outputs:
-        # Text Encodes are eval and no grad
-        with torch.no_grad():
-            train_dataset_group.cache_text_encoder_outputs(
-                (tokenizer1, tokenizer2),
-                (text_encoder1, text_encoder2),
-                accelerator.device,
-                None,
-                args.cache_text_encoder_outputs_to_disk,
-                accelerator.is_main_process,
-            )
-        accelerator.wait_for_everyone()
-
-    # prepare ControlNet
-    network = control_net_lllite.ControlNetLLLite(unet, args.cond_emb_dim, args.network_dim, args.network_dropout)
-    network.apply_to()
-
-    if args.network_weights is not None:
-        info = network.load_weights(args.network_weights)
-        accelerator.print(f"load ControlNet weights from {args.network_weights}: {info}")
-
-    if args.gradient_checkpointing:
-        unet.enable_gradient_checkpointing()
-        network.enable_gradient_checkpointing()  # may have no effect
-
-    # 学習に必要なクラスを準備する
-    accelerator.print("prepare optimizer, data loader etc.")
-
-    trainable_params = list(network.prepare_optimizer_params())
-    print(f"trainable params count: {len(trainable_params)}")
-    print(f"number of trainable parameters: {sum(p.numel() for p in trainable_params if p.requires_grad)}")
-
-    _, _, optimizer = train_util.get_optimizer(args, trainable_params)
-
-    # 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_group,
-        batch_size=1,
-        shuffle=True,
-        collate_fn=collater,
-        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 * math.ceil(
-            len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps
-        )
-        accelerator.print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
-
-    # データセット側にも学習ステップを送信
-    train_dataset_group.set_max_train_steps(args.max_train_steps)
-
-    # lr schedulerを用意する
-    lr_scheduler = train_util.get_scheduler_fix(args, optimizer, accelerator.num_processes)
-
-    # 実験的機能：勾配も含めたfp16/bf16学習を行う　モデル全体をfp16/bf16にする
-    if args.full_fp16:
-        assert (
-            args.mixed_precision == "fp16"
-        ), "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
-        accelerator.print("enable full fp16 training.")
-        unet.to(weight_dtype)
-        network.to(weight_dtype)
-    elif args.full_bf16:
-        assert (
-            args.mixed_precision == "bf16"
-        ), "full_bf16 requires mixed precision='bf16' / full_bf16を使う場合はmixed_precision='bf16'を指定してください。"
-        accelerator.print("enable full bf16 training.")
-        unet.to(weight_dtype)
-        network.to(weight_dtype)
-
-    # acceleratorがなんかよろしくやってくれるらしい
-    unet, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-        unet, network, optimizer, train_dataloader, lr_scheduler
-    )
-    network: control_net_lllite.ControlNetLLLite
-
-    # transform DDP after prepare (train_network here only)
-    unet, network = train_util.transform_models_if_DDP([unet, network])
-
-    if args.gradient_checkpointing:
-        unet.train()  # according to TI example in Diffusers, train is required -> これオリジナルのU-Netしたので本当は外せる
-    else:
-        unet.eval()
-
-    network.prepare_grad_etc()
-
-    # TextEncoderの出力をキャッシュするときにはCPUへ移動する
-    if args.cache_text_encoder_outputs:
-        # move Text Encoders for sampling images. Text Encoder doesn't work on CPU with fp16
-        text_encoder1.to("cpu", dtype=torch.float32)
-        text_encoder2.to("cpu", dtype=torch.float32)
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-    else:
-        # make sure Text Encoders are on GPU
-        text_encoder1.to(accelerator.device)
-        text_encoder2.to(accelerator.device)
-
-    if not cache_latents:
-        vae.requires_grad_(False)
-        vae.eval()
-        vae.to(accelerator.device, dtype=vae_dtype)
-
-    # 実験的機能：勾配も含めたfp16学習を行う　PyTorchにパッチを当ててfp16でのgrad scaleを有効にする
-    if args.full_fp16:
-        train_util.patch_accelerator_for_fp16_training(accelerator)
-
-    # resumeする
-    train_util.resume_from_local_or_hf_if_specified(accelerator, args)
-
-    # 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
-
-    # 学習する
-    # TODO: find a way to handle total batch size when there are multiple datasets
-    accelerator.print("running training / 学習開始")
-    accelerator.print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
-    accelerator.print(f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
-    accelerator.print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
-    accelerator.print(f"  num epochs / epoch数: {num_train_epochs}")
-    accelerator.print(f"  batch size per device / バッチサイズ: {', '.join([str(d.batch_size) for d in train_dataset_group.datasets])}")
-    # print(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
-    accelerator.print(f"  gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
-    accelerator.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
-    )
-    prepare_scheduler_for_custom_training(noise_scheduler, accelerator.device)
-    if args.zero_terminal_snr:
-        custom_train_functions.fix_noise_scheduler_betas_for_zero_terminal_snr(noise_scheduler)
-
-    if accelerator.is_main_process:
-        init_kwargs = {}
-        if args.log_tracker_config is not None:
-            init_kwargs = toml.load(args.log_tracker_config)
-        accelerator.init_trackers(
-            "lllite_control_net_train" if args.log_tracker_name is None else args.log_tracker_name, init_kwargs=init_kwargs
-        )
-
-    loss_list = []
-    loss_total = 0.0
-    del train_dataset_group
-
-    # function for saving/removing
-    def save_model(ckpt_name, unwrapped_nw, steps, epoch_no, force_sync_upload=False):
-        os.makedirs(args.output_dir, exist_ok=True)
-        ckpt_file = os.path.join(args.output_dir, ckpt_name)
-
-        accelerator.print(f"\nsaving checkpoint: {ckpt_file}")
-        sai_metadata = train_util.get_sai_model_spec(None, args, True, True, False)
-        sai_metadata["modelspec.architecture"] = sai_model_spec.ARCH_SD_XL_V1_BASE + "/control-net-lllite"
-
-        unwrapped_nw.save_weights(ckpt_file, save_dtype, sai_metadata)
-        if args.huggingface_repo_id is not None:
-            huggingface_util.upload(args, ckpt_file, "/" + ckpt_name, force_sync_upload=force_sync_upload)
-
-    def remove_model(old_ckpt_name):
-        old_ckpt_file = os.path.join(args.output_dir, old_ckpt_name)
-        if os.path.exists(old_ckpt_file):
-            accelerator.print(f"removing old checkpoint: {old_ckpt_file}")
-            os.remove(old_ckpt_file)
-
-    # training loop
-    for epoch in range(num_train_epochs):
-        accelerator.print(f"\nepoch {epoch+1}/{num_train_epochs}")
-        current_epoch.value = epoch + 1
-
-        network.on_epoch_start()  # train()
-
-        for step, batch in enumerate(train_dataloader):
-            current_step.value = global_step
-            with accelerator.accumulate(network):
-                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=vae_dtype)).latent_dist.sample()
-
-                        # NaNが含まれていれば警告を表示し0に置き換える
-                        if torch.any(torch.isnan(latents)):
-                            accelerator.print("NaN found in latents, replacing with zeros")
-                            latents = torch.where(torch.isnan(latents), torch.zeros_like(latents), latents)
-                    latents = latents * sdxl_model_util.VAE_SCALE_FACTOR
-
-                if "text_encoder_outputs1_list" not in batch or batch["text_encoder_outputs1_list"] is None:
-                    input_ids1 = batch["input_ids"]
-                    input_ids2 = batch["input_ids2"]
-                    with torch.no_grad():
-                        # Get the text embedding for conditioning
-                        input_ids1 = input_ids1.to(accelerator.device)
-                        input_ids2 = input_ids2.to(accelerator.device)
-                        encoder_hidden_states1, encoder_hidden_states2, pool2 = train_util.get_hidden_states_sdxl(
-                            args.max_token_length,
-                            input_ids1,
-                            input_ids2,
-                            tokenizer1,
-                            tokenizer2,
-                            text_encoder1,
-                            text_encoder2,
-                            None if not args.full_fp16 else weight_dtype,
-                        )
-                else:
-                    encoder_hidden_states1 = batch["text_encoder_outputs1_list"].to(accelerator.device).to(weight_dtype)
-                    encoder_hidden_states2 = batch["text_encoder_outputs2_list"].to(accelerator.device).to(weight_dtype)
-                    pool2 = batch["text_encoder_pool2_list"].to(accelerator.device).to(weight_dtype)
-
-                # get size embeddings
-                orig_size = batch["original_sizes_hw"]
-                crop_size = batch["crop_top_lefts"]
-                target_size = batch["target_sizes_hw"]
-                embs = sdxl_train_util.get_size_embeddings(orig_size, crop_size, target_size, accelerator.device).to(weight_dtype)
-
-                # concat embeddings
-                vector_embedding = torch.cat([pool2, embs], dim=1).to(weight_dtype)
-                text_embedding = torch.cat([encoder_hidden_states1, encoder_hidden_states2], dim=2).to(weight_dtype)
-
-                # Sample noise, sample a random timestep for each image, and add noise to the latents,
-                # with noise offset and/or multires noise if specified
-                noise, noisy_latents, timesteps = train_util.get_noise_noisy_latents_and_timesteps(args, noise_scheduler, latents)
-
-                noisy_latents = noisy_latents.to(weight_dtype)  # TODO check why noisy_latents is not weight_dtype
-
-                controlnet_image = batch["conditioning_images"].to(dtype=weight_dtype)
-
-                with accelerator.autocast():
-                    # conditioning imageをControlNetに渡す / pass conditioning image to ControlNet
-                    # 内部でcond_embに変換される / it will be converted to cond_emb inside
-                    network.set_cond_image(controlnet_image)
-
-                    # それらの値を使いつつ、U-Netでノイズを予測する / predict noise with U-Net using those values
-                    noise_pred = unet(noisy_latents, timesteps, text_embedding, vector_embedding)
-
-                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="none")
-                loss = loss.mean([1, 2, 3])
-
-                loss_weights = batch["loss_weights"]  # 各sampleごとのweight
-                loss = loss * loss_weights
-
-                if args.min_snr_gamma:
-                    loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma)
-                if args.scale_v_pred_loss_like_noise_pred:
-                    loss = scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler)
-                if args.v_pred_like_loss:
-                    loss = add_v_prediction_like_loss(loss, timesteps, noise_scheduler, args.v_pred_like_loss)
-
-                loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし
-
-                accelerator.backward(loss)
-                if accelerator.sync_gradients and args.max_grad_norm != 0.0:
-                    params_to_clip = network.get_trainable_params()
-                    accelerator.clip_grad_norm_(params_to_clip, 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
-
-                # sdxl_train_util.sample_images(accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
-
-                # 指定ステップごとにモデルを保存
-                if args.save_every_n_steps is not None and global_step % args.save_every_n_steps == 0:
-                    accelerator.wait_for_everyone()
-                    if accelerator.is_main_process:
-                        ckpt_name = train_util.get_step_ckpt_name(args, "." + args.save_model_as, global_step)
-                        save_model(ckpt_name, accelerator.unwrap_model(network), global_step, epoch)
-
-                        if args.save_state:
-                            train_util.save_and_remove_state_stepwise(args, accelerator, global_step)
-
-                        remove_step_no = train_util.get_remove_step_no(args, global_step)
-                        if remove_step_no is not None:
-                            remove_ckpt_name = train_util.get_step_ckpt_name(args, "." + args.save_model_as, remove_step_no)
-                            remove_model(remove_ckpt_name)
-
-            current_loss = loss.detach().item()
-            if epoch == 0:
-                loss_list.append(current_loss)
-            else:
-                loss_total -= loss_list[step]
-                loss_list[step] = current_loss
-            loss_total += current_loss
-            avr_loss = loss_total / len(loss_list)
-            logs = {"loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
-            progress_bar.set_postfix(**logs)
-
-            if args.logging_dir is not None:
-                logs = generate_step_logs(args, current_loss, avr_loss, lr_scheduler)
-                accelerator.log(logs, step=global_step)
-
-            if global_step >= args.max_train_steps:
-                break
-
-        if args.logging_dir is not None:
-            logs = {"loss/epoch": loss_total / len(loss_list)}
-            accelerator.log(logs, step=epoch + 1)
-
-        accelerator.wait_for_everyone()
-
-        # 指定エポックごとにモデルを保存
-        if args.save_every_n_epochs is not None:
-            saving = (epoch + 1) % args.save_every_n_epochs == 0 and (epoch + 1) < num_train_epochs
-            if is_main_process and saving:
-                ckpt_name = train_util.get_epoch_ckpt_name(args, "." + args.save_model_as, epoch + 1)
-                save_model(ckpt_name, accelerator.unwrap_model(network), global_step, epoch + 1)
-
-                remove_epoch_no = train_util.get_remove_epoch_no(args, epoch + 1)
-                if remove_epoch_no is not None:
-                    remove_ckpt_name = train_util.get_epoch_ckpt_name(args, "." + args.save_model_as, remove_epoch_no)
-                    remove_model(remove_ckpt_name)
-
-                if args.save_state:
-                    train_util.save_and_remove_state_on_epoch_end(args, accelerator, epoch + 1)
-
-        # self.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
-
-        # end of epoch
-
-    if is_main_process:
-        network = accelerator.unwrap_model(network)
-
-    accelerator.end_training()
-
-    if is_main_process and args.save_state:
-        train_util.save_state_on_train_end(args, accelerator)
-
-    if is_main_process:
-        ckpt_name = train_util.get_last_ckpt_name(args, "." + args.save_model_as)
-        save_model(ckpt_name, network, global_step, num_train_epochs, force_sync_upload=True)
-
-        print("model saved.")
-
-
-def setup_parser() -> argparse.ArgumentParser:
-    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_optimizer_arguments(parser)
-    config_util.add_config_arguments(parser)
-    custom_train_functions.add_custom_train_arguments(parser)
-    sdxl_train_util.add_sdxl_training_arguments(parser)
-
-    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("--cond_emb_dim", type=int, default=None, help="conditioning embedding dimension / 条件付け埋め込みの次元数")
-    parser.add_argument("--network_weights", type=str, default=None, help="pretrained weights for network / 学習するネットワークの初期重み")
-    parser.add_argument("--network_dim", type=int, default=None, help="network dimensions (rank) / モジュールの次元数")
-    parser.add_argument(
-        "--network_dropout",
-        type=float,
-        default=None,
-        help="Drops neurons out of training every step (0 or None is default behavior (no dropout), 1 would drop all neurons) / 訓練時に毎ステップでニューロンをdropする（0またはNoneはdropoutなし、1は全ニューロンをdropout）",
-    )
-    parser.add_argument(
-        "--conditioning_data_dir",
-        type=str,
-        default=None,
-        help="conditioning data directory / 条件付けデータのディレクトリ",
-    )
-    parser.add_argument(
-        "--no_half_vae",
-        action="store_true",
-        help="do not use fp16/bf16 VAE in mixed precision (use float VAE) / mixed precisionでも fp16/bf16 VAEを使わずfloat VAEを使う",
-    )
-    return parser
-
-
-if __name__ == "__main__":
-    # sdxl_original_unet.USE_REENTRANT = False
-
-    parser = setup_parser()
-
-    args = parser.parse_args()
-    args = train_util.read_config_from_file(args, parser)
-
-    train(args)

sdxl_train_control_net_lllite_alt.py

@@ -1,602 +0,0 @@
-# cond_imageをU-Netのforwardで渡すバージョンのControlNet-LLLite検証用学習コード
-# training code for ControlNet-LLLite with passing cond_image to U-Net's forward
-
-import argparse
-import gc
-import json
-import math
-import os
-import random
-import time
-from multiprocessing import Value
-from types import SimpleNamespace
-import toml
-
-from tqdm import tqdm
-import torch
-from torch.nn.parallel import DistributedDataParallel as DDP
-from accelerate.utils import set_seed
-import accelerate
-from diffusers import DDPMScheduler, ControlNetModel
-from safetensors.torch import load_file
-from library import sai_model_spec, sdxl_model_util, sdxl_original_unet, sdxl_train_util
-
-import library.model_util as model_util
-import library.train_util as train_util
-import library.config_util as config_util
-from library.config_util import (
-    ConfigSanitizer,
-    BlueprintGenerator,
-)
-import library.huggingface_util as huggingface_util
-import library.custom_train_functions as custom_train_functions
-from library.custom_train_functions import (
-    add_v_prediction_like_loss,
-    apply_snr_weight,
-    prepare_scheduler_for_custom_training,
-    pyramid_noise_like,
-    apply_noise_offset,
-    scale_v_prediction_loss_like_noise_prediction,
-)
-import networks.control_net_lllite_for_train as control_net_lllite_for_train
-
-
-# TODO 他のスクリプトと共通化する
-def generate_step_logs(args: argparse.Namespace, current_loss, avr_loss, lr_scheduler):
-    logs = {
-        "loss/current": current_loss,
-        "loss/average": avr_loss,
-        "lr": lr_scheduler.get_last_lr()[0],
-    }
-
-    if args.optimizer_type.lower().startswith("DAdapt".lower()):
-        logs["lr/d*lr"] = lr_scheduler.optimizers[-1].param_groups[0]["d"] * lr_scheduler.optimizers[-1].param_groups[0]["lr"]
-
-    return logs
-
-
-def train(args):
-    train_util.verify_training_args(args)
-    train_util.prepare_dataset_args(args, True)
-    sdxl_train_util.verify_sdxl_training_args(args)
-
-    cache_latents = args.cache_latents
-    use_user_config = args.dataset_config is not None
-
-    if args.seed is None:
-        args.seed = random.randint(0, 2**32)
-    set_seed(args.seed)
-
-    tokenizer1, tokenizer2 = sdxl_train_util.load_tokenizers(args)
-
-    # データセットを準備する
-    blueprint_generator = BlueprintGenerator(ConfigSanitizer(False, False, True, True))
-    if use_user_config:
-        print(f"Load dataset config from {args.dataset_config}")
-        user_config = config_util.load_user_config(args.dataset_config)
-        ignored = ["train_data_dir", "conditioning_data_dir"]
-        if any(getattr(args, attr) is not None for attr in ignored):
-            print(
-                "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
-                    ", ".join(ignored)
-                )
-            )
-    else:
-        user_config = {
-            "datasets": [
-                {
-                    "subsets": config_util.generate_controlnet_subsets_config_by_subdirs(
-                        args.train_data_dir,
-                        args.conditioning_data_dir,
-                        args.caption_extension,
-                    )
-                }
-            ]
-        }
-
-    blueprint = blueprint_generator.generate(user_config, args, tokenizer=[tokenizer1, tokenizer2])
-    train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
-
-    current_epoch = Value("i", 0)
-    current_step = Value("i", 0)
-    ds_for_collater = train_dataset_group if args.max_data_loader_n_workers == 0 else None
-    collater = train_util.collater_class(current_epoch, current_step, ds_for_collater)
-
-    train_dataset_group.verify_bucket_reso_steps(32)
-
-    if args.debug_dataset:
-        train_util.debug_dataset(train_dataset_group)
-        return
-    if len(train_dataset_group) == 0:
-        print(
-            "No data found. Please verify arguments (train_data_dir must be the parent of folders with images) / 画像がありません。引数指定を確認してください（train_data_dirには画像があるフォルダではなく、画像があるフォルダの親フォルダを指定する必要があります）"
-        )
-        return
-
-    if cache_latents:
-        assert (
-            train_dataset_group.is_latent_cacheable()
-        ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
-    else:
-        print("WARNING: random_crop is not supported yet for ControlNet training / ControlNetの学習ではrandom_cropはまだサポートされていません")
-
-    if args.cache_text_encoder_outputs:
-        assert (
-            train_dataset_group.is_text_encoder_output_cacheable()
-        ), "when caching Text Encoder output, either caption_dropout_rate, shuffle_caption, token_warmup_step or caption_tag_dropout_rate cannot be used / Text Encoderの出力をキャッシュするときはcaption_dropout_rate, shuffle_caption, token_warmup_step, caption_tag_dropout_rateは使えません"
-
-    # acceleratorを準備する
-    print("prepare accelerator")
-    accelerator = train_util.prepare_accelerator(args)
-    is_main_process = accelerator.is_main_process
-
-    # mixed precisionに対応した型を用意しておき適宜castする
-    weight_dtype, save_dtype = train_util.prepare_dtype(args)
-    vae_dtype = torch.float32 if args.no_half_vae else weight_dtype
-
-    # モデルを読み込む
-    (
-        load_stable_diffusion_format,
-        text_encoder1,
-        text_encoder2,
-        vae,
-        unet,
-        logit_scale,
-        ckpt_info,
-    ) = sdxl_train_util.load_target_model(args, accelerator, sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0, weight_dtype)
-
-    # 学習を準備する
-    if cache_latents:
-        vae.to(accelerator.device, dtype=vae_dtype)
-        vae.requires_grad_(False)
-        vae.eval()
-        with torch.no_grad():
-            train_dataset_group.cache_latents(
-                vae,
-                args.vae_batch_size,
-                args.cache_latents_to_disk,
-                accelerator.is_main_process,
-            )
-        vae.to("cpu")
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        gc.collect()
-
-        accelerator.wait_for_everyone()
-
-    # TextEncoderの出力をキャッシュする
-    if args.cache_text_encoder_outputs:
-        # Text Encodes are eval and no grad
-        with torch.no_grad():
-            train_dataset_group.cache_text_encoder_outputs(
-                (tokenizer1, tokenizer2),
-                (text_encoder1, text_encoder2),
-                accelerator.device,
-                None,
-                args.cache_text_encoder_outputs_to_disk,
-                accelerator.is_main_process,
-            )
-        accelerator.wait_for_everyone()
-
-    # prepare ControlNet-LLLite
-    control_net_lllite_for_train.replace_unet_linear_and_conv2d()
-
-    if args.network_weights is not None:
-        accelerator.print(f"initialize U-Net with ControlNet-LLLite")
-        with accelerate.init_empty_weights():
-            unet_lllite = control_net_lllite_for_train.SdxlUNet2DConditionModelControlNetLLLite()
-        unet_lllite.to(accelerator.device, dtype=weight_dtype)
-
-        unet_sd = unet.state_dict()
-        info = unet_lllite.load_lllite_weights(args.network_weights, unet_sd)
-        accelerator.print(f"load ControlNet-LLLite weights from {args.network_weights}: {info}")
-    else:
-        # cosumes large memory, so send to GPU before creating the LLLite model
-        accelerator.print("sending U-Net to GPU")
-        unet.to(accelerator.device, dtype=weight_dtype)
-        unet_sd = unet.state_dict()
-
-        # init LLLite weights
-        accelerator.print(f"initialize U-Net with ControlNet-LLLite")
-
-        if args.lowram:
-            with accelerate.init_on_device(accelerator.device):
-                unet_lllite = control_net_lllite_for_train.SdxlUNet2DConditionModelControlNetLLLite()
-        else:
-            unet_lllite = control_net_lllite_for_train.SdxlUNet2DConditionModelControlNetLLLite()
-        unet_lllite.to(weight_dtype)
-
-        info = unet_lllite.load_lllite_weights(None, unet_sd)
-        accelerator.print(f"init U-Net with ControlNet-LLLite weights: {info}")
-    del unet_sd, unet
-
-    unet: control_net_lllite_for_train.SdxlUNet2DConditionModelControlNetLLLite = unet_lllite
-    del unet_lllite
-
-    unet.apply_lllite(args.cond_emb_dim, args.network_dim, args.network_dropout)
-
-    # モデルに xformers とか memory efficient attention を組み込む
-    train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers, args.sdpa)
-
-    if args.gradient_checkpointing:
-        unet.enable_gradient_checkpointing()
-
-    # 学習に必要なクラスを準備する
-    accelerator.print("prepare optimizer, data loader etc.")
-
-    trainable_params = list(unet.prepare_params())
-    print(f"trainable params count: {len(trainable_params)}")
-    print(f"number of trainable parameters: {sum(p.numel() for p in trainable_params if p.requires_grad)}")
-
-    _, _, optimizer = train_util.get_optimizer(args, trainable_params)
-
-    # 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_group,
-        batch_size=1,
-        shuffle=True,
-        collate_fn=collater,
-        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 * math.ceil(
-            len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps
-        )
-        accelerator.print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
-
-    # データセット側にも学習ステップを送信
-    train_dataset_group.set_max_train_steps(args.max_train_steps)
-
-    # lr schedulerを用意する
-    lr_scheduler = train_util.get_scheduler_fix(args, optimizer, accelerator.num_processes)
-
-    # 実験的機能：勾配も含めたfp16/bf16学習を行う　モデル全体をfp16/bf16にする
-    # if args.full_fp16:
-    #     assert (
-    #         args.mixed_precision == "fp16"
-    #     ), "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
-    #     accelerator.print("enable full fp16 training.")
-    #     unet.to(weight_dtype)
-    # elif args.full_bf16:
-    #     assert (
-    #         args.mixed_precision == "bf16"
-    #     ), "full_bf16 requires mixed precision='bf16' / full_bf16を使う場合はmixed_precision='bf16'を指定してください。"
-    #     accelerator.print("enable full bf16 training.")
-    #     unet.to(weight_dtype)
-
-    unet.to(weight_dtype)
-
-    # acceleratorがなんかよろしくやってくれるらしい
-    unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(unet, optimizer, train_dataloader, lr_scheduler)
-
-    # transform DDP after prepare (train_network here only)
-    unet = train_util.transform_models_if_DDP([unet])[0]
-
-    if args.gradient_checkpointing:
-        unet.train()  # according to TI example in Diffusers, train is required -> これオリジナルのU-Netしたので本当は外せる
-    else:
-        unet.eval()
-
-    # TextEncoderの出力をキャッシュするときにはCPUへ移動する
-    if args.cache_text_encoder_outputs:
-        # move Text Encoders for sampling images. Text Encoder doesn't work on CPU with fp16
-        text_encoder1.to("cpu", dtype=torch.float32)
-        text_encoder2.to("cpu", dtype=torch.float32)
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-    else:
-        # make sure Text Encoders are on GPU
-        text_encoder1.to(accelerator.device)
-        text_encoder2.to(accelerator.device)
-
-    if not cache_latents:
-        vae.requires_grad_(False)
-        vae.eval()
-        vae.to(accelerator.device, dtype=vae_dtype)
-
-    # 実験的機能：勾配も含めたfp16学習を行う　PyTorchにパッチを当ててfp16でのgrad scaleを有効にする
-    if args.full_fp16:
-        train_util.patch_accelerator_for_fp16_training(accelerator)
-
-    # resumeする
-    train_util.resume_from_local_or_hf_if_specified(accelerator, args)
-
-    # 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
-
-    # 学習する
-    # TODO: find a way to handle total batch size when there are multiple datasets
-    accelerator.print("running training / 学習開始")
-    accelerator.print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
-    accelerator.print(f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
-    accelerator.print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
-    accelerator.print(f"  num epochs / epoch数: {num_train_epochs}")
-    accelerator.print(f"  batch size per device / バッチサイズ: {', '.join([str(d.batch_size) for d in train_dataset_group.datasets])}")
-    # print(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
-    accelerator.print(f"  gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
-    accelerator.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
-    )
-    prepare_scheduler_for_custom_training(noise_scheduler, accelerator.device)
-    if args.zero_terminal_snr:
-        custom_train_functions.fix_noise_scheduler_betas_for_zero_terminal_snr(noise_scheduler)
-
-    if accelerator.is_main_process:
-        init_kwargs = {}
-        if args.log_tracker_config is not None:
-            init_kwargs = toml.load(args.log_tracker_config)
-        accelerator.init_trackers(
-            "lllite_control_net_train" if args.log_tracker_name is None else args.log_tracker_name, init_kwargs=init_kwargs
-        )
-
-    loss_list = []
-    loss_total = 0.0
-    del train_dataset_group
-
-    # function for saving/removing
-    def save_model(
-        ckpt_name,
-        unwrapped_nw: control_net_lllite_for_train.SdxlUNet2DConditionModelControlNetLLLite,
-        steps,
-        epoch_no,
-        force_sync_upload=False,
-    ):
-        os.makedirs(args.output_dir, exist_ok=True)
-        ckpt_file = os.path.join(args.output_dir, ckpt_name)
-
-        accelerator.print(f"\nsaving checkpoint: {ckpt_file}")
-        sai_metadata = train_util.get_sai_model_spec(None, args, True, True, False)
-        sai_metadata["modelspec.architecture"] = sai_model_spec.ARCH_SD_XL_V1_BASE + "/control-net-lllite"
-
-        unwrapped_nw.save_lllite_weights(ckpt_file, save_dtype, sai_metadata)
-        if args.huggingface_repo_id is not None:
-            huggingface_util.upload(args, ckpt_file, "/" + ckpt_name, force_sync_upload=force_sync_upload)
-
-    def remove_model(old_ckpt_name):
-        old_ckpt_file = os.path.join(args.output_dir, old_ckpt_name)
-        if os.path.exists(old_ckpt_file):
-            accelerator.print(f"removing old checkpoint: {old_ckpt_file}")
-            os.remove(old_ckpt_file)
-
-    # training loop
-    for epoch in range(num_train_epochs):
-        accelerator.print(f"\nepoch {epoch+1}/{num_train_epochs}")
-        current_epoch.value = epoch + 1
-
-        for step, batch in enumerate(train_dataloader):
-            current_step.value = global_step
-            with accelerator.accumulate(unet):
-                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=vae_dtype)).latent_dist.sample()
-
-                        # NaNが含まれていれば警告を表示し0に置き換える
-                        if torch.any(torch.isnan(latents)):
-                            accelerator.print("NaN found in latents, replacing with zeros")
-                            latents = torch.where(torch.isnan(latents), torch.zeros_like(latents), latents)
-                    latents = latents * sdxl_model_util.VAE_SCALE_FACTOR
-
-                if "text_encoder_outputs1_list" not in batch or batch["text_encoder_outputs1_list"] is None:
-                    input_ids1 = batch["input_ids"]
-                    input_ids2 = batch["input_ids2"]
-                    with torch.no_grad():
-                        # Get the text embedding for conditioning
-                        input_ids1 = input_ids1.to(accelerator.device)
-                        input_ids2 = input_ids2.to(accelerator.device)
-                        encoder_hidden_states1, encoder_hidden_states2, pool2 = train_util.get_hidden_states_sdxl(
-                            args.max_token_length,
-                            input_ids1,
-                            input_ids2,
-                            tokenizer1,
-                            tokenizer2,
-                            text_encoder1,
-                            text_encoder2,
-                            None if not args.full_fp16 else weight_dtype,
-                        )
-                else:
-                    encoder_hidden_states1 = batch["text_encoder_outputs1_list"].to(accelerator.device).to(weight_dtype)
-                    encoder_hidden_states2 = batch["text_encoder_outputs2_list"].to(accelerator.device).to(weight_dtype)
-                    pool2 = batch["text_encoder_pool2_list"].to(accelerator.device).to(weight_dtype)
-
-                # get size embeddings
-                orig_size = batch["original_sizes_hw"]
-                crop_size = batch["crop_top_lefts"]
-                target_size = batch["target_sizes_hw"]
-                embs = sdxl_train_util.get_size_embeddings(orig_size, crop_size, target_size, accelerator.device).to(weight_dtype)
-
-                # concat embeddings
-                vector_embedding = torch.cat([pool2, embs], dim=1).to(weight_dtype)
-                text_embedding = torch.cat([encoder_hidden_states1, encoder_hidden_states2], dim=2).to(weight_dtype)
-
-                # Sample noise, sample a random timestep for each image, and add noise to the latents,
-                # with noise offset and/or multires noise if specified
-                noise, noisy_latents, timesteps = train_util.get_noise_noisy_latents_and_timesteps(args, noise_scheduler, latents)
-
-                noisy_latents = noisy_latents.to(weight_dtype)  # TODO check why noisy_latents is not weight_dtype
-
-                controlnet_image = batch["conditioning_images"].to(dtype=weight_dtype)
-
-                with accelerator.autocast():
-                    # conditioning imageをControlNetに渡す / pass conditioning image to ControlNet
-                    # 内部でcond_embに変換される / it will be converted to cond_emb inside
-
-                    # それらの値を使いつつ、U-Netでノイズを予測する / predict noise with U-Net using those values
-                    noise_pred = unet(noisy_latents, timesteps, text_embedding, vector_embedding, controlnet_image)
-
-                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="none")
-                loss = loss.mean([1, 2, 3])
-
-                loss_weights = batch["loss_weights"]  # 各sampleごとのweight
-                loss = loss * loss_weights
-
-                if args.min_snr_gamma:
-                    loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma)
-                if args.scale_v_pred_loss_like_noise_pred:
-                    loss = scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler)
-                if args.v_pred_like_loss:
-                    loss = add_v_prediction_like_loss(loss, timesteps, noise_scheduler, args.v_pred_like_loss)
-
-                loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし
-
-                accelerator.backward(loss)
-                if accelerator.sync_gradients and args.max_grad_norm != 0.0:
-                    params_to_clip = unet.get_trainable_params()
-                    accelerator.clip_grad_norm_(params_to_clip, 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
-
-                # sdxl_train_util.sample_images(accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
-
-                # 指定ステップごとにモデルを保存
-                if args.save_every_n_steps is not None and global_step % args.save_every_n_steps == 0:
-                    accelerator.wait_for_everyone()
-                    if accelerator.is_main_process:
-                        ckpt_name = train_util.get_step_ckpt_name(args, "." + args.save_model_as, global_step)
-                        save_model(ckpt_name, accelerator.unwrap_model(unet), global_step, epoch)
-
-                        if args.save_state:
-                            train_util.save_and_remove_state_stepwise(args, accelerator, global_step)
-
-                        remove_step_no = train_util.get_remove_step_no(args, global_step)
-                        if remove_step_no is not None:
-                            remove_ckpt_name = train_util.get_step_ckpt_name(args, "." + args.save_model_as, remove_step_no)
-                            remove_model(remove_ckpt_name)
-
-            current_loss = loss.detach().item()
-            if epoch == 0:
-                loss_list.append(current_loss)
-            else:
-                loss_total -= loss_list[step]
-                loss_list[step] = current_loss
-            loss_total += current_loss
-            avr_loss = loss_total / len(loss_list)
-            logs = {"loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
-            progress_bar.set_postfix(**logs)
-
-            if args.logging_dir is not None:
-                logs = generate_step_logs(args, current_loss, avr_loss, lr_scheduler)
-                accelerator.log(logs, step=global_step)
-
-            if global_step >= args.max_train_steps:
-                break
-
-        if args.logging_dir is not None:
-            logs = {"loss/epoch": loss_total / len(loss_list)}
-            accelerator.log(logs, step=epoch + 1)
-
-        accelerator.wait_for_everyone()
-
-        # 指定エポックごとにモデルを保存
-        if args.save_every_n_epochs is not None:
-            saving = (epoch + 1) % args.save_every_n_epochs == 0 and (epoch + 1) < num_train_epochs
-            if is_main_process and saving:
-                ckpt_name = train_util.get_epoch_ckpt_name(args, "." + args.save_model_as, epoch + 1)
-                save_model(ckpt_name, accelerator.unwrap_model(unet), global_step, epoch + 1)
-
-                remove_epoch_no = train_util.get_remove_epoch_no(args, epoch + 1)
-                if remove_epoch_no is not None:
-                    remove_ckpt_name = train_util.get_epoch_ckpt_name(args, "." + args.save_model_as, remove_epoch_no)
-                    remove_model(remove_ckpt_name)
-
-                if args.save_state:
-                    train_util.save_and_remove_state_on_epoch_end(args, accelerator, epoch + 1)
-
-        # self.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
-
-        # end of epoch
-
-    if is_main_process:
-        unet = accelerator.unwrap_model(unet)
-
-    accelerator.end_training()
-
-    if is_main_process and args.save_state:
-        train_util.save_state_on_train_end(args, accelerator)
-
-    if is_main_process:
-        ckpt_name = train_util.get_last_ckpt_name(args, "." + args.save_model_as)
-        save_model(ckpt_name, unet, global_step, num_train_epochs, force_sync_upload=True)
-
-        print("model saved.")
-
-
-def setup_parser() -> argparse.ArgumentParser:
-    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_optimizer_arguments(parser)
-    config_util.add_config_arguments(parser)
-    custom_train_functions.add_custom_train_arguments(parser)
-    sdxl_train_util.add_sdxl_training_arguments(parser)
-
-    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("--cond_emb_dim", type=int, default=None, help="conditioning embedding dimension / 条件付け埋め込みの次元数")
-    parser.add_argument("--network_weights", type=str, default=None, help="pretrained weights for network / 学習するネットワークの初期重み")
-    parser.add_argument("--network_dim", type=int, default=None, help="network dimensions (rank) / モジュールの次元数")
-    parser.add_argument(
-        "--network_dropout",
-        type=float,
-        default=None,
-        help="Drops neurons out of training every step (0 or None is default behavior (no dropout), 1 would drop all neurons) / 訓練時に毎ステップでニューロンをdropする（0またはNoneはdropoutなし、1は全ニューロンをdropout）",
-    )
-    parser.add_argument(
-        "--conditioning_data_dir",
-        type=str,
-        default=None,
-        help="conditioning data directory / 条件付けデータのディレクトリ",
-    )
-    parser.add_argument(
-        "--no_half_vae",
-        action="store_true",
-        help="do not use fp16/bf16 VAE in mixed precision (use float VAE) / mixed precisionでも fp16/bf16 VAEを使わずfloat VAEを使う",
-    )
-    return parser
-
-
-if __name__ == "__main__":
-    # sdxl_original_unet.USE_REENTRANT = False
-
-    parser = setup_parser()
-
-    args = parser.parse_args()
-    args = train_util.read_config_from_file(args, parser)
-
-    train(args)

sdxl_train_network.py

@@ -1,176 +0,0 @@
-import argparse
-import torch
-from library import sdxl_model_util, sdxl_train_util, train_util
-import train_network
-
-
-class SdxlNetworkTrainer(train_network.NetworkTrainer):
-    def __init__(self):
-        super().__init__()
-        self.vae_scale_factor = sdxl_model_util.VAE_SCALE_FACTOR
-        self.is_sdxl = True
-
-    def assert_extra_args(self, args, train_dataset_group):
-        super().assert_extra_args(args, train_dataset_group)
-        sdxl_train_util.verify_sdxl_training_args(args)
-
-        if args.cache_text_encoder_outputs:
-            assert (
-                train_dataset_group.is_text_encoder_output_cacheable()
-            ), "when caching Text Encoder output, either caption_dropout_rate, shuffle_caption, token_warmup_step or caption_tag_dropout_rate cannot be used / Text Encoderの出力をキャッシュするときはcaption_dropout_rate, shuffle_caption, token_warmup_step, caption_tag_dropout_rateは使えません"
-
-        assert (
-            args.network_train_unet_only or not args.cache_text_encoder_outputs
-        ), "network for Text Encoder cannot be trained with caching Text Encoder outputs / Text Encoderの出力をキャッシュしながらText Encoderのネットワークを学習することはできません"
-
-        train_dataset_group.verify_bucket_reso_steps(32)
-
-    def load_target_model(self, args, weight_dtype, accelerator):
-        (
-            load_stable_diffusion_format,
-            text_encoder1,
-            text_encoder2,
-            vae,
-            unet,
-            logit_scale,
-            ckpt_info,
-        ) = sdxl_train_util.load_target_model(args, accelerator, sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0, weight_dtype)
-
-        self.load_stable_diffusion_format = load_stable_diffusion_format
-        self.logit_scale = logit_scale
-        self.ckpt_info = ckpt_info
-
-        return sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0, [text_encoder1, text_encoder2], vae, unet
-
-    def load_tokenizer(self, args):
-        tokenizer = sdxl_train_util.load_tokenizers(args)
-        return tokenizer
-
-    def is_text_encoder_outputs_cached(self, args):
-        return args.cache_text_encoder_outputs
-
-    def cache_text_encoder_outputs_if_needed(
-        self, args, accelerator, unet, vae, tokenizers, text_encoders, dataset: train_util.DatasetGroup, weight_dtype
-    ):
-        if args.cache_text_encoder_outputs:
-            if not args.lowram:
-                # メモリ消費を減らす
-                print("move vae and unet to cpu to save memory")
-                org_vae_device = vae.device
-                org_unet_device = unet.device
-                vae.to("cpu")
-                unet.to("cpu")
-                if torch.cuda.is_available():
-                    torch.cuda.empty_cache()
-
-            dataset.cache_text_encoder_outputs(
-                tokenizers,
-                text_encoders,
-                accelerator.device,
-                weight_dtype,
-                args.cache_text_encoder_outputs_to_disk,
-                accelerator.is_main_process,
-            )
-
-            text_encoders[0].to("cpu", dtype=torch.float32)  # Text Encoder doesn't work with fp16 on CPU
-            text_encoders[1].to("cpu", dtype=torch.float32)
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-
-            if not args.lowram:
-                print("move vae and unet back to original device")
-                vae.to(org_vae_device)
-                unet.to(org_unet_device)
-        else:
-            # Text Encoderから毎回出力を取得するので、GPUに乗せておく
-            text_encoders[0].to(accelerator.device)
-            text_encoders[1].to(accelerator.device)
-
-    def get_text_cond(self, args, accelerator, batch, tokenizers, text_encoders, weight_dtype):
-        if "text_encoder_outputs1_list" not in batch or batch["text_encoder_outputs1_list"] is None:
-            input_ids1 = batch["input_ids"]
-            input_ids2 = batch["input_ids2"]
-            with torch.enable_grad():
-                # Get the text embedding for conditioning
-                # TODO support weighted captions
-                # if args.weighted_captions:
-                #     encoder_hidden_states = get_weighted_text_embeddings(
-                #         tokenizer,
-                #         text_encoder,
-                #         batch["captions"],
-                #         accelerator.device,
-                #         args.max_token_length // 75 if args.max_token_length else 1,
-                #         clip_skip=args.clip_skip,
-                #     )
-                # else:
-                input_ids1 = input_ids1.to(accelerator.device)
-                input_ids2 = input_ids2.to(accelerator.device)
-                encoder_hidden_states1, encoder_hidden_states2, pool2 = train_util.get_hidden_states_sdxl(
-                    args.max_token_length,
-                    input_ids1,
-                    input_ids2,
-                    tokenizers[0],
-                    tokenizers[1],
-                    text_encoders[0],
-                    text_encoders[1],
-                    None if not args.full_fp16 else weight_dtype,
-                )
-        else:
-            encoder_hidden_states1 = batch["text_encoder_outputs1_list"].to(accelerator.device).to(weight_dtype)
-            encoder_hidden_states2 = batch["text_encoder_outputs2_list"].to(accelerator.device).to(weight_dtype)
-            pool2 = batch["text_encoder_pool2_list"].to(accelerator.device).to(weight_dtype)
-
-            # # verify that the text encoder outputs are correct
-            # ehs1, ehs2, p2 = train_util.get_hidden_states_sdxl(
-            #     args.max_token_length,
-            #     batch["input_ids"].to(text_encoders[0].device),
-            #     batch["input_ids2"].to(text_encoders[0].device),
-            #     tokenizers[0],
-            #     tokenizers[1],
-            #     text_encoders[0],
-            #     text_encoders[1],
-            #     None if not args.full_fp16 else weight_dtype,
-            # )
-            # b_size = encoder_hidden_states1.shape[0]
-            # assert ((encoder_hidden_states1.to("cpu") - ehs1.to(dtype=weight_dtype)).abs().max() > 1e-2).sum() <= b_size * 2
-            # assert ((encoder_hidden_states2.to("cpu") - ehs2.to(dtype=weight_dtype)).abs().max() > 1e-2).sum() <= b_size * 2
-            # assert ((pool2.to("cpu") - p2.to(dtype=weight_dtype)).abs().max() > 1e-2).sum() <= b_size * 2
-            # print("text encoder outputs verified")
-
-        return encoder_hidden_states1, encoder_hidden_states2, pool2
-
-    def call_unet(self, args, accelerator, unet, noisy_latents, timesteps, text_conds, batch, weight_dtype):
-        noisy_latents = noisy_latents.to(weight_dtype)  # TODO check why noisy_latents is not weight_dtype
-
-        # get size embeddings
-        orig_size = batch["original_sizes_hw"]
-        crop_size = batch["crop_top_lefts"]
-        target_size = batch["target_sizes_hw"]
-        embs = sdxl_train_util.get_size_embeddings(orig_size, crop_size, target_size, accelerator.device).to(weight_dtype)
-
-        # concat embeddings
-        encoder_hidden_states1, encoder_hidden_states2, pool2 = text_conds
-        vector_embedding = torch.cat([pool2, embs], dim=1).to(weight_dtype)
-        text_embedding = torch.cat([encoder_hidden_states1, encoder_hidden_states2], dim=2).to(weight_dtype)
-
-        noise_pred = unet(noisy_latents, timesteps, text_embedding, vector_embedding)
-        return noise_pred
-
-    def sample_images(self, accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet):
-        sdxl_train_util.sample_images(accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet)
-
-
-def setup_parser() -> argparse.ArgumentParser:
-    parser = train_network.setup_parser()
-    sdxl_train_util.add_sdxl_training_arguments(parser)
-    return parser
-
-
-if __name__ == "__main__":
-    parser = setup_parser()
-
-    args = parser.parse_args()
-    args = train_util.read_config_from_file(args, parser)
-
-    trainer = SdxlNetworkTrainer()
-    trainer.train(args)

sdxl_train_textual_inversion.py

@@ -1,133 +0,0 @@
-import argparse
-import os
-
-import regex
-import torch
-import open_clip
-from library import sdxl_model_util, sdxl_train_util, train_util
-
-import train_textual_inversion
-
-
-class SdxlTextualInversionTrainer(train_textual_inversion.TextualInversionTrainer):
-    def __init__(self):
-        super().__init__()
-        self.vae_scale_factor = sdxl_model_util.VAE_SCALE_FACTOR
-        self.is_sdxl = True
-
-    def assert_extra_args(self, args, train_dataset_group):
-        super().assert_extra_args(args, train_dataset_group)
-        sdxl_train_util.verify_sdxl_training_args(args, supportTextEncoderCaching=False)
-
-        train_dataset_group.verify_bucket_reso_steps(32)
-
-    def load_target_model(self, args, weight_dtype, accelerator):
-        (
-            load_stable_diffusion_format,
-            text_encoder1,
-            text_encoder2,
-            vae,
-            unet,
-            logit_scale,
-            ckpt_info,
-        ) = sdxl_train_util.load_target_model(args, accelerator, sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0, weight_dtype)
-
-        self.load_stable_diffusion_format = load_stable_diffusion_format
-        self.logit_scale = logit_scale
-        self.ckpt_info = ckpt_info
-
-        return sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0, [text_encoder1, text_encoder2], vae, unet
-
-    def load_tokenizer(self, args):
-        tokenizer = sdxl_train_util.load_tokenizers(args)
-        return tokenizer
-
-    def get_text_cond(self, args, accelerator, batch, tokenizers, text_encoders, weight_dtype):
-        input_ids1 = batch["input_ids"]
-        input_ids2 = batch["input_ids2"]
-        with torch.enable_grad():
-            input_ids1 = input_ids1.to(accelerator.device)
-            input_ids2 = input_ids2.to(accelerator.device)
-            encoder_hidden_states1, encoder_hidden_states2, pool2 = train_util.get_hidden_states_sdxl(
-                args.max_token_length,
-                input_ids1,
-                input_ids2,
-                tokenizers[0],
-                tokenizers[1],
-                text_encoders[0],
-                text_encoders[1],
-                None if not args.full_fp16 else weight_dtype,
-            )
-        return encoder_hidden_states1, encoder_hidden_states2, pool2
-
-    def call_unet(self, args, accelerator, unet, noisy_latents, timesteps, text_conds, batch, weight_dtype):
-        noisy_latents = noisy_latents.to(weight_dtype)  # TODO check why noisy_latents is not weight_dtype
-
-        # get size embeddings
-        orig_size = batch["original_sizes_hw"]
-        crop_size = batch["crop_top_lefts"]
-        target_size = batch["target_sizes_hw"]
-        embs = sdxl_train_util.get_size_embeddings(orig_size, crop_size, target_size, accelerator.device).to(weight_dtype)
-
-        # concat embeddings
-        encoder_hidden_states1, encoder_hidden_states2, pool2 = text_conds
-        vector_embedding = torch.cat([pool2, embs], dim=1).to(weight_dtype)
-        text_embedding = torch.cat([encoder_hidden_states1, encoder_hidden_states2], dim=2).to(weight_dtype)
-
-        noise_pred = unet(noisy_latents, timesteps, text_embedding, vector_embedding)
-        return noise_pred
-
-    def sample_images(self, accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet, prompt_replacement):
-        sdxl_train_util.sample_images(
-            accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet, prompt_replacement
-        )
-
-    def save_weights(self, file, updated_embs, save_dtype, metadata):
-        state_dict = {"clip_l": updated_embs[0], "clip_g": updated_embs[1]}
-
-        if save_dtype is not None:
-            for key in list(state_dict.keys()):
-                v = state_dict[key]
-                v = v.detach().clone().to("cpu").to(save_dtype)
-                state_dict[key] = v
-
-        if os.path.splitext(file)[1] == ".safetensors":
-            from safetensors.torch import save_file
-
-            save_file(state_dict, file, metadata)
-        else:
-            torch.save(state_dict, file)
-
-    def load_weights(self, file):
-        if os.path.splitext(file)[1] == ".safetensors":
-            from safetensors.torch import load_file
-
-            data = load_file(file)
-        else:
-            data = torch.load(file, map_location="cpu")
-
-        emb_l = data.get("clip_l", None)  # ViT-L text encoder 1
-        emb_g = data.get("clip_g", None)  # BiG-G text encoder 2
-
-        assert (
-            emb_l is not None or emb_g is not None
-        ), f"weight file does not contains weights for text encoder 1 or 2 / 重みファイルにテキストエンコーダー1または2の重みが含まれていません: {file}"
-
-        return [emb_l, emb_g]
-
-
-def setup_parser() -> argparse.ArgumentParser:
-    parser = train_textual_inversion.setup_parser()
-    # don't add sdxl_train_util.add_sdxl_training_arguments(parser): because it only adds text encoder caching
-    # sdxl_train_util.add_sdxl_training_arguments(parser)
-    return parser
-
-
-if __name__ == "__main__":
-    parser = setup_parser()
-
-    args = parser.parse_args()
-    args = train_util.read_config_from_file(args, parser)
-
-    trainer = SdxlTextualInversionTrainer()
-    trainer.train(args)

tools/cache_latents.py

@@ -1,194 +0,0 @@
-# latentsのdiskへの事前キャッシュを行う / cache latents to disk
-
-import argparse
-import math
-from multiprocessing import Value
-import os
-
-from accelerate.utils import set_seed
-import torch
-from tqdm import tqdm
-
-from library import config_util
-from library import train_util
-from library import sdxl_train_util
-from library.config_util import (
-    ConfigSanitizer,
-    BlueprintGenerator,
-)
-
-
-def cache_to_disk(args: argparse.Namespace) -> None:
-    train_util.prepare_dataset_args(args, True)
-
-    # check cache latents arg
-    assert args.cache_latents_to_disk, "cache_latents_to_disk must be True / cache_latents_to_diskはTrueである必要があります"
-
-    use_dreambooth_method = args.in_json is None
-
-    if args.seed is not None:
-        set_seed(args.seed)  # 乱数系列を初期化する
-
-    # tokenizerを準備する：datasetを動かすために必要
-    if args.sdxl:
-        tokenizer1, tokenizer2 = sdxl_train_util.load_tokenizers(args)
-        tokenizers = [tokenizer1, tokenizer2]
-    else:
-        tokenizer = train_util.load_tokenizer(args)
-        tokenizers = [tokenizer]
-
-    # データセットを準備する
-    if args.dataset_class is None:
-        blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, False, True))
-        if args.dataset_config is not None:
-            print(f"Load dataset config from {args.dataset_config}")
-            user_config = config_util.load_user_config(args.dataset_config)
-            ignored = ["train_data_dir", "in_json"]
-            if any(getattr(args, attr) is not None for attr in ignored):
-                print(
-                    "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
-                        ", ".join(ignored)
-                    )
-                )
-        else:
-            if use_dreambooth_method:
-                print("Using DreamBooth method.")
-                user_config = {
-                    "datasets": [
-                        {
-                            "subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(
-                                args.train_data_dir, args.reg_data_dir
-                            )
-                        }
-                    ]
-                }
-            else:
-                print("Training with captions.")
-                user_config = {
-                    "datasets": [
-                        {
-                            "subsets": [
-                                {
-                                    "image_dir": args.train_data_dir,
-                                    "metadata_file": args.in_json,
-                                }
-                            ]
-                        }
-                    ]
-                }
-
-        blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizers)
-        train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
-    else:
-        train_dataset_group = train_util.load_arbitrary_dataset(args, tokenizers)
-
-    # datasetのcache_latentsを呼ばなければ、生の画像が返る
-
-    current_epoch = Value("i", 0)
-    current_step = Value("i", 0)
-    ds_for_collater = train_dataset_group if args.max_data_loader_n_workers == 0 else None
-    collater = train_util.collater_class(current_epoch, current_step, ds_for_collater)
-
-    # acceleratorを準備する
-    print("prepare accelerator")
-    accelerator = train_util.prepare_accelerator(args)
-
-    # mixed precisionに対応した型を用意しておき適宜castする
-    weight_dtype, _ = train_util.prepare_dtype(args)
-    vae_dtype = torch.float32 if args.no_half_vae else weight_dtype
-
-    # モデルを読み込む
-    print("load model")
-    if args.sdxl:
-        (_, _, _, vae, _, _, _) = sdxl_train_util.load_target_model(args, accelerator, "sdxl", weight_dtype)
-    else:
-        _, vae, _, _ = train_util.load_target_model(args, weight_dtype, accelerator)
-
-    if torch.__version__ >= "2.0.0": # PyTorch 2.0.0 以上対応のxformersなら以下が使える
-        vae.set_use_memory_efficient_attention_xformers(args.xformers)
-    vae.to(accelerator.device, dtype=vae_dtype)
-    vae.requires_grad_(False)
-    vae.eval()
-
-    # dataloaderを準備する
-    train_dataset_group.set_caching_mode("latents")
-
-    # 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_group,
-        batch_size=1,
-        shuffle=True,
-        collate_fn=collater,
-        num_workers=n_workers,
-        persistent_workers=args.persistent_data_loader_workers,
-    )
-
-    # acceleratorを使ってモデルを準備する：マルチGPUで使えるようになるはず
-    train_dataloader = accelerator.prepare(train_dataloader)
-
-    # データ取得のためのループ
-    for batch in tqdm(train_dataloader):
-        b_size = len(batch["images"])
-        vae_batch_size = b_size if args.vae_batch_size is None else args.vae_batch_size
-        flip_aug = batch["flip_aug"]
-        random_crop = batch["random_crop"]
-        bucket_reso = batch["bucket_reso"]
-
-        # バッチを分割して処理する
-        for i in range(0, b_size, vae_batch_size):
-            images = batch["images"][i : i + vae_batch_size]
-            absolute_paths = batch["absolute_paths"][i : i + vae_batch_size]
-            resized_sizes = batch["resized_sizes"][i : i + vae_batch_size]
-
-            image_infos = []
-            for i, (image, absolute_path, resized_size) in enumerate(zip(images, absolute_paths, resized_sizes)):
-                image_info = train_util.ImageInfo(absolute_path, 1, "dummy", False, absolute_path)
-                image_info.image = image
-                image_info.bucket_reso = bucket_reso
-                image_info.resized_size = resized_size
-                image_info.latents_npz = os.path.splitext(absolute_path)[0] + ".npz"
-
-                if args.skip_existing:
-                    if train_util.is_disk_cached_latents_is_expected(image_info.bucket_reso, image_info.latents_npz, flip_aug):
-                        print(f"Skipping {image_info.latents_npz} because it already exists.")
-                        continue
-
-                image_infos.append(image_info)
-
-            if len(image_infos) > 0:
-                train_util.cache_batch_latents(vae, True, image_infos, flip_aug, random_crop)
-
-    accelerator.wait_for_everyone()
-    accelerator.print(f"Finished caching latents for {len(train_dataset_group)} batches.")
-
-
-def setup_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser()
-
-    train_util.add_sd_models_arguments(parser)
-    train_util.add_training_arguments(parser, True)
-    train_util.add_dataset_arguments(parser, True, True, True)
-    config_util.add_config_arguments(parser)
-    parser.add_argument("--sdxl", action="store_true", help="Use SDXL model / SDXLモデルを使用する")
-    parser.add_argument(
-        "--no_half_vae",
-        action="store_true",
-        help="do not use fp16/bf16 VAE in mixed precision (use float VAE) / mixed precisionでも fp16/bf16 VAEを使わずfloat VAEを使う",
-    )
-    parser.add_argument(
-        "--skip_existing",
-        action="store_true",
-        help="skip images if npz already exists (both normal and flipped exists if flip_aug is enabled) / npzが既に存在する画像をスキップする（flip_aug有効時は通常、反転の両方が存在する画像をスキップ）",
-    )
-    return parser
-
-
-if __name__ == "__main__":
-    parser = setup_parser()
-
-    args = parser.parse_args()
-    args = train_util.read_config_from_file(args, parser)
-
-    cache_to_disk(args)

tools/cache_text_encoder_outputs.py

@@ -1,191 +0,0 @@
-# text encoder出力のdiskへの事前キャッシュを行う / cache text encoder outputs to disk in advance
-
-import argparse
-import math
-from multiprocessing import Value
-import os
-
-from accelerate.utils import set_seed
-import torch
-from tqdm import tqdm
-
-from library import config_util
-from library import train_util
-from library import sdxl_train_util
-from library.config_util import (
-    ConfigSanitizer,
-    BlueprintGenerator,
-)
-
-
-def cache_to_disk(args: argparse.Namespace) -> None:
-    train_util.prepare_dataset_args(args, True)
-
-    # check cache arg
-    assert (
-        args.cache_text_encoder_outputs_to_disk
-    ), "cache_text_encoder_outputs_to_disk must be True / cache_text_encoder_outputs_to_diskはTrueである必要があります"
-
-    # できるだけ準備はしておくが今のところSDXLのみしか動かない
-    assert (
-        args.sdxl
-    ), "cache_text_encoder_outputs_to_disk is only available for SDXL / cache_text_encoder_outputs_to_diskはSDXLのみ利用可能です"
-
-    use_dreambooth_method = args.in_json is None
-
-    if args.seed is not None:
-        set_seed(args.seed)  # 乱数系列を初期化する
-
-    # tokenizerを準備する：datasetを動かすために必要
-    if args.sdxl:
-        tokenizer1, tokenizer2 = sdxl_train_util.load_tokenizers(args)
-        tokenizers = [tokenizer1, tokenizer2]
-    else:
-        tokenizer = train_util.load_tokenizer(args)
-        tokenizers = [tokenizer]
-
-    # データセットを準備する
-    if args.dataset_class is None:
-        blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, False, True))
-        if args.dataset_config is not None:
-            print(f"Load dataset config from {args.dataset_config}")
-            user_config = config_util.load_user_config(args.dataset_config)
-            ignored = ["train_data_dir", "in_json"]
-            if any(getattr(args, attr) is not None for attr in ignored):
-                print(
-                    "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
-                        ", ".join(ignored)
-                    )
-                )
-        else:
-            if use_dreambooth_method:
-                print("Using DreamBooth method.")
-                user_config = {
-                    "datasets": [
-                        {
-                            "subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(
-                                args.train_data_dir, args.reg_data_dir
-                            )
-                        }
-                    ]
-                }
-            else:
-                print("Training with captions.")
-                user_config = {
-                    "datasets": [
-                        {
-                            "subsets": [
-                                {
-                                    "image_dir": args.train_data_dir,
-                                    "metadata_file": args.in_json,
-                                }
-                            ]
-                        }
-                    ]
-                }
-
-        blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizers)
-        train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
-    else:
-        train_dataset_group = train_util.load_arbitrary_dataset(args, tokenizers)
-
-    current_epoch = Value("i", 0)
-    current_step = Value("i", 0)
-    ds_for_collater = train_dataset_group if args.max_data_loader_n_workers == 0 else None
-    collater = train_util.collater_class(current_epoch, current_step, ds_for_collater)
-
-    # acceleratorを準備する
-    print("prepare accelerator")
-    accelerator = train_util.prepare_accelerator(args)
-
-    # mixed precisionに対応した型を用意しておき適宜castする
-    weight_dtype, _ = train_util.prepare_dtype(args)
-
-    # モデルを読み込む
-    print("load model")
-    if args.sdxl:
-        (_, text_encoder1, text_encoder2, _, _, _, _) = sdxl_train_util.load_target_model(args, accelerator, "sdxl", weight_dtype)
-        text_encoders = [text_encoder1, text_encoder2]
-    else:
-        text_encoder1, _, _, _ = train_util.load_target_model(args, weight_dtype, accelerator)
-        text_encoders = [text_encoder1]
-
-    for text_encoder in text_encoders:
-        text_encoder.to(accelerator.device, dtype=weight_dtype)
-        text_encoder.requires_grad_(False)
-        text_encoder.eval()
-
-    # dataloaderを準備する
-    train_dataset_group.set_caching_mode("text")
-
-    # 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_group,
-        batch_size=1,
-        shuffle=True,
-        collate_fn=collater,
-        num_workers=n_workers,
-        persistent_workers=args.persistent_data_loader_workers,
-    )
-
-    # acceleratorを使ってモデルを準備する：マルチGPUで使えるようになるはず
-    train_dataloader = accelerator.prepare(train_dataloader)
-
-    # データ取得のためのループ
-    for batch in tqdm(train_dataloader):
-        absolute_paths = batch["absolute_paths"]
-        input_ids1_list = batch["input_ids1_list"]
-        input_ids2_list = batch["input_ids2_list"]
-
-        image_infos = []
-        for absolute_path, input_ids1, input_ids2 in zip(absolute_paths, input_ids1_list, input_ids2_list):
-            image_info = train_util.ImageInfo(absolute_path, 1, "dummy", False, absolute_path)
-            image_info.text_encoder_outputs_npz = os.path.splitext(absolute_path)[0] + train_util.TEXT_ENCODER_OUTPUTS_CACHE_SUFFIX
-            image_info
-
-            if args.skip_existing:
-                if os.path.exists(image_info.text_encoder_outputs_npz):
-                    print(f"Skipping {image_info.text_encoder_outputs_npz} because it already exists.")
-                    continue
-                
-            image_info.input_ids1 = input_ids1
-            image_info.input_ids2 = input_ids2
-            image_infos.append(image_info)
-
-        if len(image_infos) > 0:
-            b_input_ids1 = torch.stack([image_info.input_ids1 for image_info in image_infos])
-            b_input_ids2 = torch.stack([image_info.input_ids2 for image_info in image_infos])
-            train_util.cache_batch_text_encoder_outputs(
-                image_infos, tokenizers, text_encoders, args.max_token_length, True, b_input_ids1, b_input_ids2, weight_dtype
-            )
-
-    accelerator.wait_for_everyone()
-    accelerator.print(f"Finished caching latents for {len(train_dataset_group)} batches.")
-
-
-def setup_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser()
-
-    train_util.add_sd_models_arguments(parser)
-    train_util.add_training_arguments(parser, True)
-    train_util.add_dataset_arguments(parser, True, True, True)
-    config_util.add_config_arguments(parser)
-    sdxl_train_util.add_sdxl_training_arguments(parser)
-    parser.add_argument("--sdxl", action="store_true", help="Use SDXL model / SDXLモデルを使用する")
-    parser.add_argument(
-        "--skip_existing",
-        action="store_true",
-        help="skip images if npz already exists (both normal and flipped exists if flip_aug is enabled) / npzが既に存在する画像をスキップする（flip_aug有効時は通常、反転の両方が存在する画像をスキップ）",
-    )
-    return parser
-
-
-if __name__ == "__main__":
-    parser = setup_parser()
-
-    args = parser.parse_args()
-    args = train_util.read_config_from_file(args, parser)
-
-    cache_to_disk(args)

tools/merge_models.py

@@ -1,168 +0,0 @@
-import argparse
-import os
-
-import torch
-from safetensors import safe_open
-from safetensors.torch import load_file, save_file
-from tqdm import tqdm
-
-
-def is_unet_key(key):
-    # VAE or TextEncoder, the last one is for SDXL
-    return not ("first_stage_model" in key or "cond_stage_model" in key or "conditioner." in key)
-
-
-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."),
-]
-
-
-# support for models with different text encoder keys
-def replace_text_encoder_key(key):
-    for rep_from, rep_to in TEXT_ENCODER_KEY_REPLACEMENTS:
-        if key.startswith(rep_from):
-            return True, rep_to + key[len(rep_from) :]
-    return False, key
-
-
-def merge(args):
-    if args.precision == "fp16":
-        dtype = torch.float16
-    elif args.precision == "bf16":
-        dtype = torch.bfloat16
-    else:
-        dtype = torch.float
-
-    if args.saving_precision == "fp16":
-        save_dtype = torch.float16
-    elif args.saving_precision == "bf16":
-        save_dtype = torch.bfloat16
-    else:
-        save_dtype = torch.float
-
-    # check if all models are safetensors
-    for model in args.models:
-        if not model.endswith("safetensors"):
-            print(f"Model {model} is not a safetensors model")
-            exit()
-        if not os.path.isfile(model):
-            print(f"Model {model} does not exist")
-            exit()
-
-    assert args.ratios is None or len(args.models) == len(args.ratios), "ratios must be the same length as models"
-
-    # load and merge
-    ratio = 1.0 / len(args.models)  # default
-    supplementary_key_ratios = {}  # [key] = ratio, for keys not in all models, add later
-
-    merged_sd = None
-    first_model_keys = set()  # check missing keys in other models
-    for i, model in enumerate(args.models):
-        if args.ratios is not None:
-            ratio = args.ratios[i]
-
-        if merged_sd is None:
-            # load first model
-            print(f"Loading model {model}, ratio = {ratio}...")
-            merged_sd = {}
-            with safe_open(model, framework="pt", device=args.device) as f:
-                for key in tqdm(f.keys()):
-                    value = f.get_tensor(key)
-                    _, key = replace_text_encoder_key(key)
-
-                    first_model_keys.add(key)
-
-                    if not is_unet_key(key) and args.unet_only:
-                        supplementary_key_ratios[key] = 1.0  # use first model's value for VAE or TextEncoder
-                        continue
-
-                    value = ratio * value.to(dtype)  # first model's value * ratio
-                    merged_sd[key] = value
-
-            print(f"Model has {len(merged_sd)} keys " + ("(UNet only)" if args.unet_only else ""))
-            continue
-
-        # load other models
-        print(f"Loading model {model}, ratio = {ratio}...")
-
-        with safe_open(model, framework="pt", device=args.device) as f:
-            model_keys = f.keys()
-            for key in tqdm(model_keys):
-                _, new_key = replace_text_encoder_key(key)
-                if new_key not in merged_sd:
-                    if args.show_skipped and new_key not in first_model_keys:
-                        print(f"Skip: {new_key}")
-                    continue
-
-                value = f.get_tensor(key)
-                merged_sd[new_key] = merged_sd[new_key] + ratio * value.to(dtype)
-
-            # enumerate keys not in this model
-            model_keys = set(model_keys)
-            for key in merged_sd.keys():
-                if key in model_keys:
-                    continue
-                print(f"Key {key} not in model {model}, use first model's value")
-                if key in supplementary_key_ratios:
-                    supplementary_key_ratios[key] += ratio
-                else:
-                    supplementary_key_ratios[key] = ratio
-
-    # add supplementary keys' value (including VAE and TextEncoder)
-    if len(supplementary_key_ratios) > 0:
-        print("add first model's value")
-        with safe_open(args.models[0], framework="pt", device=args.device) as f:
-            for key in tqdm(f.keys()):
-                _, new_key = replace_text_encoder_key(key)
-                if new_key not in supplementary_key_ratios:
-                    continue
-
-                if is_unet_key(new_key):  # not VAE or TextEncoder
-                    print(f"Key {new_key} not in all models, ratio = {supplementary_key_ratios[new_key]}")
-
-                value = f.get_tensor(key)  # original key
-
-                if new_key not in merged_sd:
-                    merged_sd[new_key] = supplementary_key_ratios[new_key] * value.to(dtype)
-                else:
-                    merged_sd[new_key] = merged_sd[new_key] + supplementary_key_ratios[new_key] * value.to(dtype)
-
-    # save
-    output_file = args.output
-    if not output_file.endswith(".safetensors"):
-        output_file = output_file + ".safetensors"
-
-    print(f"Saving to {output_file}...")
-
-    # convert to save_dtype
-    for k in merged_sd.keys():
-        merged_sd[k] = merged_sd[k].to(save_dtype)
-
-    save_file(merged_sd, output_file)
-
-    print("Done!")
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Merge models")
-    parser.add_argument("--models", nargs="+", type=str, help="Models to merge")
-    parser.add_argument("--output", type=str, help="Output model")
-    parser.add_argument("--ratios", nargs="+", type=float, help="Ratios of models, default is equal, total = 1.0")
-    parser.add_argument("--unet_only", action="store_true", help="Only merge unet")
-    parser.add_argument("--device", type=str, default="cpu", help="Device to use, default is cpu")
-    parser.add_argument(
-        "--precision", type=str, default="float", choices=["float", "fp16", "bf16"], help="Calculation precision, default is float"
-    )
-    parser.add_argument(
-        "--saving_precision",
-        type=str,
-        default="float",
-        choices=["float", "fp16", "bf16"],
-        help="Saving precision, default is float",
-    )
-    parser.add_argument("--show_skipped", action="store_true", help="Show skipped keys (keys not in first model)")
-
-    args = parser.parse_args()
-    merge(args)

tools/original_control_net.py

@@ -4,187 +4,176 @@ import cv2
 import torch
 from safetensors.torch import load_file
 
-from library.original_unet import UNet2DConditionModel, SampleOutput
+from diffusers import UNet2DConditionModel
+from diffusers.models.unet_2d_condition import UNet2DConditionOutput
 
 import library.model_util as model_util
 
 
 class ControlNetInfo(NamedTuple):
-    unet: Any
-    net: Any
-    prep: Any
-    weight: float
-    ratio: float
+  unet: Any
+  net: Any
+  prep: Any
+  weight: float
+  ratio: float
 
 
 class ControlNet(torch.nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
+  def __init__(self) -> None:
+    super().__init__()
 
-        # make control model
-        self.control_model = torch.nn.Module()
+    # make control model
+    self.control_model = torch.nn.Module()
 
-        dims = [320, 320, 320, 320, 640, 640, 640, 1280, 1280, 1280, 1280, 1280]
-        zero_convs = torch.nn.ModuleList()
-        for i, dim in enumerate(dims):
-            sub_list = torch.nn.ModuleList([torch.nn.Conv2d(dim, dim, 1)])
-            zero_convs.append(sub_list)
-        self.control_model.add_module("zero_convs", zero_convs)
+    dims = [320, 320, 320, 320, 640, 640, 640, 1280, 1280, 1280, 1280, 1280]
+    zero_convs = torch.nn.ModuleList()
+    for i, dim in enumerate(dims):
+      sub_list = torch.nn.ModuleList([torch.nn.Conv2d(dim, dim, 1)])
+      zero_convs.append(sub_list)
+    self.control_model.add_module("zero_convs", zero_convs)
 
-        middle_block_out = torch.nn.Conv2d(1280, 1280, 1)
-        self.control_model.add_module("middle_block_out", torch.nn.ModuleList([middle_block_out]))
+    middle_block_out = torch.nn.Conv2d(1280, 1280, 1)
+    self.control_model.add_module("middle_block_out", torch.nn.ModuleList([middle_block_out]))
 
-        dims = [16, 16, 32, 32, 96, 96, 256, 320]
-        strides = [1, 1, 2, 1, 2, 1, 2, 1]
-        prev_dim = 3
-        input_hint_block = torch.nn.Sequential()
-        for i, (dim, stride) in enumerate(zip(dims, strides)):
-            input_hint_block.append(torch.nn.Conv2d(prev_dim, dim, 3, stride, 1))
-            if i < len(dims) - 1:
-                input_hint_block.append(torch.nn.SiLU())
-            prev_dim = dim
-        self.control_model.add_module("input_hint_block", input_hint_block)
+    dims = [16, 16, 32, 32, 96, 96, 256, 320]
+    strides = [1, 1, 2, 1, 2, 1, 2, 1]
+    prev_dim = 3
+    input_hint_block = torch.nn.Sequential()
+    for i, (dim, stride) in enumerate(zip(dims, strides)):
+      input_hint_block.append(torch.nn.Conv2d(prev_dim, dim, 3, stride, 1))
+      if i < len(dims) - 1:
+        input_hint_block.append(torch.nn.SiLU())
+      prev_dim = dim
+    self.control_model.add_module("input_hint_block", input_hint_block)
 
 
 def load_control_net(v2, unet, model):
-    device = unet.device
+  device = unet.device
 
-    # control sdからキー変換しつつU-Netに対応する部分のみ取り出し、DiffusersのU-Netに読み込む
-    # state dictを読み込む
-    print(f"ControlNet: loading control SD model : {model}")
+  # control sdからキー変換しつつU-Netに対応する部分のみ取り出し、DiffusersのU-Netに読み込む
+  # state dictを読み込む
+  print(f"ControlNet: loading control SD model : {model}")
 
-    if model_util.is_safetensors(model):
-        ctrl_sd_sd = load_file(model)
-    else:
-        ctrl_sd_sd = torch.load(model, map_location="cpu")
-        ctrl_sd_sd = ctrl_sd_sd.pop("state_dict", ctrl_sd_sd)
+  if model_util.is_safetensors(model):
+    ctrl_sd_sd = load_file(model)
+  else:
+    ctrl_sd_sd = torch.load(model, map_location='cpu')
+    ctrl_sd_sd = ctrl_sd_sd.pop("state_dict", ctrl_sd_sd)
 
-    # 重みをU-Netに読み込めるようにする。ControlNetはSD版のstate dictなので、それを読み込む
-    is_difference = "difference" in ctrl_sd_sd
-    print("ControlNet: loading difference:", is_difference)
+  # 重みをU-Netに読み込めるようにする。ControlNetはSD版のstate dictなので、それを読み込む
+  is_difference = "difference" in ctrl_sd_sd
+  print("ControlNet: loading difference:", is_difference)
 
-    # ControlNetには存在しないキーがあるので、まず現在のU-NetでSD版の全keyを作っておく
-    # またTransfer Controlの元weightとなる
-    ctrl_unet_sd_sd = model_util.convert_unet_state_dict_to_sd(v2, unet.state_dict())
+  # ControlNetには存在しないキーがあるので、まず現在のU-NetでSD版の全keyを作っておく
+  # またTransfer Controlの元weightとなる
+  ctrl_unet_sd_sd = model_util.convert_unet_state_dict_to_sd(v2, unet.state_dict())
 
-    # 元のU-Netに影響しないようにコピーする。またprefixが付いていないので付ける
-    for key in list(ctrl_unet_sd_sd.keys()):
-        ctrl_unet_sd_sd["model.diffusion_model." + key] = ctrl_unet_sd_sd.pop(key).clone()
+  # 元のU-Netに影響しないようにコピーする。またprefixが付いていないので付ける
+  for key in list(ctrl_unet_sd_sd.keys()):
+    ctrl_unet_sd_sd["model.diffusion_model." + key] = ctrl_unet_sd_sd.pop(key).clone()
 
-    zero_conv_sd = {}
-    for key in list(ctrl_sd_sd.keys()):
-        if key.startswith("control_"):
-            unet_key = "model.diffusion_" + key[len("control_") :]
-            if unet_key not in ctrl_unet_sd_sd:  # zero conv
-                zero_conv_sd[key] = ctrl_sd_sd[key]
-                continue
-            if is_difference:  # Transfer Control
-                ctrl_unet_sd_sd[unet_key] += ctrl_sd_sd[key].to(device, dtype=unet.dtype)
-            else:
-                ctrl_unet_sd_sd[unet_key] = ctrl_sd_sd[key].to(device, dtype=unet.dtype)
+  zero_conv_sd = {}
+  for key in list(ctrl_sd_sd.keys()):
+    if key.startswith("control_"):
+      unet_key = "model.diffusion_" + key[len("control_"):]
+      if unet_key not in ctrl_unet_sd_sd:               # zero conv
+        zero_conv_sd[key] = ctrl_sd_sd[key]
+        continue
+      if is_difference:                                 # Transfer Control
+        ctrl_unet_sd_sd[unet_key] += ctrl_sd_sd[key].to(device, dtype=unet.dtype)
+      else:
+        ctrl_unet_sd_sd[unet_key] = ctrl_sd_sd[key].to(device, dtype=unet.dtype)
 
-    unet_config = model_util.create_unet_diffusers_config(v2)
-    ctrl_unet_du_sd = model_util.convert_ldm_unet_checkpoint(v2, ctrl_unet_sd_sd, unet_config)  # DiffUsers版ControlNetのstate dict
+  unet_config = model_util.create_unet_diffusers_config(v2)
+  ctrl_unet_du_sd = model_util.convert_ldm_unet_checkpoint(v2, ctrl_unet_sd_sd, unet_config)    # DiffUsers版ControlNetのstate dict
 
-    # ControlNetのU-Netを作成する
-    ctrl_unet = UNet2DConditionModel(**unet_config)
-    info = ctrl_unet.load_state_dict(ctrl_unet_du_sd)
-    print("ControlNet: loading Control U-Net:", info)
+  # ControlNetのU-Netを作成する
+  ctrl_unet = UNet2DConditionModel(**unet_config)
+  info = ctrl_unet.load_state_dict(ctrl_unet_du_sd)
+  print("ControlNet: loading Control U-Net:", info)
 
-    # U-Net以外のControlNetを作成する
-    # TODO support middle only
-    ctrl_net = ControlNet()
-    info = ctrl_net.load_state_dict(zero_conv_sd)
-    print("ControlNet: loading ControlNet:", info)
+  # U-Net以外のControlNetを作成する
+  # TODO support middle only
+  ctrl_net = ControlNet()
+  info = ctrl_net.load_state_dict(zero_conv_sd)
+  print("ControlNet: loading ControlNet:", info)
 
-    ctrl_unet.to(unet.device, dtype=unet.dtype)
-    ctrl_net.to(unet.device, dtype=unet.dtype)
-    return ctrl_unet, ctrl_net
+  ctrl_unet.to(unet.device, dtype=unet.dtype)
+  ctrl_net.to(unet.device, dtype=unet.dtype)
+  return ctrl_unet, ctrl_net
 
 
 def load_preprocess(prep_type: str):
-    if prep_type is None or prep_type.lower() == "none":
-        return None
-
-    if prep_type.startswith("canny"):
-        args = prep_type.split("_")
-        th1 = int(args[1]) if len(args) >= 2 else 63
-        th2 = int(args[2]) if len(args) >= 3 else 191
-
-        def canny(img):
-            img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
-            return cv2.Canny(img, th1, th2)
-
-        return canny
-
-    print("Unsupported prep type:", prep_type)
+  if prep_type is None or prep_type.lower() == "none":
     return None
 
+  if prep_type.startswith("canny"):
+    args = prep_type.split("_")
+    th1 = int(args[1]) if len(args) >= 2 else 63
+    th2 = int(args[2]) if len(args) >= 3 else 191
+
+    def canny(img):
+      img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+      return cv2.Canny(img, th1, th2)
+    return canny
+
+  print("Unsupported prep type:", prep_type)
+  return None
+
 
 def preprocess_ctrl_net_hint_image(image):
-    image = np.array(image).astype(np.float32) / 255.0
-    # ControlNetのサンプルはcv2を使っているが、読み込みはGradioなので実はRGBになっている
-    # image = image[:, :, ::-1].copy()                         # rgb to bgr
-    image = image[None].transpose(0, 3, 1, 2)  # nchw
-    image = torch.from_numpy(image)
-    return image  # 0 to 1
+  image = np.array(image).astype(np.float32) / 255.0
+  # ControlNetのサンプルはcv2を使っているが、読み込みはGradioなので実はRGBになっている
+  # image = image[:, :, ::-1].copy()                         # rgb to bgr
+  image = image[None].transpose(0, 3, 1, 2)       # nchw
+  image = torch.from_numpy(image)
+  return image                              # 0 to 1
 
 
 def get_guided_hints(control_nets: List[ControlNetInfo], num_latent_input, b_size, hints):
-    guided_hints = []
-    for i, cnet_info in enumerate(control_nets):
-        # hintは 1枚目の画像のcnet1, 1枚目の画像のcnet2, 1枚目の画像のcnet3, 2枚目の画像のcnet1, 2枚目の画像のcnet2 ... と並んでいること
-        b_hints = []
-        if len(hints) == 1:  # すべて同じ画像をhintとして使う
-            hint = hints[0]
-            if cnet_info.prep is not None:
-                hint = cnet_info.prep(hint)
-            hint = preprocess_ctrl_net_hint_image(hint)
-            b_hints = [hint for _ in range(b_size)]
-        else:
-            for bi in range(b_size):
-                hint = hints[(bi * len(control_nets) + i) % len(hints)]
-                if cnet_info.prep is not None:
-                    hint = cnet_info.prep(hint)
-                hint = preprocess_ctrl_net_hint_image(hint)
-                b_hints.append(hint)
-        b_hints = torch.cat(b_hints, dim=0)
-        b_hints = b_hints.to(cnet_info.unet.device, dtype=cnet_info.unet.dtype)
+  guided_hints = []
+  for i, cnet_info in enumerate(control_nets):
+    # hintは 1枚目の画像のcnet1, 1枚目の画像のcnet2, 1枚目の画像のcnet3, 2枚目の画像のcnet1, 2枚目の画像のcnet2 ... と並んでいること
+    b_hints = []
+    if len(hints) == 1:           # すべて同じ画像をhintとして使う
+      hint = hints[0]
+      if cnet_info.prep is not None:
+        hint = cnet_info.prep(hint)
+      hint = preprocess_ctrl_net_hint_image(hint)
+      b_hints = [hint for _ in range(b_size)]
+    else:
+      for bi in range(b_size):
+        hint = hints[(bi * len(control_nets) + i) % len(hints)]
+        if cnet_info.prep is not None:
+          hint = cnet_info.prep(hint)
+        hint = preprocess_ctrl_net_hint_image(hint)
+        b_hints.append(hint)
+    b_hints = torch.cat(b_hints, dim=0)
+    b_hints = b_hints.to(cnet_info.unet.device, dtype=cnet_info.unet.dtype)
 
-        guided_hint = cnet_info.net.control_model.input_hint_block(b_hints)
-        guided_hints.append(guided_hint)
-    return guided_hints
+    guided_hint = cnet_info.net.control_model.input_hint_block(b_hints)
+    guided_hints.append(guided_hint)
+  return guided_hints
 
 
-def call_unet_and_control_net(
-    step,
-    num_latent_input,
-    original_unet,
-    control_nets: List[ControlNetInfo],
-    guided_hints,
-    current_ratio,
-    sample,
-    timestep,
-    encoder_hidden_states,
-    encoder_hidden_states_for_control_net,
-):
-    # ControlNet
-    # 複数のControlNetの場合は、出力をマージするのではなく交互に適用する
-    cnet_cnt = len(control_nets)
-    cnet_idx = step % cnet_cnt
-    cnet_info = control_nets[cnet_idx]
+def call_unet_and_control_net(step, num_latent_input, original_unet, control_nets: List[ControlNetInfo], guided_hints, current_ratio, sample, timestep, encoder_hidden_states):
+  # ControlNet
+  # 複数のControlNetの場合は、出力をマージするのではなく交互に適用する
+  cnet_cnt = len(control_nets)
+  cnet_idx = step % cnet_cnt
+  cnet_info = control_nets[cnet_idx]
 
-    # print(current_ratio, cnet_info.prep, cnet_info.weight, cnet_info.ratio)
-    if cnet_info.ratio < current_ratio:
-        return original_unet(sample, timestep, encoder_hidden_states)
+  # print(current_ratio, cnet_info.prep, cnet_info.weight, cnet_info.ratio)
+  if cnet_info.ratio < current_ratio:
+    return original_unet(sample, timestep, encoder_hidden_states)
 
-    guided_hint = guided_hints[cnet_idx]
-    guided_hint = guided_hint.repeat((num_latent_input, 1, 1, 1))
-    outs = unet_forward(True, cnet_info.net, cnet_info.unet, guided_hint, None, sample, timestep, encoder_hidden_states_for_control_net)
-    outs = [o * cnet_info.weight for o in outs]
+  guided_hint = guided_hints[cnet_idx]
+  guided_hint = guided_hint.repeat((num_latent_input, 1, 1, 1))
+  outs = unet_forward(True, cnet_info.net, cnet_info.unet, guided_hint, None, sample, timestep, encoder_hidden_states)
+  outs = [o * cnet_info.weight for o in outs]
 
-    # U-Net
-    return unet_forward(False, cnet_info.net, original_unet, None, outs, sample, timestep, encoder_hidden_states)
+  # U-Net
+  return unet_forward(False, cnet_info.net, original_unet, None, outs, sample, timestep, encoder_hidden_states)
 
 
 """
@@ -215,123 +204,118 @@ def call_unet_and_control_net(
 """
 
 
-def unet_forward(
-    is_control_net,
-    control_net: ControlNet,
-    unet: UNet2DConditionModel,
-    guided_hint,
-    ctrl_outs,
-    sample,
-    timestep,
-    encoder_hidden_states,
-):
-    # copy from UNet2DConditionModel
-    default_overall_up_factor = 2**unet.num_upsamplers
+def unet_forward(is_control_net, control_net: ControlNet, unet: UNet2DConditionModel, guided_hint, ctrl_outs, sample, timestep, encoder_hidden_states):
+  # copy from UNet2DConditionModel
+  default_overall_up_factor = 2**unet.num_upsamplers
 
-    forward_upsample_size = False
-    upsample_size = None
+  forward_upsample_size = False
+  upsample_size = None
 
-    if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
-        print("Forward upsample size to force interpolation output size.")
-        forward_upsample_size = True
+  if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
+    print("Forward upsample size to force interpolation output size.")
+    forward_upsample_size = True
 
-    # 1. time
-    timesteps = timestep
-    if not torch.is_tensor(timesteps):
-        # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
-        # This would be a good case for the `match` statement (Python 3.10+)
-        is_mps = sample.device.type == "mps"
-        if isinstance(timestep, float):
-            dtype = torch.float32 if is_mps else torch.float64
-        else:
-            dtype = torch.int32 if is_mps else torch.int64
-        timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
-    elif len(timesteps.shape) == 0:
-        timesteps = timesteps[None].to(sample.device)
+  # 0. center input if necessary
+  if unet.config.center_input_sample:
+    sample = 2 * sample - 1.0
 
-    # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-    timesteps = timesteps.expand(sample.shape[0])
+  # 1. time
+  timesteps = timestep
+  if not torch.is_tensor(timesteps):
+    # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+    # This would be a good case for the `match` statement (Python 3.10+)
+    is_mps = sample.device.type == "mps"
+    if isinstance(timestep, float):
+      dtype = torch.float32 if is_mps else torch.float64
+    else:
+      dtype = torch.int32 if is_mps else torch.int64
+    timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+  elif len(timesteps.shape) == 0:
+    timesteps = timesteps[None].to(sample.device)
 
-    t_emb = unet.time_proj(timesteps)
+  # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+  timesteps = timesteps.expand(sample.shape[0])
 
-    # timesteps does not contain any weights and will always return f32 tensors
-    # but time_embedding might actually be running in fp16. so we need to cast here.
-    # there might be better ways to encapsulate this.
-    t_emb = t_emb.to(dtype=unet.dtype)
-    emb = unet.time_embedding(t_emb)
+  t_emb = unet.time_proj(timesteps)
 
-    outs = []  # output of ControlNet
-    zc_idx = 0
+  # timesteps does not contain any weights and will always return f32 tensors
+  # but time_embedding might actually be running in fp16. so we need to cast here.
+  # there might be better ways to encapsulate this.
+  t_emb = t_emb.to(dtype=unet.dtype)
+  emb = unet.time_embedding(t_emb)
 
-    # 2. pre-process
-    sample = unet.conv_in(sample)
+  outs = []                     # output of ControlNet
+  zc_idx = 0
+
+  # 2. pre-process
+  sample = unet.conv_in(sample)
+  if is_control_net:
+    sample += guided_hint
+    outs.append(control_net.control_model.zero_convs[zc_idx][0](sample))  # , emb, encoder_hidden_states))
+    zc_idx += 1
+
+  # 3. down
+  down_block_res_samples = (sample,)
+  for downsample_block in unet.down_blocks:
+    if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+      sample, res_samples = downsample_block(
+          hidden_states=sample,
+          temb=emb,
+          encoder_hidden_states=encoder_hidden_states,
+      )
+    else:
+      sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
     if is_control_net:
-        sample += guided_hint
-        outs.append(control_net.control_model.zero_convs[zc_idx][0](sample))  # , emb, encoder_hidden_states))
+      for rs in res_samples:
+        outs.append(control_net.control_model.zero_convs[zc_idx][0](rs))  # , emb, encoder_hidden_states))
         zc_idx += 1
 
-    # 3. down
-    down_block_res_samples = (sample,)
-    for downsample_block in unet.down_blocks:
-        if downsample_block.has_cross_attention:
-            sample, res_samples = downsample_block(
-                hidden_states=sample,
-                temb=emb,
-                encoder_hidden_states=encoder_hidden_states,
-            )
-        else:
-            sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
-        if is_control_net:
-            for rs in res_samples:
-                outs.append(control_net.control_model.zero_convs[zc_idx][0](rs))  # , emb, encoder_hidden_states))
-                zc_idx += 1
+    down_block_res_samples += res_samples
 
-        down_block_res_samples += res_samples
+  # 4. mid
+  sample = unet.mid_block(sample, emb, encoder_hidden_states=encoder_hidden_states)
+  if is_control_net:
+    outs.append(control_net.control_model.middle_block_out[0](sample))
+    return outs
 
-    # 4. mid
-    sample = unet.mid_block(sample, emb, encoder_hidden_states=encoder_hidden_states)
-    if is_control_net:
-        outs.append(control_net.control_model.middle_block_out[0](sample))
-        return outs
+  if not is_control_net:
+    sample += ctrl_outs.pop()
 
-    if not is_control_net:
-        sample += ctrl_outs.pop()
+  # 5. up
+  for i, upsample_block in enumerate(unet.up_blocks):
+    is_final_block = i == len(unet.up_blocks) - 1
 
-    # 5. up
-    for i, upsample_block in enumerate(unet.up_blocks):
-        is_final_block = i == len(unet.up_blocks) - 1
+    res_samples = down_block_res_samples[-len(upsample_block.resnets):]
+    down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
 
-        res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
-        down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+    if not is_control_net and len(ctrl_outs) > 0:
+      res_samples = list(res_samples)
+      apply_ctrl_outs = ctrl_outs[-len(res_samples):]
+      ctrl_outs = ctrl_outs[:-len(res_samples)]
+      for j in range(len(res_samples)):
+        res_samples[j] = res_samples[j] + apply_ctrl_outs[j]
+      res_samples = tuple(res_samples)
 
-        if not is_control_net and len(ctrl_outs) > 0:
-            res_samples = list(res_samples)
-            apply_ctrl_outs = ctrl_outs[-len(res_samples) :]
-            ctrl_outs = ctrl_outs[: -len(res_samples)]
-            for j in range(len(res_samples)):
-                res_samples[j] = res_samples[j] + apply_ctrl_outs[j]
-            res_samples = tuple(res_samples)
+    # if we have not reached the final block and need to forward the
+    # upsample size, we do it here
+    if not is_final_block and forward_upsample_size:
+      upsample_size = down_block_res_samples[-1].shape[2:]
 
-        # if we have not reached the final block and need to forward the
-        # upsample size, we do it here
-        if not is_final_block and forward_upsample_size:
-            upsample_size = down_block_res_samples[-1].shape[2:]
+    if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+      sample = upsample_block(
+          hidden_states=sample,
+          temb=emb,
+          res_hidden_states_tuple=res_samples,
+          encoder_hidden_states=encoder_hidden_states,
+          upsample_size=upsample_size,
+      )
+    else:
+      sample = upsample_block(
+          hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size
+      )
+  # 6. post-process
+  sample = unet.conv_norm_out(sample)
+  sample = unet.conv_act(sample)
+  sample = unet.conv_out(sample)
 
-        if upsample_block.has_cross_attention:
-            sample = upsample_block(
-                hidden_states=sample,
-                temb=emb,
-                res_hidden_states_tuple=res_samples,
-                encoder_hidden_states=encoder_hidden_states,
-                upsample_size=upsample_size,
-            )
-        else:
-            sample = upsample_block(
-                hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size
-            )
-    # 6. post-process
-    sample = unet.conv_norm_out(sample)
-    sample = unet.conv_act(sample)
-    sample = unet.conv_out(sample)
-
-    return SampleOutput(sample=sample)
+  return UNet2DConditionOutput(sample=sample)

tools/show_metadata.py

@@ -1,19 +0,0 @@
-import json
-import argparse
-from safetensors import safe_open
-
-parser = argparse.ArgumentParser()
-parser.add_argument("--model", type=str, required=True)
-args = parser.parse_args()
-
-with safe_open(args.model, framework="pt") as f:
-    metadata = f.metadata()
-
-if metadata is None:
-    print("No metadata found")
-else:
-    # metadata is json dict, but not pretty printed
-    # sort by key and pretty print
-    print(json.dumps(metadata, indent=4, sort_keys=True))
-
-    

train_controlnet.py

@@ -1,604 +0,0 @@
-import argparse
-import gc
-import json
-import math
-import os
-import random
-import time
-from multiprocessing import Value
-from types import SimpleNamespace
-import toml
-
-from tqdm import tqdm
-import torch
-from torch.nn.parallel import DistributedDataParallel as DDP
-from accelerate.utils import set_seed
-from diffusers import DDPMScheduler, ControlNetModel
-from safetensors.torch import load_file
-
-import library.model_util as model_util
-import library.train_util as train_util
-import library.config_util as config_util
-from library.config_util import (
-    ConfigSanitizer,
-    BlueprintGenerator,
-)
-import library.huggingface_util as huggingface_util
-import library.custom_train_functions as custom_train_functions
-from library.custom_train_functions import (
-    apply_snr_weight,
-    pyramid_noise_like,
-    apply_noise_offset,
-)
-
-
-# TODO 他のスクリプトと共通化する
-def generate_step_logs(args: argparse.Namespace, current_loss, avr_loss, lr_scheduler):
-    logs = {
-        "loss/current": current_loss,
-        "loss/average": avr_loss,
-        "lr": lr_scheduler.get_last_lr()[0],
-    }
-
-    if args.optimizer_type.lower().startswith("DAdapt".lower()):
-        logs["lr/d*lr"] = lr_scheduler.optimizers[-1].param_groups[0]["d"] * lr_scheduler.optimizers[-1].param_groups[0]["lr"]
-
-    return logs
-
-
-def train(args):
-    # session_id = random.randint(0, 2**32)
-    # training_started_at = time.time()
-    train_util.verify_training_args(args)
-    train_util.prepare_dataset_args(args, True)
-
-    cache_latents = args.cache_latents
-    use_user_config = args.dataset_config is not None
-
-    if args.seed is None:
-        args.seed = random.randint(0, 2**32)
-    set_seed(args.seed)
-
-    tokenizer = train_util.load_tokenizer(args)
-
-    # データセットを準備する
-    blueprint_generator = BlueprintGenerator(ConfigSanitizer(False, False, True, True))
-    if use_user_config:
-        print(f"Load dataset config from {args.dataset_config}")
-        user_config = config_util.load_user_config(args.dataset_config)
-        ignored = ["train_data_dir", "conditioning_data_dir"]
-        if any(getattr(args, attr) is not None for attr in ignored):
-            print(
-                "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
-                    ", ".join(ignored)
-                )
-            )
-    else:
-        user_config = {
-            "datasets": [
-                {
-                    "subsets": config_util.generate_controlnet_subsets_config_by_subdirs(
-                        args.train_data_dir,
-                        args.conditioning_data_dir,
-                        args.caption_extension,
-                    )
-                }
-            ]
-        }
-
-    blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
-    train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
-
-    current_epoch = Value("i", 0)
-    current_step = Value("i", 0)
-    ds_for_collater = train_dataset_group if args.max_data_loader_n_workers == 0 else None
-    collater = train_util.collater_class(current_epoch, current_step, ds_for_collater)
-
-    if args.debug_dataset:
-        train_util.debug_dataset(train_dataset_group)
-        return
-    if len(train_dataset_group) == 0:
-        print(
-            "No data found. Please verify arguments (train_data_dir must be the parent of folders with images) / 画像がありません。引数指定を確認してください（train_data_dirには画像があるフォルダではなく、画像があるフォルダの親フォルダを指定する必要があります）"
-        )
-        return
-
-    if cache_latents:
-        assert (
-            train_dataset_group.is_latent_cacheable()
-        ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
-
-    # acceleratorを準備する
-    print("prepare accelerator")
-    accelerator = train_util.prepare_accelerator(args)
-    is_main_process = accelerator.is_main_process
-
-    # mixed precisionに対応した型を用意しておき適宜castする
-    weight_dtype, save_dtype = train_util.prepare_dtype(args)
-
-    # モデルを読み込む
-    text_encoder, vae, unet, _ = train_util.load_target_model(
-        args, weight_dtype, accelerator, unet_use_linear_projection_in_v2=True
-    )
-
-    # DiffusersのControlNetが使用するデータを準備する
-    if args.v2:
-        unet.config = {
-            "act_fn": "silu",
-            "attention_head_dim": [5, 10, 20, 20],
-            "block_out_channels": [320, 640, 1280, 1280],
-            "center_input_sample": False,
-            "cross_attention_dim": 1024,
-            "down_block_types": ["CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D"],
-            "downsample_padding": 1,
-            "dual_cross_attention": False,
-            "flip_sin_to_cos": True,
-            "freq_shift": 0,
-            "in_channels": 4,
-            "layers_per_block": 2,
-            "mid_block_scale_factor": 1,
-            "norm_eps": 1e-05,
-            "norm_num_groups": 32,
-            "num_class_embeds": None,
-            "only_cross_attention": False,
-            "out_channels": 4,
-            "sample_size": 96,
-            "up_block_types": ["UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"],
-            "use_linear_projection": True,
-            "upcast_attention": True,
-            "only_cross_attention": False,
-            "downsample_padding": 1,
-            "use_linear_projection": True,
-            "class_embed_type": None,
-            "num_class_embeds": None,
-            "resnet_time_scale_shift": "default",
-            "projection_class_embeddings_input_dim": None,
-        }
-    else:
-        unet.config = {
-            "act_fn": "silu",
-            "attention_head_dim": 8,
-            "block_out_channels": [320, 640, 1280, 1280],
-            "center_input_sample": False,
-            "cross_attention_dim": 768,
-            "down_block_types": ["CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D"],
-            "downsample_padding": 1,
-            "flip_sin_to_cos": True,
-            "freq_shift": 0,
-            "in_channels": 4,
-            "layers_per_block": 2,
-            "mid_block_scale_factor": 1,
-            "norm_eps": 1e-05,
-            "norm_num_groups": 32,
-            "out_channels": 4,
-            "sample_size": 64,
-            "up_block_types": ["UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"],
-            "only_cross_attention": False,
-            "downsample_padding": 1,
-            "use_linear_projection": False,
-            "class_embed_type": None,
-            "num_class_embeds": None,
-            "upcast_attention": False,
-            "resnet_time_scale_shift": "default",
-            "projection_class_embeddings_input_dim": None,
-        }
-    unet.config = SimpleNamespace(**unet.config)
-
-    controlnet = ControlNetModel.from_unet(unet)
-
-    if args.controlnet_model_name_or_path:
-        filename = args.controlnet_model_name_or_path
-        if os.path.isfile(filename):
-            if os.path.splitext(filename)[1] == ".safetensors":
-                state_dict = load_file(filename)
-            else:
-                state_dict = torch.load(filename)
-            state_dict = model_util.convert_controlnet_state_dict_to_diffusers(state_dict)
-            controlnet.load_state_dict(state_dict)
-        elif os.path.isdir(filename):
-            controlnet = ControlNetModel.from_pretrained(filename)
-
-    # モデルに xformers とか memory efficient attention を組み込む
-    train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers, args.sdpa)
-
-    # 学習を準備する
-    if cache_latents:
-        vae.to(accelerator.device, dtype=weight_dtype)
-        vae.requires_grad_(False)
-        vae.eval()
-        with torch.no_grad():
-            train_dataset_group.cache_latents(
-                vae,
-                args.vae_batch_size,
-                args.cache_latents_to_disk,
-                accelerator.is_main_process,
-            )
-        vae.to("cpu")
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        gc.collect()
-
-        accelerator.wait_for_everyone()
-
-    if args.gradient_checkpointing:
-        controlnet.enable_gradient_checkpointing()
-
-    # 学習に必要なクラスを準備する
-    accelerator.print("prepare optimizer, data loader etc.")
-
-    trainable_params = controlnet.parameters()
-
-    _, _, optimizer = train_util.get_optimizer(args, trainable_params)
-
-    # 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_group,
-        batch_size=1,
-        shuffle=True,
-        collate_fn=collater,
-        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 * math.ceil(
-            len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps
-        )
-        accelerator.print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
-
-    # データセット側にも学習ステップを送信
-    train_dataset_group.set_max_train_steps(args.max_train_steps)
-
-    # lr schedulerを用意する
-    lr_scheduler = train_util.get_scheduler_fix(args, optimizer, accelerator.num_processes)
-
-    # 実験的機能：勾配も含めたfp16学習を行う　モデル全体をfp16にする
-    if args.full_fp16:
-        assert (
-            args.mixed_precision == "fp16"
-        ), "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
-        accelerator.print("enable full fp16 training.")
-        controlnet.to(weight_dtype)
-
-    # acceleratorがなんかよろしくやってくれるらしい
-    controlnet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-        controlnet, optimizer, train_dataloader, lr_scheduler
-    )
-
-    unet.requires_grad_(False)
-    text_encoder.requires_grad_(False)
-    unet.to(accelerator.device)
-    text_encoder.to(accelerator.device)
-
-    # transform DDP after prepare
-    controlnet = controlnet.module if isinstance(controlnet, DDP) else controlnet
-
-    controlnet.train()
-
-    if not cache_latents:
-        vae.requires_grad_(False)
-        vae.eval()
-        vae.to(accelerator.device, dtype=weight_dtype)
-
-    # 実験的機能：勾配も含めたfp16学習を行う　PyTorchにパッチを当ててfp16でのgrad scaleを有効にする
-    if args.full_fp16:
-        train_util.patch_accelerator_for_fp16_training(accelerator)
-
-    # resumeする
-    train_util.resume_from_local_or_hf_if_specified(accelerator, args)
-
-    # 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
-
-    # 学習する
-    # TODO: find a way to handle total batch size when there are multiple datasets
-    accelerator.print("running training / 学習開始")
-    accelerator.print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
-    accelerator.print(f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
-    accelerator.print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
-    accelerator.print(f"  num epochs / epoch数: {num_train_epochs}")
-    accelerator.print(f"  batch size per device / バッチサイズ: {', '.join([str(d.batch_size) for d in train_dataset_group.datasets])}")
-    # print(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
-    accelerator.print(f"  gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
-    accelerator.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:
-        init_kwargs = {}
-        if args.log_tracker_config is not None:
-            init_kwargs = toml.load(args.log_tracker_config)
-        accelerator.init_trackers("controlnet_train" if args.log_tracker_name is None else args.log_tracker_name, init_kwargs=init_kwargs)
-
-    loss_list = []
-    loss_total = 0.0
-    del train_dataset_group
-
-    # function for saving/removing
-    def save_model(ckpt_name, model, force_sync_upload=False):
-        os.makedirs(args.output_dir, exist_ok=True)
-        ckpt_file = os.path.join(args.output_dir, ckpt_name)
-
-        accelerator.print(f"\nsaving checkpoint: {ckpt_file}")
-
-        state_dict = model_util.convert_controlnet_state_dict_to_sd(model.state_dict())
-
-        if save_dtype is not None:
-            for key in list(state_dict.keys()):
-                v = state_dict[key]
-                v = v.detach().clone().to("cpu").to(save_dtype)
-                state_dict[key] = v
-
-        if os.path.splitext(ckpt_file)[1] == ".safetensors":
-            from safetensors.torch import save_file
-
-            save_file(state_dict, ckpt_file)
-        else:
-            torch.save(state_dict, ckpt_file)
-
-        if args.huggingface_repo_id is not None:
-            huggingface_util.upload(args, ckpt_file, "/" + ckpt_name, force_sync_upload=force_sync_upload)
-
-    def remove_model(old_ckpt_name):
-        old_ckpt_file = os.path.join(args.output_dir, old_ckpt_name)
-        if os.path.exists(old_ckpt_file):
-            accelerator.print(f"removing old checkpoint: {old_ckpt_file}")
-            os.remove(old_ckpt_file)
-
-    # training loop
-    for epoch in range(num_train_epochs):
-        if is_main_process:
-            accelerator.print(f"\nepoch {epoch+1}/{num_train_epochs}")
-        current_epoch.value = epoch + 1
-
-        for step, batch in enumerate(train_dataloader):
-            current_step.value = global_step
-            with accelerator.accumulate(controlnet):
-                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]
-
-                input_ids = batch["input_ids"].to(accelerator.device)
-                encoder_hidden_states = train_util.get_hidden_states(args, input_ids, tokenizer, text_encoder, weight_dtype)
-
-                # Sample noise that we'll add to the latents
-                noise = torch.randn_like(latents, device=latents.device)
-                if args.noise_offset:
-                    noise = apply_noise_offset(latents, noise, args.noise_offset, args.adaptive_noise_scale)
-                elif args.multires_noise_iterations:
-                    noise = pyramid_noise_like(
-                        noise,
-                        latents.device,
-                        args.multires_noise_iterations,
-                        args.multires_noise_discount,
-                    )
-
-                # 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)
-
-                controlnet_image = batch["conditioning_images"].to(dtype=weight_dtype)
-
-                with accelerator.autocast():
-                    down_block_res_samples, mid_block_res_sample = controlnet(
-                        noisy_latents,
-                        timesteps,
-                        encoder_hidden_states=encoder_hidden_states,
-                        controlnet_cond=controlnet_image,
-                        return_dict=False,
-                    )
-
-                    # Predict the noise residual
-                    noise_pred = unet(
-                        noisy_latents,
-                        timesteps,
-                        encoder_hidden_states,
-                        down_block_additional_residuals=[sample.to(dtype=weight_dtype) for sample in down_block_res_samples],
-                        mid_block_additional_residual=mid_block_res_sample.to(dtype=weight_dtype),
-                    ).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="none")
-                loss = loss.mean([1, 2, 3])
-
-                loss_weights = batch["loss_weights"]  # 各sampleごとのweight
-                loss = loss * loss_weights
-
-                if args.min_snr_gamma:
-                    loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma)
-
-                loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし
-
-                accelerator.backward(loss)
-                if accelerator.sync_gradients and args.max_grad_norm != 0.0:
-                    params_to_clip = controlnet.parameters()
-                    accelerator.clip_grad_norm_(params_to_clip, 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
-
-                train_util.sample_images(
-                    accelerator,
-                    args,
-                    None,
-                    global_step,
-                    accelerator.device,
-                    vae,
-                    tokenizer,
-                    text_encoder,
-                    unet,
-                    controlnet=controlnet,
-                )
-
-                # 指定ステップごとにモデルを保存
-                if args.save_every_n_steps is not None and global_step % args.save_every_n_steps == 0:
-                    accelerator.wait_for_everyone()
-                    if accelerator.is_main_process:
-                        ckpt_name = train_util.get_step_ckpt_name(args, "." + args.save_model_as, global_step)
-                        save_model(
-                            ckpt_name,
-                            accelerator.unwrap_model(controlnet),
-                        )
-
-                        if args.save_state:
-                            train_util.save_and_remove_state_stepwise(args, accelerator, global_step)
-
-                        remove_step_no = train_util.get_remove_step_no(args, global_step)
-                        if remove_step_no is not None:
-                            remove_ckpt_name = train_util.get_step_ckpt_name(args, "." + args.save_model_as, remove_step_no)
-                            remove_model(remove_ckpt_name)
-
-            current_loss = loss.detach().item()
-            if epoch == 0:
-                loss_list.append(current_loss)
-            else:
-                loss_total -= loss_list[step]
-                loss_list[step] = current_loss
-            loss_total += current_loss
-            avr_loss = loss_total / len(loss_list)
-            logs = {"loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
-            progress_bar.set_postfix(**logs)
-
-            if args.logging_dir is not None:
-                logs = generate_step_logs(args, current_loss, avr_loss, lr_scheduler)
-                accelerator.log(logs, step=global_step)
-
-            if global_step >= args.max_train_steps:
-                break
-
-        if args.logging_dir is not None:
-            logs = {"loss/epoch": loss_total / len(loss_list)}
-            accelerator.log(logs, step=epoch + 1)
-
-        accelerator.wait_for_everyone()
-
-        # 指定エポックごとにモデルを保存
-        if args.save_every_n_epochs is not None:
-            saving = (epoch + 1) % args.save_every_n_epochs == 0 and (epoch + 1) < num_train_epochs
-            if is_main_process and saving:
-                ckpt_name = train_util.get_epoch_ckpt_name(args, "." + args.save_model_as, epoch + 1)
-                save_model(ckpt_name, accelerator.unwrap_model(controlnet))
-
-                remove_epoch_no = train_util.get_remove_epoch_no(args, epoch + 1)
-                if remove_epoch_no is not None:
-                    remove_ckpt_name = train_util.get_epoch_ckpt_name(args, "." + args.save_model_as, remove_epoch_no)
-                    remove_model(remove_ckpt_name)
-
-                if args.save_state:
-                    train_util.save_and_remove_state_on_epoch_end(args, accelerator, epoch + 1)
-
-        train_util.sample_images(
-            accelerator,
-            args,
-            epoch + 1,
-            global_step,
-            accelerator.device,
-            vae,
-            tokenizer,
-            text_encoder,
-            unet,
-            controlnet=controlnet,
-        )
-
-        # end of epoch
-    if is_main_process:
-        controlnet = accelerator.unwrap_model(controlnet)
-
-    accelerator.end_training()
-
-    if is_main_process and args.save_state:
-        train_util.save_state_on_train_end(args, accelerator)
-
-    # del accelerator  # この後メモリを使うのでこれは消す→printで使うので消さずにおく
-
-    if is_main_process:
-        ckpt_name = train_util.get_last_ckpt_name(args, "." + args.save_model_as)
-        save_model(ckpt_name, controlnet, force_sync_upload=True)
-
-        print("model saved.")
-
-
-def setup_parser() -> argparse.ArgumentParser:
-    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_optimizer_arguments(parser)
-    config_util.add_config_arguments(parser)
-    custom_train_functions.add_custom_train_arguments(parser)
-
-    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(
-        "--controlnet_model_name_or_path",
-        type=str,
-        default=None,
-        help="controlnet model name or path / controlnetのモデル名またはパス",
-    )
-    parser.add_argument(
-        "--conditioning_data_dir",
-        type=str,
-        default=None,
-        help="conditioning data directory / 条件付けデータのディレクトリ",
-    )
-
-    return parser
-
-
-if __name__ == "__main__":
-    parser = setup_parser()
-
-    args = parser.parse_args()
-    args = train_util.read_config_from_file(args, parser)
-
-    train(args)

train_db.py

@@ -2,16 +2,18 @@
 # XXX dropped option: fine_tune
 
 import gc
+import time
 import argparse
 import itertools
 import math
 import os
-from multiprocessing import Value
 import toml
+from multiprocessing import Value
 
 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
@@ -46,7 +48,7 @@ def train(args):
 
     # データセットを準備する
     if args.dataset_class is None:
-        blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, False, False, True))
+        blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, False, True))
         if args.dataset_config is not None:
             print(f"Load dataset config from {args.dataset_config}")
             user_config = config_util.load_user_config(args.dataset_config)
@@ -97,7 +99,7 @@ def train(args):
             f"gradient_accumulation_stepsが{args.gradient_accumulation_steps}に設定されています。accelerateは複数モデル（U-NetおよびText Encoder）の学習時にgradient_accumulation_stepsをサポートしていないため結果は未知数です"
         )
 
-    accelerator = train_util.prepare_accelerator(args)
+    accelerator, unwrap_model = train_util.prepare_accelerator(args)
 
     # mixed precisionに対応した型を用意しておき適宜castする
     weight_dtype, save_dtype = train_util.prepare_dtype(args)
@@ -121,7 +123,7 @@ def train(args):
         use_safetensors = args.use_safetensors or ("safetensors" in args.save_model_as.lower())
 
     # モデルに xformers とか memory efficient attention を組み込む
-    train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers, args.sdpa)
+    train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
 
     # 学習を準備する
     if cache_latents:
@@ -142,7 +144,7 @@ def train(args):
     unet.requires_grad_(True)  # 念のため追加
     text_encoder.requires_grad_(train_text_encoder)
     if not train_text_encoder:
-        accelerator.print("Text Encoder is not trained.")
+        print("Text Encoder is not trained.")
 
     if args.gradient_checkpointing:
         unet.enable_gradient_checkpointing()
@@ -154,13 +156,12 @@ def train(args):
         vae.to(accelerator.device, dtype=weight_dtype)
 
     # 学習に必要なクラスを準備する
-    accelerator.print("prepare optimizer, data loader etc.")
+    print("prepare optimizer, data loader etc.")
     if train_text_encoder:
-        # wightout list, adamw8bit is crashed
-        trainable_params = list(itertools.chain(unet.parameters(), text_encoder.parameters()))
+        trainable_params = itertools.chain(unet.parameters(), text_encoder.parameters())
     else:
         trainable_params = unet.parameters()
-    
+
     _, _, optimizer = train_util.get_optimizer(args, trainable_params)
 
     # dataloaderを準備する
@@ -180,7 +181,7 @@ def train(args):
         args.max_train_steps = args.max_train_epochs * math.ceil(
             len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps
         )
-        accelerator.print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
+        print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
 
     # データセット側にも学習ステップを送信
     train_dataset_group.set_max_train_steps(args.max_train_steps)
@@ -196,7 +197,7 @@ def train(args):
         assert (
             args.mixed_precision == "fp16"
         ), "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
-        accelerator.print("enable full fp16 training.")
+        print("enable full fp16 training.")
         unet.to(weight_dtype)
         text_encoder.to(weight_dtype)
 
@@ -229,17 +230,15 @@ def train(args):
 
     # 学習する
     total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
-    accelerator.print("running training / 学習開始")
-    accelerator.print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
-    accelerator.print(f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
-    accelerator.print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
-    accelerator.print(f"  num epochs / epoch数: {num_train_epochs}")
-    accelerator.print(f"  batch size per device / バッチサイズ: {args.train_batch_size}")
-    accelerator.print(
-        f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}"
-    )
-    accelerator.print(f"  gradient ccumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
-    accelerator.print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")
+    print("running training / 学習開始")
+    print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
+    print(f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_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
@@ -248,19 +247,14 @@ def train(args):
         beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False
     )
     prepare_scheduler_for_custom_training(noise_scheduler, accelerator.device)
-    if args.zero_terminal_snr:
-        custom_train_functions.fix_noise_scheduler_betas_for_zero_terminal_snr(noise_scheduler)
 
     if accelerator.is_main_process:
-        init_kwargs = {}
-        if args.log_tracker_config is not None:
-            init_kwargs = toml.load(args.log_tracker_config)
-        accelerator.init_trackers("dreambooth" if args.log_tracker_name is None else args.log_tracker_name, init_kwargs=init_kwargs)
+        accelerator.init_trackers("dreambooth" if args.log_tracker_name is None else args.log_tracker_name)
 
     loss_list = []
     loss_total = 0.0
     for epoch in range(num_train_epochs):
-        accelerator.print(f"\nepoch {epoch+1}/{num_train_epochs}")
+        print(f"\nepoch {epoch+1}/{num_train_epochs}")
         current_epoch.value = epoch + 1
 
         # 指定したステップ数までText Encoderを学習する：epoch最初の状態
@@ -273,7 +267,7 @@ def train(args):
             current_step.value = global_step
             # 指定したステップ数でText Encoderの学習を止める
             if global_step == args.stop_text_encoder_training:
-                accelerator.print(f"stop text encoder training at step {global_step}")
+                print(f"stop text encoder training at step {global_step}")
                 if not args.gradient_checkpointing:
                     text_encoder.train(False)
                 text_encoder.requires_grad_(False)
@@ -288,6 +282,15 @@ def train(args):
                     latents = latents * 0.18215
                 b_size = latents.shape[0]
 
+                # Sample noise that we'll add to the latents
+                noise = torch.randn_like(latents, device=latents.device)
+                if args.noise_offset:
+                    noise = apply_noise_offset(latents, noise, args.noise_offset, args.adaptive_noise_scale)
+                elif args.multires_noise_iterations:
+                    noise = pyramid_noise_like(noise, latents.device, args.multires_noise_iterations, args.multires_noise_discount)
+                # elif args.perlin_noise:
+                #     noise = perlin_noise(noise, latents.device, args.perlin_noise)  # only shape of noise is used currently
+
                 # Get the text embedding for conditioning
                 with torch.set_grad_enabled(global_step < args.stop_text_encoder_training):
                     if args.weighted_captions:
@@ -305,9 +308,13 @@ def train(args):
                             args, input_ids, tokenizer, text_encoder, None if not args.full_fp16 else weight_dtype
                         )
 
-                # Sample noise, sample a random timestep for each image, and add noise to the latents,
-                # with noise offset and/or multires noise if specified
-                noise, noisy_latents, timesteps = train_util.get_noise_noisy_latents_and_timesteps(args, noise_scheduler, latents)
+                # 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
                 with accelerator.autocast():
@@ -369,17 +376,15 @@ def train(args):
                             epoch,
                             num_train_epochs,
                             global_step,
-                            accelerator.unwrap_model(text_encoder),
-                            accelerator.unwrap_model(unet),
+                            unwrap_model(text_encoder),
+                            unwrap_model(unet),
                             vae,
                         )
 
             current_loss = loss.detach().item()
             if args.logging_dir is not None:
                 logs = {"loss": current_loss, "lr": float(lr_scheduler.get_last_lr()[0])}
-                if (
-                    args.optimizer_type.lower().startswith("DAdapt".lower()) or args.optimizer_type.lower() == "Prodigy".lower()
-                ):  # tracking d*lr value
+                if args.optimizer_type.lower().startswith("DAdapt".lower()) or args.optimizer_type.lower() == "Prodigy".lower():  # tracking d*lr value
                     logs["lr/d*lr"] = (
                         lr_scheduler.optimizers[0].param_groups[0]["d"] * lr_scheduler.optimizers[0].param_groups[0]["lr"]
                     )
@@ -419,8 +424,8 @@ def train(args):
                     epoch,
                     num_train_epochs,
                     global_step,
-                    accelerator.unwrap_model(text_encoder),
-                    accelerator.unwrap_model(unet),
+                    unwrap_model(text_encoder),
+                    unwrap_model(unet),
                     vae,
                 )
 
@@ -428,8 +433,8 @@ def train(args):
 
     is_main_process = accelerator.is_main_process
     if is_main_process:
-        unet = accelerator.unwrap_model(unet)
-        text_encoder = accelerator.unwrap_model(text_encoder)
+        unet = unwrap_model(unet)
+        text_encoder = unwrap_model(text_encoder)
 
     accelerator.end_training()
 

train_textual_inversion_XTI.py

@@ -11,7 +11,6 @@ import torch
 from accelerate.utils import set_seed
 import diffusers
 from diffusers import DDPMScheduler
-import library
 
 import library.train_util as train_util
 import library.huggingface_util as huggingface_util
@@ -28,7 +27,6 @@ from library.custom_train_functions import (
     apply_noise_offset,
     scale_v_prediction_loss_like_noise_prediction,
 )
-import library.original_unet as original_unet
 from XTI_hijack import unet_forward_XTI, downblock_forward_XTI, upblock_forward_XTI
 
 imagenet_templates_small = [
@@ -109,7 +107,7 @@ def train(args):
 
     # acceleratorを準備する
     print("prepare accelerator")
-    accelerator = train_util.prepare_accelerator(args)
+    accelerator, unwrap_model = train_util.prepare_accelerator(args)
 
     # mixed precisionに対応した型を用意しておき適宜castする
     weight_dtype, save_dtype = train_util.prepare_dtype(args)
@@ -189,7 +187,7 @@ def train(args):
     print(f"create embeddings for {args.num_vectors_per_token} tokens, for {args.token_string}")
 
     # データセットを準備する
-    blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, False, False))
+    blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, False))
     if args.dataset_config is not None:
         print(f"Load dataset config from {args.dataset_config}")
         user_config = config_util.load_user_config(args.dataset_config)
@@ -234,7 +232,7 @@ def train(args):
 
     # make captions: tokenstring tokenstring1 tokenstring2 ...tokenstringn という文字列に書き換える超乱暴な実装
     if use_template:
-        print(f"use template for training captions. is object: {args.use_object_template}")
+        print("use template for training captions. is object: {args.use_object_template}")
         templates = imagenet_templates_small if args.use_object_template else imagenet_style_templates_small
         replace_to = " ".join(token_strings)
         captions = []
@@ -267,10 +265,10 @@ def train(args):
         ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
 
     # モデルに xformers とか memory efficient attention を組み込む
-    train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers, args.sdpa)
-    original_unet.UNet2DConditionModel.forward = unet_forward_XTI
-    original_unet.CrossAttnDownBlock2D.forward = downblock_forward_XTI
-    original_unet.CrossAttnUpBlock2D.forward = upblock_forward_XTI
+    train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
+    diffusers.models.UNet2DConditionModel.forward = unet_forward_XTI
+    diffusers.models.unet_2d_blocks.CrossAttnDownBlock2D.forward = downblock_forward_XTI
+    diffusers.models.unet_2d_blocks.CrossAttnUpBlock2D.forward = upblock_forward_XTI
 
     # 学習を準備する
     if cache_latents:
@@ -330,7 +328,7 @@ def train(args):
 
     index_no_updates = torch.arange(len(tokenizer)) < token_ids_XTI[0]
     # print(len(index_no_updates), torch.sum(index_no_updates))
-    orig_embeds_params = accelerator.unwrap_model(text_encoder).get_input_embeddings().weight.data.detach().clone()
+    orig_embeds_params = unwrap_model(text_encoder).get_input_embeddings().weight.data.detach().clone()
 
     # Freeze all parameters except for the token embeddings in text encoder
     text_encoder.requires_grad_(True)
@@ -384,14 +382,9 @@ def train(args):
         beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False
     )
     prepare_scheduler_for_custom_training(noise_scheduler, accelerator.device)
-    if args.zero_terminal_snr:
-        custom_train_functions.fix_noise_scheduler_betas_for_zero_terminal_snr(noise_scheduler)
 
     if accelerator.is_main_process:
-        init_kwargs = {}
-        if args.log_tracker_config is not None:
-            init_kwargs = toml.load(args.log_tracker_config)
-        accelerator.init_trackers("textual_inversion" if args.log_tracker_name is None else args.log_tracker_name, init_kwargs=init_kwargs)
+        accelerator.init_trackers("textual_inversion" if args.log_tracker_name is None else args.log_tracker_name)
 
     # function for saving/removing
     def save_model(ckpt_name, embs, steps, epoch_no, force_sync_upload=False):
@@ -440,9 +433,20 @@ def train(args):
                     ]
                 )
 
-                # Sample noise, sample a random timestep for each image, and add noise to the latents,
-                # with noise offset and/or multires noise if specified
-                noise, noisy_latents, timesteps = train_util.get_noise_noisy_latents_and_timesteps(args, noise_scheduler, latents)
+                # Sample noise that we'll add to the latents
+                noise = torch.randn_like(latents, device=latents.device)
+                if args.noise_offset:
+                    noise = apply_noise_offset(latents, noise, args.noise_offset, args.adaptive_noise_scale)
+                elif args.multires_noise_iterations:
+                    noise = pyramid_noise_like(noise, latents.device, args.multires_noise_iterations, args.multires_noise_discount)
+
+                # 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
                 with accelerator.autocast():
@@ -478,7 +482,7 @@ def train(args):
 
                 # Let's make sure we don't update any embedding weights besides the newly added token
                 with torch.no_grad():
-                    accelerator.unwrap_model(text_encoder).get_input_embeddings().weight[index_no_updates] = orig_embeds_params[
+                    unwrap_model(text_encoder).get_input_embeddings().weight[index_no_updates] = orig_embeds_params[
                         index_no_updates
                     ]
 
@@ -495,13 +499,7 @@ def train(args):
                 if args.save_every_n_steps is not None and global_step % args.save_every_n_steps == 0:
                     accelerator.wait_for_everyone()
                     if accelerator.is_main_process:
-                        updated_embs = (
-                            accelerator.unwrap_model(text_encoder)
-                            .get_input_embeddings()
-                            .weight[token_ids_XTI]
-                            .data.detach()
-                            .clone()
-                        )
+                        updated_embs = unwrap_model(text_encoder).get_input_embeddings().weight[token_ids_XTI].data.detach().clone()
 
                         ckpt_name = train_util.get_step_ckpt_name(args, "." + args.save_model_as, global_step)
                         save_model(ckpt_name, updated_embs, global_step, epoch)
@@ -517,9 +515,7 @@ def train(args):
             current_loss = loss.detach().item()
             if args.logging_dir is not None:
                 logs = {"loss": current_loss, "lr": float(lr_scheduler.get_last_lr()[0])}
-                if (
-                    args.optimizer_type.lower().startswith("DAdapt".lower()) or args.optimizer_type.lower() == "Prodigy".lower()
-                ):  # tracking d*lr value
+                if args.optimizer_type.lower().startswith("DAdapt".lower()) or args.optimizer_type.lower() == "Prodigy".lower():  # tracking d*lr value
                     logs["lr/d*lr"] = (
                         lr_scheduler.optimizers[0].param_groups[0]["d"] * lr_scheduler.optimizers[0].param_groups[0]["lr"]
                     )
@@ -539,7 +535,7 @@ def train(args):
 
         accelerator.wait_for_everyone()
 
-        updated_embs = accelerator.unwrap_model(text_encoder).get_input_embeddings().weight[token_ids_XTI].data.detach().clone()
+        updated_embs = unwrap_model(text_encoder).get_input_embeddings().weight[token_ids_XTI].data.detach().clone()
 
         if args.save_every_n_epochs is not None:
             saving = (epoch + 1) % args.save_every_n_epochs == 0 and (epoch + 1) < num_train_epochs
@@ -564,7 +560,7 @@ def train(args):
 
     is_main_process = accelerator.is_main_process
     if is_main_process:
-        text_encoder = accelerator.unwrap_model(text_encoder)
+        text_encoder = unwrap_model(text_encoder)
 
     accelerator.end_training()