multi embed training

.github/workflows/typos.yml

@@ -18,4 +18,4 @@ jobs:
       - uses: actions/checkout@v4
 
       - name: typos-action
-        uses: crate-ci/typos@v1.16.26
+        uses: crate-ci/typos@v1.16.15

README.md

@@ -249,117 +249,65 @@ ControlNet-LLLite, a novel method for ControlNet with SDXL, is added. See [docum
 
 ## Change History
 
-### Feb 24, 2024 / 2024/2/24: v0.8.4
+### Oct 9. 2023 / 2023/10/9
 
-- The log output has been improved. PR [#905](https://github.com/kohya-ss/sd-scripts/pull/905) Thanks to shirayu!
-  - The log is formatted by default. The `rich` library is required. Please see [Upgrade](#upgrade) and update the library.
-  - If `rich` is not installed, the log output will be the same as before.
-  - The following options are available in each training script:
-  - `--console_log_simple` option can be used to switch to the previous log output.
-  - `--console_log_level` option can be used to specify the log level. The default is `INFO`.
-  - `--console_log_file` option can be used to output the log to a file. The default is `None` (output to the console).
-- The sample image generation during multi-GPU training is now done with multiple GPUs. PR [#1061](https://github.com/kohya-ss/sd-scripts/pull/1061) Thanks to DKnight54!
-- The support for mps devices is improved. PR [#1054](https://github.com/kohya-ss/sd-scripts/pull/1054) Thanks to akx! If mps device exists instead of CUDA, the mps device is used automatically.
-- The `--new_conv_rank` option to specify the new rank of Conv2d is added to `networks/resize_lora.py`. PR [#1102](https://github.com/kohya-ss/sd-scripts/pull/1102) Thanks to mgz-dev!
-- An option `--highvram` to disable the optimization for environments with little VRAM is added to the training scripts. If you specify it when there is enough VRAM, the operation will be faster.
-  - Currently, only the cache part of latents is optimized.
-- The IPEX support is improved. PR [#1086](https://github.com/kohya-ss/sd-scripts/pull/1086) Thanks to Disty0!
-- Fixed a bug that `svd_merge_lora.py` crashes in some cases. PR [#1087](https://github.com/kohya-ss/sd-scripts/pull/1087) Thanks to mgz-dev!
-- DyLoRA is fixed to work with SDXL. PR [#1126](https://github.com/kohya-ss/sd-scripts/pull/1126) Thanks to tamlog06!
-- The common image generation script `gen_img.py` for SD 1/2 and SDXL is added. The basic functions are the same as the scripts for SD 1/2 and SDXL, but some new features are added.
-  - External scripts to generate prompts can be supported. It can be called with `--from_module` option. (The documentation will be added later)
-  - The normalization method after prompt weighting can be specified with `--emb_normalize_mode` option. `original` is the original method, `abs` is the normalization with the average of the absolute values, `none` is no normalization.
-- Gradual Latent Hires fix is added to each generation script. See [here](./docs/gen_img_README-ja.md#about-gradual-latent) for details.
+- `tag_images_by_wd_14_tagger.py` now supports Onnx. If you use Onnx, TensorFlow is not required anymore. [#864](https://github.com/kohya-ss/sd-scripts/pull/864) Thanks to Isotr0py!
+  - `--onnx` option is added. If you use Onnx, specify `--onnx` option.
+  - Please install Onnx and other required packages. 
+    1. Uninstall TensorFlow.
+    1. `pip install tensorboard==2.14.1` This is required for the specified version of protobuf.
+    1. `pip install protobuf==3.20.3` This is required for Onnx.
+    1. `pip install onnx==1.14.1`
+    1. `pip install onnxruntime-gpu==1.16.0` or `pip install onnxruntime==1.16.0`
+- `--append_tags` option is added to `tag_images_by_wd_14_tagger.py`. This option appends the tags to the existing tags, instead of replacing them. [#858](https://github.com/kohya-ss/sd-scripts/pull/858) Thanks to a-l-e-x-d-s-9! 
+- [OFT](https://oft.wyliu.com/) is now supported.
+  - You can use `networks.oft` for the network module in `sdxl_train_network.py`.  The usage is the same as `networks.lora`. Some options are not supported.
+  - `sdxl_gen_img.py` also supports OFT as `--network_module`. 
+  - OFT only supports SDXL currently. Because current OFT tweaks Q/K/V and O in the transformer, and SD1/2 have extremely fewer transformers than SDXL.
+  - The implementation is heavily based on laksjdjf's [OFT implementation](https://github.com/laksjdjf/sd-trainer/blob/dev/networks/lora_modules.py). Thanks to laksjdjf!
+- Other bug fixes and improvements.
 
-- ログ出力が改善されました。 PR [#905](https://github.com/kohya-ss/sd-scripts/pull/905) shirayu 氏に感謝します。
-  - デフォルトでログが成形されます。`rich` ライブラリが必要なため、[Upgrade](#upgrade) を参照し更新をお願いします。
-  - `rich` がインストールされていない場合は、従来のログ出力になります。
-  - 各学習スクリプトでは以下のオプションが有効です。
-  - `--console_log_simple` オプションで従来のログ出力に切り替えられます。
-  - `--console_log_level` でログレベルを指定できます。デフォルトは `INFO` です。
-  - `--console_log_file` でログファイルを出力できます。デフォルトは `None`（コンソールに出力） です。
-- 複数 GPU 学習時に学習中のサンプル画像生成を複数 GPU で行うようになりました。 PR [#1061](https://github.com/kohya-ss/sd-scripts/pull/1061) DKnight54 氏に感謝します。
-- mps デバイスのサポートが改善されました。 PR [#1054](https://github.com/kohya-ss/sd-scripts/pull/1054) akx 氏に感謝します。CUDA ではなく mps が存在する場合には自動的に mps デバイスを使用します。
-- `networks/resize_lora.py` に Conv2d の新しいランクを指定するオプション `--new_conv_rank` が追加されました。 PR [#1102](https://github.com/kohya-ss/sd-scripts/pull/1102) mgz-dev 氏に感謝します。
-- 学習スクリプトに VRAMが少ない環境向け最適化を無効にするオプション `--highvram` を追加しました。VRAM に余裕がある場合に指定すると動作が高速化されます。
-  - 現在は latents のキャッシュ部分のみ高速化されます。
-- IPEX サポートが改善されました。 PR [#1086](https://github.com/kohya-ss/sd-scripts/pull/1086) Disty0 氏に感謝します。
-- `svd_merge_lora.py` が場合によってエラーになる不具合が修正されました。 PR [#1087](https://github.com/kohya-ss/sd-scripts/pull/1087) mgz-dev 氏に感謝します。
-- DyLoRA が SDXL で動くよう修正されました。PR [#1126](https://github.com/kohya-ss/sd-scripts/pull/1126) tamlog06 氏に感謝します。
-- SD 1/2 および SDXL 共通の生成スクリプト `gen_img.py` を追加しました。基本的な機能は SD 1/2、SDXL 向けスクリプトと同じですが、いくつかの新機能が追加されています。
-  - プロンプトを動的に生成する外部スクリプトをサポートしました。 `--from_module` で呼び出せます。（ドキュメントはのちほど追加します）
-  - プロンプト重みづけ後の正規化方法を `--emb_normalize_mode` で指定できます。`original` は元の方法、`abs` は絶対値の平均値で正規化、`none` は正規化を行いません。
-- Gradual Latent Hires fix を各生成スクリプトに追加しました。詳細は [こちら](./docs/gen_img_README-ja.md#about-gradual-latent)。
+- `tag_images_by_wd_14_tagger.py` が Onnx をサポートしました。Onnx を使用する場合は TensorFlow は不要です。[#864](https://github.com/kohya-ss/sd-scripts/pull/864) Isotr0py氏に感謝します。
+  - Onnxを使用する場合は、`--onnx` オプションを指定してください。
+  - Onnx とその他の必要なパッケージをインストールしてください。
+    1. TensorFlow をアンインストールしてください。
+    1. `pip install tensorboard==2.14.1` protobufの指定バージョンにこれが必要。
+    1. `pip install protobuf==3.20.3` Onnxのために必要。
+    1. `pip install onnx==1.14.1`
+    1. `pip install onnxruntime-gpu==1.16.0` または `pip install onnxruntime==1.16.0`
+- `tag_images_by_wd_14_tagger.py` に `--append_tags` オプションが追加されました。このオプションを指定すると、既存のタグに上書きするのではなく、新しいタグのみが既存のタグに追加されます。 [#858](https://github.com/kohya-ss/sd-scripts/pull/858) a-l-e-x-d-s-9氏に感謝します。
+- [OFT](https://oft.wyliu.com/) をサポートしました。
+  - `sdxl_train_network.py` の`--network_module`に `networks.oft` を指定してください。使用方法は `networks.lora` と同様ですが一部のオプションは未サポートです。
+  - `sdxl_gen_img.py` でも同様に  OFT を指定できます。
+  - OFT は現在 SDXL のみサポートしています。OFT は現在 transformer の Q/K/V と O を変更しますが、SD1/2 は transformer の数が SDXL よりも極端に少ないためです。
+  - 実装は laksjdjf 氏の [OFT実装](https://github.com/laksjdjf/sd-trainer/blob/dev/networks/lora_modules.py) を多くの部分で参考にしています。laksjdjf 氏に感謝します。
+- その他のバグ修正と改善。
 
+### Oct 1. 2023 / 2023/10/1
 
-### Jan 27, 2024 / 2024/1/27: v0.8.3
+- SDXL training is now available in the main branch. The sdxl branch is merged into the main branch.
 
-- Fixed a bug that the training crashes when `--fp8_base` is specified with `--save_state`. PR [#1079](https://github.com/kohya-ss/sd-scripts/pull/1079) Thanks to feffy380!
-  - `safetensors` is updated. Please see [Upgrade](#upgrade) and update the library.
-- Fixed a bug that the training crashes when `network_multiplier` is specified with multi-GPU training. PR [#1084](https://github.com/kohya-ss/sd-scripts/pull/1084) Thanks to fireicewolf!
-- Fixed a bug that the training crashes when training ControlNet-LLLite.
+- [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.
 
-- `--fp8_base` 指定時に `--save_state` での保存がエラーになる不具合が修正されました。 PR [#1079](https://github.com/kohya-ss/sd-scripts/pull/1079) feffy380 氏に感謝します。
-  - `safetensors` がバージョンアップされていますので、[Upgrade](#upgrade) を参照し更新をお願いします。
-- 複数 GPU での学習時に `network_multiplier` を指定するとクラッシュする不具合が修正されました。 PR [#1084](https://github.com/kohya-ss/sd-scripts/pull/1084) fireicewolf 氏に感謝します。
-- ControlNet-LLLite の学習がエラーになる不具合を修正しました。 
+- `bitsandbytes` is now optional. Please install it if you want to use it. The insructions are in the later section.
 
-### Jan 23, 2024 / 2024/1/23: v0.8.2
+- `albumentations` is not required anymore.
 
-- [Experimental] The `--fp8_base` option is added to the training scripts for LoRA etc. The base model (U-Net, and Text Encoder when training modules for Text Encoder) can be trained with fp8. PR [#1057](https://github.com/kohya-ss/sd-scripts/pull/1057) Thanks to KohakuBlueleaf!
-  - Please specify `--fp8_base` in `train_network.py` or `sdxl_train_network.py`.
-  - PyTorch 2.1 or later is required.
-  - If you use xformers with PyTorch 2.1, please see [xformers repository](https://github.com/facebookresearch/xformers) and install the appropriate version according to your CUDA version.
-  - The sample image generation during training consumes a lot of memory. It is recommended to turn it off.
+- `--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.
 
-- [Experimental] The network multiplier can be specified for each dataset in the training scripts for LoRA etc.
-  - This is an experimental option and may be removed or changed in the future.
-  - For example, if you train with state A as `1.0` and state B as `-1.0`, you may be able to generate by switching between state A and B depending on the LoRA application rate.
-  - Also, if you prepare five states and train them as `0.2`, `0.4`, `0.6`, `0.8`, and `1.0`, you may be able to generate by switching the states smoothly depending on the application rate.
-  - Please specify `network_multiplier` in `[[datasets]]` in `.toml` file.
-- Some options are added to `networks/extract_lora_from_models.py` to reduce the memory usage.
-  - `--load_precision` option can be used to specify the precision when loading the model. If the model is saved in fp16, you can reduce the memory usage by specifying `--load_precision fp16` without losing precision.
-  - `--load_original_model_to` option can be used to specify the device to load the original model. `--load_tuned_model_to` option can be used to specify the device to load the derived model. The default is `cpu` for both options, but you can specify `cuda` etc. You can reduce the memory usage by loading one of them to GPU. This option is available only for SDXL.
+- Arbitrary options can be used for Diffusers' schedulers. For example `--lr_scheduler_args "lr_end=1e-8"`.
 
-- The gradient synchronization in LoRA training with multi-GPU is improved. PR [#1064](https://github.com/kohya-ss/sd-scripts/pull/1064) Thanks to KohakuBlueleaf!
-- The code for Intel IPEX support is improved. PR [#1060](https://github.com/kohya-ss/sd-scripts/pull/1060) Thanks to akx!
-- Fixed a bug in multi-GPU Textual Inversion training.
-
-- （実験的）　LoRA等の学習スクリプトで、ベースモデル（U-Net、および Text Encoder のモジュール学習時は Text Encoder も）の重みを fp8 にして学習するオプションが追加されました。 PR [#1057](https://github.com/kohya-ss/sd-scripts/pull/1057) KohakuBlueleaf 氏に感謝します。
-  - `train_network.py` または `sdxl_train_network.py` で `--fp8_base` を指定してください。
-  - PyTorch 2.1 以降が必要です。
-  - PyTorch 2.1 で xformers を使用する場合は、[xformers のリポジトリ](https://github.com/facebookresearch/xformers) を参照し、CUDA バージョンに応じて適切なバージョンをインストールしてください。
-  - 学習中のサンプル画像生成はメモリを大量に消費するため、オフにすることをお勧めします。
-- (実験的)　LoRA 等の学習で、データセットごとに異なるネットワーク適用率を指定できるようになりました。 
-  - 実験的オプションのため、将来的に削除または仕様変更される可能性があります。
-  - たとえば状態 A を `1.0`、状態 B を `-1.0` として学習すると、LoRA の適用率に応じて状態 A と B を切り替えつつ生成できるかもしれません。
-  - また、五段階の状態を用意し、それぞれ `0.2`、`0.4`、`0.6`、`0.8`、`1.0` として学習すると、適用率でなめらかに状態を切り替えて生成できるかもしれません。 
-  - `.toml` ファイルで `[[datasets]]` に `network_multiplier` を指定してください。
-- `networks/extract_lora_from_models.py` に使用メモリ量を削減するいくつかのオプションを追加しました。 
-  - `--load_precision` で読み込み時の精度を指定できます。モデルが fp16 で保存されている場合は `--load_precision fp16` を指定して精度を変えずにメモリ量を削減できます。
-  - `--load_original_model_to` で元モデルを読み込むデバイスを、`--load_tuned_model_to` で派生モデルを読み込むデバイスを指定できます。デフォルトは両方とも `cpu` ですがそれぞれ `cuda` 等を指定できます。片方を GPU に読み込むことでメモリ量を削減できます。SDXL の場合のみ有効です。
-- マルチ GPU での LoRA 等の学習時に勾配の同期が改善されました。 PR [#1064](https://github.com/kohya-ss/sd-scripts/pull/1064) KohakuBlueleaf 氏に感謝します。
-- Intel IPEX サポートのコードが改善されました。PR [#1060](https://github.com/kohya-ss/sd-scripts/pull/1060) akx 氏に感謝します。
-- マルチ GPU での Textual Inversion 学習の不具合を修正しました。
-
-- `.toml` example for network multiplier / ネットワーク適用率の `.toml` の記述例
-
-```toml
-[general]
-[[datasets]]
-resolution = 512
-batch_size = 8
-network_multiplier = 1.0
-
-... subset settings ...
-
-[[datasets]]
-resolution = 512
-batch_size = 8
-network_multiplier = -1.0
-
-... subset settings ...
-```
+- 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.
+- JPEG XL is supported. [#786](https://github.com/kohya-ss/sd-scripts/pull/786) 
+- 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`.
+- Intel ARC support with IPEX is added. [#825](https://github.com/kohya-ss/sd-scripts/pull/825)
+- Other bug fixes and improvements.
 
 
 Please read [Releases](https://github.com/kohya-ss/sd-scripts/releases) for recent updates.

XTI_hijack.py

@@ -1,7 +1,11 @@
 import torch
-from library.device_utils import init_ipex
-init_ipex()
-
+try:
+    import intel_extension_for_pytorch as ipex
+    if torch.xpu.is_available():
+        from library.ipex import ipex_init
+        ipex_init()
+except Exception:
+    pass
 from typing import Union, List, Optional, Dict, Any, Tuple
 from diffusers.models.unet_2d_condition import UNet2DConditionOutput
 

docs/gen_img_README-ja.md

@@ -452,36 +452,3 @@ python gen_img_diffusers.py --ckpt wd-v1-3-full-pruned-half.ckpt
 
 - `--network_show_meta` : 追加ネットワークのメタデータを表示します。
 
-
---- 
-
-# About Gradual Latent
-
-Gradual Latent is a Hires fix that gradually increases the size of the latent.  `gen_img.py`, `sdxl_gen_img.py`, and `gen_img_diffusers.py` have the following options.
-
-- `--gradual_latent_timesteps`: Specifies the timestep to start increasing the size of the latent. The default is None, which means Gradual Latent is not used. Please try around 750 at first.
-- `--gradual_latent_ratio`: Specifies the initial size of the latent. The default is 0.5, which means it starts with half the default latent size.
-- `--gradual_latent_ratio_step`: Specifies the ratio to increase the size of the latent. The default is 0.125, which means the latent size is gradually increased to 0.625, 0.75, 0.875, 1.0.
-- `--gradual_latent_ratio_every_n_steps`: Specifies the interval to increase the size of the latent. The default is 3, which means the latent size is increased every 3 steps.
-
-Each option can also be specified with prompt options, `--glt`, `--glr`, `--gls`, `--gle`.
-
-__Please specify `euler_a` for the sampler.__ Because the source code of the sampler is modified. It will not work with other samplers.
-
-It is more effective with SD 1.5. It is quite subtle with SDXL.
-
-# Gradual Latent について
-
-latentのサイズを徐々に大きくしていくHires fixです。`gen_img.py` 、``sdxl_gen_img.py` 、`gen_img_diffusers.py` に以下のオプションが追加されています。
-
-- `--gradual_latent_timesteps` : latentのサイズを大きくし始めるタイムステップを指定します。デフォルトは None で、Gradual Latentを使用しません。750 くらいから始めてみてください。
-- `--gradual_latent_ratio` : latentの初期サイズを指定します。デフォルトは 0.5 で、デフォルトの latent サイズの半分のサイズから始めます。
-- `--gradual_latent_ratio_step`: latentのサイズを大きくする割合を指定します。デフォルトは 0.125 で、latentのサイズを 0.625, 0.75, 0.875, 1.0 と徐々に大きくします。
-- `--gradual_latent_ratio_every_n_steps`: latentのサイズを大きくする間隔を指定します。デフォルトは 3 で、3ステップごとに latent のサイズを大きくします。
-
-それぞれのオプションは、プロンプトオプション、`--glt`、`--glr`、`--gls`、`--gle` でも指定できます。
-
-サンプラーに手を加えているため、__サンプラーに `euler_a` を指定してください。__ 他のサンプラーでは動作しません。
-
-SD 1.5 のほうが効果があります。SDXL ではかなり微妙です。
-

docs/train_README-ja.md

@@ -374,10 +374,6 @@ classがひとつで対象が複数の場合、正則化画像フォルダはひ
     
     サンプル出力するステップ数またはエポック数を指定します。この数ごとにサンプル出力します。両方指定するとエポック数が優先されます。
 
-- `--sample_at_first`
-    
-    学習開始前にサンプル出力します。学習前との比較ができます。
-
 - `--sample_prompts`
 
     サンプル出力用プロンプトのファイルを指定します。

fine_tune.py

@@ -2,27 +2,24 @@
 # XXX dropped option: hypernetwork training
 
 import argparse
+import gc
 import math
 import os
 from multiprocessing import Value
 import toml
 
 from tqdm import tqdm
-
 import torch
-from library.device_utils import init_ipex, clean_memory_on_device
-init_ipex()
-
+try:
+    import intel_extension_for_pytorch as ipex
+    if torch.xpu.is_available():
+        from library.ipex import ipex_init
+        ipex_init()
+except Exception:
+    pass
 from accelerate.utils import set_seed
 from diffusers import DDPMScheduler
 
-from library.utils import setup_logging, add_logging_arguments
-
-setup_logging()
-import logging
-
-logger = logging.getLogger(__name__)
-
 import library.train_util as train_util
 import library.config_util as config_util
 from library.config_util import (
@@ -35,14 +32,12 @@ from library.custom_train_functions import (
     get_weighted_text_embeddings,
     prepare_scheduler_for_custom_training,
     scale_v_prediction_loss_like_noise_prediction,
-    apply_debiased_estimation,
 )
 
 
 def train(args):
     train_util.verify_training_args(args)
     train_util.prepare_dataset_args(args, True)
-    setup_logging(args, reset=True)
 
     cache_latents = args.cache_latents
 
@@ -55,11 +50,11 @@ def train(args):
     if args.dataset_class is None:
         blueprint_generator = BlueprintGenerator(ConfigSanitizer(False, True, False, True))
         if args.dataset_config is not None:
-            logger.info(f"Load dataset config from {args.dataset_config}")
+            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):
-                logger.warning(
+                print(
                     "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
                         ", ".join(ignored)
                     )
@@ -92,7 +87,7 @@ def train(args):
         train_util.debug_dataset(train_dataset_group)
         return
     if len(train_dataset_group) == 0:
-        logger.error(
+        print(
             "No data found. Please verify the metadata file and train_data_dir option. / 画像がありません。メタデータおよびtrain_data_dirオプションを確認してください。"
         )
         return
@@ -103,7 +98,7 @@ def train(args):
         ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
 
     # acceleratorを準備する
-    logger.info("prepare accelerator")
+    print("prepare accelerator")
     accelerator = train_util.prepare_accelerator(args)
 
     # mixed precisionに対応した型を用意しておき適宜castする
@@ -164,7 +159,9 @@ def train(args):
         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")
-        clean_memory_on_device(accelerator.device)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
 
         accelerator.wait_for_everyone()
 
@@ -195,24 +192,18 @@ def train(args):
 
     for m in training_models:
         m.requires_grad_(True)
-
-    trainable_params = []
-    if args.learning_rate_te is None or not args.train_text_encoder:
-        for m in training_models:
-            trainable_params.extend(m.parameters())
-    else:
-        trainable_params = [
-            {"params": list(unet.parameters()), "lr": args.learning_rate},
-            {"params": list(text_encoder.parameters()), "lr": args.learning_rate_te},
-        ]
+    params = []
+    for m in training_models:
+        params.extend(m.parameters())
+    params_to_optimize = params
 
     # 学習に必要なクラスを準備する
     accelerator.print("prepare optimizer, data loader etc.")
-    _, _, optimizer = train_util.get_optimizer(args, trainable_params=trainable_params)
+    _, _, optimizer = train_util.get_optimizer(args, trainable_params=params_to_optimize)
 
     # dataloaderを準備する
-    # DataLoaderのプロセス数：0 は persistent_workers が使えないので注意
-    n_workers = min(args.max_data_loader_n_workers, os.cpu_count())  # cpu_count or max_data_loader_n_workers
+    # 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,
@@ -227,9 +218,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}"
-        )
+        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)
@@ -254,6 +243,9 @@ def train(args):
     else:
         unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(unet, optimizer, train_dataloader, lr_scheduler)
 
+    # transform DDP after prepare
+    text_encoder, unet = train_util.transform_if_model_is_DDP(text_encoder, unet)
+
     # 実験的機能：勾配も含めたfp16学習を行う　PyTorchにパッチを当ててfp16でのgrad scaleを有効にする
     if args.full_fp16:
         train_util.patch_accelerator_for_fp16_training(accelerator)
@@ -292,16 +284,10 @@ def train(args):
 
     if accelerator.is_main_process:
         init_kwargs = {}
-        if args.wandb_run_name:
-            init_kwargs["wandb"] = {"name": args.wandb_run_name}
         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 --sample_at_first
-    train_util.sample_images(accelerator, args, 0, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
-
-    loss_recorder = train_util.LossRecorder()
     for epoch in range(num_train_epochs):
         accelerator.print(f"\nepoch {epoch+1}/{num_train_epochs}")
         current_epoch.value = epoch + 1
@@ -309,6 +295,7 @@ def train(args):
         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]):  # 複数モデルに対応していない模様だがとりあえずこうしておく
@@ -352,17 +339,15 @@ def train(args):
                 else:
                     target = noise
 
-                if args.min_snr_gamma or args.scale_v_pred_loss_like_noise_pred or args.debiased_estimation_loss:
+                if args.min_snr_gamma or args.scale_v_pred_loss_like_noise_pred:
                     # 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, args.v_parameterization)
+                        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.debiased_estimation_loss:
-                        loss = apply_debiased_estimation(loss, timesteps, noise_scheduler)
 
                     loss = loss.mean()  # mean over batch dimension
                 else:
@@ -411,20 +396,26 @@ def train(args):
 
             current_loss = loss.detach().item()  # 平均なのでbatch sizeは関係ないはず
             if args.logging_dir is not None:
-                logs = {"loss": current_loss}
-                train_util.append_lr_to_logs(logs, lr_scheduler, args.optimizer_type, including_unet=True)
+                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
+                    logs["lr/d*lr"] = (
+                        lr_scheduler.optimizers[0].param_groups[0]["d"] * lr_scheduler.optimizers[0].param_groups[0]["lr"]
+                    )
                 accelerator.log(logs, step=global_step)
 
-            loss_recorder.add(epoch=epoch, step=step, loss=current_loss)
-            avr_loss: float = loss_recorder.moving_average
-            logs = {"avr_loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
+            # 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_recorder.moving_average}
+            logs = {"loss/epoch": loss_total / len(train_dataloader)}
             accelerator.log(logs, step=epoch + 1)
 
         accelerator.wait_for_everyone()
@@ -467,13 +458,12 @@ def train(args):
         train_util.save_sd_model_on_train_end(
             args, src_path, save_stable_diffusion_format, use_safetensors, save_dtype, epoch, global_step, text_encoder, unet, vae
         )
-        logger.info("model saved.")
+        print("model saved.")
 
 
 def setup_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser()
 
-    add_logging_arguments(parser)
     train_util.add_sd_models_arguments(parser)
     train_util.add_dataset_arguments(parser, False, True, True)
     train_util.add_training_arguments(parser, False)
@@ -482,16 +472,8 @@ def setup_parser() -> argparse.ArgumentParser:
     config_util.add_config_arguments(parser)
     custom_train_functions.add_custom_train_arguments(parser)
 
-    parser.add_argument(
-        "--diffusers_xformers", action="store_true", help="use xformers by diffusers / Diffusersでxformersを使用する"
-    )
+    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(
-        "--learning_rate_te",
-        type=float,
-        default=None,
-        help="learning rate for text encoder, default is same as unet / Text Encoderの学習率、デフォルトはunetと同じ",
-    )
 
     return parser
 

finetune/blip/blip.py

@@ -21,10 +21,6 @@ import torch.nn.functional as F
 import os
 from urllib.parse import urlparse
 from timm.models.hub import download_cached_file
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
 
 class BLIP_Base(nn.Module):
     def __init__(self,                 
@@ -239,6 +235,6 @@ def load_checkpoint(model,url_or_filename):
                 del state_dict[key]
     
     msg = model.load_state_dict(state_dict,strict=False)
-    logger.info('load checkpoint from %s'%url_or_filename)  
+    print('load checkpoint from %s'%url_or_filename)  
     return model,msg
     

finetune/clean_captions_and_tags.py

@@ -8,10 +8,6 @@ import json
 import re
 
 from tqdm import tqdm
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
 
 PATTERN_HAIR_LENGTH = re.compile(r', (long|short|medium) hair, ')
 PATTERN_HAIR_CUT = re.compile(r', (bob|hime) cut, ')
@@ -40,13 +36,13 @@ def clean_tags(image_key, tags):
   tokens = tags.split(", rating")
   if len(tokens) == 1:
     # WD14 taggerのときはこちらになるのでメッセージは出さない
-    # logger.info("no rating:")
-    # logger.info(f"{image_key} {tags}")
+    # print("no rating:")
+    # print(f"{image_key} {tags}")
     pass
   else:
     if len(tokens) > 2:
-      logger.info("multiple ratings:")
-      logger.info(f"{image_key} {tags}")
+      print("multiple ratings:")
+      print(f"{image_key} {tags}")
     tags = tokens[0]
 
   tags = ", " + tags.replace(", ", ", , ") + ", "     # カンマ付きで検索をするための身も蓋もない対策
@@ -128,43 +124,43 @@ def clean_caption(caption):
 
 def main(args):
   if os.path.exists(args.in_json):
-    logger.info(f"loading existing metadata: {args.in_json}")
+    print(f"loading existing metadata: {args.in_json}")
     with open(args.in_json, "rt", encoding='utf-8') as f:
       metadata = json.load(f)
   else:
-    logger.error("no metadata / メタデータファイルがありません")
+    print("no metadata / メタデータファイルがありません")
     return
 
-  logger.info("cleaning captions and tags.")
+  print("cleaning captions and tags.")
   image_keys = list(metadata.keys())
   for image_key in tqdm(image_keys):
     tags = metadata[image_key].get('tags')
     if tags is None:
-      logger.error(f"image does not have tags / メタデータにタグがありません: {image_key}")
+      print(f"image does not have tags / メタデータにタグがありません: {image_key}")
     else:
       org = tags
       tags = clean_tags(image_key, tags)
       metadata[image_key]['tags'] = tags
       if args.debug and org != tags:
-        logger.info("FROM: " + org)
-        logger.info("TO:   " + tags)
+        print("FROM: " + org)
+        print("TO:   " + tags)
 
     caption = metadata[image_key].get('caption')
     if caption is None:
-      logger.error(f"image does not have caption / メタデータにキャプションがありません: {image_key}")
+      print(f"image does not have caption / メタデータにキャプションがありません: {image_key}")
     else:
       org = caption
       caption = clean_caption(caption)
       metadata[image_key]['caption'] = caption
       if args.debug and org != caption:
-        logger.info("FROM: " + org)
-        logger.info("TO:   " + caption)
+        print("FROM: " + org)
+        print("TO:   " + caption)
 
   # metadataを書き出して終わり
-  logger.info(f"writing metadata: {args.out_json}")
+  print(f"writing metadata: {args.out_json}")
   with open(args.out_json, "wt", encoding='utf-8') as f:
     json.dump(metadata, f, indent=2)
-  logger.info("done!")
+  print("done!")
 
 
 def setup_parser() -> argparse.ArgumentParser:
@@ -182,10 +178,10 @@ if __name__ == '__main__':
 
   args, unknown = parser.parse_known_args()
   if len(unknown) == 1:
-    logger.warning("WARNING: train_data_dir argument is removed. This script will not work with three arguments in future. Please specify two arguments: in_json and out_json.")
-    logger.warning("All captions and tags in the metadata are processed.")
-    logger.warning("警告: train_data_dir引数は不要になりました。将来的には三つの引数を指定すると動かなくなる予定です。読み込み元のメタデータと書き出し先の二つの引数だけ指定してください。")
-    logger.warning("メタデータ内のすべてのキャプションとタグが処理されます。")
+    print("WARNING: train_data_dir argument is removed. This script will not work with three arguments in future. Please specify two arguments: in_json and out_json.")
+    print("All captions and tags in the metadata are processed.")
+    print("警告: train_data_dir引数は不要になりました。将来的には三つの引数を指定すると動かなくなる予定です。読み込み元のメタデータと書き出し先の二つの引数だけ指定してください。")
+    print("メタデータ内のすべてのキャプションとタグが処理されます。")
     args.in_json = args.out_json
     args.out_json = unknown[0]
   elif len(unknown) > 0:

finetune/make_captions.py

@@ -9,22 +9,14 @@ from pathlib import Path
 from PIL import Image
 from tqdm import tqdm
 import numpy as np
-
 import torch
-from library.device_utils import init_ipex, get_preferred_device
-init_ipex()
-
 from torchvision import transforms
 from torchvision.transforms.functional import InterpolationMode
 sys.path.append(os.path.dirname(__file__))
-from blip.blip import blip_decoder, is_url
+from blip.blip import blip_decoder
 import library.train_util as train_util
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
 
-DEVICE = get_preferred_device()
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
 IMAGE_SIZE = 384
@@ -55,7 +47,7 @@ class ImageLoadingTransformDataset(torch.utils.data.Dataset):
             # convert to tensor temporarily so dataloader will accept it
             tensor = IMAGE_TRANSFORM(image)
         except Exception as e:
-            logger.error(f"Could not load image path / 画像を読み込めません: {img_path}, error: {e}")
+            print(f"Could not load image path / 画像を読み込めません: {img_path}, error: {e}")
             return None
 
         return (tensor, img_path)
@@ -82,21 +74,19 @@ def main(args):
         args.train_data_dir = os.path.abspath(args.train_data_dir)  # convert to absolute path
 
         cwd = os.getcwd()
-        logger.info(f"Current Working Directory is: {cwd}")
+        print("Current Working Directory is: ", cwd)
         os.chdir("finetune")
-        if not is_url(args.caption_weights) and not os.path.isfile(args.caption_weights):
-            args.caption_weights = os.path.join("..", args.caption_weights)
 
-    logger.info(f"load images from {args.train_data_dir}")
+    print(f"load images from {args.train_data_dir}")
     train_data_dir_path = Path(args.train_data_dir)
     image_paths = train_util.glob_images_pathlib(train_data_dir_path, args.recursive)
-    logger.info(f"found {len(image_paths)} images.")
+    print(f"found {len(image_paths)} images.")
 
-    logger.info(f"loading BLIP caption: {args.caption_weights}")
+    print(f"loading BLIP caption: {args.caption_weights}")
     model = blip_decoder(pretrained=args.caption_weights, image_size=IMAGE_SIZE, vit="large", med_config="./blip/med_config.json")
     model.eval()
     model = model.to(DEVICE)
-    logger.info("BLIP loaded")
+    print("BLIP loaded")
 
     # captioningする
     def run_batch(path_imgs):
@@ -116,7 +106,7 @@ def main(args):
             with open(os.path.splitext(image_path)[0] + args.caption_extension, "wt", encoding="utf-8") as f:
                 f.write(caption + "\n")
                 if args.debug:
-                    logger.info(f'{image_path} {caption}')
+                    print(image_path, caption)
 
     # 読み込みの高速化のためにDataLoaderを使うオプション
     if args.max_data_loader_n_workers is not None:
@@ -146,7 +136,7 @@ def main(args):
                         raw_image = raw_image.convert("RGB")
                     img_tensor = IMAGE_TRANSFORM(raw_image)
                 except Exception as e:
-                    logger.error(f"Could not load image path / 画像を読み込めません: {image_path}, error: {e}")
+                    print(f"Could not load image path / 画像を読み込めません: {image_path}, error: {e}")
                     continue
 
             b_imgs.append((image_path, img_tensor))
@@ -156,7 +146,7 @@ def main(args):
     if len(b_imgs) > 0:
         run_batch(b_imgs)
 
-    logger.info("done!")
+    print("done!")
 
 
 def setup_parser() -> argparse.ArgumentParser:

finetune/make_captions_by_git.py

@@ -5,19 +5,12 @@ import re
 from pathlib import Path
 from PIL import Image
 from tqdm import tqdm
-
 import torch
-from library.device_utils import init_ipex, get_preferred_device
-init_ipex()
-
 from transformers import AutoProcessor, AutoModelForCausalLM
 from transformers.generation.utils import GenerationMixin
 
 import library.train_util as train_util
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -42,8 +35,8 @@ def remove_words(captions, debug):
         for pat in PATTERN_REPLACE:
             cap = pat.sub("", cap)
         if debug and cap != caption:
-            logger.info(caption)
-            logger.info(cap)
+            print(caption)
+            print(cap)
         removed_caps.append(cap)
     return removed_caps
 
@@ -59,9 +52,6 @@ def collate_fn_remove_corrupted(batch):
 
 
 def main(args):
-    r"""
-    transformers 4.30.2で、バッチサイズ>1でも動くようになったので、以下コメントアウト
-
     # GITにバッチサイズが1より大きくても動くようにパッチを当てる: transformers 4.26.0用
     org_prepare_input_ids_for_generation = GenerationMixin._prepare_input_ids_for_generation
     curr_batch_size = [args.batch_size]  # ループの最後で件数がbatch_size未満になるので入れ替えられるように
@@ -75,24 +65,23 @@ def main(args):
         return input_ids
 
     GenerationMixin._prepare_input_ids_for_generation = _prepare_input_ids_for_generation_patch
-    """
 
-    logger.info(f"load images from {args.train_data_dir}")
+    print(f"load images from {args.train_data_dir}")
     train_data_dir_path = Path(args.train_data_dir)
     image_paths = train_util.glob_images_pathlib(train_data_dir_path, args.recursive)
-    logger.info(f"found {len(image_paths)} images.")
+    print(f"found {len(image_paths)} images.")
 
     # できればcacheに依存せず明示的にダウンロードしたい
-    logger.info(f"loading GIT: {args.model_id}")
+    print(f"loading GIT: {args.model_id}")
     git_processor = AutoProcessor.from_pretrained(args.model_id)
     git_model = AutoModelForCausalLM.from_pretrained(args.model_id).to(DEVICE)
-    logger.info("GIT loaded")
+    print("GIT loaded")
 
     # captioningする
     def run_batch(path_imgs):
         imgs = [im for _, im in path_imgs]
 
-        # curr_batch_size[0] = len(path_imgs)
+        curr_batch_size[0] = len(path_imgs)
         inputs = git_processor(images=imgs, return_tensors="pt").to(DEVICE)  # 画像はpil形式
         generated_ids = git_model.generate(pixel_values=inputs.pixel_values, max_length=args.max_length)
         captions = git_processor.batch_decode(generated_ids, skip_special_tokens=True)
@@ -104,7 +93,7 @@ def main(args):
             with open(os.path.splitext(image_path)[0] + args.caption_extension, "wt", encoding="utf-8") as f:
                 f.write(caption + "\n")
                 if args.debug:
-                    logger.info(f"{image_path} {caption}")
+                    print(image_path, caption)
 
     # 読み込みの高速化のためにDataLoaderを使うオプション
     if args.max_data_loader_n_workers is not None:
@@ -133,7 +122,7 @@ def main(args):
                     if image.mode != "RGB":
                         image = image.convert("RGB")
                 except Exception as e:
-                    logger.error(f"Could not load image path / 画像を読み込めません: {image_path}, error: {e}")
+                    print(f"Could not load image path / 画像を読み込めません: {image_path}, error: {e}")
                     continue
 
             b_imgs.append((image_path, image))
@@ -144,7 +133,7 @@ def main(args):
     if len(b_imgs) > 0:
         run_batch(b_imgs)
 
-    logger.info("done!")
+    print("done!")
 
 
 def setup_parser() -> argparse.ArgumentParser:

finetune/merge_captions_to_metadata.py

@@ -5,30 +5,26 @@ from typing import List
 from tqdm import tqdm
 import library.train_util as train_util
 import os
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
 
 def main(args):
   assert not args.recursive or (args.recursive and args.full_path), "recursive requires full_path / recursiveはfull_pathと同時に指定してください"
 
   train_data_dir_path = Path(args.train_data_dir)
   image_paths: List[Path] = train_util.glob_images_pathlib(train_data_dir_path, args.recursive)
-  logger.info(f"found {len(image_paths)} images.")
+  print(f"found {len(image_paths)} images.")
 
   if args.in_json is None and Path(args.out_json).is_file():
     args.in_json = args.out_json
 
   if args.in_json is not None:
-    logger.info(f"loading existing metadata: {args.in_json}")
+    print(f"loading existing metadata: {args.in_json}")
     metadata = json.loads(Path(args.in_json).read_text(encoding='utf-8'))
-    logger.warning("captions for existing images will be overwritten / 既存の画像のキャプションは上書きされます")
+    print("captions for existing images will be overwritten / 既存の画像のキャプションは上書きされます")
   else:
-    logger.info("new metadata will be created / 新しいメタデータファイルが作成されます")
+    print("new metadata will be created / 新しいメタデータファイルが作成されます")
     metadata = {}
 
-  logger.info("merge caption texts to metadata json.")
+  print("merge caption texts to metadata json.")
   for image_path in tqdm(image_paths):
     caption_path = image_path.with_suffix(args.caption_extension)
     caption = caption_path.read_text(encoding='utf-8').strip()
@@ -42,12 +38,12 @@ def main(args):
 
     metadata[image_key]['caption'] = caption
     if args.debug:
-      logger.info(f"{image_key} {caption}")
+      print(image_key, caption)
 
   # metadataを書き出して終わり
-  logger.info(f"writing metadata: {args.out_json}")
+  print(f"writing metadata: {args.out_json}")
   Path(args.out_json).write_text(json.dumps(metadata, indent=2), encoding='utf-8')
-  logger.info("done!")
+  print("done!")
 
 
 def setup_parser() -> argparse.ArgumentParser:

finetune/merge_dd_tags_to_metadata.py

@@ -5,30 +5,26 @@ from typing import List
 from tqdm import tqdm
 import library.train_util as train_util
 import os
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
 
 def main(args):
   assert not args.recursive or (args.recursive and args.full_path), "recursive requires full_path / recursiveはfull_pathと同時に指定してください"
 
   train_data_dir_path = Path(args.train_data_dir)
   image_paths: List[Path] = train_util.glob_images_pathlib(train_data_dir_path, args.recursive)
-  logger.info(f"found {len(image_paths)} images.")
+  print(f"found {len(image_paths)} images.")
 
   if args.in_json is None and Path(args.out_json).is_file():
     args.in_json = args.out_json
 
   if args.in_json is not None:
-    logger.info(f"loading existing metadata: {args.in_json}")
+    print(f"loading existing metadata: {args.in_json}")
     metadata = json.loads(Path(args.in_json).read_text(encoding='utf-8'))
-    logger.warning("tags data for existing images will be overwritten / 既存の画像のタグは上書きされます")
+    print("tags data for existing images will be overwritten / 既存の画像のタグは上書きされます")
   else:
-    logger.info("new metadata will be created / 新しいメタデータファイルが作成されます")
+    print("new metadata will be created / 新しいメタデータファイルが作成されます")
     metadata = {}
 
-  logger.info("merge tags to metadata json.")
+  print("merge tags to metadata json.")
   for image_path in tqdm(image_paths):
     tags_path = image_path.with_suffix(args.caption_extension)
     tags = tags_path.read_text(encoding='utf-8').strip()
@@ -42,13 +38,13 @@ def main(args):
 
     metadata[image_key]['tags'] = tags
     if args.debug:
-      logger.info(f"{image_key} {tags}")
+      print(image_key, tags)
 
   # metadataを書き出して終わり
-  logger.info(f"writing metadata: {args.out_json}")
+  print(f"writing metadata: {args.out_json}")
   Path(args.out_json).write_text(json.dumps(metadata, indent=2), encoding='utf-8')
 
-  logger.info("done!")
+  print("done!")
 
 
 def setup_parser() -> argparse.ArgumentParser:

finetune/prepare_buckets_latents.py

@@ -8,21 +8,13 @@ from tqdm import tqdm
 import numpy as np
 from PIL import Image
 import cv2
-
 import torch
-from library.device_utils import init_ipex, get_preferred_device
-init_ipex()
-
 from torchvision import transforms
 
 import library.model_util as model_util
 import library.train_util as train_util
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
 
-DEVICE = get_preferred_device()
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 IMAGE_TRANSFORMS = transforms.Compose(
     [
@@ -59,22 +51,22 @@ def get_npz_filename(data_dir, image_key, is_full_path, recursive):
 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:
-        logger.warning(f"resolution of buckets in training time is a multiple of 8 / 学習時の各bucketの解像度は8単位になります")
+        print(f"resolution of buckets in training time is a multiple of 8 / 学習時の各bucketの解像度は8単位になります")
     if args.bucket_reso_steps % 32 > 0:
-        logger.warning(
+        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)]
-    logger.info(f"found {len(image_paths)} images.")
+    print(f"found {len(image_paths)} images.")
 
     if os.path.exists(args.in_json):
-        logger.info(f"loading existing metadata: {args.in_json}")
+        print(f"loading existing metadata: {args.in_json}")
         with open(args.in_json, "rt", encoding="utf-8") as f:
             metadata = json.load(f)
     else:
-        logger.error(f"no metadata / メタデータファイルがありません: {args.in_json}")
+        print(f"no metadata / メタデータファイルがありません: {args.in_json}")
         return
 
     weight_dtype = torch.float32
@@ -97,7 +89,7 @@ def main(args):
     if not args.bucket_no_upscale:
         bucket_manager.make_buckets()
     else:
-        logger.warning(
+        print(
             "min_bucket_reso and max_bucket_reso are ignored if bucket_no_upscale is set, because bucket reso is defined by image size automatically / bucket_no_upscaleが指定された場合は、bucketの解像度は画像サイズから自動計算されるため、min_bucket_resoとmax_bucket_resoは無視されます"
         )
 
@@ -138,7 +130,7 @@ def main(args):
                 if image.mode != "RGB":
                     image = image.convert("RGB")
             except Exception as e:
-                logger.error(f"Could not load image path / 画像を読み込めません: {image_path}, error: {e}")
+                print(f"Could not load image path / 画像を読み込めません: {image_path}, error: {e}")
                 continue
 
         image_key = image_path if args.full_path else os.path.splitext(os.path.basename(image_path))[0]
@@ -191,15 +183,15 @@ def main(args):
     for i, reso in enumerate(bucket_manager.resos):
         count = bucket_counts.get(reso, 0)
         if count > 0:
-            logger.info(f"bucket {i} {reso}: {count}")
+            print(f"bucket {i} {reso}: {count}")
     img_ar_errors = np.array(img_ar_errors)
-    logger.info(f"mean ar error: {np.mean(img_ar_errors)}")
+    print(f"mean ar error: {np.mean(img_ar_errors)}")
 
     # metadataを書き出して終わり
-    logger.info(f"writing metadata: {args.out_json}")
+    print(f"writing metadata: {args.out_json}")
     with open(args.out_json, "wt", encoding="utf-8") as f:
         json.dump(metadata, f, indent=2)
-    logger.info("done!")
+    print("done!")
 
 
 def setup_parser() -> argparse.ArgumentParser:
@@ -223,7 +215,7 @@ def setup_parser() -> argparse.ArgumentParser:
         help="max resolution in fine tuning (width,height) / fine tuning時の最大画像サイズ 「幅,高さ」（使用メモリ量に関係します）",
     )
     parser.add_argument("--min_bucket_reso", type=int, default=256, help="minimum resolution for buckets / bucketの最小解像度")
-    parser.add_argument("--max_bucket_reso", type=int, default=1024, help="maximum resolution for buckets / bucketの最大解像度")
+    parser.add_argument("--max_bucket_reso", type=int, default=1024, help="maximum resolution for buckets / bucketの最小解像度")
     parser.add_argument(
         "--bucket_reso_steps",
         type=int,

finetune/tag_images_by_wd14_tagger.py

@@ -11,10 +11,6 @@ from PIL import Image
 from tqdm import tqdm
 
 import library.train_util as train_util
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
 
 # from wd14 tagger
 IMAGE_SIZE = 448
@@ -62,7 +58,7 @@ class ImageLoadingPrepDataset(torch.utils.data.Dataset):
             image = preprocess_image(image)
             tensor = torch.tensor(image)
         except Exception as e:
-            logger.error(f"Could not load image path / 画像を読み込めません: {img_path}, error: {e}")
+            print(f"Could not load image path / 画像を読み込めません: {img_path}, error: {e}")
             return None
 
         return (tensor, img_path)
@@ -83,7 +79,7 @@ def main(args):
     # depreacatedの警告が出るけどなくなったらその時
     # https://github.com/toriato/stable-diffusion-webui-wd14-tagger/issues/22
     if not os.path.exists(args.model_dir) or args.force_download:
-        logger.info(f"downloading wd14 tagger model from hf_hub. id: {args.repo_id}")
+        print(f"downloading wd14 tagger model from hf_hub. id: {args.repo_id}")
         files = FILES
         if args.onnx:
             files += FILES_ONNX
@@ -99,7 +95,7 @@ def main(args):
                 force_filename=file,
             )
     else:
-        logger.info("using existing wd14 tagger model")
+        print("using existing wd14 tagger model")
 
     # 画像を読み込む
     if args.onnx:
@@ -107,8 +103,8 @@ def main(args):
         import onnxruntime as ort
 
         onnx_path = f"{args.model_dir}/model.onnx"
-        logger.info("Running wd14 tagger with onnx")
-        logger.info(f"loading onnx model: {onnx_path}")
+        print("Running wd14 tagger with onnx")
+        print(f"loading onnx model: {onnx_path}")
 
         if not os.path.exists(onnx_path):
             raise Exception(
@@ -125,7 +121,7 @@ def main(args):
 
         if args.batch_size != batch_size and type(batch_size) != str:
             # some rebatch model may use 'N' as dynamic axes
-            logger.warning(
+            print(
                 f"Batch size {args.batch_size} doesn't match onnx model batch size {batch_size}, use model batch size {batch_size}"
             )
             args.batch_size = batch_size
@@ -160,13 +156,11 @@ def main(args):
 
     train_data_dir_path = Path(args.train_data_dir)
     image_paths = train_util.glob_images_pathlib(train_data_dir_path, args.recursive)
-    logger.info(f"found {len(image_paths)} images.")
+    print(f"found {len(image_paths)} images.")
 
     tag_freq = {}
 
-    caption_separator = args.caption_separator
-    stripped_caption_separator = caption_separator.strip()
-    undesired_tags = set(args.undesired_tags.split(stripped_caption_separator))
+    undesired_tags = set(args.undesired_tags.split(","))
 
     def run_batch(path_imgs):
         imgs = np.array([im for _, im in path_imgs])
@@ -200,7 +194,7 @@ def main(args):
 
                     if tag_name not in undesired_tags:
                         tag_freq[tag_name] = tag_freq.get(tag_name, 0) + 1
-                        general_tag_text += caption_separator + tag_name
+                        general_tag_text += ", " + tag_name
                         combined_tags.append(tag_name)
                 elif i >= len(general_tags) and p >= args.character_threshold:
                     tag_name = character_tags[i - len(general_tags)]
@@ -209,18 +203,18 @@ def main(args):
 
                     if tag_name not in undesired_tags:
                         tag_freq[tag_name] = tag_freq.get(tag_name, 0) + 1
-                        character_tag_text += caption_separator + tag_name
+                        character_tag_text += ", " + tag_name
                         combined_tags.append(tag_name)
 
             # 先頭のカンマを取る
             if len(general_tag_text) > 0:
-                general_tag_text = general_tag_text[len(caption_separator) :]
+                general_tag_text = general_tag_text[2:]
             if len(character_tag_text) > 0:
-                character_tag_text = character_tag_text[len(caption_separator) :]
+                character_tag_text = character_tag_text[2:]
 
             caption_file = os.path.splitext(image_path)[0] + args.caption_extension
 
-            tag_text = caption_separator.join(combined_tags)
+            tag_text = ", ".join(combined_tags)
 
             if args.append_tags:
                 # Check if file exists
@@ -230,21 +224,18 @@ def main(args):
                         existing_content = f.read().strip("\n")  # Remove newlines
 
                     # Split the content into tags and store them in a list
-                    existing_tags = [tag.strip() for tag in existing_content.split(stripped_caption_separator) if tag.strip()]
+                    existing_tags = [tag.strip() for tag in existing_content.split(",") if tag.strip()]
 
                     # Check and remove repeating tags in tag_text
                     new_tags = [tag for tag in combined_tags if tag not in existing_tags]
 
                     # Create new tag_text
-                    tag_text = caption_separator.join(existing_tags + new_tags)
+                    tag_text = ", ".join(existing_tags + new_tags)
 
             with open(caption_file, "wt", encoding="utf-8") as f:
                 f.write(tag_text + "\n")
                 if args.debug:
-                    logger.info("")
-                    logger.info(f"{image_path}:")
-                    logger.info(f"\tCharacter tags: {character_tag_text}")
-                    logger.info(f"\tGeneral tags: {general_tag_text}")
+                    print(f"\n{image_path}:\n  Character tags: {character_tag_text}\n  General tags: {general_tag_text}")
 
     # 読み込みの高速化のためにDataLoaderを使うオプション
     if args.max_data_loader_n_workers is not None:
@@ -276,7 +267,7 @@ def main(args):
                         image = image.convert("RGB")
                     image = preprocess_image(image)
                 except Exception as e:
-                    logger.error(f"Could not load image path / 画像を読み込めません: {image_path}, error: {e}")
+                    print(f"Could not load image path / 画像を読み込めません: {image_path}, error: {e}")
                     continue
             b_imgs.append((image_path, image))
 
@@ -291,11 +282,11 @@ def main(args):
 
     if args.frequency_tags:
         sorted_tags = sorted(tag_freq.items(), key=lambda x: x[1], reverse=True)
-        print("Tag frequencies:")
+        print("\nTag frequencies:")
         for tag, freq in sorted_tags:
             print(f"{tag}: {freq}")
 
-    logger.info("done!")
+    print("done!")
 
 
 def setup_parser() -> argparse.ArgumentParser:
@@ -359,12 +350,6 @@ def setup_parser() -> argparse.ArgumentParser:
     parser.add_argument("--frequency_tags", action="store_true", help="Show frequency of tags for images / 画像ごとのタグの出現頻度を表示する")
     parser.add_argument("--onnx", action="store_true", help="use onnx model for inference / onnxモデルを推論に使用する")
     parser.add_argument("--append_tags", action="store_true", help="Append captions instead of overwriting / 上書きではなくキャプションを追記する")
-    parser.add_argument(
-        "--caption_separator",
-        type=str,
-        default=", ",
-        help="Separator for captions, include space if needed / キャプションの区切り文字、必要ならスペースを含めてください",
-    )
 
     return parser
 

library/custom_train_functions.py

@@ -3,10 +3,7 @@ import argparse
 import random
 import re
 from typing import List, Optional, Union
-from .utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 def prepare_scheduler_for_custom_training(noise_scheduler, device):
     if hasattr(noise_scheduler, "all_snr"):
@@ -24,7 +21,7 @@ def prepare_scheduler_for_custom_training(noise_scheduler, device):
 
 def fix_noise_scheduler_betas_for_zero_terminal_snr(noise_scheduler):
     # fix beta: zero terminal SNR
-    logger.info(f"fix noise scheduler betas: https://arxiv.org/abs/2305.08891")
+    print(f"fix noise scheduler betas: https://arxiv.org/abs/2305.08891")
 
     def enforce_zero_terminal_snr(betas):
         # Convert betas to alphas_bar_sqrt
@@ -52,21 +49,18 @@ def fix_noise_scheduler_betas_for_zero_terminal_snr(noise_scheduler):
     alphas = 1.0 - betas
     alphas_cumprod = torch.cumprod(alphas, dim=0)
 
-    # logger.info(f"original: {noise_scheduler.betas}")
-    # logger.info(f"fixed: {betas}")
+    # 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, v_prediction=False):
+def apply_snr_weight(loss, timesteps, noise_scheduler, gamma):
     snr = torch.stack([noise_scheduler.all_snr[t] for t in timesteps])
-    min_snr_gamma = torch.minimum(snr, torch.full_like(snr, gamma))
-    if v_prediction:
-        snr_weight = torch.div(min_snr_gamma, snr+1).float().to(loss.device)
-    else:
-        snr_weight = torch.div(min_snr_gamma, snr).float().to(loss.device)
+    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
     loss = loss * snr_weight
     return loss
 
@@ -82,22 +76,16 @@ def get_snr_scale(timesteps, noise_scheduler):
     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
-    # logger.info(f"timesteps: {timesteps}, snr_t: {snr_t}, scale: {scale}")
+    # 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)
-    # logger.info(f"add v-prediction like loss: {v_pred_like_loss}, scale: {scale}, loss: {loss}, time: {timesteps}")
+    # 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
     return loss
 
-def apply_debiased_estimation(loss, 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
-    weight = 1/torch.sqrt(snr_t)
-    loss = weight * loss
-    return loss
 
 # TODO train_utilと分散しているのでどちらかに寄せる
 
@@ -120,11 +108,6 @@ def add_custom_train_arguments(parser: argparse.ArgumentParser, support_weighted
         default=None,
         help="add v-prediction like loss multiplied by this value / v-prediction lossをこの値をかけたものをlossに加算する",
     )
-    parser.add_argument(
-        "--debiased_estimation_loss",
-        action="store_true",
-        help="debiased estimation loss / debiased estimation loss",
-    )
     if support_weighted_captions:
         parser.add_argument(
             "--weighted_captions",
@@ -271,7 +254,7 @@ def get_prompts_with_weights(tokenizer, prompt: List[str], max_length: int):
         tokens.append(text_token)
         weights.append(text_weight)
     if truncated:
-        logger.warning("Prompt was truncated. Try to shorten the prompt or increase max_embeddings_multiples")
+        print("Prompt was truncated. Try to shorten the prompt or increase max_embeddings_multiples")
     return tokens, weights
 
 

library/device_utils.py

@@ -1,84 +0,0 @@
-import functools
-import gc
-
-import torch
-
-try:
-    HAS_CUDA = torch.cuda.is_available()
-except Exception:
-    HAS_CUDA = False
-
-try:
-    HAS_MPS = torch.backends.mps.is_available()
-except Exception:
-    HAS_MPS = False
-
-try:
-    import intel_extension_for_pytorch as ipex  # noqa
-
-    HAS_XPU = torch.xpu.is_available()
-except Exception:
-    HAS_XPU = False
-
-
-def clean_memory():
-    gc.collect()
-    if HAS_CUDA:
-        torch.cuda.empty_cache()
-    if HAS_XPU:
-        torch.xpu.empty_cache()
-    if HAS_MPS:
-        torch.mps.empty_cache()
-
-
-def clean_memory_on_device(device: torch.device):
-    r"""
-    Clean memory on the specified device, will be called from training scripts.
-    """
-    gc.collect()
-
-    # device may "cuda" or "cuda:0", so we need to check the type of device
-    if device.type == "cuda":
-        torch.cuda.empty_cache()
-    if device.type == "xpu":
-        torch.xpu.empty_cache()
-    if device.type == "mps":
-        torch.mps.empty_cache()
-
-
-@functools.lru_cache(maxsize=None)
-def get_preferred_device() -> torch.device:
-    r"""
-    Do not call this function from training scripts. Use accelerator.device instead.
-    """
-    if HAS_CUDA:
-        device = torch.device("cuda")
-    elif HAS_XPU:
-        device = torch.device("xpu")
-    elif HAS_MPS:
-        device = torch.device("mps")
-    else:
-        device = torch.device("cpu")
-    print(f"get_preferred_device() -> {device}")
-    return device
-
-
-def init_ipex():
-    """
-    Apply IPEX to CUDA hijacks using `library.ipex.ipex_init`.
-
-    This function should run right after importing torch and before doing anything else.
-
-    If IPEX is not available, this function does nothing.
-    """
-    try:
-        if HAS_XPU:
-            from library.ipex import ipex_init
-
-            is_initialized, error_message = ipex_init()
-            if not is_initialized:
-                print("failed to initialize ipex:", error_message)
-        else:
-            return
-    except Exception as e:
-        print("failed to initialize ipex:", e)

library/huggingface_util.py

@@ -4,10 +4,7 @@ from pathlib import Path
 import argparse
 import os
 from library.utils import fire_in_thread
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 def exists_repo(repo_id: str, repo_type: str, revision: str = "main", token: str = None):
     api = HfApi(
@@ -36,9 +33,9 @@ def upload(
         try:
             api.create_repo(repo_id=repo_id, repo_type=repo_type, private=private)
         except Exception as e:  # とりあえずRepositoryNotFoundErrorは確認したが他にあると困るので
-            logger.error("===========================================")
-            logger.error(f"failed to create HuggingFace repo / HuggingFaceのリポジトリの作成に失敗しました : {e}")
-            logger.error("===========================================")
+            print("===========================================")
+            print(f"failed to create HuggingFace repo / HuggingFaceのリポジトリの作成に失敗しました : {e}")
+            print("===========================================")
 
     is_folder = (type(src) == str and os.path.isdir(src)) or (isinstance(src, Path) and src.is_dir())
 
@@ -59,9 +56,9 @@ def upload(
                     path_in_repo=path_in_repo,
                 )
         except Exception as e:  # RuntimeErrorを確認済みだが他にあると困るので
-            logger.error("===========================================")
-            logger.error(f"failed to upload to HuggingFace / HuggingFaceへのアップロードに失敗しました : {e}")
-            logger.error("===========================================")
+            print("===========================================")
+            print(f"failed to upload to HuggingFace / HuggingFaceへのアップロードに失敗しました : {e}")
+            print("===========================================")
 
     if args.async_upload and not force_sync_upload:
         fire_in_thread(uploader)

library/ipex/__init__.py

@@ -4,176 +4,172 @@ import contextlib
 import torch
 import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import
 from .hijacks import ipex_hijacks
+from .attention import attention_init
 
 # pylint: disable=protected-access, missing-function-docstring, line-too-long
 
 def ipex_init(): # pylint: disable=too-many-statements
     try:
-        if hasattr(torch, "cuda") and hasattr(torch.cuda, "is_xpu_hijacked") and torch.cuda.is_xpu_hijacked:
-            return True, "Skipping IPEX hijack"
-        else:
-            # Replace cuda with xpu:
-            torch.cuda.current_device = torch.xpu.current_device
-            torch.cuda.current_stream = torch.xpu.current_stream
-            torch.cuda.device = torch.xpu.device
-            torch.cuda.device_count = torch.xpu.device_count
-            torch.cuda.device_of = torch.xpu.device_of
-            torch.cuda.get_device_name = torch.xpu.get_device_name
-            torch.cuda.get_device_properties = torch.xpu.get_device_properties
-            torch.cuda.init = torch.xpu.init
-            torch.cuda.is_available = torch.xpu.is_available
-            torch.cuda.is_initialized = torch.xpu.is_initialized
-            torch.cuda.is_current_stream_capturing = lambda: False
-            torch.cuda.set_device = torch.xpu.set_device
-            torch.cuda.stream = torch.xpu.stream
-            torch.cuda.synchronize = torch.xpu.synchronize
-            torch.cuda.Event = torch.xpu.Event
-            torch.cuda.Stream = torch.xpu.Stream
-            torch.cuda.FloatTensor = torch.xpu.FloatTensor
-            torch.Tensor.cuda = torch.Tensor.xpu
-            torch.Tensor.is_cuda = torch.Tensor.is_xpu
-            torch.UntypedStorage.cuda = torch.UntypedStorage.xpu
-            torch.cuda._initialization_lock = torch.xpu.lazy_init._initialization_lock
-            torch.cuda._initialized = torch.xpu.lazy_init._initialized
-            torch.cuda._lazy_seed_tracker = torch.xpu.lazy_init._lazy_seed_tracker
-            torch.cuda._queued_calls = torch.xpu.lazy_init._queued_calls
-            torch.cuda._tls = torch.xpu.lazy_init._tls
-            torch.cuda.threading = torch.xpu.lazy_init.threading
-            torch.cuda.traceback = torch.xpu.lazy_init.traceback
-            torch.cuda.Optional = torch.xpu.Optional
-            torch.cuda.__cached__ = torch.xpu.__cached__
-            torch.cuda.__loader__ = torch.xpu.__loader__
-            torch.cuda.ComplexFloatStorage = torch.xpu.ComplexFloatStorage
-            torch.cuda.Tuple = torch.xpu.Tuple
-            torch.cuda.streams = torch.xpu.streams
-            torch.cuda._lazy_new = torch.xpu._lazy_new
-            torch.cuda.FloatStorage = torch.xpu.FloatStorage
-            torch.cuda.Any = torch.xpu.Any
-            torch.cuda.__doc__ = torch.xpu.__doc__
-            torch.cuda.default_generators = torch.xpu.default_generators
-            torch.cuda.HalfTensor = torch.xpu.HalfTensor
-            torch.cuda._get_device_index = torch.xpu._get_device_index
-            torch.cuda.__path__ = torch.xpu.__path__
-            torch.cuda.Device = torch.xpu.Device
-            torch.cuda.IntTensor = torch.xpu.IntTensor
-            torch.cuda.ByteStorage = torch.xpu.ByteStorage
-            torch.cuda.set_stream = torch.xpu.set_stream
-            torch.cuda.BoolStorage = torch.xpu.BoolStorage
-            torch.cuda.os = torch.xpu.os
-            torch.cuda.torch = torch.xpu.torch
-            torch.cuda.BFloat16Storage = torch.xpu.BFloat16Storage
-            torch.cuda.Union = torch.xpu.Union
-            torch.cuda.DoubleTensor = torch.xpu.DoubleTensor
-            torch.cuda.ShortTensor = torch.xpu.ShortTensor
-            torch.cuda.LongTensor = torch.xpu.LongTensor
-            torch.cuda.IntStorage = torch.xpu.IntStorage
-            torch.cuda.LongStorage = torch.xpu.LongStorage
-            torch.cuda.__annotations__ = torch.xpu.__annotations__
-            torch.cuda.__package__ = torch.xpu.__package__
-            torch.cuda.__builtins__ = torch.xpu.__builtins__
-            torch.cuda.CharTensor = torch.xpu.CharTensor
-            torch.cuda.List = torch.xpu.List
-            torch.cuda._lazy_init = torch.xpu._lazy_init
-            torch.cuda.BFloat16Tensor = torch.xpu.BFloat16Tensor
-            torch.cuda.DoubleStorage = torch.xpu.DoubleStorage
-            torch.cuda.ByteTensor = torch.xpu.ByteTensor
-            torch.cuda.StreamContext = torch.xpu.StreamContext
-            torch.cuda.ComplexDoubleStorage = torch.xpu.ComplexDoubleStorage
-            torch.cuda.ShortStorage = torch.xpu.ShortStorage
-            torch.cuda._lazy_call = torch.xpu._lazy_call
-            torch.cuda.HalfStorage = torch.xpu.HalfStorage
-            torch.cuda.random = torch.xpu.random
-            torch.cuda._device = torch.xpu._device
-            torch.cuda.classproperty = torch.xpu.classproperty
-            torch.cuda.__name__ = torch.xpu.__name__
-            torch.cuda._device_t = torch.xpu._device_t
-            torch.cuda.warnings = torch.xpu.warnings
-            torch.cuda.__spec__ = torch.xpu.__spec__
-            torch.cuda.BoolTensor = torch.xpu.BoolTensor
-            torch.cuda.CharStorage = torch.xpu.CharStorage
-            torch.cuda.__file__ = torch.xpu.__file__
-            torch.cuda._is_in_bad_fork = torch.xpu.lazy_init._is_in_bad_fork
-            # torch.cuda.is_current_stream_capturing = torch.xpu.is_current_stream_capturing
+        #Replace cuda with xpu:
+        torch.cuda.current_device = torch.xpu.current_device
+        torch.cuda.current_stream = torch.xpu.current_stream
+        torch.cuda.device = torch.xpu.device
+        torch.cuda.device_count = torch.xpu.device_count
+        torch.cuda.device_of = torch.xpu.device_of
+        torch.cuda.get_device_name = torch.xpu.get_device_name
+        torch.cuda.get_device_properties = torch.xpu.get_device_properties
+        torch.cuda.init = torch.xpu.init
+        torch.cuda.is_available = torch.xpu.is_available
+        torch.cuda.is_initialized = torch.xpu.is_initialized
+        torch.cuda.is_current_stream_capturing = lambda: False
+        torch.cuda.set_device = torch.xpu.set_device
+        torch.cuda.stream = torch.xpu.stream
+        torch.cuda.synchronize = torch.xpu.synchronize
+        torch.cuda.Event = torch.xpu.Event
+        torch.cuda.Stream = torch.xpu.Stream
+        torch.cuda.FloatTensor = torch.xpu.FloatTensor
+        torch.Tensor.cuda = torch.Tensor.xpu
+        torch.Tensor.is_cuda = torch.Tensor.is_xpu
+        torch.cuda._initialization_lock = torch.xpu.lazy_init._initialization_lock
+        torch.cuda._initialized = torch.xpu.lazy_init._initialized
+        torch.cuda._lazy_seed_tracker = torch.xpu.lazy_init._lazy_seed_tracker
+        torch.cuda._queued_calls = torch.xpu.lazy_init._queued_calls
+        torch.cuda._tls = torch.xpu.lazy_init._tls
+        torch.cuda.threading = torch.xpu.lazy_init.threading
+        torch.cuda.traceback = torch.xpu.lazy_init.traceback
+        torch.cuda.Optional = torch.xpu.Optional
+        torch.cuda.__cached__ = torch.xpu.__cached__
+        torch.cuda.__loader__ = torch.xpu.__loader__
+        torch.cuda.ComplexFloatStorage = torch.xpu.ComplexFloatStorage
+        torch.cuda.Tuple = torch.xpu.Tuple
+        torch.cuda.streams = torch.xpu.streams
+        torch.cuda._lazy_new = torch.xpu._lazy_new
+        torch.cuda.FloatStorage = torch.xpu.FloatStorage
+        torch.cuda.Any = torch.xpu.Any
+        torch.cuda.__doc__ = torch.xpu.__doc__
+        torch.cuda.default_generators = torch.xpu.default_generators
+        torch.cuda.HalfTensor = torch.xpu.HalfTensor
+        torch.cuda._get_device_index = torch.xpu._get_device_index
+        torch.cuda.__path__ = torch.xpu.__path__
+        torch.cuda.Device = torch.xpu.Device
+        torch.cuda.IntTensor = torch.xpu.IntTensor
+        torch.cuda.ByteStorage = torch.xpu.ByteStorage
+        torch.cuda.set_stream = torch.xpu.set_stream
+        torch.cuda.BoolStorage = torch.xpu.BoolStorage
+        torch.cuda.os = torch.xpu.os
+        torch.cuda.torch = torch.xpu.torch
+        torch.cuda.BFloat16Storage = torch.xpu.BFloat16Storage
+        torch.cuda.Union = torch.xpu.Union
+        torch.cuda.DoubleTensor = torch.xpu.DoubleTensor
+        torch.cuda.ShortTensor = torch.xpu.ShortTensor
+        torch.cuda.LongTensor = torch.xpu.LongTensor
+        torch.cuda.IntStorage = torch.xpu.IntStorage
+        torch.cuda.LongStorage = torch.xpu.LongStorage
+        torch.cuda.__annotations__ = torch.xpu.__annotations__
+        torch.cuda.__package__ = torch.xpu.__package__
+        torch.cuda.__builtins__ = torch.xpu.__builtins__
+        torch.cuda.CharTensor = torch.xpu.CharTensor
+        torch.cuda.List = torch.xpu.List
+        torch.cuda._lazy_init = torch.xpu._lazy_init
+        torch.cuda.BFloat16Tensor = torch.xpu.BFloat16Tensor
+        torch.cuda.DoubleStorage = torch.xpu.DoubleStorage
+        torch.cuda.ByteTensor = torch.xpu.ByteTensor
+        torch.cuda.StreamContext = torch.xpu.StreamContext
+        torch.cuda.ComplexDoubleStorage = torch.xpu.ComplexDoubleStorage
+        torch.cuda.ShortStorage = torch.xpu.ShortStorage
+        torch.cuda._lazy_call = torch.xpu._lazy_call
+        torch.cuda.HalfStorage = torch.xpu.HalfStorage
+        torch.cuda.random = torch.xpu.random
+        torch.cuda._device = torch.xpu._device
+        torch.cuda.classproperty = torch.xpu.classproperty
+        torch.cuda.__name__ = torch.xpu.__name__
+        torch.cuda._device_t = torch.xpu._device_t
+        torch.cuda.warnings = torch.xpu.warnings
+        torch.cuda.__spec__ = torch.xpu.__spec__
+        torch.cuda.BoolTensor = torch.xpu.BoolTensor
+        torch.cuda.CharStorage = torch.xpu.CharStorage
+        torch.cuda.__file__ = torch.xpu.__file__
+        torch.cuda._is_in_bad_fork = torch.xpu.lazy_init._is_in_bad_fork
+        #torch.cuda.is_current_stream_capturing = torch.xpu.is_current_stream_capturing
 
-            # Memory:
-            torch.cuda.memory = torch.xpu.memory
-            if 'linux' in sys.platform and "WSL2" in os.popen("uname -a").read():
-                torch.xpu.empty_cache = lambda: None
-            torch.cuda.empty_cache = torch.xpu.empty_cache
-            torch.cuda.memory_stats = torch.xpu.memory_stats
-            torch.cuda.memory_summary = torch.xpu.memory_summary
-            torch.cuda.memory_snapshot = torch.xpu.memory_snapshot
-            torch.cuda.memory_allocated = torch.xpu.memory_allocated
-            torch.cuda.max_memory_allocated = torch.xpu.max_memory_allocated
-            torch.cuda.memory_reserved = torch.xpu.memory_reserved
-            torch.cuda.memory_cached = torch.xpu.memory_reserved
-            torch.cuda.max_memory_reserved = torch.xpu.max_memory_reserved
-            torch.cuda.max_memory_cached = torch.xpu.max_memory_reserved
-            torch.cuda.reset_peak_memory_stats = torch.xpu.reset_peak_memory_stats
-            torch.cuda.reset_max_memory_cached = torch.xpu.reset_peak_memory_stats
-            torch.cuda.reset_max_memory_allocated = torch.xpu.reset_peak_memory_stats
-            torch.cuda.memory_stats_as_nested_dict = torch.xpu.memory_stats_as_nested_dict
-            torch.cuda.reset_accumulated_memory_stats = torch.xpu.reset_accumulated_memory_stats
+        #Memory:
+        torch.cuda.memory = torch.xpu.memory
+        if 'linux' in sys.platform and "WSL2" in os.popen("uname -a").read():
+            torch.xpu.empty_cache = lambda: None
+        torch.cuda.empty_cache = torch.xpu.empty_cache
+        torch.cuda.memory_stats = torch.xpu.memory_stats
+        torch.cuda.memory_summary = torch.xpu.memory_summary
+        torch.cuda.memory_snapshot = torch.xpu.memory_snapshot
+        torch.cuda.memory_allocated = torch.xpu.memory_allocated
+        torch.cuda.max_memory_allocated = torch.xpu.max_memory_allocated
+        torch.cuda.memory_reserved = torch.xpu.memory_reserved
+        torch.cuda.memory_cached = torch.xpu.memory_reserved
+        torch.cuda.max_memory_reserved = torch.xpu.max_memory_reserved
+        torch.cuda.max_memory_cached = torch.xpu.max_memory_reserved
+        torch.cuda.reset_peak_memory_stats = torch.xpu.reset_peak_memory_stats
+        torch.cuda.reset_max_memory_cached = torch.xpu.reset_peak_memory_stats
+        torch.cuda.reset_max_memory_allocated = torch.xpu.reset_peak_memory_stats
+        torch.cuda.memory_stats_as_nested_dict = torch.xpu.memory_stats_as_nested_dict
+        torch.cuda.reset_accumulated_memory_stats = torch.xpu.reset_accumulated_memory_stats
 
-            # RNG:
-            torch.cuda.get_rng_state = torch.xpu.get_rng_state
-            torch.cuda.get_rng_state_all = torch.xpu.get_rng_state_all
-            torch.cuda.set_rng_state = torch.xpu.set_rng_state
-            torch.cuda.set_rng_state_all = torch.xpu.set_rng_state_all
-            torch.cuda.manual_seed = torch.xpu.manual_seed
-            torch.cuda.manual_seed_all = torch.xpu.manual_seed_all
-            torch.cuda.seed = torch.xpu.seed
-            torch.cuda.seed_all = torch.xpu.seed_all
-            torch.cuda.initial_seed = torch.xpu.initial_seed
-
-            # AMP:
-            torch.cuda.amp = torch.xpu.amp
-            torch.is_autocast_enabled = torch.xpu.is_autocast_xpu_enabled
-            torch.get_autocast_gpu_dtype = torch.xpu.get_autocast_xpu_dtype
-
-            if not hasattr(torch.cuda.amp, "common"):
-                torch.cuda.amp.common = contextlib.nullcontext()
-            torch.cuda.amp.common.amp_definitely_not_available = lambda: False
+        #RNG:
+        torch.cuda.get_rng_state = torch.xpu.get_rng_state
+        torch.cuda.get_rng_state_all = torch.xpu.get_rng_state_all
+        torch.cuda.set_rng_state = torch.xpu.set_rng_state
+        torch.cuda.set_rng_state_all = torch.xpu.set_rng_state_all
+        torch.cuda.manual_seed = torch.xpu.manual_seed
+        torch.cuda.manual_seed_all = torch.xpu.manual_seed_all
+        torch.cuda.seed = torch.xpu.seed
+        torch.cuda.seed_all = torch.xpu.seed_all
+        torch.cuda.initial_seed = torch.xpu.initial_seed
 
+        #AMP:
+        torch.cuda.amp = torch.xpu.amp
+        if not hasattr(torch.cuda.amp, "common"):
+            torch.cuda.amp.common = contextlib.nullcontext()
+        torch.cuda.amp.common.amp_definitely_not_available = lambda: False
+        try:
+            torch.cuda.amp.GradScaler = torch.xpu.amp.GradScaler
+        except Exception: # pylint: disable=broad-exception-caught
             try:
+                from .gradscaler import gradscaler_init # pylint: disable=import-outside-toplevel, import-error
+                gradscaler_init()
                 torch.cuda.amp.GradScaler = torch.xpu.amp.GradScaler
             except Exception: # pylint: disable=broad-exception-caught
-                try:
-                    from .gradscaler import gradscaler_init # pylint: disable=import-outside-toplevel, import-error
-                    gradscaler_init()
-                    torch.cuda.amp.GradScaler = torch.xpu.amp.GradScaler
-                except Exception: # pylint: disable=broad-exception-caught
-                    torch.cuda.amp.GradScaler = ipex.cpu.autocast._grad_scaler.GradScaler
+                torch.cuda.amp.GradScaler = ipex.cpu.autocast._grad_scaler.GradScaler
 
-            # C
-            torch._C._cuda_getCurrentRawStream = ipex._C._getCurrentStream
-            ipex._C._DeviceProperties.multi_processor_count = ipex._C._DeviceProperties.gpu_eu_count
-            ipex._C._DeviceProperties.major = 2023
-            ipex._C._DeviceProperties.minor = 2
+        #C
+        torch._C._cuda_getCurrentRawStream = ipex._C._getCurrentStream
+        ipex._C._DeviceProperties.major = 2023
+        ipex._C._DeviceProperties.minor = 2
 
-            # Fix functions with ipex:
-            torch.cuda.mem_get_info = lambda device=None: [(torch.xpu.get_device_properties(device).total_memory - torch.xpu.memory_reserved(device)), torch.xpu.get_device_properties(device).total_memory]
-            torch._utils._get_available_device_type = lambda: "xpu"
-            torch.has_cuda = True
-            torch.cuda.has_half = True
-            torch.cuda.is_bf16_supported = lambda *args, **kwargs: True
-            torch.cuda.is_fp16_supported = lambda *args, **kwargs: True
-            torch.backends.cuda.is_built = lambda *args, **kwargs: True
-            torch.version.cuda = "12.1"
-            torch.cuda.get_device_capability = lambda *args, **kwargs: [12,1]
-            torch.cuda.get_device_properties.major = 12
-            torch.cuda.get_device_properties.minor = 1
-            torch.cuda.ipc_collect = lambda *args, **kwargs: None
-            torch.cuda.utilization = lambda *args, **kwargs: 0
+        #Fix functions with ipex:
+        torch.cuda.mem_get_info = lambda device=None: [(torch.xpu.get_device_properties(device).total_memory - torch.xpu.memory_reserved(device)), torch.xpu.get_device_properties(device).total_memory]
+        torch._utils._get_available_device_type = lambda: "xpu"
+        torch.has_cuda = True
+        torch.cuda.has_half = True
+        torch.cuda.is_bf16_supported = lambda *args, **kwargs: True
+        torch.cuda.is_fp16_supported = lambda *args, **kwargs: True
+        torch.version.cuda = "11.7"
+        torch.cuda.get_device_capability = lambda *args, **kwargs: [11,7]
+        torch.cuda.get_device_properties.major = 11
+        torch.cuda.get_device_properties.minor = 7
+        torch.cuda.ipc_collect = lambda *args, **kwargs: None
+        torch.cuda.utilization = lambda *args, **kwargs: 0
+        if hasattr(torch.xpu, 'getDeviceIdListForCard'):
+            torch.cuda.getDeviceIdListForCard = torch.xpu.getDeviceIdListForCard
+            torch.cuda.get_device_id_list_per_card = torch.xpu.getDeviceIdListForCard
+        else:
+            torch.cuda.getDeviceIdListForCard = torch.xpu.get_device_id_list_per_card
+            torch.cuda.get_device_id_list_per_card = torch.xpu.get_device_id_list_per_card
 
-            ipex_hijacks()
-            if not torch.xpu.has_fp64_dtype() or os.environ.get('IPEX_FORCE_ATTENTION_SLICE', None) is not None:
-                try:
-                    from .diffusers import ipex_diffusers
-                    ipex_diffusers()
-                except Exception: # pylint: disable=broad-exception-caught
-                    pass
-            torch.cuda.is_xpu_hijacked = True
+        ipex_hijacks()
+        attention_init()
+        try:
+            from .diffusers import ipex_diffusers
+            ipex_diffusers()
+        except Exception: # pylint: disable=broad-exception-caught
+            pass
     except Exception as e:
         return False, e
     return True, None

library/ipex/attention.py

@@ -1,98 +1,45 @@
-import os
 import torch
 import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import
-from functools import cache
 
 # pylint: disable=protected-access, missing-function-docstring, line-too-long
 
-# ARC GPUs can't allocate more than 4GB to a single block so we slice the attetion layers
-
-sdpa_slice_trigger_rate = float(os.environ.get('IPEX_SDPA_SLICE_TRIGGER_RATE', 4))
-attention_slice_rate = float(os.environ.get('IPEX_ATTENTION_SLICE_RATE', 4))
-
-# Find something divisible with the input_tokens
-@cache
-def find_slice_size(slice_size, slice_block_size):
-    while (slice_size * slice_block_size) > attention_slice_rate:
-        slice_size = slice_size // 2
-        if slice_size <= 1:
-            slice_size = 1
-            break
-    return slice_size
-
-# Find slice sizes for SDPA
-@cache
-def find_sdpa_slice_sizes(query_shape, query_element_size):
-    if len(query_shape) == 3:
-        batch_size_attention, query_tokens, shape_three = query_shape
-        shape_four = 1
-    else:
-        batch_size_attention, query_tokens, shape_three, shape_four = query_shape
-
-    slice_block_size = query_tokens * shape_three * shape_four / 1024 / 1024 * query_element_size
-    block_size = batch_size_attention * slice_block_size
-
-    split_slice_size = batch_size_attention
-    split_2_slice_size = query_tokens
-    split_3_slice_size = shape_three
-
-    do_split = False
-    do_split_2 = False
-    do_split_3 = False
-
-    if block_size > sdpa_slice_trigger_rate:
-        do_split = True
-        split_slice_size = find_slice_size(split_slice_size, slice_block_size)
-        if split_slice_size * slice_block_size > attention_slice_rate:
-            slice_2_block_size = split_slice_size * shape_three * shape_four / 1024 / 1024 * query_element_size
-            do_split_2 = True
-            split_2_slice_size = find_slice_size(split_2_slice_size, slice_2_block_size)
-            if split_2_slice_size * slice_2_block_size > attention_slice_rate:
-                slice_3_block_size = split_slice_size * split_2_slice_size * shape_four / 1024 / 1024 * query_element_size
-                do_split_3 = True
-                split_3_slice_size = find_slice_size(split_3_slice_size, slice_3_block_size)
-
-    return do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size
-
-# Find slice sizes for BMM
-@cache
-def find_bmm_slice_sizes(input_shape, input_element_size, mat2_shape):
-    batch_size_attention, input_tokens, mat2_atten_shape = input_shape[0], input_shape[1], mat2_shape[2]
-    slice_block_size = input_tokens * mat2_atten_shape / 1024 / 1024 * input_element_size
-    block_size = batch_size_attention * slice_block_size
-
-    split_slice_size = batch_size_attention
-    split_2_slice_size = input_tokens
-    split_3_slice_size = mat2_atten_shape
-
-    do_split = False
-    do_split_2 = False
-    do_split_3 = False
-
-    if block_size > attention_slice_rate:
-        do_split = True
-        split_slice_size = find_slice_size(split_slice_size, slice_block_size)
-        if split_slice_size * slice_block_size > attention_slice_rate:
-            slice_2_block_size = split_slice_size * mat2_atten_shape / 1024 / 1024 * input_element_size
-            do_split_2 = True
-            split_2_slice_size = find_slice_size(split_2_slice_size, slice_2_block_size)
-            if split_2_slice_size * slice_2_block_size > attention_slice_rate:
-                slice_3_block_size = split_slice_size * split_2_slice_size / 1024 / 1024 * input_element_size
-                do_split_3 = True
-                split_3_slice_size = find_slice_size(split_3_slice_size, slice_3_block_size)
-
-    return do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size
-
-
 original_torch_bmm = torch.bmm
-def torch_bmm_32_bit(input, mat2, *, out=None):
-    if input.device.type != "xpu":
-        return original_torch_bmm(input, mat2, out=out)
-    do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size = find_bmm_slice_sizes(input.shape, input.element_size(), mat2.shape)
+def torch_bmm(input, mat2, *, out=None):
+    if input.dtype != mat2.dtype:
+        mat2 = mat2.to(input.dtype)
+
+    #ARC GPUs can't allocate more than 4GB to a single block, Slice it:
+    batch_size_attention, input_tokens, mat2_shape = input.shape[0], input.shape[1], mat2.shape[2]
+    block_multiply = input.element_size()
+    slice_block_size = input_tokens * mat2_shape / 1024 / 1024 * block_multiply
+    block_size = batch_size_attention * slice_block_size
+
+    split_slice_size = batch_size_attention
+    if block_size > 4:
+        do_split = True
+        #Find something divisible with the input_tokens
+        while (split_slice_size * slice_block_size) > 4:
+            split_slice_size = split_slice_size // 2
+            if split_slice_size <= 1:
+                split_slice_size = 1
+                break
+    else:
+        do_split = False
+
+    split_2_slice_size = input_tokens
+    if split_slice_size * slice_block_size > 4:
+        slice_block_size2 = split_slice_size * mat2_shape / 1024 / 1024 * block_multiply
+        do_split_2 = True
+        #Find something divisible with the input_tokens
+        while (split_2_slice_size * slice_block_size2) > 4:
+            split_2_slice_size = split_2_slice_size // 2
+            if split_2_slice_size <= 1:
+                split_2_slice_size = 1
+                break
+    else:
+        do_split_2 = False
 
-    # Slice BMM
     if do_split:
-        batch_size_attention, input_tokens, mat2_atten_shape = input.shape[0], input.shape[1], mat2.shape[2]
         hidden_states = torch.zeros(input.shape[0], input.shape[1], mat2.shape[2], device=input.device, dtype=input.dtype)
         for i in range(batch_size_attention // split_slice_size):
             start_idx = i * split_slice_size
@@ -101,21 +48,11 @@ def torch_bmm_32_bit(input, mat2, *, out=None):
                 for i2 in range(input_tokens // split_2_slice_size): # pylint: disable=invalid-name
                     start_idx_2 = i2 * split_2_slice_size
                     end_idx_2 = (i2 + 1) * split_2_slice_size
-                    if do_split_3:
-                        for i3 in range(mat2_atten_shape // split_3_slice_size): # pylint: disable=invalid-name
-                            start_idx_3 = i3 * split_3_slice_size
-                            end_idx_3 = (i3 + 1) * split_3_slice_size
-                            hidden_states[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] = original_torch_bmm(
-                                input[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3],
-                                mat2[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3],
-                                out=out
-                            )
-                    else:
-                        hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = original_torch_bmm(
-                            input[start_idx:end_idx, start_idx_2:end_idx_2],
-                            mat2[start_idx:end_idx, start_idx_2:end_idx_2],
-                            out=out
-                        )
+                    hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = original_torch_bmm(
+                        input[start_idx:end_idx, start_idx_2:end_idx_2],
+                        mat2[start_idx:end_idx, start_idx_2:end_idx_2],
+                        out=out
+                    )
             else:
                 hidden_states[start_idx:end_idx] = original_torch_bmm(
                     input[start_idx:end_idx],
@@ -124,18 +61,49 @@ def torch_bmm_32_bit(input, mat2, *, out=None):
                 )
     else:
         return original_torch_bmm(input, mat2, out=out)
-    torch.xpu.synchronize(input.device)
     return hidden_states
 
 original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention
-def scaled_dot_product_attention_32_bit(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False):
-    if query.device.type != "xpu":
-        return original_scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal)
-    do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size = find_sdpa_slice_sizes(query.shape, query.element_size())
+def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False):
+    #ARC GPUs can't allocate more than 4GB to a single block, Slice it:
+    if len(query.shape) == 3:
+        batch_size_attention, query_tokens, shape_four = query.shape
+        shape_one = 1
+        no_shape_one = True
+    else:
+        shape_one, batch_size_attention, query_tokens, shape_four = query.shape
+        no_shape_one = False
+
+    block_multiply = query.element_size()
+    slice_block_size = shape_one * query_tokens * shape_four / 1024 / 1024 * block_multiply
+    block_size = batch_size_attention * slice_block_size
+
+    split_slice_size = batch_size_attention
+    if block_size > 4:
+        do_split = True
+        #Find something divisible with the shape_one
+        while (split_slice_size * slice_block_size) > 4:
+            split_slice_size = split_slice_size // 2
+            if split_slice_size <= 1:
+                split_slice_size = 1
+                break
+    else:
+        do_split = False
+
+    split_2_slice_size = query_tokens
+    if split_slice_size * slice_block_size > 4:
+        slice_block_size2 = shape_one * split_slice_size * shape_four / 1024 / 1024 * block_multiply
+        do_split_2 = True
+        #Find something divisible with the batch_size_attention
+        while (split_2_slice_size * slice_block_size2) > 4:
+            split_2_slice_size = split_2_slice_size // 2
+            if split_2_slice_size <= 1:
+                split_2_slice_size = 1
+                break
+    else:
+        do_split_2 = False
 
-    # Slice SDPA
     if do_split:
-        batch_size_attention, query_tokens, shape_three = query.shape[0], query.shape[1], query.shape[2]
         hidden_states = torch.zeros(query.shape, device=query.device, dtype=query.dtype)
         for i in range(batch_size_attention // split_slice_size):
             start_idx = i * split_slice_size
@@ -144,18 +112,7 @@ def scaled_dot_product_attention_32_bit(query, key, value, attn_mask=None, dropo
                 for i2 in range(query_tokens // split_2_slice_size): # pylint: disable=invalid-name
                     start_idx_2 = i2 * split_2_slice_size
                     end_idx_2 = (i2 + 1) * split_2_slice_size
-                    if do_split_3:
-                        for i3 in range(shape_three // split_3_slice_size): # pylint: disable=invalid-name
-                            start_idx_3 = i3 * split_3_slice_size
-                            end_idx_3 = (i3 + 1) * split_3_slice_size
-                            hidden_states[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] = original_scaled_dot_product_attention(
-                                query[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3],
-                                key[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3],
-                                value[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3],
-                                attn_mask=attn_mask[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] if attn_mask is not None else attn_mask,
-                                dropout_p=dropout_p, is_causal=is_causal
-                            )
-                    else:
+                    if no_shape_one:
                         hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = original_scaled_dot_product_attention(
                             query[start_idx:end_idx, start_idx_2:end_idx_2],
                             key[start_idx:end_idx, start_idx_2:end_idx_2],
@@ -163,15 +120,38 @@ def scaled_dot_product_attention_32_bit(query, key, value, attn_mask=None, dropo
                             attn_mask=attn_mask[start_idx:end_idx, start_idx_2:end_idx_2] if attn_mask is not None else attn_mask,
                             dropout_p=dropout_p, is_causal=is_causal
                         )
+                    else:
+                        hidden_states[:, start_idx:end_idx, start_idx_2:end_idx_2] = original_scaled_dot_product_attention(
+                            query[:, start_idx:end_idx, start_idx_2:end_idx_2],
+                            key[:, start_idx:end_idx, start_idx_2:end_idx_2],
+                            value[:, start_idx:end_idx, start_idx_2:end_idx_2],
+                            attn_mask=attn_mask[:, start_idx:end_idx, start_idx_2:end_idx_2] if attn_mask is not None else attn_mask,
+                            dropout_p=dropout_p, is_causal=is_causal
+                        )
             else:
-                hidden_states[start_idx:end_idx] = original_scaled_dot_product_attention(
-                    query[start_idx:end_idx],
-                    key[start_idx:end_idx],
-                    value[start_idx:end_idx],
-                    attn_mask=attn_mask[start_idx:end_idx] if attn_mask is not None else attn_mask,
-                    dropout_p=dropout_p, is_causal=is_causal
-                )
+                if no_shape_one:
+                    hidden_states[start_idx:end_idx] = original_scaled_dot_product_attention(
+                        query[start_idx:end_idx],
+                        key[start_idx:end_idx],
+                        value[start_idx:end_idx],
+                        attn_mask=attn_mask[start_idx:end_idx] if attn_mask is not None else attn_mask,
+                        dropout_p=dropout_p, is_causal=is_causal
+                    )
+                else:
+                    hidden_states[:, start_idx:end_idx] = original_scaled_dot_product_attention(
+                        query[:, start_idx:end_idx],
+                        key[:, start_idx:end_idx],
+                        value[:, start_idx:end_idx],
+                        attn_mask=attn_mask[:, start_idx:end_idx] if attn_mask is not None else attn_mask,
+                        dropout_p=dropout_p, is_causal=is_causal
+                    )
     else:
-        return original_scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal)
-    torch.xpu.synchronize(query.device)
+        return original_scaled_dot_product_attention(
+            query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal
+        )
     return hidden_states
+
+def attention_init():
+    #ARC GPUs can't allocate more than 4GB to a single block:
+    torch.bmm = torch_bmm
+    torch.nn.functional.scaled_dot_product_attention = scaled_dot_product_attention

library/ipex/diffusers.py

@@ -1,62 +1,10 @@
-import os
 import torch
 import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import
-import diffusers #0.24.0 # pylint: disable=import-error
+import diffusers #0.21.1 # pylint: disable=import-error
 from diffusers.models.attention_processor import Attention
-from diffusers.utils import USE_PEFT_BACKEND
-from functools import cache
 
 # pylint: disable=protected-access, missing-function-docstring, line-too-long
 
-attention_slice_rate = float(os.environ.get('IPEX_ATTENTION_SLICE_RATE', 4))
-
-@cache
-def find_slice_size(slice_size, slice_block_size):
-    while (slice_size * slice_block_size) > attention_slice_rate:
-        slice_size = slice_size // 2
-        if slice_size <= 1:
-            slice_size = 1
-            break
-    return slice_size
-
-@cache
-def find_attention_slice_sizes(query_shape, query_element_size, query_device_type, slice_size=None):
-    if len(query_shape) == 3:
-        batch_size_attention, query_tokens, shape_three = query_shape
-        shape_four = 1
-    else:
-        batch_size_attention, query_tokens, shape_three, shape_four = query_shape
-    if slice_size is not None:
-        batch_size_attention = slice_size
-
-    slice_block_size = query_tokens * shape_three * shape_four / 1024 / 1024 * query_element_size
-    block_size = batch_size_attention * slice_block_size
-
-    split_slice_size = batch_size_attention
-    split_2_slice_size = query_tokens
-    split_3_slice_size = shape_three
-
-    do_split = False
-    do_split_2 = False
-    do_split_3 = False
-
-    if query_device_type != "xpu":
-        return do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size
-
-    if block_size > attention_slice_rate:
-        do_split = True
-        split_slice_size = find_slice_size(split_slice_size, slice_block_size)
-        if split_slice_size * slice_block_size > attention_slice_rate:
-            slice_2_block_size = split_slice_size * shape_three * shape_four / 1024 / 1024 * query_element_size
-            do_split_2 = True
-            split_2_slice_size = find_slice_size(split_2_slice_size, slice_2_block_size)
-            if split_2_slice_size * slice_2_block_size > attention_slice_rate:
-                slice_3_block_size = split_slice_size * split_2_slice_size * shape_four / 1024 / 1024 * query_element_size
-                do_split_3 = True
-                split_3_slice_size = find_slice_size(split_3_slice_size, slice_3_block_size)
-
-    return do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size
-
 class SlicedAttnProcessor: # pylint: disable=too-few-public-methods
     r"""
     Processor for implementing sliced attention.
@@ -70,9 +18,7 @@ class SlicedAttnProcessor: # pylint: disable=too-few-public-methods
     def __init__(self, slice_size):
         self.slice_size = slice_size
 
-    def __call__(self, attn: Attention, hidden_states: torch.FloatTensor,
-    encoder_hidden_states=None, attention_mask=None) -> torch.FloatTensor: # pylint: disable=too-many-statements, too-many-locals, too-many-branches
-
+    def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, attention_mask=None): # pylint: disable=too-many-statements, too-many-locals, too-many-branches
         residual = hidden_states
 
         input_ndim = hidden_states.ndim
@@ -108,62 +54,49 @@ class SlicedAttnProcessor: # pylint: disable=too-few-public-methods
             (batch_size_attention, query_tokens, dim // attn.heads), device=query.device, dtype=query.dtype
         )
 
-        ####################################################################
-        # ARC GPUs can't allocate more than 4GB to a single block, Slice it:
-        _, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size = find_attention_slice_sizes(query.shape, query.element_size(), query.device.type, slice_size=self.slice_size)
+        #ARC GPUs can't allocate more than 4GB to a single block, Slice it:
+        block_multiply = query.element_size()
+        slice_block_size = self.slice_size * shape_three / 1024 / 1024 * block_multiply
+        block_size = query_tokens * slice_block_size
+        split_2_slice_size = query_tokens
+        if block_size > 4:
+            do_split_2 = True
+            #Find something divisible with the query_tokens
+            while (split_2_slice_size * slice_block_size) > 4:
+                split_2_slice_size = split_2_slice_size // 2
+                if split_2_slice_size <= 1:
+                    split_2_slice_size = 1
+                    break
+        else:
+            do_split_2 = False
+
+        for i in range(batch_size_attention // self.slice_size):
+            start_idx = i * self.slice_size
+            end_idx = (i + 1) * self.slice_size
 
-        for i in range(batch_size_attention // split_slice_size):
-            start_idx = i * split_slice_size
-            end_idx = (i + 1) * split_slice_size
             if do_split_2:
                 for i2 in range(query_tokens // split_2_slice_size): # pylint: disable=invalid-name
                     start_idx_2 = i2 * split_2_slice_size
                     end_idx_2 = (i2 + 1) * split_2_slice_size
-                    if do_split_3:
-                        for i3 in range(shape_three // split_3_slice_size): # pylint: disable=invalid-name
-                            start_idx_3 = i3 * split_3_slice_size
-                            end_idx_3 = (i3 + 1) * split_3_slice_size
 
-                            query_slice = query[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3]
-                            key_slice = key[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3]
-                            attn_mask_slice = attention_mask[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] if attention_mask is not None else None
+                    query_slice = query[start_idx:end_idx, start_idx_2:end_idx_2]
+                    key_slice = key[start_idx:end_idx, start_idx_2:end_idx_2]
+                    attn_mask_slice = attention_mask[start_idx:end_idx, start_idx_2:end_idx_2] if attention_mask is not None else None
 
-                            attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice)
-                            del query_slice
-                            del key_slice
-                            del attn_mask_slice
-                            attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3])
+                    attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice)
+                    attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx, start_idx_2:end_idx_2])
 
-                            hidden_states[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] = attn_slice
-                            del attn_slice
-                    else:
-                        query_slice = query[start_idx:end_idx, start_idx_2:end_idx_2]
-                        key_slice = key[start_idx:end_idx, start_idx_2:end_idx_2]
-                        attn_mask_slice = attention_mask[start_idx:end_idx, start_idx_2:end_idx_2] if attention_mask is not None else None
-
-                        attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice)
-                        del query_slice
-                        del key_slice
-                        del attn_mask_slice
-                        attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx, start_idx_2:end_idx_2])
-
-                        hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = attn_slice
-                        del attn_slice
-                torch.xpu.synchronize(query.device)
+                    hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = attn_slice
             else:
                 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)
-                del query_slice
-                del key_slice
-                del attn_mask_slice
+
                 attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx])
 
                 hidden_states[start_idx:end_idx] = attn_slice
-                del attn_slice
-        ####################################################################
 
         hidden_states = attn.batch_to_head_dim(hidden_states)
 
@@ -182,131 +115,6 @@ class SlicedAttnProcessor: # pylint: disable=too-few-public-methods
 
         return hidden_states
 
-
-class AttnProcessor:
-    r"""
-    Default processor for performing attention-related computations.
-    """
-
-    def __call__(self, attn: Attention, hidden_states: torch.FloatTensor,
-    encoder_hidden_states=None, attention_mask=None,
-    temb=None, scale: float = 1.0) -> torch.Tensor: # pylint: disable=too-many-statements, too-many-locals, too-many-branches
-
-        residual = hidden_states
-
-        args = () if USE_PEFT_BACKEND else (scale,)
-
-        if attn.spatial_norm is not None:
-            hidden_states = attn.spatial_norm(hidden_states, temb)
-
-        input_ndim = hidden_states.ndim
-
-        if input_ndim == 4:
-            batch_size, channel, height, width = hidden_states.shape
-            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
-
-        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)
-
-        if attn.group_norm is not None:
-            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
-
-        query = attn.to_q(hidden_states, *args)
-
-        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)
-
-        key = attn.to_k(encoder_hidden_states, *args)
-        value = attn.to_v(encoder_hidden_states, *args)
-
-        query = attn.head_to_batch_dim(query)
-        key = attn.head_to_batch_dim(key)
-        value = attn.head_to_batch_dim(value)
-
-        ####################################################################
-        # ARC GPUs can't allocate more than 4GB to a single block, Slice it:
-        batch_size_attention, query_tokens, shape_three = query.shape[0], query.shape[1], query.shape[2]
-        hidden_states = torch.zeros(query.shape, device=query.device, dtype=query.dtype)
-        do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size = find_attention_slice_sizes(query.shape, query.element_size(), query.device.type)
-
-        if do_split:
-            for i in range(batch_size_attention // split_slice_size):
-                start_idx = i * split_slice_size
-                end_idx = (i + 1) * split_slice_size
-                if do_split_2:
-                    for i2 in range(query_tokens // split_2_slice_size): # pylint: disable=invalid-name
-                        start_idx_2 = i2 * split_2_slice_size
-                        end_idx_2 = (i2 + 1) * split_2_slice_size
-                        if do_split_3:
-                            for i3 in range(shape_three // split_3_slice_size): # pylint: disable=invalid-name
-                                start_idx_3 = i3 * split_3_slice_size
-                                end_idx_3 = (i3 + 1) * split_3_slice_size
-
-                                query_slice = query[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3]
-                                key_slice = key[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3]
-                                attn_mask_slice = attention_mask[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] if attention_mask is not None else None
-
-                                attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice)
-                                del query_slice
-                                del key_slice
-                                del attn_mask_slice
-                                attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3])
-
-                                hidden_states[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] = attn_slice
-                                del attn_slice
-                        else:
-                            query_slice = query[start_idx:end_idx, start_idx_2:end_idx_2]
-                            key_slice = key[start_idx:end_idx, start_idx_2:end_idx_2]
-                            attn_mask_slice = attention_mask[start_idx:end_idx, start_idx_2:end_idx_2] if attention_mask is not None else None
-
-                            attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice)
-                            del query_slice
-                            del key_slice
-                            del attn_mask_slice
-                            attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx, start_idx_2:end_idx_2])
-
-                            hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = attn_slice
-                            del attn_slice
-                else:
-                    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)
-                    del query_slice
-                    del key_slice
-                    del attn_mask_slice
-                    attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx])
-
-                    hidden_states[start_idx:end_idx] = attn_slice
-                    del attn_slice
-            torch.xpu.synchronize(query.device)
-        else:
-            attention_probs = attn.get_attention_scores(query, key, attention_mask)
-            hidden_states = torch.bmm(attention_probs, value)
-        ####################################################################
-        hidden_states = attn.batch_to_head_dim(hidden_states)
-
-        # linear proj
-        hidden_states = attn.to_out[0](hidden_states, *args)
-        # dropout
-        hidden_states = attn.to_out[1](hidden_states)
-
-        if input_ndim == 4:
-            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
-
-        if attn.residual_connection:
-            hidden_states = hidden_states + residual
-
-        hidden_states = hidden_states / attn.rescale_output_factor
-
-        return hidden_states
-
 def ipex_diffusers():
     #ARC GPUs can't allocate more than 4GB to a single block:
     diffusers.models.attention_processor.SlicedAttnProcessor = SlicedAttnProcessor
-    diffusers.models.attention_processor.AttnProcessor = AttnProcessor

library/ipex/gradscaler.py

@@ -5,7 +5,6 @@ import intel_extension_for_pytorch._C as core # pylint: disable=import-error, un
 
 # pylint: disable=protected-access, missing-function-docstring, line-too-long
 
-device_supports_fp64 = torch.xpu.has_fp64_dtype()
 OptState = ipex.cpu.autocast._grad_scaler.OptState
 _MultiDeviceReplicator = ipex.cpu.autocast._grad_scaler._MultiDeviceReplicator
 _refresh_per_optimizer_state = ipex.cpu.autocast._grad_scaler._refresh_per_optimizer_state
@@ -97,10 +96,7 @@ def unscale_(self, optimizer):
 
     # FP32 division can be imprecise for certain compile options, so we carry out the reciprocal in FP64.
     assert self._scale is not None
-    if device_supports_fp64:
-        inv_scale = self._scale.double().reciprocal().float()
-    else:
-        inv_scale = self._scale.to("cpu").double().reciprocal().float().to(self._scale.device)
+    inv_scale = self._scale.to("cpu").double().reciprocal().float().to(self._scale.device)
     found_inf = torch.full(
         (1,), 0.0, dtype=torch.float32, device=self._scale.device
     )

library/ipex/hijacks.py

@@ -1,26 +1,75 @@
-import os
-from functools import wraps
-from contextlib import nullcontext
+import contextlib
+import importlib
 import torch
 import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import
-import numpy as np
-
-device_supports_fp64 = torch.xpu.has_fp64_dtype()
 
 # pylint: disable=protected-access, missing-function-docstring, line-too-long, unnecessary-lambda, no-else-return
 
+class CondFunc: # pylint: disable=missing-class-docstring
+    def __new__(cls, orig_func, sub_func, cond_func):
+        self = super(CondFunc, cls).__new__(cls)
+        if isinstance(orig_func, str):
+            func_path = orig_func.split('.')
+            for i in range(len(func_path)-1, -1, -1):
+                try:
+                    resolved_obj = importlib.import_module('.'.join(func_path[:i]))
+                    break
+                except ImportError:
+                    pass
+            for attr_name in func_path[i:-1]:
+                resolved_obj = getattr(resolved_obj, attr_name)
+            orig_func = getattr(resolved_obj, func_path[-1])
+            setattr(resolved_obj, func_path[-1], lambda *args, **kwargs: self(*args, **kwargs))
+        self.__init__(orig_func, sub_func, cond_func)
+        return lambda *args, **kwargs: self(*args, **kwargs)
+    def __init__(self, orig_func, sub_func, cond_func):
+        self.__orig_func = orig_func
+        self.__sub_func = sub_func
+        self.__cond_func = cond_func
+    def __call__(self, *args, **kwargs):
+        if not self.__cond_func or self.__cond_func(self.__orig_func, *args, **kwargs):
+            return self.__sub_func(self.__orig_func, *args, **kwargs)
+        else:
+            return self.__orig_func(*args, **kwargs)
+
+_utils = torch.utils.data._utils
+def _shutdown_workers(self):
+    if torch.utils.data._utils is None or torch.utils.data._utils.python_exit_status is True or torch.utils.data._utils.python_exit_status is None:
+        return
+    if hasattr(self, "_shutdown") and not self._shutdown:
+        self._shutdown = True
+        try:
+            if hasattr(self, '_pin_memory_thread'):
+                self._pin_memory_thread_done_event.set()
+                self._worker_result_queue.put((None, None))
+                self._pin_memory_thread.join()
+                self._worker_result_queue.cancel_join_thread()
+                self._worker_result_queue.close()
+            self._workers_done_event.set()
+            for worker_id in range(len(self._workers)):
+                if self._persistent_workers or self._workers_status[worker_id]:
+                    self._mark_worker_as_unavailable(worker_id, shutdown=True)
+            for w in self._workers: # pylint: disable=invalid-name
+                w.join(timeout=torch.utils.data._utils.MP_STATUS_CHECK_INTERVAL)
+            for q in self._index_queues: # pylint: disable=invalid-name
+                q.cancel_join_thread()
+                q.close()
+        finally:
+            if self._worker_pids_set:
+                torch.utils.data._utils.signal_handling._remove_worker_pids(id(self))
+                self._worker_pids_set = False
+            for w in self._workers: # pylint: disable=invalid-name
+                if w.is_alive():
+                    w.terminate()
+
 class DummyDataParallel(torch.nn.Module): # pylint: disable=missing-class-docstring, unused-argument, too-few-public-methods
     def __new__(cls, module, device_ids=None, output_device=None, dim=0): # pylint: disable=unused-argument
         if isinstance(device_ids, list) and len(device_ids) > 1:
-            logger.error("IPEX backend doesn't support DataParallel on multiple XPU devices")
+            print("IPEX backend doesn't support DataParallel on multiple XPU devices")
         return module.to("xpu")
 
 def return_null_context(*args, **kwargs): # pylint: disable=unused-argument
-    return nullcontext()
-
-@property
-def is_cuda(self):
-    return self.device.type == 'xpu' or self.device.type == 'cuda'
+    return contextlib.nullcontext()
 
 def check_device(device):
     return bool((isinstance(device, torch.device) and device.type == "cuda") or (isinstance(device, str) and "cuda" in device) or isinstance(device, int))
@@ -28,19 +77,26 @@ def check_device(device):
 def return_xpu(device):
     return f"xpu:{device.split(':')[-1]}" if isinstance(device, str) and ":" in device else f"xpu:{device}" if isinstance(device, int) else torch.device("xpu") if isinstance(device, torch.device) else "xpu"
 
+def ipex_no_cuda(orig_func, *args, **kwargs):
+    torch.cuda.is_available = lambda: False
+    orig_func(*args, **kwargs)
+    torch.cuda.is_available = torch.xpu.is_available
 
-# Autocast
-original_autocast_init = torch.amp.autocast_mode.autocast.__init__
-@wraps(torch.amp.autocast_mode.autocast.__init__)
-def autocast_init(self, device_type, dtype=None, enabled=True, cache_enabled=None):
-    if device_type == "cuda":
-        return original_autocast_init(self, device_type="xpu", dtype=dtype, enabled=enabled, cache_enabled=cache_enabled)
+original_autocast = torch.autocast
+def ipex_autocast(*args, **kwargs):
+    if len(args) > 0 and args[0] == "cuda":
+        return original_autocast("xpu", *args[1:], **kwargs)
     else:
-        return original_autocast_init(self, device_type=device_type, dtype=dtype, enabled=enabled, cache_enabled=cache_enabled)
+        return original_autocast(*args, **kwargs)
+
+original_torch_cat = torch.cat
+def torch_cat(tensor, *args, **kwargs):
+    if len(tensor) == 3 and (tensor[0].dtype != tensor[1].dtype or tensor[2].dtype != tensor[1].dtype):
+        return original_torch_cat([tensor[0].to(tensor[1].dtype), tensor[1], tensor[2].to(tensor[1].dtype)], *args, **kwargs)
+    else:
+        return original_torch_cat(tensor, *args, **kwargs)
 
-# Latent Antialias CPU Offload:
 original_interpolate = torch.nn.functional.interpolate
-@wraps(torch.nn.functional.interpolate)
 def interpolate(tensor, size=None, scale_factor=None, mode='nearest', align_corners=None, recompute_scale_factor=None, antialias=False): # pylint: disable=too-many-arguments
     if antialias or align_corners is not None:
         return_device = tensor.device
@@ -51,248 +107,90 @@ def interpolate(tensor, size=None, scale_factor=None, mode='nearest', align_corn
         return original_interpolate(tensor, size=size, scale_factor=scale_factor, mode=mode,
         align_corners=align_corners, recompute_scale_factor=recompute_scale_factor, antialias=antialias)
 
-
-# Diffusers Float64 (Alchemist GPUs doesn't support 64 bit):
-original_from_numpy = torch.from_numpy
-@wraps(torch.from_numpy)
-def from_numpy(ndarray):
-    if ndarray.dtype == float:
-        return original_from_numpy(ndarray.astype('float32'))
+original_linalg_solve = torch.linalg.solve
+def linalg_solve(A, B, *args, **kwargs): # pylint: disable=invalid-name
+    if A.device != torch.device("cpu") or B.device != torch.device("cpu"):
+        return_device = A.device
+        return original_linalg_solve(A.to("cpu"), B.to("cpu"), *args, **kwargs).to(return_device)
     else:
-        return original_from_numpy(ndarray)
+        return original_linalg_solve(A, B, *args, **kwargs)
 
-original_as_tensor = torch.as_tensor
-@wraps(torch.as_tensor)
-def as_tensor(data, dtype=None, device=None):
-    if check_device(device):
-        device = return_xpu(device)
-    if isinstance(data, np.ndarray) and data.dtype == float and not (
-        (isinstance(device, torch.device) and device.type == "cpu") or (isinstance(device, str) and "cpu" in device)):
-        return original_as_tensor(data, dtype=torch.float32, device=device)
-    else:
-        return original_as_tensor(data, dtype=dtype, device=device)
-
-
-if device_supports_fp64 and os.environ.get('IPEX_FORCE_ATTENTION_SLICE', None) is None:
-    original_torch_bmm = torch.bmm
-    original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention
-else:
-    # 32 bit attention workarounds for Alchemist:
-    try:
-        from .attention import torch_bmm_32_bit as original_torch_bmm
-        from .attention import scaled_dot_product_attention_32_bit as original_scaled_dot_product_attention
-    except Exception: # pylint: disable=broad-exception-caught
-        original_torch_bmm = torch.bmm
-        original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention
-
-
-# Data Type Errors:
-@wraps(torch.bmm)
-def torch_bmm(input, mat2, *, out=None):
-    if input.dtype != mat2.dtype:
-        mat2 = mat2.to(input.dtype)
-    return original_torch_bmm(input, mat2, out=out)
-
-@wraps(torch.nn.functional.scaled_dot_product_attention)
-def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False):
-    if query.dtype != key.dtype:
-        key = key.to(dtype=query.dtype)
-    if query.dtype != value.dtype:
-        value = value.to(dtype=query.dtype)
-    if attn_mask is not None and query.dtype != attn_mask.dtype:
-        attn_mask = attn_mask.to(dtype=query.dtype)
-    return original_scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal)
-
-# A1111 FP16
-original_functional_group_norm = torch.nn.functional.group_norm
-@wraps(torch.nn.functional.group_norm)
-def functional_group_norm(input, num_groups, weight=None, bias=None, eps=1e-05):
-    if weight is not None and input.dtype != weight.data.dtype:
-        input = input.to(dtype=weight.data.dtype)
-    if bias is not None and weight is not None and bias.data.dtype != weight.data.dtype:
-        bias.data = bias.data.to(dtype=weight.data.dtype)
-    return original_functional_group_norm(input, num_groups, weight=weight, bias=bias, eps=eps)
-
-# A1111 BF16
-original_functional_layer_norm = torch.nn.functional.layer_norm
-@wraps(torch.nn.functional.layer_norm)
-def functional_layer_norm(input, normalized_shape, weight=None, bias=None, eps=1e-05):
-    if weight is not None and input.dtype != weight.data.dtype:
-        input = input.to(dtype=weight.data.dtype)
-    if bias is not None and weight is not None and bias.data.dtype != weight.data.dtype:
-        bias.data = bias.data.to(dtype=weight.data.dtype)
-    return original_functional_layer_norm(input, normalized_shape, weight=weight, bias=bias, eps=eps)
-
-# Training
-original_functional_linear = torch.nn.functional.linear
-@wraps(torch.nn.functional.linear)
-def functional_linear(input, weight, bias=None):
-    if input.dtype != weight.data.dtype:
-        input = input.to(dtype=weight.data.dtype)
-    if bias is not None and bias.data.dtype != weight.data.dtype:
-        bias.data = bias.data.to(dtype=weight.data.dtype)
-    return original_functional_linear(input, weight, bias=bias)
-
-original_functional_conv2d = torch.nn.functional.conv2d
-@wraps(torch.nn.functional.conv2d)
-def functional_conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
-    if input.dtype != weight.data.dtype:
-        input = input.to(dtype=weight.data.dtype)
-    if bias is not None and bias.data.dtype != weight.data.dtype:
-        bias.data = bias.data.to(dtype=weight.data.dtype)
-    return original_functional_conv2d(input, weight, bias=bias, stride=stride, padding=padding, dilation=dilation, groups=groups)
-
-# A1111 Embedding BF16
-original_torch_cat = torch.cat
-@wraps(torch.cat)
-def torch_cat(tensor, *args, **kwargs):
-    if len(tensor) == 3 and (tensor[0].dtype != tensor[1].dtype or tensor[2].dtype != tensor[1].dtype):
-        return original_torch_cat([tensor[0].to(tensor[1].dtype), tensor[1], tensor[2].to(tensor[1].dtype)], *args, **kwargs)
-    else:
-        return original_torch_cat(tensor, *args, **kwargs)
-
-# SwinIR BF16:
-original_functional_pad = torch.nn.functional.pad
-@wraps(torch.nn.functional.pad)
-def functional_pad(input, pad, mode='constant', value=None):
-    if mode == 'reflect' and input.dtype == torch.bfloat16:
-        return original_functional_pad(input.to(torch.float32), pad, mode=mode, value=value).to(dtype=torch.bfloat16)
-    else:
-        return original_functional_pad(input, pad, mode=mode, value=value)
-
-
-original_torch_tensor = torch.tensor
-@wraps(torch.tensor)
-def torch_tensor(data, *args, dtype=None, device=None, **kwargs):
-    if check_device(device):
-        device = return_xpu(device)
-    if not device_supports_fp64:
-        if (isinstance(device, torch.device) and device.type == "xpu") or (isinstance(device, str) and "xpu" in device):
-            if dtype == torch.float64:
-                dtype = torch.float32
-            elif dtype is None and (hasattr(data, "dtype") and (data.dtype == torch.float64 or data.dtype == float)):
-                dtype = torch.float32
-    return original_torch_tensor(data, *args, dtype=dtype, device=device, **kwargs)
-
-original_Tensor_to = torch.Tensor.to
-@wraps(torch.Tensor.to)
-def Tensor_to(self, device=None, *args, **kwargs):
-    if check_device(device):
-        return original_Tensor_to(self, return_xpu(device), *args, **kwargs)
-    else:
-        return original_Tensor_to(self, device, *args, **kwargs)
-
-original_Tensor_cuda = torch.Tensor.cuda
-@wraps(torch.Tensor.cuda)
-def Tensor_cuda(self, device=None, *args, **kwargs):
-    if check_device(device):
-        return original_Tensor_cuda(self, return_xpu(device), *args, **kwargs)
-    else:
-        return original_Tensor_cuda(self, device, *args, **kwargs)
-
-original_UntypedStorage_init = torch.UntypedStorage.__init__
-@wraps(torch.UntypedStorage.__init__)
-def UntypedStorage_init(*args, device=None, **kwargs):
-    if check_device(device):
-        return original_UntypedStorage_init(*args, device=return_xpu(device), **kwargs)
-    else:
-        return original_UntypedStorage_init(*args, device=device, **kwargs)
-
-original_UntypedStorage_cuda = torch.UntypedStorage.cuda
-@wraps(torch.UntypedStorage.cuda)
-def UntypedStorage_cuda(self, device=None, *args, **kwargs):
-    if check_device(device):
-        return original_UntypedStorage_cuda(self, return_xpu(device), *args, **kwargs)
-    else:
-        return original_UntypedStorage_cuda(self, device, *args, **kwargs)
-
-original_torch_empty = torch.empty
-@wraps(torch.empty)
-def torch_empty(*args, device=None, **kwargs):
-    if check_device(device):
-        return original_torch_empty(*args, device=return_xpu(device), **kwargs)
-    else:
-        return original_torch_empty(*args, device=device, **kwargs)
-
-original_torch_randn = torch.randn
-@wraps(torch.randn)
-def torch_randn(*args, device=None, **kwargs):
-    if check_device(device):
-        return original_torch_randn(*args, device=return_xpu(device), **kwargs)
-    else:
-        return original_torch_randn(*args, device=device, **kwargs)
-
-original_torch_ones = torch.ones
-@wraps(torch.ones)
-def torch_ones(*args, device=None, **kwargs):
-    if check_device(device):
-        return original_torch_ones(*args, device=return_xpu(device), **kwargs)
-    else:
-        return original_torch_ones(*args, device=device, **kwargs)
-
-original_torch_zeros = torch.zeros
-@wraps(torch.zeros)
-def torch_zeros(*args, device=None, **kwargs):
-    if check_device(device):
-        return original_torch_zeros(*args, device=return_xpu(device), **kwargs)
-    else:
-        return original_torch_zeros(*args, device=device, **kwargs)
-
-original_torch_linspace = torch.linspace
-@wraps(torch.linspace)
-def torch_linspace(*args, device=None, **kwargs):
-    if check_device(device):
-        return original_torch_linspace(*args, device=return_xpu(device), **kwargs)
-    else:
-        return original_torch_linspace(*args, device=device, **kwargs)
-
-original_torch_Generator = torch.Generator
-@wraps(torch.Generator)
-def torch_Generator(device=None):
-    if check_device(device):
-        return original_torch_Generator(return_xpu(device))
-    else:
-        return original_torch_Generator(device)
-
-original_torch_load = torch.load
-@wraps(torch.load)
-def torch_load(f, map_location=None, pickle_module=None, *, weights_only=False, mmap=None, **kwargs):
-    if check_device(map_location):
-        return original_torch_load(f, map_location=return_xpu(map_location), pickle_module=pickle_module, weights_only=weights_only, mmap=mmap, **kwargs)
-    else:
-        return original_torch_load(f, map_location=map_location, pickle_module=pickle_module, weights_only=weights_only, mmap=mmap, **kwargs)
-
-
-# Hijack Functions:
 def ipex_hijacks():
-    torch.tensor = torch_tensor
-    torch.Tensor.to = Tensor_to
-    torch.Tensor.cuda = Tensor_cuda
-    torch.UntypedStorage.__init__ = UntypedStorage_init
-    torch.UntypedStorage.cuda = UntypedStorage_cuda
-    torch.empty = torch_empty
-    torch.randn = torch_randn
-    torch.ones = torch_ones
-    torch.zeros = torch_zeros
-    torch.linspace = torch_linspace
-    torch.Generator = torch_Generator
-    torch.load = torch_load
+    CondFunc('torch.Tensor.to',
+        lambda orig_func, self, device=None, *args, **kwargs: orig_func(self, return_xpu(device), *args, **kwargs),
+        lambda orig_func, self, device=None, *args, **kwargs: check_device(device))
+    CondFunc('torch.Tensor.cuda',
+        lambda orig_func, self, device=None, *args, **kwargs: orig_func(self, return_xpu(device), *args, **kwargs),
+        lambda orig_func, self, device=None, *args, **kwargs: check_device(device))
+    CondFunc('torch.empty',
+        lambda orig_func, *args, device=None, **kwargs: orig_func(*args, device=return_xpu(device), **kwargs),
+        lambda orig_func, *args, device=None, **kwargs: check_device(device))
+    CondFunc('torch.load',
+        lambda orig_func, *args, map_location=None, **kwargs: orig_func(*args, return_xpu(map_location), **kwargs),
+        lambda orig_func, *args, map_location=None, **kwargs: map_location is None or check_device(map_location))
+    CondFunc('torch.randn',
+        lambda orig_func, *args, device=None, **kwargs: orig_func(*args, device=return_xpu(device), **kwargs),
+        lambda orig_func, *args, device=None, **kwargs: check_device(device))
+    CondFunc('torch.ones',
+        lambda orig_func, *args, device=None, **kwargs: orig_func(*args, device=return_xpu(device), **kwargs),
+        lambda orig_func, *args, device=None, **kwargs: check_device(device))
+    CondFunc('torch.zeros',
+        lambda orig_func, *args, device=None, **kwargs: orig_func(*args, device=return_xpu(device), **kwargs),
+        lambda orig_func, *args, device=None, **kwargs: check_device(device))
+    CondFunc('torch.tensor',
+        lambda orig_func, *args, device=None, **kwargs: orig_func(*args, device=return_xpu(device), **kwargs),
+        lambda orig_func, *args, device=None, **kwargs: check_device(device))
+    CondFunc('torch.linspace',
+        lambda orig_func, *args, device=None, **kwargs: orig_func(*args, device=return_xpu(device), **kwargs),
+        lambda orig_func, *args, device=None, **kwargs: check_device(device))
 
-    torch.backends.cuda.sdp_kernel = return_null_context
+    CondFunc('torch.Generator',
+        lambda orig_func, device=None: torch.xpu.Generator(device),
+        lambda orig_func, device=None: device is not None and device != torch.device("cpu") and device != "cpu")
+
+    CondFunc('torch.batch_norm',
+        lambda orig_func, input, weight, bias, *args, **kwargs: orig_func(input,
+        weight if weight is not None else torch.ones(input.size()[1], device=input.device),
+        bias if bias is not None else torch.zeros(input.size()[1], device=input.device), *args, **kwargs),
+        lambda orig_func, input, *args, **kwargs: input.device != torch.device("cpu"))
+    CondFunc('torch.instance_norm',
+        lambda orig_func, input, weight, bias, *args, **kwargs: orig_func(input,
+        weight if weight is not None else torch.ones(input.size()[1], device=input.device),
+        bias if bias is not None else torch.zeros(input.size()[1], device=input.device), *args, **kwargs),
+        lambda orig_func, input, *args, **kwargs: input.device != torch.device("cpu"))
+
+    #Functions with dtype errors:
+    CondFunc('torch.nn.modules.GroupNorm.forward',
+        lambda orig_func, self, input: orig_func(self, input.to(self.weight.data.dtype)),
+        lambda orig_func, self, input: input.dtype != self.weight.data.dtype)
+    CondFunc('torch.nn.modules.linear.Linear.forward',
+        lambda orig_func, self, input: orig_func(self, input.to(self.weight.data.dtype)),
+        lambda orig_func, self, input: input.dtype != self.weight.data.dtype)
+    CondFunc('torch.nn.modules.conv.Conv2d.forward',
+        lambda orig_func, self, input: orig_func(self, input.to(self.weight.data.dtype)),
+        lambda orig_func, self, input: input.dtype != self.weight.data.dtype)
+    CondFunc('torch.nn.functional.layer_norm',
+        lambda orig_func, input, normalized_shape=None, weight=None, *args, **kwargs:
+        orig_func(input.to(weight.data.dtype), normalized_shape, weight, *args, **kwargs),
+        lambda orig_func, input, normalized_shape=None, weight=None, *args, **kwargs:
+        weight is not None and input.dtype != weight.data.dtype)
+
+    #Diffusers Float64 (ARC GPUs doesn't support double or Float64):
+    if not torch.xpu.has_fp64_dtype():
+        CondFunc('torch.from_numpy',
+        lambda orig_func, ndarray: orig_func(ndarray.astype('float32')),
+        lambda orig_func, ndarray: ndarray.dtype == float)
+
+    #Broken functions when torch.cuda.is_available is True:
+    CondFunc('torch.utils.data.dataloader._BaseDataLoaderIter.__init__',
+        lambda orig_func, *args, **kwargs: ipex_no_cuda(orig_func, *args, **kwargs),
+        lambda orig_func, *args, **kwargs: True)
+
+    #Functions that make compile mad with CondFunc:
+    torch.utils.data.dataloader._MultiProcessingDataLoaderIter._shutdown_workers = _shutdown_workers
     torch.nn.DataParallel = DummyDataParallel
-    torch.UntypedStorage.is_cuda = is_cuda
-    torch.amp.autocast_mode.autocast.__init__ = autocast_init
-
-    torch.nn.functional.scaled_dot_product_attention = scaled_dot_product_attention
-    torch.nn.functional.group_norm = functional_group_norm
-    torch.nn.functional.layer_norm = functional_layer_norm
-    torch.nn.functional.linear = functional_linear
-    torch.nn.functional.conv2d = functional_conv2d
-    torch.nn.functional.interpolate = interpolate
-    torch.nn.functional.pad = functional_pad
-
-    torch.bmm = torch_bmm
+    torch.autocast = ipex_autocast
     torch.cat = torch_cat
-    if not device_supports_fp64:
-        torch.from_numpy = from_numpy
-        torch.as_tensor = as_tensor
+    torch.linalg.solve = linalg_solve
+    torch.nn.functional.interpolate = interpolate
+    torch.backends.cuda.sdp_kernel = return_null_context

library/lpw_stable_diffusion.py

@@ -9,7 +9,7 @@ import numpy as np
 import PIL.Image
 import torch
 from packaging import version
-from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
+from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
 
 import diffusers
 from diffusers import SchedulerMixin, StableDiffusionPipeline
@@ -17,6 +17,7 @@ from diffusers.models import AutoencoderKL, UNet2DConditionModel
 from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput, StableDiffusionSafetyChecker
 from diffusers.utils import logging
 
+
 try:
     from diffusers.utils import PIL_INTERPOLATION
 except ImportError:
@@ -519,7 +520,6 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
         safety_checker: StableDiffusionSafetyChecker,
         feature_extractor: CLIPFeatureExtractor,
         requires_safety_checker: bool = True,
-        image_encoder: CLIPVisionModelWithProjection = None,
         clip_skip: int = 1,
     ):
         super().__init__(
@@ -531,11 +531,32 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
             safety_checker=safety_checker,
             feature_extractor=feature_extractor,
             requires_safety_checker=requires_safety_checker,
-            image_encoder=image_encoder,
         )
-        self.custom_clip_skip = clip_skip
+        self.clip_skip = clip_skip
         self.__init__additional__()
 
+    # else:
+    #     def __init__(
+    #         self,
+    #         vae: AutoencoderKL,
+    #         text_encoder: CLIPTextModel,
+    #         tokenizer: CLIPTokenizer,
+    #         unet: UNet2DConditionModel,
+    #         scheduler: SchedulerMixin,
+    #         safety_checker: StableDiffusionSafetyChecker,
+    #         feature_extractor: CLIPFeatureExtractor,
+    #     ):
+    #         super().__init__(
+    #             vae=vae,
+    #             text_encoder=text_encoder,
+    #             tokenizer=tokenizer,
+    #             unet=unet,
+    #             scheduler=scheduler,
+    #             safety_checker=safety_checker,
+    #             feature_extractor=feature_extractor,
+    #         )
+    #         self.__init__additional__()
+
     def __init__additional__(self):
         if not hasattr(self, "vae_scale_factor"):
             setattr(self, "vae_scale_factor", 2 ** (len(self.vae.config.block_out_channels) - 1))
@@ -603,7 +624,7 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
             prompt=prompt,
             uncond_prompt=negative_prompt if do_classifier_free_guidance else None,
             max_embeddings_multiples=max_embeddings_multiples,
-            clip_skip=self.custom_clip_skip,
+            clip_skip=self.clip_skip,
         )
         bs_embed, seq_len, _ = text_embeddings.shape
         text_embeddings = text_embeddings.repeat(1, num_images_per_prompt, 1)
@@ -625,7 +646,7 @@ class StableDiffusionLongPromptWeightingPipeline(StableDiffusionPipeline):
             raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")
 
         if height % 8 != 0 or width % 8 != 0:
-            logger.info(f'{height} {width}')
+            print(height, width)
             raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
 
         if (callback_steps is None) or (

library/model_util.py

@@ -3,20 +3,19 @@
 
 import math
 import os
-
 import torch
-from library.device_utils import init_ipex
-init_ipex()
-
+try:
+    import intel_extension_for_pytorch as ipex
+    if torch.xpu.is_available():
+        from library.ipex import ipex_init
+        ipex_init()
+except Exception:
+    pass
 import diffusers
 from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig, logging
 from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline  # , UNet2DConditionModel
 from safetensors.torch import load_file, save_file
 from library.original_unet import UNet2DConditionModel
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
 
 # DiffUsers版StableDiffusionのモデルパラメータ
 NUM_TRAIN_TIMESTEPS = 1000
@@ -572,9 +571,9 @@ def convert_ldm_clip_checkpoint_v1(checkpoint):
         if key.startswith("cond_stage_model.transformer"):
             text_model_dict[key[len("cond_stage_model.transformer.") :]] = checkpoint[key]
 
-    # remove position_ids for newer transformer, which causes error :(
-    if "text_model.embeddings.position_ids" in text_model_dict:
-        text_model_dict.pop("text_model.embeddings.position_ids")
+    # 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
 
     return text_model_dict
 
@@ -948,7 +947,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:
-                # logger.info(f"Reshaping {k} for SD format: shape {v.shape} -> {v.shape} x 1 x 1")
+                # print(f"Reshaping {k} for SD format: shape {v.shape} -> {v.shape} x 1 x 1")
                 new_state_dict[k] = reshape_weight_for_sd(v)
 
     return new_state_dict
@@ -1006,7 +1005,7 @@ def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, device="cpu", dt
 
     unet = UNet2DConditionModel(**unet_config).to(device)
     info = unet.load_state_dict(converted_unet_checkpoint)
-    logger.info(f"loading u-net: {info}")
+    print("loading u-net:", info)
 
     # Convert the VAE model.
     vae_config = create_vae_diffusers_config()
@@ -1014,7 +1013,7 @@ def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, device="cpu", dt
 
     vae = AutoencoderKL(**vae_config).to(device)
     info = vae.load_state_dict(converted_vae_checkpoint)
-    logger.info(f"loading vae: {info}")
+    print("loading vae:", info)
 
     # convert text_model
     if v2:
@@ -1048,7 +1047,7 @@ def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, device="cpu", dt
         # 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()
-        # logger.info(f"config: {text_model.config}")
+        # print(f"config: {text_model.config}")
         cfg = CLIPTextConfig(
             vocab_size=49408,
             hidden_size=768,
@@ -1071,7 +1070,7 @@ def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, device="cpu", dt
         )
         text_model = CLIPTextModel._from_config(cfg)
         info = text_model.load_state_dict(converted_text_encoder_checkpoint)
-    logger.info(f"loading text encoder: {info}")
+    print("loading text encoder:", info)
 
     return text_model, vae, unet
 
@@ -1146,7 +1145,7 @@ def convert_text_encoder_state_dict_to_sd_v2(checkpoint, make_dummy_weights=Fals
 
     # 最後の層などを捏造するか
     if make_dummy_weights:
-        logger.info("make dummy weights for resblock.23, text_projection and logit scale.")
+        print("make dummy weights for resblock.23, text_projection and logit scale.")
         keys = list(new_sd.keys())
         for key in keys:
             if key.startswith("transformer.resblocks.22."):
@@ -1243,13 +1242,8 @@ def save_diffusers_checkpoint(v2, output_dir, text_encoder, unet, pretrained_mod
     if vae is None:
         vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae")
 
-    # original U-Net cannot be saved, so we need to convert it to the Diffusers version
-    # TODO this consumes a lot of memory
-    diffusers_unet = diffusers.UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, subfolder="unet")
-    diffusers_unet.load_state_dict(unet.state_dict())
-
     pipeline = StableDiffusionPipeline(
-        unet=diffusers_unet,
+        unet=unet,
         text_encoder=text_encoder,
         vae=vae,
         scheduler=scheduler,
@@ -1265,14 +1259,14 @@ VAE_PREFIX = "first_stage_model."
 
 
 def load_vae(vae_id, dtype):
-    logger.info(f"load VAE: {vae_id}")
+    print(f"load VAE: {vae_id}")
     if os.path.isdir(vae_id) or not os.path.isfile(vae_id):
         # Diffusers local/remote
         try:
             vae = AutoencoderKL.from_pretrained(vae_id, subfolder=None, torch_dtype=dtype)
         except EnvironmentError as e:
-            logger.error(f"exception occurs in loading vae: {e}")
-            logger.error("retry with subfolder='vae'")
+            print(f"exception occurs in loading vae: {e}")
+            print("retry with subfolder='vae'")
             vae = AutoencoderKL.from_pretrained(vae_id, subfolder="vae", torch_dtype=dtype)
         return vae
 
@@ -1313,19 +1307,19 @@ def load_vae(vae_id, dtype):
 
 def make_bucket_resolutions(max_reso, min_size=256, max_size=1024, divisible=64):
     max_width, max_height = max_reso
-    max_area = max_width * max_height
+    max_area = (max_width // divisible) * (max_height // divisible)
 
     resos = set()
 
-    width = int(math.sqrt(max_area) // divisible) * divisible
-    resos.add((width, width))
+    size = int(math.sqrt(max_area)) * divisible
+    resos.add((size, size))
 
-    width = min_size
-    while width <= max_size:
-        height = min(max_size, int((max_area // width) // divisible) * divisible)
-        if height >= min_size:
-            resos.add((width, height))
-            resos.add((height, width))
+    size = min_size
+    while size <= max_size:
+        width = size
+        height = min(max_size, (max_area // (width // divisible)) * divisible)
+        resos.add((width, height))
+        resos.add((height, width))
 
         # # make additional resos
         # if width >= height and width - divisible >= min_size:
@@ -1335,7 +1329,7 @@ def make_bucket_resolutions(max_reso, min_size=256, max_size=1024, divisible=64)
         #   resos.add((width, height - divisible))
         #   resos.add((height - divisible, width))
 
-        width += divisible
+        size += divisible
 
     resos = list(resos)
     resos.sort()
@@ -1344,13 +1338,13 @@ def make_bucket_resolutions(max_reso, min_size=256, max_size=1024, divisible=64)
 
 if __name__ == "__main__":
     resos = make_bucket_resolutions((512, 768))
-    logger.info(f"{len(resos)}")
-    logger.info(f"{resos}")
+    print(len(resos))
+    print(resos)
     aspect_ratios = [w / h for w, h in resos]
-    logger.info(f"{aspect_ratios}")
+    print(aspect_ratios)
 
     ars = set()
     for ar in aspect_ratios:
         if ar in ars:
-            logger.error(f"error! duplicate ar: {ar}")
+            print("error! duplicate ar:", ar)
         ars.add(ar)

library/original_unet.py

@@ -113,10 +113,6 @@ import torch
 from torch import nn
 from torch.nn import functional as F
 from einops import rearrange
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
 
 BLOCK_OUT_CHANNELS: Tuple[int] = (320, 640, 1280, 1280)
 TIMESTEP_INPUT_DIM = BLOCK_OUT_CHANNELS[0]
@@ -365,23 +361,6 @@ def get_timestep_embedding(
     return emb
 
 
-# Deep Shrink: We do not common this function, because minimize dependencies.
-def resize_like(x, target, mode="bicubic", align_corners=False):
-    org_dtype = x.dtype
-    if org_dtype == torch.bfloat16:
-        x = x.to(torch.float32)
-
-    if x.shape[-2:] != target.shape[-2:]:
-        if mode == "nearest":
-            x = F.interpolate(x, size=target.shape[-2:], mode=mode)
-        else:
-            x = F.interpolate(x, size=target.shape[-2:], mode=mode, align_corners=align_corners)
-
-    if org_dtype == torch.bfloat16:
-        x = x.to(org_dtype)
-    return x
-
-
 class SampleOutput:
     def __init__(self, sample):
         self.sample = sample
@@ -590,9 +569,6 @@ class CrossAttention(nn.Module):
         self.use_memory_efficient_attention_mem_eff = False
         self.use_sdpa = False
 
-        # Attention processor
-        self.processor = None
-
     def set_use_memory_efficient_attention(self, xformers, mem_eff):
         self.use_memory_efficient_attention_xformers = xformers
         self.use_memory_efficient_attention_mem_eff = mem_eff
@@ -614,28 +590,7 @@ class CrossAttention(nn.Module):
         tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size)
         return tensor
 
-    def set_processor(self):
-        return self.processor
-
-    def get_processor(self):
-        return self.processor
-
-    def forward(self, hidden_states, context=None, mask=None, **kwargs):
-        if self.processor is not None:
-            (
-                hidden_states,
-                encoder_hidden_states,
-                attention_mask,
-            ) = translate_attention_names_from_diffusers(
-                hidden_states=hidden_states, context=context, mask=mask, **kwargs
-            )
-            return self.processor(
-                attn=self,
-                hidden_states=hidden_states,
-                encoder_hidden_states=context,
-                attention_mask=mask,
-                **kwargs
-            )
+    def forward(self, hidden_states, context=None, mask=None):
         if self.use_memory_efficient_attention_xformers:
             return self.forward_memory_efficient_xformers(hidden_states, context, mask)
         if self.use_memory_efficient_attention_mem_eff:
@@ -748,21 +703,6 @@ class CrossAttention(nn.Module):
         out = self.to_out[0](out)
         return out
 
-def translate_attention_names_from_diffusers(
-    hidden_states: torch.FloatTensor,
-    context: Optional[torch.FloatTensor] = None,
-    mask: Optional[torch.FloatTensor] = None,
-    # HF naming
-    encoder_hidden_states: Optional[torch.FloatTensor] = None,
-    attention_mask: Optional[torch.FloatTensor] = None
-):
-    # translate from hugging face diffusers
-    context = context if context is not None else encoder_hidden_states
-
-    # translate from hugging face diffusers
-    mask = mask if mask is not None else attention_mask
-
-    return hidden_states, context, mask
 
 # feedforward
 class GEGLU(nn.Module):
@@ -1190,7 +1130,6 @@ class UpBlock2D(nn.Module):
             # pop res hidden states
             res_hidden_states = res_hidden_states_tuple[-1]
             res_hidden_states_tuple = res_hidden_states_tuple[:-1]
-
             hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
 
             if self.training and self.gradient_checkpointing:
@@ -1266,9 +1205,9 @@ class CrossAttnUpBlock2D(nn.Module):
         for attn in self.attentions:
             attn.set_use_memory_efficient_attention(xformers, mem_eff)
 
-    def set_use_sdpa(self, sdpa):
+    def set_use_sdpa(self, spda):
         for attn in self.attentions:
-            attn.set_use_sdpa(sdpa)
+            attn.set_use_sdpa(spda)
 
     def forward(
         self,
@@ -1282,7 +1221,6 @@ class CrossAttnUpBlock2D(nn.Module):
             # pop res hidden states
             res_hidden_states = res_hidden_states_tuple[-1]
             res_hidden_states_tuple = res_hidden_states_tuple[:-1]
-
             hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
 
             if self.training and self.gradient_checkpointing:
@@ -1384,7 +1322,7 @@ class UNet2DConditionModel(nn.Module):
     ):
         super().__init__()
         assert sample_size is not None, "sample_size must be specified"
-        logger.info(
+        print(
             f"UNet2DConditionModel: {sample_size}, {attention_head_dim}, {cross_attention_dim}, {use_linear_projection}, {upcast_attention}"
         )
 
@@ -1393,7 +1331,7 @@ class UNet2DConditionModel(nn.Module):
         self.out_channels = OUT_CHANNELS
 
         self.sample_size = sample_size
-        self.prepare_config(sample_size=sample_size)
+        self.prepare_config()
 
         # state_dictの書式が変わるのでmoduleの持ち方は変えられない
 
@@ -1480,8 +1418,8 @@ class UNet2DConditionModel(nn.Module):
         self.conv_out = nn.Conv2d(BLOCK_OUT_CHANNELS[0], OUT_CHANNELS, kernel_size=3, padding=1)
 
     # region diffusers compatibility
-    def prepare_config(self, *args, **kwargs):
-        self.config = SimpleNamespace(**kwargs)
+    def prepare_config(self):
+        self.config = SimpleNamespace()
 
     @property
     def dtype(self) -> torch.dtype:
@@ -1518,7 +1456,7 @@ class UNet2DConditionModel(nn.Module):
     def set_gradient_checkpointing(self, value=False):
         modules = self.down_blocks + [self.mid_block] + self.up_blocks
         for module in modules:
-            logger.info(f"{module.__class__.__name__} {module.gradient_checkpointing} -> {value}")
+            print(module.__class__.__name__, module.gradient_checkpointing, "->", value)
             module.gradient_checkpointing = value
 
     # endregion
@@ -1581,6 +1519,7 @@ class UNet2DConditionModel(nn.Module):
         # 2. pre-process
         sample = self.conv_in(sample)
 
+        # 3. down
         down_block_res_samples = (sample,)
         for downsample_block in self.down_blocks:
             # downblockはforwardで必ずencoder_hidden_statesを受け取るようにしても良さそうだけど、
@@ -1665,255 +1604,3 @@ class UNet2DConditionModel(nn.Module):
         timesteps = timesteps.expand(sample.shape[0])
 
         return timesteps
-
-
-class InferUNet2DConditionModel:
-    def __init__(self, original_unet: UNet2DConditionModel):
-        self.delegate = original_unet
-
-        # override original model's forward method: because forward is not called by `__call__`
-        # overriding `__call__` is not enough, because nn.Module.forward has a special handling
-        self.delegate.forward = self.forward
-
-        # override original model's up blocks' forward method
-        for up_block in self.delegate.up_blocks:
-            if up_block.__class__.__name__ == "UpBlock2D":
-
-                def resnet_wrapper(func, block):
-                    def forward(*args, **kwargs):
-                        return func(block, *args, **kwargs)
-
-                    return forward
-
-                up_block.forward = resnet_wrapper(self.up_block_forward, up_block)
-
-            elif up_block.__class__.__name__ == "CrossAttnUpBlock2D":
-
-                def cross_attn_up_wrapper(func, block):
-                    def forward(*args, **kwargs):
-                        return func(block, *args, **kwargs)
-
-                    return forward
-
-                up_block.forward = cross_attn_up_wrapper(self.cross_attn_up_block_forward, up_block)
-
-        # Deep Shrink
-        self.ds_depth_1 = None
-        self.ds_depth_2 = None
-        self.ds_timesteps_1 = None
-        self.ds_timesteps_2 = None
-        self.ds_ratio = None
-
-    # call original model's methods
-    def __getattr__(self, name):
-        return getattr(self.delegate, name)
-
-    def __call__(self, *args, **kwargs):
-        return self.delegate(*args, **kwargs)
-
-    def set_deep_shrink(self, ds_depth_1, ds_timesteps_1=650, ds_depth_2=None, ds_timesteps_2=None, ds_ratio=0.5):
-        if ds_depth_1 is None:
-            logger.info("Deep Shrink is disabled.")
-            self.ds_depth_1 = None
-            self.ds_timesteps_1 = None
-            self.ds_depth_2 = None
-            self.ds_timesteps_2 = None
-            self.ds_ratio = None
-        else:
-            logger.info(
-                f"Deep Shrink is enabled: [depth={ds_depth_1}/{ds_depth_2}, timesteps={ds_timesteps_1}/{ds_timesteps_2}, ratio={ds_ratio}]"
-            )
-            self.ds_depth_1 = ds_depth_1
-            self.ds_timesteps_1 = ds_timesteps_1
-            self.ds_depth_2 = ds_depth_2 if ds_depth_2 is not None else -1
-            self.ds_timesteps_2 = ds_timesteps_2 if ds_timesteps_2 is not None else 1000
-            self.ds_ratio = ds_ratio
-
-    def up_block_forward(self, _self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None):
-        for resnet in _self.resnets:
-            # pop res hidden states
-            res_hidden_states = res_hidden_states_tuple[-1]
-            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
-
-            # Deep Shrink
-            if res_hidden_states.shape[-2:] != hidden_states.shape[-2:]:
-                hidden_states = resize_like(hidden_states, res_hidden_states)
-
-            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
-            hidden_states = resnet(hidden_states, temb)
-
-        if _self.upsamplers is not None:
-            for upsampler in _self.upsamplers:
-                hidden_states = upsampler(hidden_states, upsample_size)
-
-        return hidden_states
-
-    def cross_attn_up_block_forward(
-        self,
-        _self,
-        hidden_states,
-        res_hidden_states_tuple,
-        temb=None,
-        encoder_hidden_states=None,
-        upsample_size=None,
-    ):
-        for resnet, attn in zip(_self.resnets, _self.attentions):
-            # pop res hidden states
-            res_hidden_states = res_hidden_states_tuple[-1]
-            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
-
-            # Deep Shrink
-            if res_hidden_states.shape[-2:] != hidden_states.shape[-2:]:
-                hidden_states = resize_like(hidden_states, res_hidden_states)
-
-            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
-            hidden_states = resnet(hidden_states, temb)
-            hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample
-
-        if _self.upsamplers is not None:
-            for upsampler in _self.upsamplers:
-                hidden_states = upsampler(hidden_states, upsample_size)
-
-        return hidden_states
-
-    def forward(
-        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,
-        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
-        mid_block_additional_residual: Optional[torch.Tensor] = None,
-    ) -> Union[Dict, Tuple]:
-        r"""
-        current implementation is a copy of `UNet2DConditionModel.forward()` with Deep Shrink.
-        """
-
-        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.
-
-        Returns:
-            `SampleOutput` or `tuple`:
-            `SampleOutput` if `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is the sample tensor.
-        """
-
-        _self = self.delegate
-
-        # 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
-
-        # 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
-
-        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
-        timesteps = _self.handle_unusual_timesteps(sample, 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)
-
-        # 2. pre-process
-        sample = _self.conv_in(sample)
-
-        down_block_res_samples = (sample,)
-        for depth, downsample_block in enumerate(_self.down_blocks):
-            # Deep Shrink
-            if self.ds_depth_1 is not None:
-                if (depth == self.ds_depth_1 and timesteps[0] >= self.ds_timesteps_1) or (
-                    self.ds_depth_2 is not None
-                    and depth == self.ds_depth_2
-                    and timesteps[0] < self.ds_timesteps_1
-                    and timesteps[0] >= self.ds_timesteps_2
-                ):
-                    org_dtype = sample.dtype
-                    if org_dtype == torch.bfloat16:
-                        sample = sample.to(torch.float32)
-                    sample = F.interpolate(sample, scale_factor=self.ds_ratio, mode="bicubic", align_corners=False).to(org_dtype)
-
-            # 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,
-                )
-            else:
-                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
-
-            down_block_res_samples += res_samples
-
-        # skip connectionにControlNetの出力を追加する
-        if down_block_additional_residuals is not None:
-            down_block_res_samples = list(down_block_res_samples)
-            for i in range(len(down_block_res_samples)):
-                down_block_res_samples[i] += down_block_additional_residuals[i]
-            down_block_res_samples = tuple(down_block_res_samples)
-
-        # 4. mid
-        sample = _self.mid_block(sample, emb, encoder_hidden_states=encoder_hidden_states)
-
-        # ControlNetの出力を追加する
-        if mid_block_additional_residual is not None:
-            sample += mid_block_additional_residual
-
-        # 5. up
-        for i, upsample_block in enumerate(_self.up_blocks):
-            is_final_block = i == len(_self.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)]  # skip connection
-
-            # 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:]
-
-            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 = _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)

library/sai_model_spec.py

@@ -5,10 +5,6 @@ from io import BytesIO
 import os
 from typing import List, Optional, Tuple, Union
 import safetensors
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
 
 r"""
 # Metadata Example
@@ -235,7 +231,7 @@ def build_metadata(
     # # 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()]):
-        logger.error(f"Internal error: some metadata values are None: {metadata}")
+        print(f"Internal error: some metadata values are None: {metadata}")
     
     return metadata
 

library/sdxl_lpw_stable_diffusion.py

@@ -923,11 +923,7 @@ class SdxlStableDiffusionLongPromptWeightingPipeline:
             if up1 is not None:
                 uncond_pool = up1
 
-        unet_dtype = self.unet.dtype
-        dtype = unet_dtype
-        if hasattr(dtype, "itemsize") and dtype.itemsize == 1:  # fp8
-            dtype = torch.float16
-            self.unet.to(dtype)
+        dtype = self.unet.dtype
 
         # 4. Preprocess image and mask
         if isinstance(image, PIL.Image.Image):
@@ -1032,7 +1028,6 @@ class SdxlStableDiffusionLongPromptWeightingPipeline:
                 if is_cancelled_callback is not None and is_cancelled_callback():
                     return None
 
-        self.unet.to(unet_dtype)
         return latents
 
     def latents_to_image(self, latents):

library/sdxl_model_util.py

@@ -7,10 +7,7 @@ from typing import List
 from diffusers import AutoencoderKL, EulerDiscreteScheduler, UNet2DConditionModel
 from library import model_util
 from library import sdxl_original_unet
-from .utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 VAE_SCALE_FACTOR = 0.13025
 MODEL_VERSION_SDXL_BASE_V1_0 = "sdxl_base_v1-0"
@@ -103,7 +100,7 @@ def convert_sdxl_text_encoder_2_checkpoint(checkpoint, max_length):
             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: position_ids is not used in newer transformers
+            key = None  # remove this key: make position_ids by ourselves
         return key
 
     keys = list(checkpoint.keys())
@@ -129,15 +126,13 @@ def convert_sdxl_text_encoder_2_checkpoint(checkpoint, max_length):
             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)
 
-    # temporary workaround for text_projection.weight.weight for Playground-v2
-    if "text_projection.weight.weight" in new_sd:
-        logger.info("convert_sdxl_text_encoder_2_checkpoint: convert text_projection.weight.weight to text_projection.weight")
-        new_sd["text_projection.weight"] = new_sd["text_projection.weight.weight"]
-        del new_sd["text_projection.weight.weight"]
-
     return new_sd, logit_scale
 
 
@@ -189,20 +184,20 @@ def load_models_from_sdxl_checkpoint(model_version, ckpt_path, map_location, dty
         checkpoint = None
 
     # U-Net
-    logger.info("building U-Net")
+    print("building U-Net")
     with init_empty_weights():
         unet = sdxl_original_unet.SdxlUNet2DConditionModel()
 
-    logger.info("loading U-Net from checkpoint")
+    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)
-    logger.info(f"U-Net: {info}")
+    print("U-Net: ", info)
 
     # Text Encoders
-    logger.info("building text encoders")
+    print("building text encoders")
 
     # Text Encoder 1 is same to Stability AI's SDXL
     text_model1_cfg = CLIPTextConfig(
@@ -255,7 +250,7 @@ def load_models_from_sdxl_checkpoint(model_version, ckpt_path, map_location, dty
     with init_empty_weights():
         text_model2 = CLIPTextModelWithProjection(text_model2_cfg)
 
-    logger.info("loading text encoders from checkpoint")
+    print("loading text encoders from checkpoint")
     te1_sd = {}
     te2_sd = {}
     for k in list(state_dict.keys()):
@@ -263,28 +258,28 @@ def load_models_from_sdxl_checkpoint(model_version, ckpt_path, map_location, dty
             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)
-
-    # 最新の transformers では position_ids を含むとエラーになるので削除 / remove position_ids for latest transformers
-    if "text_model.embeddings.position_ids" in te1_sd:
-        te1_sd.pop("text_model.embeddings.position_ids")
+    
+    # 一部の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
-    logger.info(f"text encoder 1: {info1}")
+    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
-    logger.info(f"text encoder 2: {info2}")
+    print("text encoder 2:", info2)
 
     # prepare vae
-    logger.info("building VAE")
+    print("building VAE")
     vae_config = model_util.create_vae_diffusers_config()
     with init_empty_weights():
         vae = AutoencoderKL(**vae_config)
 
-    logger.info("loading VAE from checkpoint")
+    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)
-    logger.info(f"VAE: {info}")
+    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

library/sdxl_original_unet.py

@@ -24,16 +24,13 @@
 
 import math
 from types import SimpleNamespace
-from typing import Any, Optional
+from typing import Optional
 import torch
 import torch.utils.checkpoint
 from torch import nn
 from torch.nn import functional as F
 from einops import rearrange
-from .utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 IN_CHANNELS: int = 4
 OUT_CHANNELS: int = 4
@@ -269,23 +266,6 @@ def get_timestep_embedding(
     return emb
 
 
-# Deep Shrink: We do not common this function, because minimize dependencies.
-def resize_like(x, target, mode="bicubic", align_corners=False):
-    org_dtype = x.dtype
-    if org_dtype == torch.bfloat16:
-        x = x.to(torch.float32)
-
-    if x.shape[-2:] != target.shape[-2:]:
-        if mode == "nearest":
-            x = F.interpolate(x, size=target.shape[-2:], mode=mode)
-        else:
-            x = F.interpolate(x, size=target.shape[-2:], mode=mode, align_corners=align_corners)
-
-    if org_dtype == torch.bfloat16:
-        x = x.to(org_dtype)
-    return x
-
-
 class GroupNorm32(nn.GroupNorm):
     def forward(self, x):
         if self.weight.dtype != torch.float32:
@@ -335,7 +315,7 @@ class ResnetBlock2D(nn.Module):
 
     def forward(self, x, emb):
         if self.training and self.gradient_checkpointing:
-            # logger.info("ResnetBlock2D: gradient_checkpointing")
+            # print("ResnetBlock2D: gradient_checkpointing")
 
             def create_custom_forward(func):
                 def custom_forward(*inputs):
@@ -369,7 +349,7 @@ class Downsample2D(nn.Module):
 
     def forward(self, hidden_states):
         if self.training and self.gradient_checkpointing:
-            # logger.info("Downsample2D: gradient_checkpointing")
+            # print("Downsample2D: gradient_checkpointing")
 
             def create_custom_forward(func):
                 def custom_forward(*inputs):
@@ -656,7 +636,7 @@ class BasicTransformerBlock(nn.Module):
 
     def forward(self, hidden_states, context=None, timestep=None):
         if self.training and self.gradient_checkpointing:
-            # logger.info("BasicTransformerBlock: checkpointing")
+            # print("BasicTransformerBlock: checkpointing")
 
             def create_custom_forward(func):
                 def custom_forward(*inputs):
@@ -799,7 +779,7 @@ class Upsample2D(nn.Module):
 
     def forward(self, hidden_states, output_size=None):
         if self.training and self.gradient_checkpointing:
-            # logger.info("Upsample2D: gradient_checkpointing")
+            # print("Upsample2D: gradient_checkpointing")
 
             def create_custom_forward(func):
                 def custom_forward(*inputs):
@@ -1049,7 +1029,7 @@ class SdxlUNet2DConditionModel(nn.Module):
         for block in blocks:
             for module in block:
                 if hasattr(module, "set_use_memory_efficient_attention"):
-                    # logger.info(module.__class__.__name__)
+                    # print(module.__class__.__name__)
                     module.set_use_memory_efficient_attention(xformers, mem_eff)
 
     def set_use_sdpa(self, sdpa: bool) -> None:
@@ -1064,7 +1044,7 @@ class SdxlUNet2DConditionModel(nn.Module):
         for block in blocks:
             for module in block.modules():
                 if hasattr(module, "gradient_checkpointing"):
-                    # logger.info(f{module.__class__.__name__} {module.gradient_checkpointing} -> {value}")
+                    # print(module.__class__.__name__, module.gradient_checkpointing, "->", value)
                     module.gradient_checkpointing = value
 
     # endregion
@@ -1086,7 +1066,7 @@ class SdxlUNet2DConditionModel(nn.Module):
         def call_module(module, h, emb, context):
             x = h
             for layer in module:
-                # logger.info(layer.__class__.__name__, x.dtype, emb.dtype, context.dtype if context is not None else None)
+                # print(layer.__class__.__name__, x.dtype, emb.dtype, context.dtype if context is not None else None)
                 if isinstance(layer, ResnetBlock2D):
                     x = layer(x, emb)
                 elif isinstance(layer, Transformer2DModel):
@@ -1097,7 +1077,6 @@ class SdxlUNet2DConditionModel(nn.Module):
 
         # h = x.type(self.dtype)
         h = x
-
         for module in self.input_blocks:
             h = call_module(module, h, emb, context)
             hs.append(h)
@@ -1114,125 +1093,10 @@ class SdxlUNet2DConditionModel(nn.Module):
         return h
 
 
-class InferSdxlUNet2DConditionModel:
-    def __init__(self, original_unet: SdxlUNet2DConditionModel, **kwargs):
-        self.delegate = original_unet
-
-        # override original model's forward method: because forward is not called by `__call__`
-        # overriding `__call__` is not enough, because nn.Module.forward has a special handling
-        self.delegate.forward = self.forward
-
-        # Deep Shrink
-        self.ds_depth_1 = None
-        self.ds_depth_2 = None
-        self.ds_timesteps_1 = None
-        self.ds_timesteps_2 = None
-        self.ds_ratio = None
-
-    # call original model's methods
-    def __getattr__(self, name):
-        return getattr(self.delegate, name)
-    
-    def __call__(self, *args, **kwargs):
-        return self.delegate(*args, **kwargs)
-
-    def set_deep_shrink(self, ds_depth_1, ds_timesteps_1=650, ds_depth_2=None, ds_timesteps_2=None, ds_ratio=0.5):
-        if ds_depth_1 is None:
-            logger.info("Deep Shrink is disabled.")
-            self.ds_depth_1 = None
-            self.ds_timesteps_1 = None
-            self.ds_depth_2 = None
-            self.ds_timesteps_2 = None
-            self.ds_ratio = None
-        else:
-            logger.info(
-                f"Deep Shrink is enabled: [depth={ds_depth_1}/{ds_depth_2}, timesteps={ds_timesteps_1}/{ds_timesteps_2}, ratio={ds_ratio}]"
-            )
-            self.ds_depth_1 = ds_depth_1
-            self.ds_timesteps_1 = ds_timesteps_1
-            self.ds_depth_2 = ds_depth_2 if ds_depth_2 is not None else -1
-            self.ds_timesteps_2 = ds_timesteps_2 if ds_timesteps_2 is not None else 1000
-            self.ds_ratio = ds_ratio
-
-    def forward(self, x, timesteps=None, context=None, y=None, **kwargs):
-        r"""
-        current implementation is a copy of `SdxlUNet2DConditionModel.forward()` with Deep Shrink.
-        """
-        _self = self.delegate
-
-        # broadcast timesteps to batch dimension
-        timesteps = timesteps.expand(x.shape[0])
-
-        hs = []
-        t_emb = get_timestep_embedding(timesteps, _self.model_channels)  # , repeat_only=False)
-        t_emb = t_emb.to(x.dtype)
-        emb = _self.time_embed(t_emb)
-
-        assert x.shape[0] == y.shape[0], f"batch size mismatch: {x.shape[0]} != {y.shape[0]}"
-        assert x.dtype == y.dtype, f"dtype mismatch: {x.dtype} != {y.dtype}"
-        # assert x.dtype == _self.dtype
-        emb = emb + _self.label_emb(y)
-
-        def call_module(module, h, emb, context):
-            x = h
-            for layer in module:
-                # print(layer.__class__.__name__, x.dtype, emb.dtype, context.dtype if context is not None else None)
-                if isinstance(layer, ResnetBlock2D):
-                    x = layer(x, emb)
-                elif isinstance(layer, Transformer2DModel):
-                    x = layer(x, context)
-                else:
-                    x = layer(x)
-            return x
-
-        # h = x.type(self.dtype)
-        h = x
-
-        for depth, module in enumerate(_self.input_blocks):
-            # Deep Shrink
-            if self.ds_depth_1 is not None:
-                if (depth == self.ds_depth_1 and timesteps[0] >= self.ds_timesteps_1) or (
-                    self.ds_depth_2 is not None
-                    and depth == self.ds_depth_2
-                    and timesteps[0] < self.ds_timesteps_1
-                    and timesteps[0] >= self.ds_timesteps_2
-                ):
-                    # print("downsample", h.shape, self.ds_ratio)
-                    org_dtype = h.dtype
-                    if org_dtype == torch.bfloat16:
-                        h = h.to(torch.float32)
-                    h = F.interpolate(h, scale_factor=self.ds_ratio, mode="bicubic", align_corners=False).to(org_dtype)
-
-            h = call_module(module, h, emb, context)
-            hs.append(h)
-
-        h = call_module(_self.middle_block, h, emb, context)
-
-        for module in _self.output_blocks:
-            # Deep Shrink
-            if self.ds_depth_1 is not None:
-                if hs[-1].shape[-2:] != h.shape[-2:]:
-                    # print("upsample", h.shape, hs[-1].shape)
-                    h = resize_like(h, hs[-1])
-
-            h = torch.cat([h, hs.pop()], dim=1)
-            h = call_module(module, h, emb, context)
-
-        # Deep Shrink: in case of depth 0
-        if self.ds_depth_1 == 0 and h.shape[-2:] != x.shape[-2:]:
-            # print("upsample", h.shape, x.shape)
-            h = resize_like(h, x)
-
-        h = h.type(x.dtype)
-        h = call_module(_self.out, h, emb, context)
-
-        return h
-
-
 if __name__ == "__main__":
     import time
 
-    logger.info("create unet")
+    print("create unet")
     unet = SdxlUNet2DConditionModel()
 
     unet.to("cuda")
@@ -1241,7 +1105,7 @@ if __name__ == "__main__":
     unet.train()
 
     # 使用メモリ量確認用の疑似学習ループ
-    logger.info("preparing optimizer")
+    print("preparing optimizer")
 
     # optimizer = torch.optim.SGD(unet.parameters(), lr=1e-3, nesterov=True, momentum=0.9) # not working
 
@@ -1256,12 +1120,12 @@ if __name__ == "__main__":
 
     scaler = torch.cuda.amp.GradScaler(enabled=True)
 
-    logger.info("start training")
+    print("start training")
     steps = 10
     batch_size = 1
 
     for step in range(steps):
-        logger.info(f"step {step}")
+        print(f"step {step}")
         if step == 1:
             time_start = time.perf_counter()
 
@@ -1281,4 +1145,4 @@ if __name__ == "__main__":
         optimizer.zero_grad(set_to_none=True)
 
     time_end = time.perf_counter()
-    logger.info(f"elapsed time: {time_end - time_start} [sec] for last {steps - 1} steps")
+    print(f"elapsed time: {time_end - time_start} [sec] for last {steps - 1} steps")

library/sdxl_train_util.py

@@ -1,21 +1,14 @@
 import argparse
+import gc
 import math
 import os
 from typing import Optional
-
 import torch
-from library.device_utils import init_ipex, clean_memory_on_device
-init_ipex()
-
 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
-from .utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
 
 TOKENIZER1_PATH = "openai/clip-vit-large-patch14"
 TOKENIZER2_PATH = "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"
@@ -28,7 +21,7 @@ def load_target_model(args, accelerator, model_version: str, weight_dtype):
     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:
-            logger.info(f"loading model for process {accelerator.state.local_process_index}/{accelerator.state.num_processes}")
+            print(f"loading model for process {accelerator.state.local_process_index}/{accelerator.state.num_processes}")
 
             (
                 load_stable_diffusion_format,
@@ -54,9 +47,12 @@ def load_target_model(args, accelerator, model_version: str, weight_dtype):
                 unet.to(accelerator.device)
                 vae.to(accelerator.device)
 
-            clean_memory_on_device(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
 
 
@@ -68,7 +64,7 @@ def _load_target_model(
     load_stable_diffusion_format = os.path.isfile(name_or_path)  # determine SD or Diffusers
 
     if load_stable_diffusion_format:
-        logger.info(f"load StableDiffusion checkpoint: {name_or_path}")
+        print(f"load StableDiffusion checkpoint: {name_or_path}")
         (
             text_encoder1,
             text_encoder2,
@@ -82,7 +78,7 @@ def _load_target_model(
         from diffusers import StableDiffusionXLPipeline
 
         variant = "fp16" if weight_dtype == torch.float16 else None
-        logger.info(f"load Diffusers pretrained models: {name_or_path}, variant={variant}")
+        print(f"load Diffusers pretrained models: {name_or_path}, variant={variant}")
         try:
             try:
                 pipe = StableDiffusionXLPipeline.from_pretrained(
@@ -90,12 +86,12 @@ def _load_target_model(
                 )
             except EnvironmentError as ex:
                 if variant is not None:
-                    logger.info("try to load fp32 model")
+                    print("try to load fp32 model")
                     pipe = StableDiffusionXLPipeline.from_pretrained(name_or_path, variant=None, tokenizer=None)
                 else:
                     raise ex
         except EnvironmentError as ex:
-            logger.error(
+            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
@@ -118,7 +114,7 @@ def _load_target_model(
         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)
-        logger.info("U-Net converted to original U-Net")
+        print("U-Net converted to original U-Net")
 
         logit_scale = None
         ckpt_info = None
@@ -126,13 +122,13 @@ def _load_target_model(
     # VAEを読み込む
     if vae_path is not None:
         vae = model_util.load_vae(vae_path, weight_dtype)
-        logger.info("additional VAE loaded")
+        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):
-    logger.info("prepare tokenizers")
+    print("prepare tokenizers")
 
     original_paths = [TOKENIZER1_PATH, TOKENIZER2_PATH]
     tokeniers = []
@@ -141,14 +137,14 @@ def load_tokenizers(args: argparse.Namespace):
         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):
-                logger.info(f"load tokenizer from cache: {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):
-            logger.info(f"save Tokenizer to cache: {local_tokenizer_path}")
+            print(f"save Tokenizer to cache: {local_tokenizer_path}")
             tokenizer.save_pretrained(local_tokenizer_path)
 
         if i == 1:
@@ -157,7 +153,7 @@ def load_tokenizers(args: argparse.Namespace):
         tokeniers.append(tokenizer)
 
     if hasattr(args, "max_token_length") and args.max_token_length is not None:
-        logger.info(f"update token length: {args.max_token_length}")
+        print(f"update token length: {args.max_token_length}")
 
     return tokeniers
 
@@ -338,23 +334,23 @@ def add_sdxl_training_arguments(parser: argparse.ArgumentParser):
 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:
-        logger.warning("v_parameterization will be unexpected / SDXL学習ではv_parameterizationは想定外の動作になります")
+        print("v_parameterization will be unexpected / SDXL学習ではv_parameterizationは想定外の動作になります")
 
     if args.clip_skip is not None:
-        logger.warning("clip_skip will be unexpected / SDXL学習ではclip_skipは動作しません")
+        print("clip_skip will be unexpected / SDXL学習ではclip_skipは動作しません")
 
     # if args.multires_noise_iterations:
-    #     logger.info(
+    #     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:
-    #         logger.info(
+    #         print(
     #             f"Warning: SDXL has been trained with noise_offset={DEFAULT_NOISE_OFFSET} / SDXLはnoise_offset={DEFAULT_NOISE_OFFSET}で学習されています"
     #         )
-    #     logger.info(f"noise_offset is set to {args.noise_offset} / noise_offsetが{args.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
@@ -363,7 +359,7 @@ def verify_sdxl_training_args(args: argparse.Namespace, supportTextEncoderCachin
     if supportTextEncoderCaching:
         if args.cache_text_encoder_outputs_to_disk and not args.cache_text_encoder_outputs:
             args.cache_text_encoder_outputs = True
-            logger.warning(
+            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が有効になりました"
             )

library/slicing_vae.py

@@ -26,10 +26,7 @@ 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 .utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 def slice_h(x, num_slices):
     # slice with pad 1 both sides: to eliminate side effect of padding of conv2d
@@ -65,7 +62,7 @@ def cat_h(sliced):
     return x
 
 
-def resblock_forward(_self, num_slices, input_tensor, temb, **kwargs):
+def resblock_forward(_self, num_slices, input_tensor, temb):
     assert _self.upsample is None and _self.downsample is None
     assert _self.norm1.num_groups == _self.norm2.num_groups
     assert temb is None
@@ -92,7 +89,7 @@ def resblock_forward(_self, num_slices, input_tensor, temb, **kwargs):
     #     sliced_tensor = torch.chunk(x, num_div, dim=1)
     #     sliced_weight = torch.chunk(norm.weight, num_div, dim=0)
     #     sliced_bias = torch.chunk(norm.bias, num_div, dim=0)
-    #     logger.info(sliced_tensor[0].shape, num_div, sliced_weight[0].shape, sliced_bias[0].shape)
+    #     print(sliced_tensor[0].shape, num_div, sliced_weight[0].shape, sliced_bias[0].shape)
     #     normed_tensor = []
     #     for i in range(num_div):
     #         n = torch.group_norm(sliced_tensor[i], norm.num_groups, sliced_weight[i], sliced_bias[i], norm.eps)
@@ -246,7 +243,7 @@ class SlicingEncoder(nn.Module):
 
         self.num_slices = num_slices
         div = num_slices / (2 ** (len(self.down_blocks) - 1))  # 深い層はそこまで分割しなくていいので適宜減らす
-        # logger.info(f"initial divisor: {div}")
+        # print(f"initial divisor: {div}")
         if div >= 2:
             div = int(div)
             for resnet in self.mid_block.resnets:
@@ -256,11 +253,11 @@ class SlicingEncoder(nn.Module):
         for i, down_block in enumerate(self.down_blocks[::-1]):
             if div >= 2:
                 div = int(div)
-                # logger.info(f"down block: {i} divisor: {div}")
+                # print(f"down block: {i} divisor: {div}")
                 for resnet in down_block.resnets:
                     resnet.forward = wrapper(resblock_forward, resnet, div)
                 if down_block.downsamplers is not None:
-                    # logger.info("has downsample")
+                    # print("has downsample")
                     for downsample in down_block.downsamplers:
                         downsample.forward = wrapper(self.downsample_forward, downsample, div * 2)
             div *= 2
@@ -310,7 +307,7 @@ class SlicingEncoder(nn.Module):
     def downsample_forward(self, _self, num_slices, hidden_states):
         assert hidden_states.shape[1] == _self.channels
         assert _self.use_conv and _self.padding == 0
-        logger.info(f"downsample forward {num_slices} {hidden_states.shape}")
+        print("downsample forward", num_slices, hidden_states.shape)
 
         org_device = hidden_states.device
         cpu_device = torch.device("cpu")
@@ -353,7 +350,7 @@ class SlicingEncoder(nn.Module):
                 hidden_states = torch.cat([hidden_states, x], dim=2)
 
         hidden_states = hidden_states.to(org_device)
-        # logger.info(f"downsample forward done {hidden_states.shape}")
+        # print("downsample forward done", hidden_states.shape)
         return hidden_states
 
 
@@ -429,7 +426,7 @@ class SlicingDecoder(nn.Module):
 
         self.num_slices = num_slices
         div = num_slices / (2 ** (len(self.up_blocks) - 1))
-        logger.info(f"initial divisor: {div}")
+        print(f"initial divisor: {div}")
         if div >= 2:
             div = int(div)
             for resnet in self.mid_block.resnets:
@@ -439,11 +436,11 @@ class SlicingDecoder(nn.Module):
         for i, up_block in enumerate(self.up_blocks):
             if div >= 2:
                 div = int(div)
-                # logger.info(f"up block: {i} divisor: {div}")
+                # print(f"up block: {i} divisor: {div}")
                 for resnet in up_block.resnets:
                     resnet.forward = wrapper(resblock_forward, resnet, div)
                 if up_block.upsamplers is not None:
-                    # logger.info("has upsample")
+                    # print("has upsample")
                     for upsample in up_block.upsamplers:
                         upsample.forward = wrapper(self.upsample_forward, upsample, div * 2)
             div *= 2
@@ -531,7 +528,7 @@ class SlicingDecoder(nn.Module):
             del x
 
         hidden_states = torch.cat(sliced, dim=2)
-        # logger.info(f"us hidden_states {hidden_states.shape}")
+        # print("us hidden_states", hidden_states.shape)
         del sliced
 
         hidden_states = hidden_states.to(org_device)

library/utils.py

@@ -1,266 +1,6 @@
-import logging
-import sys
 import threading
-import torch
-from torchvision import transforms
 from typing import *
-from diffusers import EulerAncestralDiscreteScheduler
-import diffusers.schedulers.scheduling_euler_ancestral_discrete
-from diffusers.schedulers.scheduling_euler_ancestral_discrete import EulerAncestralDiscreteSchedulerOutput
 
 
 def fire_in_thread(f, *args, **kwargs):
-    threading.Thread(target=f, args=args, kwargs=kwargs).start()
-
-
-def add_logging_arguments(parser):
-    parser.add_argument(
-        "--console_log_level",
-        type=str,
-        default=None,
-        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
-        help="Set the logging level, default is INFO / ログレベルを設定する。デフォルトはINFO",
-    )
-    parser.add_argument(
-        "--console_log_file",
-        type=str,
-        default=None,
-        help="Log to a file instead of stderr / 標準エラー出力ではなくファイルにログを出力する",
-    )
-    parser.add_argument("--console_log_simple", action="store_true", help="Simple log output / シンプルなログ出力")
-
-
-def setup_logging(args=None, log_level=None, reset=False):
-    if logging.root.handlers:
-        if reset:
-            # remove all handlers
-            for handler in logging.root.handlers[:]:
-                logging.root.removeHandler(handler)
-        else:
-            return
-
-    # log_level can be set by the caller or by the args, the caller has priority. If not set, use INFO
-    if log_level is None and args is not None:
-        log_level = args.console_log_level
-    if log_level is None:
-        log_level = "INFO"
-    log_level = getattr(logging, log_level)
-
-    msg_init = None
-    if args is not None and args.console_log_file:
-        handler = logging.FileHandler(args.console_log_file, mode="w")
-    else:
-        handler = None
-        if not args or not args.console_log_simple:
-            try:
-                from rich.logging import RichHandler
-                from rich.console import Console
-                from rich.logging import RichHandler
-
-                handler = RichHandler(console=Console(stderr=True))
-            except ImportError:
-                # print("rich is not installed, using basic logging")
-                msg_init = "rich is not installed, using basic logging"
-
-        if handler is None:
-            handler = logging.StreamHandler(sys.stdout)  # same as print
-            handler.propagate = False
-
-    formatter = logging.Formatter(
-        fmt="%(message)s",
-        datefmt="%Y-%m-%d %H:%M:%S",
-    )
-    handler.setFormatter(formatter)
-    logging.root.setLevel(log_level)
-    logging.root.addHandler(handler)
-
-    if msg_init is not None:
-        logger = logging.getLogger(__name__)
-        logger.info(msg_init)
-
-
-
-# TODO make inf_utils.py
-
-
-# region Gradual Latent hires fix
-
-
-class GradualLatent:
-    def __init__(
-        self,
-        ratio,
-        start_timesteps,
-        every_n_steps,
-        ratio_step,
-        s_noise=1.0,
-        gaussian_blur_ksize=None,
-        gaussian_blur_sigma=0.5,
-        gaussian_blur_strength=0.5,
-        unsharp_target_x=True,
-    ):
-        self.ratio = ratio
-        self.start_timesteps = start_timesteps
-        self.every_n_steps = every_n_steps
-        self.ratio_step = ratio_step
-        self.s_noise = s_noise
-        self.gaussian_blur_ksize = gaussian_blur_ksize
-        self.gaussian_blur_sigma = gaussian_blur_sigma
-        self.gaussian_blur_strength = gaussian_blur_strength
-        self.unsharp_target_x = unsharp_target_x
-
-    def __str__(self) -> str:
-        return (
-            f"GradualLatent(ratio={self.ratio}, start_timesteps={self.start_timesteps}, "
-            + f"every_n_steps={self.every_n_steps}, ratio_step={self.ratio_step}, s_noise={self.s_noise}, "
-            + f"gaussian_blur_ksize={self.gaussian_blur_ksize}, gaussian_blur_sigma={self.gaussian_blur_sigma}, gaussian_blur_strength={self.gaussian_blur_strength}, "
-            + f"unsharp_target_x={self.unsharp_target_x})"
-        )
-
-    def apply_unshark_mask(self, x: torch.Tensor):
-        if self.gaussian_blur_ksize is None:
-            return x
-        blurred = transforms.functional.gaussian_blur(x, self.gaussian_blur_ksize, self.gaussian_blur_sigma)
-        # mask = torch.sigmoid((x - blurred) * self.gaussian_blur_strength)
-        mask = (x - blurred) * self.gaussian_blur_strength
-        sharpened = x + mask
-        return sharpened
-
-    def interpolate(self, x: torch.Tensor, resized_size, unsharp=True):
-        org_dtype = x.dtype
-        if org_dtype == torch.bfloat16:
-            x = x.float()
-
-        x = torch.nn.functional.interpolate(x, size=resized_size, mode="bicubic", align_corners=False).to(dtype=org_dtype)
-
-        # apply unsharp mask / アンシャープマスクを適用する
-        if unsharp and self.gaussian_blur_ksize:
-            x = self.apply_unshark_mask(x)
-
-        return x
-
-
-class EulerAncestralDiscreteSchedulerGL(EulerAncestralDiscreteScheduler):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.resized_size = None
-        self.gradual_latent = None
-
-    def set_gradual_latent_params(self, size, gradual_latent: GradualLatent):
-        self.resized_size = size
-        self.gradual_latent = gradual_latent
-
-    def step(
-        self,
-        model_output: torch.FloatTensor,
-        timestep: Union[float, torch.FloatTensor],
-        sample: torch.FloatTensor,
-        generator: Optional[torch.Generator] = None,
-        return_dict: bool = True,
-    ) -> Union[EulerAncestralDiscreteSchedulerOutput, Tuple]:
-        """
-        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
-        process from the learned model outputs (most often the predicted noise).
-
-        Args:
-            model_output (`torch.FloatTensor`):
-                The direct output from learned diffusion model.
-            timestep (`float`):
-                The current discrete timestep in the diffusion chain.
-            sample (`torch.FloatTensor`):
-                A current instance of a sample created by the diffusion process.
-            generator (`torch.Generator`, *optional*):
-                A random number generator.
-            return_dict (`bool`):
-                Whether or not to return a
-                [`~schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteSchedulerOutput`] or tuple.
-
-        Returns:
-            [`~schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteSchedulerOutput`] or `tuple`:
-                If return_dict is `True`,
-                [`~schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteSchedulerOutput`] is returned,
-                otherwise a tuple is returned where the first element is the sample tensor.
-
-        """
-
-        if isinstance(timestep, int) or isinstance(timestep, torch.IntTensor) or isinstance(timestep, torch.LongTensor):
-            raise ValueError(
-                (
-                    "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
-                    " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
-                    " one of the `scheduler.timesteps` as a timestep."
-                ),
-            )
-
-        if not self.is_scale_input_called:
-            # logger.warning(
-            print(
-                "The `scale_model_input` function should be called before `step` to ensure correct denoising. "
-                "See `StableDiffusionPipeline` for a usage example."
-            )
-
-        if self.step_index is None:
-            self._init_step_index(timestep)
-
-        sigma = self.sigmas[self.step_index]
-
-        # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
-        if self.config.prediction_type == "epsilon":
-            pred_original_sample = sample - sigma * model_output
-        elif self.config.prediction_type == "v_prediction":
-            # * c_out + input * c_skip
-            pred_original_sample = model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1))
-        elif self.config.prediction_type == "sample":
-            raise NotImplementedError("prediction_type not implemented yet: sample")
-        else:
-            raise ValueError(f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`")
-
-        sigma_from = self.sigmas[self.step_index]
-        sigma_to = self.sigmas[self.step_index + 1]
-        sigma_up = (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5
-        sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5
-
-        # 2. Convert to an ODE derivative
-        derivative = (sample - pred_original_sample) / sigma
-
-        dt = sigma_down - sigma
-
-        device = model_output.device
-        if self.resized_size is None:
-            prev_sample = sample + derivative * dt
-
-            noise = diffusers.schedulers.scheduling_euler_ancestral_discrete.randn_tensor(
-                model_output.shape, dtype=model_output.dtype, device=device, generator=generator
-            )
-            s_noise = 1.0
-        else:
-            print("resized_size", self.resized_size, "model_output.shape", model_output.shape, "sample.shape", sample.shape)
-            s_noise = self.gradual_latent.s_noise
-
-            if self.gradual_latent.unsharp_target_x:
-                prev_sample = sample + derivative * dt
-                prev_sample = self.gradual_latent.interpolate(prev_sample, self.resized_size)
-            else:
-                sample = self.gradual_latent.interpolate(sample, self.resized_size)
-                derivative = self.gradual_latent.interpolate(derivative, self.resized_size, unsharp=False)
-                prev_sample = sample + derivative * dt
-
-            noise = diffusers.schedulers.scheduling_euler_ancestral_discrete.randn_tensor(
-                (model_output.shape[0], model_output.shape[1], self.resized_size[0], self.resized_size[1]),
-                dtype=model_output.dtype,
-                device=device,
-                generator=generator,
-            )
-
-        prev_sample = prev_sample + noise * sigma_up * s_noise
-
-        # upon completion increase step index by one
-        self._step_index += 1
-
-        if not return_dict:
-            return (prev_sample,)
-
-        return EulerAncestralDiscreteSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)
-
-
-# endregion
+    threading.Thread(target=f, args=args, kwargs=kwargs).start()

networks/check_lora_weights.py

@@ -2,13 +2,10 @@ import argparse
 import os
 import torch
 from safetensors.torch import load_file
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 def main(file):
-    logger.info(f"loading: {file}")
+    print(f"loading: {file}")
     if os.path.splitext(file)[1] == ".safetensors":
         sd = load_file(file)
     else:

networks/control_net_lllite.py

@@ -2,10 +2,7 @@ import os
 from typing import Optional, List, Type
 import torch
 from library import sdxl_original_unet
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 # input_blocksに適用するかどうか / if True, input_blocks are not applied
 SKIP_INPUT_BLOCKS = False
@@ -128,7 +125,7 @@ class LLLiteModule(torch.nn.Module):
             return
 
         # timestepごとに呼ばれないので、あらかじめ計算しておく / it is not called for each timestep, so calculate it in advance
-        # logger.info(f"C {self.lllite_name}, cond_image.shape={cond_image.shape}")
+        # 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
@@ -158,7 +155,7 @@ class LLLiteModule(torch.nn.Module):
             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
-        # logger.info(f"C {self.lllite_name}, x.shape={x.shape}, cx.shape={cx.shape}")
+        # print(f"C {self.lllite_name}, x.shape={x.shape}, cx.shape={cx.shape}")
 
         # downで入力の次元数を削減し、conditioning image embeddingと結合する
         # 加算ではなくchannel方向に結合することで、うまいこと混ぜてくれることを期待している
@@ -289,7 +286,7 @@ class ControlNetLLLite(torch.nn.Module):
 
         # create module instances
         self.unet_modules: List[LLLiteModule] = create_modules(unet, target_modules, LLLiteModule)
-        logger.info(f"create ControlNet LLLite for U-Net: {len(self.unet_modules)} modules.")
+        print(f"create ControlNet LLLite for U-Net: {len(self.unet_modules)} modules.")
 
     def forward(self, x):
         return x  # dummy
@@ -322,7 +319,7 @@ class ControlNetLLLite(torch.nn.Module):
         return info
 
     def apply_to(self):
-        logger.info("applying LLLite for U-Net...")
+        print("applying LLLite for U-Net...")
         for module in self.unet_modules:
             module.apply_to()
             self.add_module(module.lllite_name, module)
@@ -377,19 +374,19 @@ if __name__ == "__main__":
     # sdxl_original_unet.USE_REENTRANT = False
 
     # test shape etc
-    logger.info("create unet")
+    print("create unet")
     unet = sdxl_original_unet.SdxlUNet2DConditionModel()
     unet.to("cuda").to(torch.float16)
 
-    logger.info("create ControlNet-LLLite")
+    print("create ControlNet-LLLite")
     control_net = ControlNetLLLite(unet, 32, 64)
     control_net.apply_to()
     control_net.to("cuda")
 
-    logger.info(control_net)
+    print(control_net)
 
-    # logger.info number of parameters
-    logger.info(f"number of parameters {sum(p.numel() for p in control_net.parameters() if p.requires_grad)}")
+    # print number of parameters
+    print("number of parameters", sum(p.numel() for p in control_net.parameters() if p.requires_grad))
 
     input()
 
@@ -401,12 +398,12 @@ if __name__ == "__main__":
 
     # # visualize
     # import torchviz
-    # logger.info("run visualize")
+    # print("run visualize")
     # controlnet.set_control(conditioning_image)
     # output = unet(x, t, ctx, y)
-    # logger.info("make_dot")
+    # print("make_dot")
     # image = torchviz.make_dot(output, params=dict(controlnet.named_parameters()))
-    # logger.info("render")
+    # print("render")
     # image.format = "svg" # "png"
     # image.render("NeuralNet") # すごく時間がかかるので注意 / be careful because it takes a long time
     # input()
@@ -417,12 +414,12 @@ if __name__ == "__main__":
 
     scaler = torch.cuda.amp.GradScaler(enabled=True)
 
-    logger.info("start training")
+    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):
-        logger.info(f"step {step}")
+        print(f"step {step}")
 
         batch_size = 1
         conditioning_image = torch.rand(batch_size, 3, 1024, 1024).cuda() * 2.0 - 1.0
@@ -442,7 +439,7 @@ if __name__ == "__main__":
         scaler.step(optimizer)
         scaler.update()
         optimizer.zero_grad(set_to_none=True)
-        logger.info(f"{sample_param}")
+        print(sample_param)
 
     # from safetensors.torch import save_file
 

networks/control_net_lllite_for_train.py

@@ -6,10 +6,7 @@ import re
 from typing import Optional, List, Type
 import torch
 from library import sdxl_original_unet
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 # input_blocksに適用するかどうか / if True, input_blocks are not applied
 SKIP_INPUT_BLOCKS = False
@@ -273,7 +270,7 @@ class SdxlUNet2DConditionModelControlNetLLLite(sdxl_original_unet.SdxlUNet2DCond
 
         # create module instances
         self.lllite_modules = apply_to_modules(self, target_modules)
-        logger.info(f"enable ControlNet LLLite for U-Net: {len(self.lllite_modules)} modules.")
+        print(f"enable ControlNet LLLite for U-Net: {len(self.lllite_modules)} modules.")
 
     # def prepare_optimizer_params(self):
     def prepare_params(self):
@@ -284,8 +281,8 @@ class SdxlUNet2DConditionModelControlNetLLLite(sdxl_original_unet.SdxlUNet2DCond
                 train_params.append(p)
             else:
                 non_train_params.append(p)
-        logger.info(f"count of trainable parameters: {len(train_params)}")
-        logger.info(f"count of non-trainable parameters: {len(non_train_params)}")
+        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)
@@ -391,7 +388,7 @@ class SdxlUNet2DConditionModelControlNetLLLite(sdxl_original_unet.SdxlUNet2DCond
             matches = pattern.findall(module_name)
             if matches is not None:
                 for m in matches:
-                    logger.info(f"{module_name} {m}")
+                    print(module_name, m)
                     module_name = module_name.replace(m, m.replace("_", "@"))
             module_name = module_name.replace("_", ".")
             module_name = module_name.replace("@", "_")
@@ -410,7 +407,7 @@ class SdxlUNet2DConditionModelControlNetLLLite(sdxl_original_unet.SdxlUNet2DCond
 
 
 def replace_unet_linear_and_conv2d():
-    logger.info("replace torch.nn.Linear and torch.nn.Conv2d to LLLiteLinear and LLLiteConv2d in U-Net")
+    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
 
@@ -422,10 +419,10 @@ if __name__ == "__main__":
     replace_unet_linear_and_conv2d()
 
     # test shape etc
-    logger.info("create unet")
+    print("create unet")
     unet = SdxlUNet2DConditionModelControlNetLLLite()
 
-    logger.info("enable ControlNet-LLLite")
+    print("enable ControlNet-LLLite")
     unet.apply_lllite(32, 64, None, False, 1.0)
     unet.to("cuda")  # .to(torch.float16)
 
@@ -442,14 +439,14 @@ if __name__ == "__main__":
     #         unet_sd[converted_key] = model_sd[key]
 
     # info = unet.load_lllite_weights("r:/lllite_from_unet.safetensors", unet_sd)
-    # logger.info(info)
+    # print(info)
 
-    # logger.info(unet)
+    # print(unet)
 
-    # logger.info number of parameters
+    # print number of parameters
     params = unet.prepare_params()
-    logger.info(f"number of parameters {sum(p.numel() for p in params)}")
-    # logger.info("type any key to continue")
+    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)
@@ -458,12 +455,12 @@ if __name__ == "__main__":
 
     # # visualize
     # import torchviz
-    # logger.info("run visualize")
+    # print("run visualize")
     # controlnet.set_control(conditioning_image)
     # output = unet(x, t, ctx, y)
-    # logger.info("make_dot")
+    # print("make_dot")
     # image = torchviz.make_dot(output, params=dict(controlnet.named_parameters()))
-    # logger.info("render")
+    # print("render")
     # image.format = "svg" # "png"
     # image.render("NeuralNet") # すごく時間がかかるので注意 / be careful because it takes a long time
     # input()
@@ -474,13 +471,13 @@ if __name__ == "__main__":
 
     scaler = torch.cuda.amp.GradScaler(enabled=True)
 
-    logger.info("start training")
+    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):
-        logger.info(f"step {step}")
+        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()
@@ -497,9 +494,9 @@ if __name__ == "__main__":
         scaler.step(optimizer)
         scaler.update()
         optimizer.zero_grad(set_to_none=True)
-        logger.info(sample_param)
+        print(sample_param)
 
     # from safetensors.torch import save_file
 
-    # logger.info("save weights")
+    # print("save weights")
     # unet.save_lllite_weights("r:/lllite_from_unet.safetensors", torch.float16, None)

networks/dylora.py

@@ -12,15 +12,10 @@
 import math
 import os
 import random
-from typing import Dict, List, Optional, Tuple, Type, Union
-from diffusers import AutoencoderKL
-from transformers import CLIPTextModel
+from typing import List, Tuple, Union
 import torch
 from torch import nn
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 class DyLoRAModule(torch.nn.Module):
     """
@@ -170,15 +165,7 @@ class DyLoRAModule(torch.nn.Module):
         super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
 
 
-def create_network(
-    multiplier: float,
-    network_dim: Optional[int],
-    network_alpha: Optional[float],
-    vae: AutoencoderKL,
-    text_encoder: Union[CLIPTextModel, List[CLIPTextModel]],
-    unet,
-    **kwargs,
-):
+def create_network(multiplier, network_dim, network_alpha, vae, text_encoder, unet, **kwargs):
     if network_dim is None:
         network_dim = 4  # default
     if network_alpha is None:
@@ -195,7 +182,6 @@ def create_network(
             conv_alpha = 1.0
         else:
             conv_alpha = float(conv_alpha)
-            
     if unit is not None:
         unit = int(unit)
     else:
@@ -237,7 +223,7 @@ def create_network_from_weights(multiplier, file, vae, text_encoder, unet, weigh
         elif "lora_down" in key:
             dim = value.size()[0]
             modules_dim[lora_name] = dim
-            # logger.info(f"{lora_name} {value.size()} {dim}")
+            # print(lora_name, value.size(), dim)
 
     # support old LoRA without alpha
     for key in modules_dim.keys():
@@ -281,11 +267,11 @@ class DyLoRANetwork(torch.nn.Module):
         self.apply_to_conv = apply_to_conv
 
         if modules_dim is not None:
-            logger.info("create LoRA network from weights")
+            print(f"create LoRA network from weights")
         else:
-            logger.info(f"create LoRA network. base dim (rank): {lora_dim}, alpha: {alpha}, unit: {unit}")
+            print(f"create LoRA network. base dim (rank): {lora_dim}, alpha: {alpha}, unit: {unit}")
             if self.apply_to_conv:
-                logger.info("apply LoRA to Conv2d with kernel size (3,3).")
+                print(f"apply LoRA to Conv2d with kernel size (3,3).")
 
         # create module instances
         def create_modules(is_unet, root_module: torch.nn.Module, target_replace_modules) -> List[DyLoRAModule]:
@@ -320,23 +306,9 @@ class DyLoRANetwork(torch.nn.Module):
                             lora = module_class(lora_name, child_module, self.multiplier, dim, alpha, unit)
                             loras.append(lora)
             return loras
-        
-        text_encoders = text_encoder if type(text_encoder) == list else [text_encoder]
-        
-        self.text_encoder_loras = []
-        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 = create_modules(False, text_encoder, DyLoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE)
-            self.text_encoder_loras.extend(text_encoder_loras)
 
-        # self.text_encoder_loras = create_modules(False, text_encoder, DyLoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE)
-        logger.info(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
+        self.text_encoder_loras = create_modules(False, text_encoder, DyLoRANetwork.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
         target_modules = DyLoRANetwork.UNET_TARGET_REPLACE_MODULE
@@ -344,7 +316,7 @@ class DyLoRANetwork(torch.nn.Module):
             target_modules += DyLoRANetwork.UNET_TARGET_REPLACE_MODULE_CONV2D_3X3
 
         self.unet_loras = create_modules(True, unet, target_modules)
-        logger.info(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
+        print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
 
     def set_multiplier(self, multiplier):
         self.multiplier = multiplier
@@ -364,12 +336,12 @@ class DyLoRANetwork(torch.nn.Module):
 
     def apply_to(self, text_encoder, unet, apply_text_encoder=True, apply_unet=True):
         if apply_text_encoder:
-            logger.info("enable LoRA for text encoder")
+            print("enable LoRA for text encoder")
         else:
             self.text_encoder_loras = []
 
         if apply_unet:
-            logger.info("enable LoRA for U-Net")
+            print("enable LoRA for U-Net")
         else:
             self.unet_loras = []
 
@@ -387,12 +359,12 @@ class DyLoRANetwork(torch.nn.Module):
                 apply_unet = True
 
         if apply_text_encoder:
-            logger.info("enable LoRA for text encoder")
+            print("enable LoRA for text encoder")
         else:
             self.text_encoder_loras = []
 
         if apply_unet:
-            logger.info("enable LoRA for U-Net")
+            print("enable LoRA for U-Net")
         else:
             self.unet_loras = []
 
@@ -403,7 +375,7 @@ class DyLoRANetwork(torch.nn.Module):
                     sd_for_lora[key[len(lora.lora_name) + 1 :]] = weights_sd[key]
             lora.merge_to(sd_for_lora, dtype, device)
 
-        logger.info(f"weights are merged")
+        print(f"weights are merged")
     """
 
     def prepare_optimizer_params(self, text_encoder_lr, unet_lr, default_lr):

networks/extract_lora_from_dylora.py

@@ -10,10 +10,7 @@ from safetensors.torch import load_file, save_file, safe_open
 from tqdm import tqdm
 from library import train_util, model_util
 import numpy as np
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 def load_state_dict(file_name):
     if model_util.is_safetensors(file_name):
@@ -43,13 +40,13 @@ def split_lora_model(lora_sd, unit):
             rank = value.size()[0]
             if rank > max_rank:
                 max_rank = rank
-    logger.info(f"Max rank: {max_rank}")
+    print(f"Max rank: {max_rank}")
 
     rank = unit
     split_models = []
     new_alpha = None
     while rank < max_rank:
-        logger.info(f"Splitting rank {rank}")
+        print(f"Splitting rank {rank}")
         new_sd = {}
         for key, value in lora_sd.items():
             if "lora_down" in key:
@@ -60,7 +57,7 @@ def split_lora_model(lora_sd, unit):
                 # なぜかscaleするとおかしくなる……
                 # this_rank = lora_sd[key.replace("alpha", "lora_down.weight")].size()[0]
                 # scale = math.sqrt(this_rank / rank)  # rank is > unit
-                # logger.info(key, value.size(), this_rank, rank, value, scale)
+                # print(key, value.size(), this_rank, rank, value, scale)
                 # new_alpha = value * scale  # always same
                 # new_sd[key] = new_alpha
                 new_sd[key] = value
@@ -72,10 +69,10 @@ def split_lora_model(lora_sd, unit):
 
 
 def split(args):
-    logger.info("loading Model...")
+    print("loading Model...")
     lora_sd, metadata = load_state_dict(args.model)
 
-    logger.info("Splitting Model...")
+    print("Splitting Model...")
     original_rank, split_models = split_lora_model(lora_sd, args.unit)
 
     comment = metadata.get("ss_training_comment", "")
@@ -97,7 +94,7 @@ def split(args):
         filename, ext = os.path.splitext(args.save_to)
         model_file_name = filename + f"-{new_rank:04d}{ext}"
 
-        logger.info(f"saving model to: {model_file_name}")
+        print(f"saving model to: {model_file_name}")
         save_to_file(model_file_name, state_dict, new_metadata)
 
 

networks/extract_lora_from_models.py

@@ -11,13 +11,10 @@ from safetensors.torch import load_file, save_file
 from tqdm import tqdm
 from library import sai_model_spec, model_util, sdxl_model_util
 import lora
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
 
-# CLAMP_QUANTILE = 0.99
-# MIN_DIFF = 1e-1
+
+CLAMP_QUANTILE = 0.99
+MIN_DIFF = 1e-1
 
 
 def save_to_file(file_name, model, state_dict, dtype):
@@ -32,24 +29,7 @@ def save_to_file(file_name, model, state_dict, dtype):
         torch.save(model, file_name)
 
 
-def svd(
-    model_org=None,
-    model_tuned=None,
-    save_to=None,
-    dim=4,
-    v2=None,
-    sdxl=None,
-    conv_dim=None,
-    v_parameterization=None,
-    device=None,
-    save_precision=None,
-    clamp_quantile=0.99,
-    min_diff=0.01,
-    no_metadata=False,
-    load_precision=None,
-    load_original_model_to=None,
-    load_tuned_model_to=None,
-):
+def svd(args):
     def str_to_dtype(p):
         if p == "float":
             return torch.float
@@ -59,65 +39,44 @@ def svd(
             return torch.bfloat16
         return None
 
-    assert v2 != sdxl or (not v2 and not sdxl), "v2 and sdxl cannot be specified at the same time / v2とsdxlは同時に指定できません"
-    if v_parameterization is None:
-        v_parameterization = v2
+    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
 
-    load_dtype = str_to_dtype(load_precision) if load_precision else None
-    save_dtype = str_to_dtype(save_precision)
-    work_device = "cpu"
+    save_dtype = str_to_dtype(args.save_precision)
 
     # load models
-    if not sdxl:
-        logger.info(f"loading original SD model : {model_org}")
-        text_encoder_o, _, unet_o = model_util.load_models_from_stable_diffusion_checkpoint(v2, model_org)
+    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]
-        if load_dtype is not None:
-            text_encoder_o = text_encoder_o.to(load_dtype)
-            unet_o = unet_o.to(load_dtype)
-
-        logger.info(f"loading tuned SD model : {model_tuned}")
-        text_encoder_t, _, unet_t = model_util.load_models_from_stable_diffusion_checkpoint(v2, model_tuned)
+        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]
-        if load_dtype is not None:
-            text_encoder_t = text_encoder_t.to(load_dtype)
-            unet_t = unet_t.to(load_dtype)
-
-        model_version = model_util.get_model_version_str_for_sd1_sd2(v2, v_parameterization)
+        model_version = model_util.get_model_version_str_for_sd1_sd2(args.v2, args.v_parameterization)
     else:
-        device_org = load_original_model_to if load_original_model_to else "cpu"
-        device_tuned = load_tuned_model_to if load_tuned_model_to else "cpu"
-
-        logger.info(f"loading original SDXL model : {model_org}")
+        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, model_org, device_org
+            sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0, args.model_org, "cpu"
         )
         text_encoders_o = [text_encoder_o1, text_encoder_o2]
-        if load_dtype is not None:
-            text_encoder_o1 = text_encoder_o1.to(load_dtype)
-            text_encoder_o2 = text_encoder_o2.to(load_dtype)
-            unet_o = unet_o.to(load_dtype)
-
-        logger.info(f"loading original SDXL model : {model_tuned}")
+        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, model_tuned, device_tuned
+            sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0, args.model_tuned, "cpu"
         )
         text_encoders_t = [text_encoder_t1, text_encoder_t2]
-        if load_dtype is not None:
-            text_encoder_t1 = text_encoder_t1.to(load_dtype)
-            text_encoder_t2 = text_encoder_t2.to(load_dtype)
-            unet_t = unet_t.to(load_dtype)
-
         model_version = sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0
 
     # create LoRA network to extract weights: Use dim (rank) as alpha
-    if conv_dim is None:
+    if args.conv_dim is None:
         kwargs = {}
     else:
-        kwargs = {"conv_dim": conv_dim, "conv_alpha": conv_dim}
+        kwargs = {"conv_dim": args.conv_dim, "conv_alpha": args.conv_dim}
 
-    lora_network_o = lora.create_network(1.0, dim, dim, None, text_encoders_o, unet_o, **kwargs)
-    lora_network_t = lora.create_network(1.0, dim, dim, None, text_encoders_t, unet_t, **kwargs)
+    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ベース） "
@@ -129,66 +88,50 @@ def svd(
         lora_name = lora_o.lora_name
         module_o = lora_o.org_module
         module_t = lora_t.org_module
-        diff = module_t.weight.to(work_device) - module_o.weight.to(work_device)
-
-        # clear weight to save memory
-        module_o.weight = None
-        module_t.weight = None
+        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:
+        if not text_encoder_different and torch.max(torch.abs(diff)) > MIN_DIFF:
             text_encoder_different = True
-            logger.info(f"Text encoder is different. {torch.max(torch.abs(diff))} > {min_diff}")
+            print(f"Text encoder is different. {torch.max(torch.abs(diff))} > {MIN_DIFF}")
 
+        diff = diff.float()
         diffs[lora_name] = diff
 
-    # clear target Text Encoder to save memory
-    for text_encoder in text_encoders_t:
-        del text_encoder
-
     if not text_encoder_different:
-        logger.warning("Text encoder is same. Extract U-Net only.")
+        print("Text encoder is same. Extract U-Net only.")
         lora_network_o.text_encoder_loras = []
-        diffs = {}  # clear diffs
+        diffs = {}
 
     for i, (lora_o, lora_t) in enumerate(zip(lora_network_o.unet_loras, lora_network_t.unet_loras)):
         lora_name = lora_o.lora_name
         module_o = lora_o.org_module
         module_t = lora_t.org_module
-        diff = module_t.weight.to(work_device) - module_o.weight.to(work_device)
+        diff = module_t.weight - module_o.weight
+        diff = diff.float()
 
-        # clear weight to save memory
-        module_o.weight = None
-        module_t.weight = None
+        if args.device:
+            diff = diff.to(args.device)
 
         diffs[lora_name] = diff
 
-    # clear LoRA network, target U-Net to save memory
-    del lora_network_o
-    del lora_network_t
-    del unet_t
-
     # make LoRA with svd
-    logger.info("calculating by svd")
+    print("calculating by svd")
     lora_weights = {}
     with torch.no_grad():
         for lora_name, mat in tqdm(list(diffs.items())):
-            if args.device:
-                mat = mat.to(args.device)
-            mat = mat.to(torch.float)  # calc by float
-
-            # if conv_dim is None, diffs do not include LoRAs for conv2d-3x3
+            # 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)
 
-            rank = dim if not conv2d_3x3 or conv_dim is None else conv_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 device:
-                mat = mat.to(device)
+            if args.device:
+                mat = mat.to(args.device)
 
-            # logger.info(lora_name, mat.size(), mat.device, rank, in_dim, out_dim)
+            # 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
 
             if conv2d:
@@ -206,7 +149,7 @@ def svd(
             Vh = Vh[:rank, :]
 
             dist = torch.cat([U.flatten(), Vh.flatten()])
-            hi_val = torch.quantile(dist, clamp_quantile)
+            hi_val = torch.quantile(dist, CLAMP_QUANTILE)
             low_val = -hi_val
 
             U = U.clamp(low_val, hi_val)
@@ -216,8 +159,8 @@ def svd(
                 U = U.reshape(out_dim, rank, 1, 1)
                 Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1])
 
-            U = U.to(work_device, dtype=save_dtype).contiguous()
-            Vh = Vh.to(work_device, dtype=save_dtype).contiguous()
+            U = U.to("cpu").contiguous()
+            Vh = Vh.to("cpu").contiguous()
 
             lora_weights[lora_name] = (U, Vh)
 
@@ -233,34 +176,36 @@ def svd(
     lora_network_save.apply_to(text_encoders_o, unet_o)  # create internal module references for state_dict
 
     info = lora_network_save.load_state_dict(lora_sd)
-    logger.info(f"Loading extracted LoRA weights: {info}")
+    print(f"Loading extracted LoRA weights: {info}")
 
-    dir_name = os.path.dirname(save_to)
+    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 conv_dim is not None:
-        net_kwargs["conv_dim"] = str(conv_dim)
-        net_kwargs["conv_alpha"] = str(float(conv_dim))
+    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(v2),
+        "ss_v2": str(args.v2),
         "ss_base_model_version": model_version,
         "ss_network_module": "networks.lora",
-        "ss_network_dim": str(dim),
-        "ss_network_alpha": str(float(dim)),
+        "ss_network_dim": str(args.dim),
+        "ss_network_alpha": str(args.dim),
         "ss_network_args": json.dumps(net_kwargs),
     }
 
-    if not no_metadata:
-        title = os.path.splitext(os.path.basename(save_to))[0]
-        sai_metadata = sai_model_spec.build_metadata(None, v2, v_parameterization, sdxl, True, False, time.time(), title=title)
+    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, args.sdxl, True, False, time.time(), title=title
+        )
         metadata.update(sai_metadata)
 
-    lora_network_save.save_weights(save_to, save_dtype, metadata)
-    logger.info(f"LoRA weights are saved to: {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:
@@ -268,20 +213,13 @@ def setup_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",
-        action="store_true",
+        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(
-        "--load_precision",
-        type=str,
-        default=None,
-        choices=[None, "float", "fp16", "bf16"],
-        help="precision in loading, model default if omitted / 読み込み時に精度を変更して読み込む、省略時はモデルファイルによる"
-    )
     parser.add_argument(
         "--save_precision",
         type=str,
@@ -293,22 +231,16 @@ def setup_parser() -> argparse.ArgumentParser:
         "--model_org",
         type=str,
         default=None,
-        required=True,
         help="Stable Diffusion original model: ckpt or safetensors file / 元モデル、ckptまたはsafetensors",
     )
     parser.add_argument(
         "--model_tuned",
         type=str,
         default=None,
-        required=True,
         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,
-        required=True,
-        help="destination file name: ckpt or safetensors file / 保存先のファイル名、ckptまたはsafetensors",
+        "--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(
@@ -318,37 +250,12 @@ def setup_parser() -> argparse.ArgumentParser:
         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(
-        "--clamp_quantile",
-        type=float,
-        default=0.99,
-        help="Quantile clamping value, float, (0-1). Default = 0.99 / 値をクランプするための分位点、float、(0-1)。デフォルトは0.99",
-    )
-    parser.add_argument(
-        "--min_diff",
-        type=float,
-        default=0.01,
-        help="Minimum difference between finetuned model and base to consider them different enough to extract, float, (0-1). Default = 0.01 /"
-        + "LoRAを抽出するために元モデルと派生モデルの差分の最小値、float、(0-1)。デフォルトは0.01",
-    )
     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.add_argument(
-        "--load_original_model_to",
-        type=str,
-        default=None,
-        help="location to load original model, cpu or cuda, cuda:0, etc, default is cpu, only for SDXL / 元モデル読み込み先、cpuまたはcuda、cuda:0など、省略時はcpu、SDXLのみ有効",
-    )
-    parser.add_argument(
-        "--load_tuned_model_to",
-        type=str,
-        default=None,
-        help="location to load tuned model, cpu or cuda, cuda:0, etc, default is cpu, only for SDXL / 派生モデル読み込み先、cpuまたはcuda、cuda:0など、省略時はcpu、SDXLのみ有効",
-    )
 
     return parser
 
@@ -357,4 +264,4 @@ if __name__ == "__main__":
     parser = setup_parser()
 
     args = parser.parse_args()
-    svd(**vars(args))
+    svd(args)

networks/lora.py

@@ -11,12 +11,7 @@ from transformers import CLIPTextModel
 import numpy as np
 import torch
 import re
-from library.utils import setup_logging
 
-setup_logging()
-import logging
-
-logger = logging.getLogger(__name__)
 
 RE_UPDOWN = re.compile(r"(up|down)_blocks_(\d+)_(resnets|upsamplers|downsamplers|attentions)_(\d+)_")
 
@@ -51,7 +46,7 @@ class LoRAModule(torch.nn.Module):
         # if limit_rank:
         #   self.lora_dim = min(lora_dim, in_dim, out_dim)
         #   if self.lora_dim != lora_dim:
-        #     logger.info(f"{lora_name} dim (rank) is changed to: {self.lora_dim}")
+        #     print(f"{lora_name} dim (rank) is changed to: {self.lora_dim}")
         # else:
         self.lora_dim = lora_dim
 
@@ -182,7 +177,7 @@ class LoRAInfModule(LoRAModule):
         else:
             # conv2d 3x3
             conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
-            # logger.info(conved.size(), weight.size(), module.stride, module.padding)
+            # print(conved.size(), weight.size(), module.stride, module.padding)
             weight = weight + self.multiplier * conved * self.scale
 
         # set weight to org_module
@@ -221,7 +216,7 @@ class LoRAInfModule(LoRAModule):
         self.region_mask = None
 
     def default_forward(self, x):
-        # logger.info(f"default_forward {self.lora_name} {x.size()}")
+        # print("default_forward", self.lora_name, x.size())
         return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
 
     def forward(self, x):
@@ -250,8 +245,7 @@ class LoRAInfModule(LoRAModule):
         if mask is None:
             # raise ValueError(f"mask is None for resolution {area}")
             # emb_layers in SDXL doesn't have mask
-            # if "emb" not in self.lora_name:
-            #     print(f"mask is None for resolution {self.lora_name}, {area}, {x.size()}")
+            # 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
         if len(x.size()) != 4:
@@ -269,8 +263,6 @@ class LoRAInfModule(LoRAModule):
         lx = self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
         mask = self.get_mask_for_x(lx)
         # print("regional", self.lora_name, self.network.sub_prompt_index, lx.size(), mask.size())
-        # if mask.ndim > lx.ndim:  # in some resolution, lx is 2d and mask is 3d (the reason is not checked)
-        #     mask = mask.squeeze(-1)
         lx = lx * mask
 
         x = self.org_forward(x)
@@ -299,7 +291,7 @@ class LoRAInfModule(LoRAModule):
         if has_real_uncond:
             query[-self.network.batch_size :] = x[-self.network.batch_size :]
 
-        # logger.info(f"postp_to_q {self.lora_name} {x.size()} {query.size()} {self.network.num_sub_prompts}")
+        # print("postp_to_q", self.lora_name, x.size(), query.size(), self.network.num_sub_prompts)
         return query
 
     def sub_prompt_forward(self, x):
@@ -314,7 +306,7 @@ class LoRAInfModule(LoRAModule):
         lx = x[emb_idx :: self.network.num_sub_prompts]
         lx = self.lora_up(self.lora_down(lx)) * self.multiplier * self.scale
 
-        # logger.info(f"sub_prompt_forward {self.lora_name} {x.size()} {lx.size()} {emb_idx}")
+        # print("sub_prompt_forward", self.lora_name, x.size(), lx.size(), emb_idx)
 
         x = self.org_forward(x)
         x[emb_idx :: self.network.num_sub_prompts] += lx
@@ -322,7 +314,7 @@ class LoRAInfModule(LoRAModule):
         return x
 
     def to_out_forward(self, x):
-        # logger.info(f"to_out_forward {self.lora_name} {x.size()} {self.network.is_last_network}")
+        # print("to_out_forward", self.lora_name, x.size(), self.network.is_last_network)
 
         if self.network.is_last_network:
             masks = [None] * self.network.num_sub_prompts
@@ -340,7 +332,7 @@ class LoRAInfModule(LoRAModule):
             )
             self.network.shared[self.lora_name] = (lx, masks)
 
-        # logger.info(f"to_out_forward {lx.size()} {lx1.size()} {self.network.sub_prompt_index} {self.network.num_sub_prompts}")
+        # print("to_out_forward", lx.size(), lx1.size(), self.network.sub_prompt_index, self.network.num_sub_prompts)
         lx[self.network.sub_prompt_index :: self.network.num_sub_prompts] += lx1
         masks[self.network.sub_prompt_index] = self.get_mask_for_x(lx1)
 
@@ -359,7 +351,7 @@ class LoRAInfModule(LoRAModule):
         if has_real_uncond:
             out[-self.network.batch_size :] = x[-self.network.batch_size :]  # real_uncond
 
-        # logger.info(f"to_out_forward {self.lora_name} {self.network.sub_prompt_index} {self.network.num_sub_prompts}")
+        # 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:
@@ -382,7 +374,7 @@ class LoRAInfModule(LoRAModule):
             x1 = x1 + lx1
             out[self.network.batch_size + i] = x1
 
-        # logger.info(f"to_out_forward {x.size()} {out.size()} {has_real_uncond}")
+        # print("to_out_forward", x.size(), out.size(), has_real_uncond)
         return out
 
 
@@ -519,9 +511,7 @@ def get_block_dims_and_alphas(
             len(block_dims) == num_total_blocks
         ), f"block_dims must have {num_total_blocks} elements / block_dimsは{num_total_blocks}個指定してください"
     else:
-        logger.warning(
-            f"block_dims is not specified. all dims are set to {network_dim} / block_dimsが指定されていません。すべてのdimは{network_dim}になります"
-        )
+        print(f"block_dims is not specified. all dims are set to {network_dim} / block_dimsが指定されていません。すべてのdimは{network_dim}になります")
         block_dims = [network_dim] * num_total_blocks
 
     if block_alphas is not None:
@@ -530,7 +520,7 @@ def get_block_dims_and_alphas(
             len(block_alphas) == num_total_blocks
         ), f"block_alphas must have {num_total_blocks} elements / block_alphasは{num_total_blocks}個指定してください"
     else:
-        logger.warning(
+        print(
             f"block_alphas is not specified. all alphas are set to {network_alpha} / block_alphasが指定されていません。すべてのalphaは{network_alpha}になります"
         )
         block_alphas = [network_alpha] * num_total_blocks
@@ -550,13 +540,13 @@ def get_block_dims_and_alphas(
         else:
             if conv_alpha is None:
                 conv_alpha = 1.0
-            logger.warning(
+            print(
                 f"conv_block_alphas is not specified. all alphas are set to {conv_alpha} / conv_block_alphasが指定されていません。すべてのalphaは{conv_alpha}になります"
             )
             conv_block_alphas = [conv_alpha] * num_total_blocks
     else:
         if conv_dim is not None:
-            logger.warning(
+            print(
                 f"conv_dim/alpha for all blocks are set to {conv_dim} and {conv_alpha} / すべてのブロックのconv_dimとalphaは{conv_dim}および{conv_alpha}になります"
             )
             conv_block_dims = [conv_dim] * num_total_blocks
@@ -596,7 +586,7 @@ def get_block_lr_weight(
         elif name == "zeros":
             return [0.0 + base_lr] * max_len
         else:
-            logger.error(
+            print(
                 "Unknown lr_weight argument %s is used. Valid arguments:  / 不明なlr_weightの引数 %s が使われました。有効な引数:\n\tcosine, sine, linear, reverse_linear, zeros"
                 % (name)
             )
@@ -608,14 +598,14 @@ def get_block_lr_weight(
         up_lr_weight = get_list(up_lr_weight)
 
     if (up_lr_weight != None and len(up_lr_weight) > max_len) or (down_lr_weight != None and len(down_lr_weight) > max_len):
-        logger.warning("down_weight or up_weight is too long. Parameters after %d-th are ignored." % max_len)
-        logger.warning("down_weightもしくはup_weightが長すぎます。%d個目以降のパラメータは無視されます。" % max_len)
+        print("down_weight or up_weight is too long. Parameters after %d-th are ignored." % max_len)
+        print("down_weightもしくはup_weightが長すぎます。%d個目以降のパラメータは無視されます。" % max_len)
         up_lr_weight = up_lr_weight[:max_len]
         down_lr_weight = down_lr_weight[:max_len]
 
     if (up_lr_weight != None and len(up_lr_weight) < max_len) or (down_lr_weight != None and len(down_lr_weight) < max_len):
-        logger.warning("down_weight or up_weight is too short. Parameters after %d-th are filled with 1." % max_len)
-        logger.warning("down_weightもしくはup_weightが短すぎます。%d個目までの不足したパラメータは1で補われます。" % max_len)
+        print("down_weight or up_weight is too short. Parameters after %d-th are filled with 1." % max_len)
+        print("down_weightもしくはup_weightが短すぎます。%d個目までの不足したパラメータは1で補われます。" % max_len)
 
         if down_lr_weight != None and len(down_lr_weight) < max_len:
             down_lr_weight = down_lr_weight + [1.0] * (max_len - len(down_lr_weight))
@@ -623,24 +613,24 @@ def get_block_lr_weight(
             up_lr_weight = up_lr_weight + [1.0] * (max_len - len(up_lr_weight))
 
     if (up_lr_weight != None) or (mid_lr_weight != None) or (down_lr_weight != None):
-        logger.info("apply block learning rate / 階層別学習率を適用します。")
+        print("apply block learning rate / 階層別学習率を適用します。")
         if down_lr_weight != None:
             down_lr_weight = [w if w > zero_threshold else 0 for w in down_lr_weight]
-            logger.info(f"down_lr_weight (shallower -> deeper, 浅い層->深い層): {down_lr_weight}")
+            print("down_lr_weight (shallower -> deeper, 浅い層->深い層):", down_lr_weight)
         else:
-            logger.info("down_lr_weight: all 1.0, すべて1.0")
+            print("down_lr_weight: all 1.0, すべて1.0")
 
         if mid_lr_weight != None:
             mid_lr_weight = mid_lr_weight if mid_lr_weight > zero_threshold else 0
-            logger.info(f"mid_lr_weight: {mid_lr_weight}")
+            print("mid_lr_weight:", mid_lr_weight)
         else:
-            logger.info("mid_lr_weight: 1.0")
+            print("mid_lr_weight: 1.0")
 
         if up_lr_weight != None:
             up_lr_weight = [w if w > zero_threshold else 0 for w in up_lr_weight]
-            logger.info(f"up_lr_weight (deeper -> shallower, 深い層->浅い層): {up_lr_weight}")
+            print("up_lr_weight (deeper -> shallower, 深い層->浅い層):", up_lr_weight)
         else:
-            logger.info("up_lr_weight: all 1.0, すべて1.0")
+            print("up_lr_weight: all 1.0, すべて1.0")
 
     return down_lr_weight, mid_lr_weight, up_lr_weight
 
@@ -721,7 +711,7 @@ def create_network_from_weights(multiplier, file, vae, text_encoder, unet, weigh
         elif "lora_down" in key:
             dim = value.size()[0]
             modules_dim[lora_name] = dim
-            # logger.info(lora_name, value.size(), dim)
+            # print(lora_name, value.size(), dim)
 
     # support old LoRA without alpha
     for key in modules_dim.keys():
@@ -796,26 +786,20 @@ class LoRANetwork(torch.nn.Module):
         self.module_dropout = module_dropout
 
         if modules_dim is not None:
-            logger.info(f"create LoRA network from weights")
+            print(f"create LoRA network from weights")
         elif block_dims is not None:
-            logger.info(f"create LoRA network from block_dims")
-            logger.info(
-                f"neuron dropout: p={self.dropout}, rank dropout: p={self.rank_dropout}, module dropout: p={self.module_dropout}"
-            )
-            logger.info(f"block_dims: {block_dims}")
-            logger.info(f"block_alphas: {block_alphas}")
+            print(f"create LoRA network from block_dims")
+            print(f"neuron dropout: p={self.dropout}, rank dropout: p={self.rank_dropout}, module dropout: p={self.module_dropout}")
+            print(f"block_dims: {block_dims}")
+            print(f"block_alphas: {block_alphas}")
             if conv_block_dims is not None:
-                logger.info(f"conv_block_dims: {conv_block_dims}")
-                logger.info(f"conv_block_alphas: {conv_block_alphas}")
+                print(f"conv_block_dims: {conv_block_dims}")
+                print(f"conv_block_alphas: {conv_block_alphas}")
         else:
-            logger.info(f"create LoRA network. base dim (rank): {lora_dim}, alpha: {alpha}")
-            logger.info(
-                f"neuron dropout: p={self.dropout}, rank dropout: p={self.rank_dropout}, module dropout: p={self.module_dropout}"
-            )
+            print(f"create LoRA network. base dim (rank): {lora_dim}, alpha: {alpha}")
+            print(f"neuron dropout: p={self.dropout}, rank dropout: p={self.rank_dropout}, module dropout: p={self.module_dropout}")
             if self.conv_lora_dim is not None:
-                logger.info(
-                    f"apply LoRA to Conv2d with kernel size (3,3). dim (rank): {self.conv_lora_dim}, alpha: {self.conv_alpha}"
-                )
+                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(
@@ -900,15 +884,15 @@ class LoRANetwork(torch.nn.Module):
         for i, text_encoder in enumerate(text_encoders):
             if len(text_encoders) > 1:
                 index = i + 1
-                logger.info(f"create LoRA for Text Encoder {index}:")
+                print(f"create LoRA for Text Encoder {index}:")
             else:
                 index = None
-                logger.info(f"create LoRA for Text Encoder:")
+                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
-        logger.info(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
+        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
         target_modules = LoRANetwork.UNET_TARGET_REPLACE_MODULE
@@ -916,15 +900,15 @@ class LoRANetwork(torch.nn.Module):
             target_modules += LoRANetwork.UNET_TARGET_REPLACE_MODULE_CONV2D_3X3
 
         self.unet_loras, skipped_un = create_modules(True, None, unet, target_modules)
-        logger.info(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
+        print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
 
         skipped = skipped_te + skipped_un
         if varbose and len(skipped) > 0:
-            logger.warning(
+            print(
                 f"because block_lr_weight is 0 or dim (rank) is 0, {len(skipped)} LoRA modules are skipped / block_lr_weightまたはdim (rank)が0の為、次の{len(skipped)}個のLoRAモジュールはスキップされます:"
             )
             for name in skipped:
-                logger.info(f"\t{name}")
+                print(f"\t{name}")
 
         self.up_lr_weight: List[float] = None
         self.down_lr_weight: List[float] = None
@@ -942,10 +926,6 @@ class LoRANetwork(torch.nn.Module):
         for lora in self.text_encoder_loras + self.unet_loras:
             lora.multiplier = self.multiplier
 
-    def set_enabled(self, is_enabled):
-        for lora in self.text_encoder_loras + self.unet_loras:
-            lora.enabled = is_enabled
-
     def load_weights(self, file):
         if os.path.splitext(file)[1] == ".safetensors":
             from safetensors.torch import load_file
@@ -959,12 +939,12 @@ class LoRANetwork(torch.nn.Module):
 
     def apply_to(self, text_encoder, unet, apply_text_encoder=True, apply_unet=True):
         if apply_text_encoder:
-            logger.info("enable LoRA for text encoder")
+            print("enable LoRA for text encoder")
         else:
             self.text_encoder_loras = []
 
         if apply_unet:
-            logger.info("enable LoRA for U-Net")
+            print("enable LoRA for U-Net")
         else:
             self.unet_loras = []
 
@@ -986,12 +966,12 @@ class LoRANetwork(torch.nn.Module):
                 apply_unet = True
 
         if apply_text_encoder:
-            logger.info("enable LoRA for text encoder")
+            print("enable LoRA for text encoder")
         else:
             self.text_encoder_loras = []
 
         if apply_unet:
-            logger.info("enable LoRA for U-Net")
+            print("enable LoRA for U-Net")
         else:
             self.unet_loras = []
 
@@ -1002,7 +982,7 @@ class LoRANetwork(torch.nn.Module):
                     sd_for_lora[key[len(lora.lora_name) + 1 :]] = weights_sd[key]
             lora.merge_to(sd_for_lora, dtype, device)
 
-        logger.info(f"weights are merged")
+        print(f"weights are merged")
 
     # 層別学習率用に層ごとの学習率に対する倍率を定義する　引数の順番が逆だがとりあえず気にしない
     def set_block_lr_weight(
@@ -1133,7 +1113,7 @@ class LoRANetwork(torch.nn.Module):
         for lora in self.text_encoder_loras + self.unet_loras:
             lora.set_network(self)
 
-    def set_current_generation(self, batch_size, num_sub_prompts, width, height, shared, ds_ratio=None):
+    def set_current_generation(self, batch_size, num_sub_prompts, width, height, shared):
         self.batch_size = batch_size
         self.num_sub_prompts = num_sub_prompts
         self.current_size = (height, width)
@@ -1148,7 +1128,7 @@ class LoRANetwork(torch.nn.Module):
         device = ref_weight.device
 
         def resize_add(mh, mw):
-            # logger.info(mh, mw, mh * mw)
+            # print(mh, mw, mh * mw)
             m = torch.nn.functional.interpolate(mask, (mh, mw), mode="bilinear")  # doesn't work in bf16
             m = m.to(device, dtype=dtype)
             mask_dic[mh * mw] = m
@@ -1159,13 +1139,6 @@ class LoRANetwork(torch.nn.Module):
             resize_add(h, w)
             if h % 2 == 1 or w % 2 == 1:  # add extra shape if h/w is not divisible by 2
                 resize_add(h + h % 2, w + w % 2)
-
-            # deep shrink
-            if ds_ratio is not None:
-                hd = int(h * ds_ratio)
-                wd = int(w * ds_ratio)
-                resize_add(hd, wd)
-
             h = (h + 1) // 2
             w = (w + 1) // 2
 

networks/lora_diffusers.py

@@ -9,15 +9,8 @@ from diffusers import UNet2DConditionModel
 import numpy as np
 from tqdm import tqdm
 from transformers import CLIPTextModel
-
 import torch
-from library.device_utils import init_ipex, get_preferred_device
-init_ipex()
 
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
 
 def make_unet_conversion_map() -> Dict[str, str]:
     unet_conversion_map_layer = []
@@ -255,7 +248,7 @@ def create_network_from_weights(
         elif "lora_down" in key:
             dim = value.size()[0]
             modules_dim[lora_name] = dim
-            # logger.info(f"{lora_name} {value.size()} {dim}")
+            # print(lora_name, value.size(), dim)
 
     # support old LoRA without alpha
     for key in modules_dim.keys():
@@ -298,12 +291,12 @@ class LoRANetwork(torch.nn.Module):
         super().__init__()
         self.multiplier = multiplier
 
-        logger.info("create LoRA network from weights")
+        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:
-            logger.info(f"converted {converted} Stability AI's U-Net LoRA modules to Diffusers (SDXL)")
+            print(f"converted {converted} Stability AI's U-Net LoRA modules to Diffusers (SDXL)")
 
         # create module instances
         def create_modules(
@@ -338,7 +331,7 @@ class LoRANetwork(torch.nn.Module):
                             lora_name = lora_name.replace(".", "_")
 
                             if lora_name not in modules_dim:
-                                # logger.info(f"skipped {lora_name} (not found in modules_dim)")
+                                # print(f"skipped {lora_name} (not found in modules_dim)")
                                 skipped.append(lora_name)
                                 continue
 
@@ -369,18 +362,18 @@ class LoRANetwork(torch.nn.Module):
             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
-        logger.info(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
+        print(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
         if len(skipped_te) > 0:
-            logger.warning(f"skipped {len(skipped_te)} modules because of missing weight for text encoder.")
+            print(f"skipped {len(skipped_te)} modules because of missing weight for text encoder.")
 
         # 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)
-        logger.info(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
+        print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
         if len(skipped_un) > 0:
-            logger.warning(f"skipped {len(skipped_un)} modules because of missing weight for U-Net.")
+            print(f"skipped {len(skipped_un)} modules because of missing weight for U-Net.")
 
         # assertion
         names = set()
@@ -427,11 +420,11 @@ class LoRANetwork(torch.nn.Module):
 
     def apply_to(self, multiplier=1.0, apply_text_encoder=True, apply_unet=True):
         if apply_text_encoder:
-            logger.info("enable LoRA for text encoder")
+            print("enable LoRA for text encoder")
             for lora in self.text_encoder_loras:
                 lora.apply_to(multiplier)
         if apply_unet:
-            logger.info("enable LoRA for U-Net")
+            print("enable LoRA for U-Net")
             for lora in self.unet_loras:
                 lora.apply_to(multiplier)
 
@@ -440,16 +433,16 @@ class LoRANetwork(torch.nn.Module):
             lora.unapply_to()
 
     def merge_to(self, multiplier=1.0):
-        logger.info("merge LoRA weights to original weights")
+        print("merge LoRA weights to original weights")
         for lora in tqdm(self.text_encoder_loras + self.unet_loras):
             lora.merge_to(multiplier)
-        logger.info(f"weights are merged")
+        print(f"weights are merged")
 
     def restore_from(self, multiplier=1.0):
-        logger.info("restore LoRA weights from original weights")
+        print("restore LoRA weights from original weights")
         for lora in tqdm(self.text_encoder_loras + self.unet_loras):
             lora.restore_from(multiplier)
-        logger.info(f"weights are restored")
+        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
@@ -470,7 +463,7 @@ class LoRANetwork(torch.nn.Module):
         my_state_dict = self.state_dict()
         for key in state_dict.keys():
             if state_dict[key].size() != my_state_dict[key].size():
-                # logger.info(f"convert {key} from {state_dict[key].size()} to {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)
@@ -483,7 +476,7 @@ if __name__ == "__main__":
     from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
     import torch
 
-    device = get_preferred_device()
+    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")
@@ -497,7 +490,7 @@ if __name__ == "__main__":
     image_prefix = args.model_id.replace("/", "_") + "_"
 
     # load Diffusers model
-    logger.info(f"load model from {args.model_id}")
+    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
@@ -510,7 +503,7 @@ if __name__ == "__main__":
     text_encoders = [pipe.text_encoder, pipe.text_encoder_2] if args.sdxl else [pipe.text_encoder]
 
     # load LoRA weights
-    logger.info(f"load LoRA weights from {args.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
 
@@ -519,10 +512,10 @@ if __name__ == "__main__":
         lora_sd = torch.load(args.lora_weights)
 
     # create by LoRA weights and load weights
-    logger.info(f"create LoRA network")
+    print(f"create LoRA network")
     lora_network: LoRANetwork = create_network_from_weights(text_encoders, pipe.unet, lora_sd, multiplier=1.0)
 
-    logger.info(f"load LoRA network weights")
+    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
@@ -551,34 +544,34 @@ if __name__ == "__main__":
         random.seed(seed)
 
     # create image with original weights
-    logger.info(f"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
-    logger.info(f"apply LoRA network to the model")
+    print(f"apply LoRA network to the model")
     lora_network.apply_to(multiplier=1.0)
 
-    logger.info(f"create image with applied LoRA")
+    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
-    logger.info(f"unapply LoRA network to the model")
+    print(f"unapply LoRA network to the model")
     lora_network.unapply_to()
 
-    logger.info(f"create image with unapplied LoRA")
+    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)
-    logger.info(f"merge LoRA network to the model")
+    print(f"merge LoRA network to the model")
     lora_network.merge_to(multiplier=1.0)
 
-    logger.info(f"create image with LoRA")
+    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")
@@ -586,31 +579,31 @@ if __name__ == "__main__":
     # restore (unmerge) LoRA weights: numerically unstable
     # マージされた重みを元に戻す。計算誤差のため、元の重みと完全に一致しないことがあるかもしれない
     # 保存したstate_dictから元の重みを復元するのが確実
-    logger.info(f"restore (unmerge) LoRA weights")
+    print(f"restore (unmerge) LoRA weights")
     lora_network.restore_from(multiplier=1.0)
 
-    logger.info(f"create image without LoRA")
+    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
-    logger.info(f"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)
 
-    logger.info(f"create image with restored original weights")
+    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
-    logger.info(f"merge LoRA weights with convenience function")
+    print(f"merge LoRA weights with convenience function")
     merge_lora_weights(pipe, lora_sd, multiplier=1.0)
 
-    logger.info(f"create image with merged LoRA weights")
+    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/lora_fa.py

@@ -14,10 +14,7 @@ from transformers import CLIPTextModel
 import numpy as np
 import torch
 import re
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 RE_UPDOWN = re.compile(r"(up|down)_blocks_(\d+)_(resnets|upsamplers|downsamplers|attentions)_(\d+)_")
 
@@ -52,7 +49,7 @@ class LoRAModule(torch.nn.Module):
         # if limit_rank:
         #   self.lora_dim = min(lora_dim, in_dim, out_dim)
         #   if self.lora_dim != lora_dim:
-        #     logger.info(f"{lora_name} dim (rank) is changed to: {self.lora_dim}")
+        #     print(f"{lora_name} dim (rank) is changed to: {self.lora_dim}")
         # else:
         self.lora_dim = lora_dim
 
@@ -200,7 +197,7 @@ class LoRAInfModule(LoRAModule):
         else:
             # conv2d 3x3
             conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
-            # logger.info(conved.size(), weight.size(), module.stride, module.padding)
+            # print(conved.size(), weight.size(), module.stride, module.padding)
             weight = weight + self.multiplier * conved * self.scale
 
         # set weight to org_module
@@ -239,7 +236,7 @@ class LoRAInfModule(LoRAModule):
         self.region_mask = None
 
     def default_forward(self, x):
-        # logger.info("default_forward", self.lora_name, x.size())
+        # print("default_forward", self.lora_name, x.size())
         return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
 
     def forward(self, x):
@@ -281,7 +278,7 @@ class LoRAInfModule(LoRAModule):
         # apply mask for LoRA result
         lx = self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
         mask = self.get_mask_for_x(lx)
-        # logger.info("regional", self.lora_name, self.network.sub_prompt_index, lx.size(), mask.size())
+        # print("regional", self.lora_name, self.network.sub_prompt_index, lx.size(), mask.size())
         lx = lx * mask
 
         x = self.org_forward(x)
@@ -310,7 +307,7 @@ class LoRAInfModule(LoRAModule):
         if has_real_uncond:
             query[-self.network.batch_size :] = x[-self.network.batch_size :]
 
-        # logger.info("postp_to_q", self.lora_name, x.size(), query.size(), self.network.num_sub_prompts)
+        # print("postp_to_q", self.lora_name, x.size(), query.size(), self.network.num_sub_prompts)
         return query
 
     def sub_prompt_forward(self, x):
@@ -325,7 +322,7 @@ class LoRAInfModule(LoRAModule):
         lx = x[emb_idx :: self.network.num_sub_prompts]
         lx = self.lora_up(self.lora_down(lx)) * self.multiplier * self.scale
 
-        # logger.info("sub_prompt_forward", self.lora_name, x.size(), lx.size(), emb_idx)
+        # print("sub_prompt_forward", self.lora_name, x.size(), lx.size(), emb_idx)
 
         x = self.org_forward(x)
         x[emb_idx :: self.network.num_sub_prompts] += lx
@@ -333,7 +330,7 @@ class LoRAInfModule(LoRAModule):
         return x
 
     def to_out_forward(self, x):
-        # logger.info("to_out_forward", self.lora_name, x.size(), self.network.is_last_network)
+        # print("to_out_forward", self.lora_name, x.size(), self.network.is_last_network)
 
         if self.network.is_last_network:
             masks = [None] * self.network.num_sub_prompts
@@ -351,7 +348,7 @@ class LoRAInfModule(LoRAModule):
             )
             self.network.shared[self.lora_name] = (lx, masks)
 
-        # logger.info("to_out_forward", lx.size(), lx1.size(), self.network.sub_prompt_index, self.network.num_sub_prompts)
+        # print("to_out_forward", lx.size(), lx1.size(), self.network.sub_prompt_index, self.network.num_sub_prompts)
         lx[self.network.sub_prompt_index :: self.network.num_sub_prompts] += lx1
         masks[self.network.sub_prompt_index] = self.get_mask_for_x(lx1)
 
@@ -370,7 +367,7 @@ class LoRAInfModule(LoRAModule):
         if has_real_uncond:
             out[-self.network.batch_size :] = x[-self.network.batch_size :]  # real_uncond
 
-        # logger.info("to_out_forward", self.lora_name, self.network.sub_prompt_index, self.network.num_sub_prompts)
+        # print("to_out_forward", self.lora_name, self.network.sub_prompt_index, self.network.num_sub_prompts)
         # for i in range(len(masks)):
         #     if masks[i] is None:
         #         masks[i] = torch.zeros_like(masks[-1])
@@ -392,7 +389,7 @@ class LoRAInfModule(LoRAModule):
             x1 = x1 + lx1
             out[self.network.batch_size + i] = x1
 
-        # logger.info("to_out_forward", x.size(), out.size(), has_real_uncond)
+        # print("to_out_forward", x.size(), out.size(), has_real_uncond)
         return out
 
 
@@ -529,7 +526,7 @@ def get_block_dims_and_alphas(
             len(block_dims) == num_total_blocks
         ), f"block_dims must have {num_total_blocks} elements / block_dimsは{num_total_blocks}個指定してください"
     else:
-        logger.warning(f"block_dims is not specified. all dims are set to {network_dim} / block_dimsが指定されていません。すべてのdimは{network_dim}になります")
+        print(f"block_dims is not specified. all dims are set to {network_dim} / block_dimsが指定されていません。すべてのdimは{network_dim}になります")
         block_dims = [network_dim] * num_total_blocks
 
     if block_alphas is not None:
@@ -538,7 +535,7 @@ def get_block_dims_and_alphas(
             len(block_alphas) == num_total_blocks
         ), f"block_alphas must have {num_total_blocks} elements / block_alphasは{num_total_blocks}個指定してください"
     else:
-        logger.warning(
+        print(
             f"block_alphas is not specified. all alphas are set to {network_alpha} / block_alphasが指定されていません。すべてのalphaは{network_alpha}になります"
         )
         block_alphas = [network_alpha] * num_total_blocks
@@ -558,13 +555,13 @@ def get_block_dims_and_alphas(
         else:
             if conv_alpha is None:
                 conv_alpha = 1.0
-            logger.warning(
+            print(
                 f"conv_block_alphas is not specified. all alphas are set to {conv_alpha} / conv_block_alphasが指定されていません。すべてのalphaは{conv_alpha}になります"
             )
             conv_block_alphas = [conv_alpha] * num_total_blocks
     else:
         if conv_dim is not None:
-            logger.warning(
+            print(
                 f"conv_dim/alpha for all blocks are set to {conv_dim} and {conv_alpha} / すべてのブロックのconv_dimとalphaは{conv_dim}および{conv_alpha}になります"
             )
             conv_block_dims = [conv_dim] * num_total_blocks
@@ -604,7 +601,7 @@ def get_block_lr_weight(
         elif name == "zeros":
             return [0.0 + base_lr] * max_len
         else:
-            logger.error(
+            print(
                 "Unknown lr_weight argument %s is used. Valid arguments:  / 不明なlr_weightの引数 %s が使われました。有効な引数:\n\tcosine, sine, linear, reverse_linear, zeros"
                 % (name)
             )
@@ -616,14 +613,14 @@ def get_block_lr_weight(
         up_lr_weight = get_list(up_lr_weight)
 
     if (up_lr_weight != None and len(up_lr_weight) > max_len) or (down_lr_weight != None and len(down_lr_weight) > max_len):
-        logger.warning("down_weight or up_weight is too long. Parameters after %d-th are ignored." % max_len)
-        logger.warning("down_weightもしくはup_weightが長すぎます。%d個目以降のパラメータは無視されます。" % max_len)
+        print("down_weight or up_weight is too long. Parameters after %d-th are ignored." % max_len)
+        print("down_weightもしくはup_weightが長すぎます。%d個目以降のパラメータは無視されます。" % max_len)
         up_lr_weight = up_lr_weight[:max_len]
         down_lr_weight = down_lr_weight[:max_len]
 
     if (up_lr_weight != None and len(up_lr_weight) < max_len) or (down_lr_weight != None and len(down_lr_weight) < max_len):
-        logger.warning("down_weight or up_weight is too short. Parameters after %d-th are filled with 1." % max_len)
-        logger.warning("down_weightもしくはup_weightが短すぎます。%d個目までの不足したパラメータは1で補われます。" % max_len)
+        print("down_weight or up_weight is too short. Parameters after %d-th are filled with 1." % max_len)
+        print("down_weightもしくはup_weightが短すぎます。%d個目までの不足したパラメータは1で補われます。" % max_len)
 
         if down_lr_weight != None and len(down_lr_weight) < max_len:
             down_lr_weight = down_lr_weight + [1.0] * (max_len - len(down_lr_weight))
@@ -631,24 +628,24 @@ def get_block_lr_weight(
             up_lr_weight = up_lr_weight + [1.0] * (max_len - len(up_lr_weight))
 
     if (up_lr_weight != None) or (mid_lr_weight != None) or (down_lr_weight != None):
-        logger.info("apply block learning rate / 階層別学習率を適用します。")
+        print("apply block learning rate / 階層別学習率を適用します。")
         if down_lr_weight != None:
             down_lr_weight = [w if w > zero_threshold else 0 for w in down_lr_weight]
-            logger.info(f"down_lr_weight (shallower -> deeper, 浅い層->深い層): {down_lr_weight}")
+            print("down_lr_weight (shallower -> deeper, 浅い層->深い層):", down_lr_weight)
         else:
-            logger.info("down_lr_weight: all 1.0, すべて1.0")
+            print("down_lr_weight: all 1.0, すべて1.0")
 
         if mid_lr_weight != None:
             mid_lr_weight = mid_lr_weight if mid_lr_weight > zero_threshold else 0
-            logger.info(f"mid_lr_weight: {mid_lr_weight}")
+            print("mid_lr_weight:", mid_lr_weight)
         else:
-            logger.info("mid_lr_weight: 1.0")
+            print("mid_lr_weight: 1.0")
 
         if up_lr_weight != None:
             up_lr_weight = [w if w > zero_threshold else 0 for w in up_lr_weight]
-            logger.info(f"up_lr_weight (deeper -> shallower, 深い層->浅い層): {up_lr_weight}")
+            print("up_lr_weight (deeper -> shallower, 深い層->浅い層):", up_lr_weight)
         else:
-            logger.info("up_lr_weight: all 1.0, すべて1.0")
+            print("up_lr_weight: all 1.0, すべて1.0")
 
     return down_lr_weight, mid_lr_weight, up_lr_weight
 
@@ -729,7 +726,7 @@ def create_network_from_weights(multiplier, file, vae, text_encoder, unet, weigh
         elif "lora_down" in key:
             dim = value.size()[0]
             modules_dim[lora_name] = dim
-            # logger.info(lora_name, value.size(), dim)
+            # print(lora_name, value.size(), dim)
 
     # support old LoRA without alpha
     for key in modules_dim.keys():
@@ -804,20 +801,20 @@ class LoRANetwork(torch.nn.Module):
         self.module_dropout = module_dropout
 
         if modules_dim is not None:
-            logger.info(f"create LoRA network from weights")
+            print(f"create LoRA network from weights")
         elif block_dims is not None:
-            logger.info(f"create LoRA network from block_dims")
-            logger.info(f"neuron dropout: p={self.dropout}, rank dropout: p={self.rank_dropout}, module dropout: p={self.module_dropout}")
-            logger.info(f"block_dims: {block_dims}")
-            logger.info(f"block_alphas: {block_alphas}")
+            print(f"create LoRA network from block_dims")
+            print(f"neuron dropout: p={self.dropout}, rank dropout: p={self.rank_dropout}, module dropout: p={self.module_dropout}")
+            print(f"block_dims: {block_dims}")
+            print(f"block_alphas: {block_alphas}")
             if conv_block_dims is not None:
-                logger.info(f"conv_block_dims: {conv_block_dims}")
-                logger.info(f"conv_block_alphas: {conv_block_alphas}")
+                print(f"conv_block_dims: {conv_block_dims}")
+                print(f"conv_block_alphas: {conv_block_alphas}")
         else:
-            logger.info(f"create LoRA network. base dim (rank): {lora_dim}, alpha: {alpha}")
-            logger.info(f"neuron dropout: p={self.dropout}, rank dropout: p={self.rank_dropout}, module dropout: p={self.module_dropout}")
+            print(f"create LoRA network. base dim (rank): {lora_dim}, alpha: {alpha}")
+            print(f"neuron dropout: p={self.dropout}, rank dropout: p={self.rank_dropout}, module dropout: p={self.module_dropout}")
             if self.conv_lora_dim is not None:
-                logger.info(f"apply LoRA to Conv2d with kernel size (3,3). dim (rank): {self.conv_lora_dim}, alpha: {self.conv_alpha}")
+                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(
@@ -902,15 +899,15 @@ class LoRANetwork(torch.nn.Module):
         for i, text_encoder in enumerate(text_encoders):
             if len(text_encoders) > 1:
                 index = i + 1
-                logger.info(f"create LoRA for Text Encoder {index}:")
+                print(f"create LoRA for Text Encoder {index}:")
             else:
                 index = None
-                logger.info(f"create LoRA for Text Encoder:")
+                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
-        logger.info(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
+        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
         target_modules = LoRANetwork.UNET_TARGET_REPLACE_MODULE
@@ -918,15 +915,15 @@ class LoRANetwork(torch.nn.Module):
             target_modules += LoRANetwork.UNET_TARGET_REPLACE_MODULE_CONV2D_3X3
 
         self.unet_loras, skipped_un = create_modules(True, None, unet, target_modules)
-        logger.info(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
+        print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
 
         skipped = skipped_te + skipped_un
         if varbose and len(skipped) > 0:
-            logger.warning(
+            print(
                 f"because block_lr_weight is 0 or dim (rank) is 0, {len(skipped)} LoRA modules are skipped / block_lr_weightまたはdim (rank)が0の為、次の{len(skipped)}個のLoRAモジュールはスキップされます:"
             )
             for name in skipped:
-                logger.info(f"\t{name}")
+                print(f"\t{name}")
 
         self.up_lr_weight: List[float] = None
         self.down_lr_weight: List[float] = None
@@ -957,12 +954,12 @@ class LoRANetwork(torch.nn.Module):
 
     def apply_to(self, text_encoder, unet, apply_text_encoder=True, apply_unet=True):
         if apply_text_encoder:
-            logger.info("enable LoRA for text encoder")
+            print("enable LoRA for text encoder")
         else:
             self.text_encoder_loras = []
 
         if apply_unet:
-            logger.info("enable LoRA for U-Net")
+            print("enable LoRA for U-Net")
         else:
             self.unet_loras = []
 
@@ -984,12 +981,12 @@ class LoRANetwork(torch.nn.Module):
                 apply_unet = True
 
         if apply_text_encoder:
-            logger.info("enable LoRA for text encoder")
+            print("enable LoRA for text encoder")
         else:
             self.text_encoder_loras = []
 
         if apply_unet:
-            logger.info("enable LoRA for U-Net")
+            print("enable LoRA for U-Net")
         else:
             self.unet_loras = []
 
@@ -1000,7 +997,7 @@ class LoRANetwork(torch.nn.Module):
                     sd_for_lora[key[len(lora.lora_name) + 1 :]] = weights_sd[key]
             lora.merge_to(sd_for_lora, dtype, device)
 
-        logger.info(f"weights are merged")
+        print(f"weights are merged")
 
     # 層別学習率用に層ごとの学習率に対する倍率を定義する　引数の順番が逆だがとりあえず気にしない
     def set_block_lr_weight(
@@ -1147,7 +1144,7 @@ class LoRANetwork(torch.nn.Module):
         device = ref_weight.device
 
         def resize_add(mh, mw):
-            # logger.info(mh, mw, mh * mw)
+            # print(mh, mw, mh * mw)
             m = torch.nn.functional.interpolate(mask, (mh, mw), mode="bilinear")  # doesn't work in bf16
             m = m.to(device, dtype=dtype)
             mask_dic[mh * mw] = m

networks/lora_interrogator.py

@@ -5,34 +5,27 @@ from library import model_util
 import library.train_util as train_util
 import argparse
 from transformers import CLIPTokenizer
-
 import torch
-from library.device_utils import init_ipex, get_preferred_device
-init_ipex()
 
 import library.model_util as model_util
 import lora
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
 
 TOKENIZER_PATH = "openai/clip-vit-large-patch14"
 V2_STABLE_DIFFUSION_PATH = "stabilityai/stable-diffusion-2"     # ここからtokenizerだけ使う
 
-DEVICE = get_preferred_device()
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 
 def interrogate(args):
   weights_dtype = torch.float16
 
   # いろいろ準備する
-  logger.info(f"loading SD model: {args.sd_model}")
+  print(f"loading SD model: {args.sd_model}")
   args.pretrained_model_name_or_path = args.sd_model
   args.vae = None
   text_encoder, vae, unet, _ = train_util._load_target_model(args,weights_dtype, DEVICE)
 
-  logger.info(f"loading LoRA: {args.model}")
+  print(f"loading LoRA: {args.model}")
   network, weights_sd = lora.create_network_from_weights(1.0, args.model, vae, text_encoder, unet)
 
   # text encoder向けの重みがあるかチェックする：本当はlora側でやるのがいい
@@ -42,11 +35,11 @@ def interrogate(args):
       has_te_weight = True
       break
   if not has_te_weight:
-    logger.error("This LoRA does not have modules for Text Encoder, cannot interrogate / このLoRAはText Encoder向けのモジュールがないため調査できません")
+    print("This LoRA does not have modules for Text Encoder, cannot interrogate / このLoRAはText Encoder向けのモジュールがないため調査できません")
     return
   del vae
 
-  logger.info("loading tokenizer")
+  print("loading tokenizer")
   if args.v2:
     tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(V2_STABLE_DIFFUSION_PATH, subfolder="tokenizer")
   else:
@@ -60,7 +53,7 @@ def interrogate(args):
   # トークンをひとつひとつ当たっていく
   token_id_start = 0
   token_id_end = max(tokenizer.all_special_ids)
-  logger.info(f"interrogate tokens are: {token_id_start} to {token_id_end}")
+  print(f"interrogate tokens are: {token_id_start} to {token_id_end}")
 
   def get_all_embeddings(text_encoder):
     embs = []
@@ -86,24 +79,24 @@ def interrogate(args):
         embs.extend(encoder_hidden_states)
     return torch.stack(embs)
 
-  logger.info("get original text encoder embeddings.")
+  print("get original text encoder embeddings.")
   orig_embs = get_all_embeddings(text_encoder)
 
   network.apply_to(text_encoder, unet, True, len(network.unet_loras) > 0)
   info = network.load_state_dict(weights_sd, strict=False)
-  logger.info(f"Loading LoRA weights: {info}")
+  print(f"Loading LoRA weights: {info}")
 
   network.to(DEVICE, dtype=weights_dtype)
   network.eval()
 
   del unet
 
-  logger.info("You can ignore warning messages start with '_IncompatibleKeys' (LoRA model does not have alpha because trained by older script) / '_IncompatibleKeys'の警告は無視して構いません（以前のスクリプトで学習されたLoRAモデルのためalphaの定義がありません）")
-  logger.info("get text encoder embeddings with lora.")
+  print("You can ignore warning messages start with '_IncompatibleKeys' (LoRA model does not have alpha because trained by older script) / '_IncompatibleKeys'の警告は無視して構いません（以前のスクリプトで学習されたLoRAモデルのためalphaの定義がありません）")
+  print("get text encoder embeddings with lora.")
   lora_embs = get_all_embeddings(text_encoder)
 
   # 比べる：とりあえず単純に差分の絶対値で
-  logger.info("comparing...")
+  print("comparing...")
   diffs = {}
   for i, (orig_emb, lora_emb) in enumerate(zip(orig_embs, tqdm(lora_embs))):
     diff = torch.mean(torch.abs(orig_emb - lora_emb))

networks/merge_lora.py

@@ -7,10 +7,7 @@ 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
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 def load_state_dict(file_name, dtype):
     if os.path.splitext(file_name)[1] == ".safetensors":
@@ -64,10 +61,10 @@ def merge_to_sd_model(text_encoder, unet, models, ratios, merge_dtype):
                         name_to_module[lora_name] = child_module
 
     for model, ratio in zip(models, ratios):
-        logger.info(f"loading: {model}")
+        print(f"loading: {model}")
         lora_sd, _ = load_state_dict(model, merge_dtype)
 
-        logger.info(f"merging...")
+        print(f"merging...")
         for key in lora_sd.keys():
             if "lora_down" in key:
                 up_key = key.replace("lora_down", "lora_up")
@@ -76,10 +73,10 @@ def merge_to_sd_model(text_encoder, unet, models, ratios, merge_dtype):
                 # 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:
-                    logger.info(f"no module found for LoRA weight: {key}")
+                    print(f"no module found for LoRA weight: {key}")
                     continue
                 module = name_to_module[module_name]
-                # logger.info(f"apply {key} to {module}")
+                # print(f"apply {key} to {module}")
 
                 down_weight = lora_sd[key]
                 up_weight = lora_sd[up_key]
@@ -107,7 +104,7 @@ def merge_to_sd_model(text_encoder, unet, models, ratios, merge_dtype):
                 else:
                     # conv2d 3x3
                     conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
-                    # logger.info(conved.size(), weight.size(), module.stride, module.padding)
+                    # print(conved.size(), weight.size(), module.stride, module.padding)
                     weight = weight + ratio * conved * scale
 
                 module.weight = torch.nn.Parameter(weight)
@@ -121,7 +118,7 @@ def merge_lora_models(models, ratios, merge_dtype, concat=False, shuffle=False):
     v2 = None
     base_model = None
     for model, ratio in zip(models, ratios):
-        logger.info(f"loading: {model}")
+        print(f"loading: {model}")
         lora_sd, lora_metadata = load_state_dict(model, merge_dtype)
 
         if lora_metadata is not None:
@@ -154,10 +151,10 @@ def merge_lora_models(models, ratios, merge_dtype, concat=False, shuffle=False):
                 if lora_module_name not in base_alphas:
                     base_alphas[lora_module_name] = alpha
 
-        logger.info(f"dim: {list(set(dims.values()))}, alpha: {list(set(alphas.values()))}")
+        print(f"dim: {list(set(dims.values()))}, alpha: {list(set(alphas.values()))}")
 
         # merge
-        logger.info(f"merging...")
+        print(f"merging...")
         for key in lora_sd.keys():
             if "alpha" in key:
                 continue
@@ -199,8 +196,8 @@ def merge_lora_models(models, ratios, merge_dtype, concat=False, shuffle=False):
             merged_sd[key_down] = merged_sd[key_down][perm]
             merged_sd[key_up] = merged_sd[key_up][:,perm]
 
-    logger.info("merged model")
-    logger.info(f"dim: {list(set(base_dims.values()))}, alpha: {list(set(base_alphas.values()))}")
+    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()))
@@ -242,7 +239,7 @@ def merge(args):
         save_dtype = merge_dtype
 
     if args.sd_model is not None:
-        logger.info(f"loading SD model: {args.sd_model}")
+        print(f"loading SD model: {args.sd_model}")
 
         text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.sd_model)
 
@@ -267,18 +264,18 @@ def merge(args):
             )
             if args.v2:
                 # TODO read sai modelspec
-                logger.warning(
+                print(
                     "Cannot determine if model is for v-prediction, so save metadata as v-prediction / modelがv-prediction用か否か不明なため、仮にv-prediction用としてmetadataを保存します"
                 )
 
-        logger.info(f"saving SD model to: {args.save_to}")
+        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
         )
     else:
         state_dict, metadata, v2 = merge_lora_models(args.models, args.ratios, merge_dtype, args.concat, args.shuffle)
 
-        logger.info(f"calculating hashes and creating metadata...")
+        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
@@ -292,12 +289,12 @@ def merge(args):
             )
             if v2:
                 # TODO read sai modelspec
-                logger.warning(
+                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)
 
-        logger.info(f"saving model to: {args.save_to}")
+        print(f"saving model to: {args.save_to}")
         save_to_file(args.save_to, state_dict, state_dict, save_dtype, metadata)
 
 

networks/merge_lora_old.py

@@ -6,10 +6,7 @@ import torch
 from safetensors.torch import load_file, save_file
 import library.model_util as model_util
 import lora
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 def load_state_dict(file_name, dtype):
   if os.path.splitext(file_name)[1] == '.safetensors':
@@ -57,10 +54,10 @@ def merge_to_sd_model(text_encoder, unet, models, ratios, merge_dtype):
             name_to_module[lora_name] = child_module
 
   for model, ratio in zip(models, ratios):
-    logger.info(f"loading: {model}")
+    print(f"loading: {model}")
     lora_sd = load_state_dict(model, merge_dtype)
 
-    logger.info(f"merging...")
+    print(f"merging...")
     for key in lora_sd.keys():
       if "lora_down" in key:
         up_key = key.replace("lora_down", "lora_up")
@@ -69,10 +66,10 @@ def merge_to_sd_model(text_encoder, unet, models, ratios, merge_dtype):
         # 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:
-          logger.info(f"no module found for LoRA weight: {key}")
+          print(f"no module found for LoRA weight: {key}")
           continue
         module = name_to_module[module_name]
-        # logger.info(f"apply {key} to {module}")
+        # print(f"apply {key} to {module}")
 
         down_weight = lora_sd[key]
         up_weight = lora_sd[up_key]
@@ -99,10 +96,10 @@ def merge_lora_models(models, ratios, merge_dtype):
   alpha = None
   dim = None
   for model, ratio in zip(models, ratios):
-    logger.info(f"loading: {model}")
+    print(f"loading: {model}")
     lora_sd = load_state_dict(model, merge_dtype)
 
-    logger.info(f"merging...")
+    print(f"merging...")
     for key in lora_sd.keys():
       if 'alpha' in key:
         if key in merged_sd:
@@ -120,7 +117,7 @@ def merge_lora_models(models, ratios, merge_dtype):
             dim = lora_sd[key].size()[0]
           merged_sd[key] = lora_sd[key] * ratio
 
-  logger.info(f"dim (rank): {dim}, alpha: {alpha}")
+  print(f"dim (rank): {dim}, alpha: {alpha}")
   if alpha is None:
     alpha = dim
 
@@ -145,21 +142,19 @@ def merge(args):
     save_dtype = merge_dtype
 
   if args.sd_model is not None:
-    logger.info(f"loading SD model: {args.sd_model}")
+    print(f"loading SD model: {args.sd_model}")
 
     text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.sd_model)
 
     merge_to_sd_model(text_encoder, unet, args.models, args.ratios, merge_dtype)
 
-    logger.info("")
-    logger.info(f"saving SD model to: {args.save_to}")
+    print(f"\nsaving SD model to: {args.save_to}")
     model_util.save_stable_diffusion_checkpoint(args.v2, args.save_to, text_encoder, unet,
                                                 args.sd_model, 0, 0, save_dtype, vae)
   else:
     state_dict, _, _ = merge_lora_models(args.models, args.ratios, merge_dtype)
 
-    logger.info(f"")
-    logger.info(f"saving model to: {args.save_to}")
+    print(f"\nsaving model to: {args.save_to}")
     save_to_file(args.save_to, state_dict, state_dict, save_dtype)
 
 

networks/oft.py

@@ -8,10 +8,7 @@ from transformers import CLIPTextModel
 import numpy as np
 import torch
 import re
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 RE_UPDOWN = re.compile(r"(up|down)_blocks_(\d+)_(resnets|upsamplers|downsamplers|attentions)_(\d+)_")
 
@@ -240,7 +237,7 @@ class OFTNetwork(torch.nn.Module):
         self.dim = dim
         self.alpha = alpha
 
-        logger.info(
+        print(
             f"create OFT network. num blocks: {self.dim}, constraint: {self.alpha}, multiplier: {self.multiplier}, enable_conv: {enable_conv}"
         )
 
@@ -261,7 +258,7 @@ class OFTNetwork(torch.nn.Module):
                         if is_linear or is_conv2d_1x1 or (is_conv2d and enable_conv):
                             oft_name = prefix + "." + name + "." + child_name
                             oft_name = oft_name.replace(".", "_")
-                            # logger.info(oft_name)
+                            # print(oft_name)
 
                             oft = module_class(
                                 oft_name,
@@ -282,7 +279,7 @@ class OFTNetwork(torch.nn.Module):
             target_modules += OFTNetwork.UNET_TARGET_REPLACE_MODULE_CONV2D_3X3
 
         self.unet_ofts: List[OFTModule] = create_modules(unet, target_modules)
-        logger.info(f"create OFT for U-Net: {len(self.unet_ofts)} modules.")
+        print(f"create OFT for U-Net: {len(self.unet_ofts)} modules.")
 
         # assertion
         names = set()
@@ -319,7 +316,7 @@ class OFTNetwork(torch.nn.Module):
 
     # TODO refactor to common function with apply_to
     def merge_to(self, text_encoder, unet, weights_sd, dtype, device):
-        logger.info("enable OFT for U-Net")
+        print("enable OFT for U-Net")
 
         for oft in self.unet_ofts:
             sd_for_lora = {}
@@ -329,7 +326,7 @@ class OFTNetwork(torch.nn.Module):
             oft.load_state_dict(sd_for_lora, False)
             oft.merge_to()
 
-        logger.info(f"weights are merged")
+        print(f"weights are merged")
 
     # 二つのText Encoderに別々の学習率を設定できるようにするといいかも
     def prepare_optimizer_params(self, text_encoder_lr, unet_lr, default_lr):
@@ -341,11 +338,11 @@ class OFTNetwork(torch.nn.Module):
             for oft in ofts:
                 params.extend(oft.parameters())
 
-            # logger.info num of params
+            # print num of params
             num_params = 0
             for p in params:
                 num_params += p.numel()
-            logger.info(f"OFT params: {num_params}")
+            print(f"OFT params: {num_params}")
             return params
 
         param_data = {"params": enumerate_params(self.unet_ofts)}

networks/resize_lora.py

@@ -2,91 +2,80 @@
 # This code is based off the extract_lora_from_models.py file which is based on https://github.com/cloneofsimo/lora/blob/develop/lora_diffusion/cli_svd.py
 # Thanks to cloneofsimo
 
-import os
 import argparse
 import torch
 from safetensors.torch import load_file, save_file, safe_open
 from tqdm import tqdm
+from library import train_util, model_util
 import numpy as np
 
-from library import train_util
-from library import model_util
-from library.utils import setup_logging
-
-setup_logging()
-import logging
-
-logger = logging.getLogger(__name__)
-
 MIN_SV = 1e-6
 
 # Model save and load functions
 
-
 def load_state_dict(file_name, dtype):
-    if model_util.is_safetensors(file_name):
-        sd = load_file(file_name)
-        with safe_open(file_name, framework="pt") as f:
-            metadata = f.metadata()
-    else:
-        sd = torch.load(file_name, map_location="cpu")
-        metadata = None
+  if model_util.is_safetensors(file_name):
+    sd = load_file(file_name)
+    with safe_open(file_name, framework="pt") as f:
+      metadata = f.metadata()
+  else:
+    sd = torch.load(file_name, map_location='cpu')
+    metadata = None
 
-    for key in list(sd.keys()):
-        if type(sd[key]) == torch.Tensor:
-            sd[key] = sd[key].to(dtype)
+  for key in list(sd.keys()):
+    if type(sd[key]) == torch.Tensor:
+      sd[key] = sd[key].to(dtype)
 
-    return sd, metadata
+  return sd, metadata
 
 
-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, model, state_dict, dtype, metadata):
+  if dtype is not None:
+    for key in list(state_dict.keys()):
+      if type(state_dict[key]) == torch.Tensor:
+        state_dict[key] = state_dict[key].to(dtype)
 
-    if model_util.is_safetensors(file_name):
-        save_file(state_dict, file_name, metadata)
-    else:
-        torch.save(state_dict, file_name)
+  if model_util.is_safetensors(file_name):
+    save_file(model, file_name, metadata)
+  else:
+    torch.save(model, file_name)
 
 
 # Indexing functions
 
-
 def index_sv_cumulative(S, target):
-    original_sum = float(torch.sum(S))
-    cumulative_sums = torch.cumsum(S, dim=0) / original_sum
-    index = int(torch.searchsorted(cumulative_sums, target)) + 1
-    index = max(1, min(index, len(S) - 1))
+  original_sum = float(torch.sum(S))
+  cumulative_sums = torch.cumsum(S, dim=0)/original_sum
+  index = int(torch.searchsorted(cumulative_sums, target)) + 1
+  index = max(1, min(index, len(S)-1))
 
-    return index
+  return index
 
 
 def index_sv_fro(S, target):
-    S_squared = S.pow(2)
-    S_fro_sq = float(torch.sum(S_squared))
-    sum_S_squared = torch.cumsum(S_squared, dim=0) / S_fro_sq
-    index = int(torch.searchsorted(sum_S_squared, target**2)) + 1
-    index = max(1, min(index, len(S) - 1))
+  S_squared = S.pow(2)
+  s_fro_sq = float(torch.sum(S_squared))
+  sum_S_squared = torch.cumsum(S_squared, dim=0)/s_fro_sq
+  index = int(torch.searchsorted(sum_S_squared, target**2)) + 1
+  index = max(1, min(index, len(S)-1))
 
-    return index
+  return index
 
 
 def index_sv_ratio(S, target):
-    max_sv = S[0]
-    min_sv = max_sv / target
-    index = int(torch.sum(S > min_sv).item())
-    index = max(1, min(index, len(S) - 1))
+  max_sv = S[0]
+  min_sv = max_sv/target
+  index = int(torch.sum(S > min_sv).item())
+  index = max(1, min(index, len(S)-1))
 
-    return index
+  return index
 
 
 # Modified from Kohaku-blueleaf's extract/merge functions
 def extract_conv(weight, lora_rank, dynamic_method, dynamic_param, device, scale=1):
     out_size, in_size, kernel_size, _ = weight.size()
     U, S, Vh = torch.linalg.svd(weight.reshape(out_size, -1).to(device))
-
+    
     param_dict = rank_resize(S, lora_rank, dynamic_method, dynamic_param, scale)
     lora_rank = param_dict["new_rank"]
 
@@ -103,17 +92,17 @@ def extract_conv(weight, lora_rank, dynamic_method, dynamic_param, device, scale
 
 def extract_linear(weight, lora_rank, dynamic_method, dynamic_param, device, scale=1):
     out_size, in_size = weight.size()
-
+    
     U, S, Vh = torch.linalg.svd(weight.to(device))
-
+    
     param_dict = rank_resize(S, lora_rank, dynamic_method, dynamic_param, scale)
     lora_rank = param_dict["new_rank"]
-
+    
     U = U[:, :lora_rank]
     S = S[:lora_rank]
     U = U @ torch.diag(S)
     Vh = Vh[:lora_rank, :]
-
+    
     param_dict["lora_down"] = Vh.reshape(lora_rank, in_size).cpu()
     param_dict["lora_up"] = U.reshape(out_size, lora_rank).cpu()
     del U, S, Vh, weight
@@ -124,7 +113,7 @@ def merge_conv(lora_down, lora_up, device):
     in_rank, in_size, kernel_size, k_ = lora_down.shape
     out_size, out_rank, _, _ = lora_up.shape
     assert in_rank == out_rank and kernel_size == k_, f"rank {in_rank} {out_rank} or kernel {kernel_size} {k_} mismatch"
-
+    
     lora_down = lora_down.to(device)
     lora_up = lora_up.to(device)
 
@@ -138,274 +127,236 @@ def merge_linear(lora_down, lora_up, device):
     in_rank, in_size = lora_down.shape
     out_size, out_rank = lora_up.shape
     assert in_rank == out_rank, f"rank {in_rank} {out_rank} mismatch"
-
+    
     lora_down = lora_down.to(device)
     lora_up = lora_up.to(device)
-
+    
     weight = lora_up @ lora_down
     del lora_up, lora_down
     return weight
-
+  
 
 # Calculate new rank
 
-
 def rank_resize(S, rank, dynamic_method, dynamic_param, scale=1):
     param_dict = {}
 
-    if dynamic_method == "sv_ratio":
+    if dynamic_method=="sv_ratio":
         # Calculate new dim and alpha based off ratio
         new_rank = index_sv_ratio(S, dynamic_param) + 1
-        new_alpha = float(scale * new_rank)
+        new_alpha = float(scale*new_rank)
 
-    elif dynamic_method == "sv_cumulative":
+    elif dynamic_method=="sv_cumulative":
         # Calculate new dim and alpha based off cumulative sum
         new_rank = index_sv_cumulative(S, dynamic_param) + 1
-        new_alpha = float(scale * new_rank)
+        new_alpha = float(scale*new_rank)
 
-    elif dynamic_method == "sv_fro":
+    elif dynamic_method=="sv_fro":
         # Calculate new dim and alpha based off sqrt sum of squares
         new_rank = index_sv_fro(S, dynamic_param) + 1
-        new_alpha = float(scale * new_rank)
+        new_alpha = float(scale*new_rank)
     else:
         new_rank = rank
-        new_alpha = float(scale * new_rank)
+        new_alpha = float(scale*new_rank)
 
-    if S[0] <= MIN_SV:  # Zero matrix, set dim to 1
+    
+    if S[0] <= MIN_SV: # Zero matrix, set dim to 1
         new_rank = 1
-        new_alpha = float(scale * new_rank)
-    elif new_rank > rank:  # cap max rank at rank
+        new_alpha = float(scale*new_rank)
+    elif new_rank > rank: # cap max rank at rank
         new_rank = rank
-        new_alpha = float(scale * new_rank)
+        new_alpha = float(scale*new_rank)
+
 
     # Calculate resize info
     s_sum = torch.sum(torch.abs(S))
     s_rank = torch.sum(torch.abs(S[:new_rank]))
-
+    
     S_squared = S.pow(2)
     s_fro = torch.sqrt(torch.sum(S_squared))
     s_red_fro = torch.sqrt(torch.sum(S_squared[:new_rank]))
-    fro_percent = float(s_red_fro / s_fro)
+    fro_percent = float(s_red_fro/s_fro)
 
     param_dict["new_rank"] = new_rank
     param_dict["new_alpha"] = new_alpha
-    param_dict["sum_retained"] = (s_rank) / s_sum
+    param_dict["sum_retained"] = (s_rank)/s_sum
     param_dict["fro_retained"] = fro_percent
-    param_dict["max_ratio"] = S[0] / S[new_rank - 1]
+    param_dict["max_ratio"] = S[0]/S[new_rank - 1]
 
     return param_dict
 
 
-def resize_lora_model(lora_sd, new_rank, new_conv_rank, save_dtype, device, dynamic_method, dynamic_param, verbose):
-    network_alpha = None
-    network_dim = None
-    verbose_str = "\n"
-    fro_list = []
+def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dynamic_param, verbose):
+  network_alpha = None
+  network_dim = None
+  verbose_str = "\n"
+  fro_list = []
 
-    # Extract loaded lora dim and alpha
-    for key, value in lora_sd.items():
-        if network_alpha is None and "alpha" in key:
-            network_alpha = value
-        if network_dim is None and "lora_down" in key and len(value.size()) == 2:
-            network_dim = value.size()[0]
-        if network_alpha is not None and network_dim is not None:
-            break
-        if network_alpha is None:
-            network_alpha = network_dim
+  # Extract loaded lora dim and alpha
+  for key, value in lora_sd.items():
+    if network_alpha is None and 'alpha' in key:
+      network_alpha = value
+    if network_dim is None and 'lora_down' in key and len(value.size()) == 2:
+      network_dim = value.size()[0]
+    if network_alpha is not None and network_dim is not None:
+      break
+    if network_alpha is None:
+      network_alpha = network_dim
 
-    scale = network_alpha / network_dim
+  scale = network_alpha/network_dim
 
-    if dynamic_method:
-        logger.info(
-            f"Dynamically determining new alphas and dims based off {dynamic_method}: {dynamic_param}, max rank is {new_rank}"
-        )
+  if dynamic_method:
+    print(f"Dynamically determining new alphas and dims based off {dynamic_method}: {dynamic_param}, max rank is {new_rank}")
 
-    lora_down_weight = None
-    lora_up_weight = None
+  lora_down_weight = None
+  lora_up_weight = None
 
-    o_lora_sd = lora_sd.copy()
-    block_down_name = None
-    block_up_name = None
+  o_lora_sd = lora_sd.copy()
+  block_down_name = None
+  block_up_name = None
 
-    with torch.no_grad():
-        for key, value in tqdm(lora_sd.items()):
-            weight_name = None
-            if "lora_down" in key:
-                block_down_name = key.rsplit(".lora_down", 1)[0]
-                weight_name = key.rsplit(".", 1)[-1]
-                lora_down_weight = value
-            else:
-                continue
+  with torch.no_grad():
+    for key, value in tqdm(lora_sd.items()):
+      weight_name = None
+      if 'lora_down' in key:
+        block_down_name = key.rsplit('.lora_down', 1)[0]
+        weight_name = key.rsplit(".", 1)[-1]
+        lora_down_weight = value
+      else:
+        continue
 
-            # find corresponding lora_up and alpha
-            block_up_name = block_down_name
-            lora_up_weight = lora_sd.get(block_up_name + ".lora_up." + weight_name, None)
-            lora_alpha = lora_sd.get(block_down_name + ".alpha", None)
+      # find corresponding lora_up and alpha
+      block_up_name = block_down_name
+      lora_up_weight = lora_sd.get(block_up_name + '.lora_up.' + weight_name, None)
+      lora_alpha = lora_sd.get(block_down_name + '.alpha', None)
 
-            weights_loaded = lora_down_weight is not None and lora_up_weight is not None
+      weights_loaded = (lora_down_weight is not None and lora_up_weight is not None)
 
-            if weights_loaded:
+      if weights_loaded:
 
-                conv2d = len(lora_down_weight.size()) == 4
-                if lora_alpha is None:
-                    scale = 1.0
-                else:
-                    scale = lora_alpha / lora_down_weight.size()[0]
+        conv2d = (len(lora_down_weight.size()) == 4)
+        if lora_alpha is None:
+          scale = 1.0
+        else:
+          scale = lora_alpha/lora_down_weight.size()[0]
 
-                if conv2d:
-                    full_weight_matrix = merge_conv(lora_down_weight, lora_up_weight, device)
-                    param_dict = extract_conv(full_weight_matrix, new_conv_rank, dynamic_method, dynamic_param, device, scale)
-                else:
-                    full_weight_matrix = merge_linear(lora_down_weight, lora_up_weight, device)
-                    param_dict = extract_linear(full_weight_matrix, new_rank, dynamic_method, dynamic_param, device, scale)
+        if conv2d:
+          full_weight_matrix = merge_conv(lora_down_weight, lora_up_weight, device)
+          param_dict = extract_conv(full_weight_matrix, new_rank, dynamic_method, dynamic_param, device, scale)
+        else:
+          full_weight_matrix = merge_linear(lora_down_weight, lora_up_weight, device)
+          param_dict = extract_linear(full_weight_matrix, new_rank, dynamic_method, dynamic_param, device, scale)
 
-                if verbose:
-                    max_ratio = param_dict["max_ratio"]
-                    sum_retained = param_dict["sum_retained"]
-                    fro_retained = param_dict["fro_retained"]
-                    if not np.isnan(fro_retained):
-                        fro_list.append(float(fro_retained))
+        if verbose:
+          max_ratio = param_dict['max_ratio']
+          sum_retained = param_dict['sum_retained']
+          fro_retained = param_dict['fro_retained']
+          if not np.isnan(fro_retained):
+            fro_list.append(float(fro_retained))
 
-                    verbose_str += f"{block_down_name:75} | "
-                    verbose_str += (
-                        f"sum(S) retained: {sum_retained:.1%}, fro retained: {fro_retained:.1%}, max(S) ratio: {max_ratio:0.1f}"
-                    )
+          verbose_str+=f"{block_down_name:75} | "
+          verbose_str+=f"sum(S) retained: {sum_retained:.1%}, fro retained: {fro_retained:.1%}, max(S) ratio: {max_ratio:0.1f}"
 
-                if verbose and dynamic_method:
-                    verbose_str += f", dynamic | dim: {param_dict['new_rank']}, alpha: {param_dict['new_alpha']}\n"
-                else:
-                    verbose_str += "\n"
+        if verbose and dynamic_method:
+          verbose_str+=f", dynamic | dim: {param_dict['new_rank']}, alpha: {param_dict['new_alpha']}\n"
+        else:
+          verbose_str+=f"\n"
 
-                new_alpha = param_dict["new_alpha"]
-                o_lora_sd[block_down_name + "." + "lora_down.weight"] = param_dict["lora_down"].to(save_dtype).contiguous()
-                o_lora_sd[block_up_name + "." + "lora_up.weight"] = param_dict["lora_up"].to(save_dtype).contiguous()
-                o_lora_sd[block_up_name + "." "alpha"] = torch.tensor(param_dict["new_alpha"]).to(save_dtype)
+        new_alpha = param_dict['new_alpha']
+        o_lora_sd[block_down_name + "." + "lora_down.weight"] = param_dict["lora_down"].to(save_dtype).contiguous()
+        o_lora_sd[block_up_name + "." + "lora_up.weight"] = param_dict["lora_up"].to(save_dtype).contiguous()
+        o_lora_sd[block_up_name + "." "alpha"] = torch.tensor(param_dict['new_alpha']).to(save_dtype)
 
-                block_down_name = None
-                block_up_name = None
-                lora_down_weight = None
-                lora_up_weight = None
-                weights_loaded = False
-                del param_dict
+        block_down_name = None
+        block_up_name = None
+        lora_down_weight = None
+        lora_up_weight = None
+        weights_loaded = False
+        del param_dict
 
-    if verbose:
-        print(verbose_str)
-        print(f"Average Frobenius norm retention: {np.mean(fro_list):.2%} | std: {np.std(fro_list):0.3f}")
-    logger.info("resizing complete")
-    return o_lora_sd, network_dim, new_alpha
+  if verbose:
+    print(verbose_str)
+
+    print(f"Average Frobenius norm retention: {np.mean(fro_list):.2%} | std: {np.std(fro_list):0.3f}")
+  print("resizing complete")
+  return o_lora_sd, network_dim, new_alpha
 
 
 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.")
+  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.")
 
-    args.new_conv_rank = args.new_conv_rank if args.new_conv_rank is not None else args.new_rank
+    
+  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
+  if args.dynamic_method and not args.dynamic_param:
+    raise Exception("If using dynamic_method, then dynamic_param is required")
 
-    if args.dynamic_method and not args.dynamic_param:
-        raise Exception("If using dynamic_method, then dynamic_param is required")
+  merge_dtype = str_to_dtype('float')  # matmul method above only seems to work in float32
+  save_dtype = str_to_dtype(args.save_precision)
+  if save_dtype is None:
+    save_dtype = merge_dtype
 
-    merge_dtype = str_to_dtype("float")  # matmul method above only seems to work in float32
-    save_dtype = str_to_dtype(args.save_precision)
-    if save_dtype is None:
-        save_dtype = merge_dtype
+  print("loading Model...")
+  lora_sd, metadata = load_state_dict(args.model, merge_dtype)
 
-    logger.info("loading Model...")
-    lora_sd, metadata = load_state_dict(args.model, merge_dtype)
+  print("Resizing Lora...")
+  state_dict, old_dim, new_alpha = resize_lora_model(lora_sd, args.new_rank, save_dtype, args.device, args.dynamic_method, args.dynamic_param, args.verbose)
 
-    logger.info("Resizing Lora...")
-    state_dict, old_dim, new_alpha = resize_lora_model(
-        lora_sd, args.new_rank, args.new_conv_rank, save_dtype, args.device, args.dynamic_method, args.dynamic_param, args.verbose
-    )
+  # update metadata
+  if metadata is None:
+    metadata = {}
 
-    # update metadata
-    if metadata is None:
-        metadata = {}
+  comment = metadata.get("ss_training_comment", "")
 
-    comment = metadata.get("ss_training_comment", "")
+  if not args.dynamic_method:
+    metadata["ss_training_comment"] = f"dimension is resized from {old_dim} to {args.new_rank}; {comment}"
+    metadata["ss_network_dim"] = str(args.new_rank)
+    metadata["ss_network_alpha"] = str(new_alpha)
+  else:
+    metadata["ss_training_comment"] = f"Dynamic resize with {args.dynamic_method}: {args.dynamic_param} from {old_dim}; {comment}"
+    metadata["ss_network_dim"] = 'Dynamic'
+    metadata["ss_network_alpha"] = 'Dynamic'
 
-    if not args.dynamic_method:
-        conv_desc = "" if args.new_rank == args.new_conv_rank else f" (conv: {args.new_conv_rank})"
-        metadata["ss_training_comment"] = f"dimension is resized from {old_dim} to {args.new_rank}{conv_desc}; {comment}"
-        metadata["ss_network_dim"] = str(args.new_rank)
-        metadata["ss_network_alpha"] = str(new_alpha)
-    else:
-        metadata["ss_training_comment"] = (
-            f"Dynamic resize with {args.dynamic_method}: {args.dynamic_param} from {old_dim}; {comment}"
-        )
-        metadata["ss_network_dim"] = "Dynamic"
-        metadata["ss_network_alpha"] = "Dynamic"
+  model_hash, legacy_hash = train_util.precalculate_safetensors_hashes(state_dict, metadata)
+  metadata["sshs_model_hash"] = model_hash
+  metadata["sshs_legacy_hash"] = legacy_hash
 
-    model_hash, legacy_hash = train_util.precalculate_safetensors_hashes(state_dict, metadata)
-    metadata["sshs_model_hash"] = model_hash
-    metadata["sshs_legacy_hash"] = legacy_hash
-
-    logger.info(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, state_dict, save_dtype, metadata)
 
 
 def setup_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser()
+  parser = argparse.ArgumentParser()
 
-    parser.add_argument(
-        "--save_precision",
-        type=str,
-        default=None,
-        choices=[None, "float", "fp16", "bf16"],
-        help="precision in saving, float if omitted / 保存時の精度、未指定時はfloat",
-    )
-    parser.add_argument("--new_rank", type=int, default=4, help="Specify rank of output LoRA / 出力するLoRAのrank (dim)")
-    parser.add_argument(
-        "--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(
-        "--save_to",
-        type=str,
-        default=None,
-        help="destination file name: ckpt or safetensors file / 保存先のファイル名、ckptまたはsafetensors",
-    )
-    parser.add_argument(
-        "--model",
-        type=str,
-        default=None,
-        help="LoRA model to resize at to new rank: ckpt or safetensors file / 読み込むLoRAモデル、ckptまたはsafetensors",
-    )
-    parser.add_argument(
-        "--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う"
-    )
-    parser.add_argument(
-        "--verbose", action="store_true", help="Display verbose resizing information / rank変更時の詳細情報を出力する"
-    )
-    parser.add_argument(
-        "--dynamic_method",
-        type=str,
-        default=None,
-        choices=[None, "sv_ratio", "sv_fro", "sv_cumulative"],
-        help="Specify dynamic resizing method, --new_rank is used as a hard limit for max rank",
-    )
-    parser.add_argument("--dynamic_param", type=float, default=None, help="Specify target for dynamic reduction")
-
-    return parser
+  parser.add_argument("--save_precision", type=str, default=None,
+                      choices=[None, "float", "fp16", "bf16"], help="precision in saving, float if omitted / 保存時の精度、未指定時はfloat")
+  parser.add_argument("--new_rank", type=int, default=4,
+                      help="Specify rank of output LoRA / 出力するLoRAのrank (dim)")
+  parser.add_argument("--save_to", type=str, default=None,
+                      help="destination file name: ckpt or safetensors file / 保存先のファイル名、ckptまたはsafetensors")
+  parser.add_argument("--model", type=str, default=None,
+                      help="LoRA model to resize at to new rank: ckpt or safetensors file / 読み込むLoRAモデル、ckptまたはsafetensors")
+  parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")
+  parser.add_argument("--verbose", action="store_true", 
+                      help="Display verbose resizing information / rank変更時の詳細情報を出力する")
+  parser.add_argument("--dynamic_method", type=str, default=None, choices=[None, "sv_ratio", "sv_fro", "sv_cumulative"],
+                      help="Specify dynamic resizing method, --new_rank is used as a hard limit for max rank")
+  parser.add_argument("--dynamic_param", type=float, default=None,
+                      help="Specify target for dynamic reduction")
+       
+  return parser
 
 
-if __name__ == "__main__":
-    parser = setup_parser()
+if __name__ == '__main__':
+  parser = setup_parser()
 
-    args = parser.parse_args()
-    resize(args)
+  args = parser.parse_args()
+  resize(args)

networks/sdxl_merge_lora.py

@@ -8,10 +8,7 @@ from tqdm import tqdm
 from library import sai_model_spec, sdxl_model_util, train_util
 import library.model_util as model_util
 import lora
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 def load_state_dict(file_name, dtype):
     if os.path.splitext(file_name)[1] == ".safetensors":
@@ -69,10 +66,10 @@ def merge_to_sd_model(text_encoder1, text_encoder2, unet, models, ratios, merge_
                         name_to_module[lora_name] = child_module
 
     for model, ratio in zip(models, ratios):
-        logger.info(f"loading: {model}")
+        print(f"loading: {model}")
         lora_sd, _ = load_state_dict(model, merge_dtype)
 
-        logger.info(f"merging...")
+        print(f"merging...")
         for key in tqdm(lora_sd.keys()):
             if "lora_down" in key:
                 up_key = key.replace("lora_down", "lora_up")
@@ -81,10 +78,10 @@ def merge_to_sd_model(text_encoder1, text_encoder2, unet, models, ratios, merge_
                 # 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:
-                    logger.info(f"no module found for LoRA weight: {key}")
+                    print(f"no module found for LoRA weight: {key}")
                     continue
                 module = name_to_module[module_name]
-                # logger.info(f"apply {key} to {module}")
+                # print(f"apply {key} to {module}")
 
                 down_weight = lora_sd[key]
                 up_weight = lora_sd[up_key]
@@ -95,7 +92,7 @@ def merge_to_sd_model(text_encoder1, text_encoder2, unet, models, ratios, merge_
 
                 # W <- W + U * D
                 weight = module.weight
-                # logger.info(module_name, down_weight.size(), up_weight.size())
+                # print(module_name, down_weight.size(), up_weight.size())
                 if len(weight.size()) == 2:
                     # linear
                     weight = weight + ratio * (up_weight @ down_weight) * scale
@@ -110,7 +107,7 @@ def merge_to_sd_model(text_encoder1, text_encoder2, unet, models, ratios, merge_
                 else:
                     # conv2d 3x3
                     conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
-                    # logger.info(conved.size(), weight.size(), module.stride, module.padding)
+                    # print(conved.size(), weight.size(), module.stride, module.padding)
                     weight = weight + ratio * conved * scale
 
                 module.weight = torch.nn.Parameter(weight)
@@ -124,7 +121,7 @@ def merge_lora_models(models, ratios, merge_dtype, concat=False, shuffle=False):
     v2 = None
     base_model = None
     for model, ratio in zip(models, ratios):
-        logger.info(f"loading: {model}")
+        print(f"loading: {model}")
         lora_sd, lora_metadata = load_state_dict(model, merge_dtype)
 
         if lora_metadata is not None:
@@ -157,10 +154,10 @@ def merge_lora_models(models, ratios, merge_dtype, concat=False, shuffle=False):
                 if lora_module_name not in base_alphas:
                     base_alphas[lora_module_name] = alpha
 
-        logger.info(f"dim: {list(set(dims.values()))}, alpha: {list(set(alphas.values()))}")
+        print(f"dim: {list(set(dims.values()))}, alpha: {list(set(alphas.values()))}")
 
         # merge
-        logger.info(f"merging...")
+        print(f"merging...")
         for key in tqdm(lora_sd.keys()):
             if "alpha" in key:
                 continue
@@ -203,8 +200,8 @@ def merge_lora_models(models, ratios, merge_dtype, concat=False, shuffle=False):
             merged_sd[key_down] = merged_sd[key_down][perm]
             merged_sd[key_up] = merged_sd[key_up][:,perm]
 
-    logger.info("merged model")
-    logger.info(f"dim: {list(set(base_dims.values()))}, alpha: {list(set(base_alphas.values()))}")
+    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()))
@@ -246,7 +243,7 @@ def merge(args):
         save_dtype = merge_dtype
 
     if args.sd_model is not None:
-        logger.info(f"loading SD model: {args.sd_model}")
+        print(f"loading SD model: {args.sd_model}")
 
         (
             text_model1,
@@ -268,14 +265,14 @@ def merge(args):
                 None, False, False, True, False, False, time.time(), title=title, merged_from=merged_from
             )
 
-        logger.info(f"saving SD model to: {args.save_to}")
+        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, args.concat, args.shuffle)
 
-        logger.info(f"calculating hashes and creating metadata...")
+        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
@@ -289,7 +286,7 @@ def merge(args):
             )
             metadata.update(sai_metadata)
 
-        logger.info(f"saving model to: {args.save_to}")
+        print(f"saving model to: {args.save_to}")
         save_to_file(args.save_to, state_dict, state_dict, save_dtype, metadata)
 
 

networks/svd_merge_lora.py

@@ -1,3 +1,4 @@
+import math
 import argparse
 import os
 import time
@@ -7,10 +8,7 @@ from tqdm import tqdm
 from library import sai_model_spec, train_util
 import library.model_util as model_util
 import lora
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 CLAMP_QUANTILE = 0.99
 
@@ -43,12 +41,12 @@ def save_to_file(file_name, state_dict, dtype, metadata):
 
 
 def merge_lora_models(models, ratios, new_rank, new_conv_rank, device, merge_dtype):
-    logger.info(f"new rank: {new_rank}, new conv rank: {new_conv_rank}")
+    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):
-        logger.info(f"loading: {model}")
+        print(f"loading: {model}")
         lora_sd, lora_metadata = load_state_dict(model, merge_dtype)
 
         if lora_metadata is not None:
@@ -58,7 +56,7 @@ def merge_lora_models(models, ratios, new_rank, new_conv_rank, device, merge_dty
                 base_model = lora_metadata.get(train_util.SS_METADATA_KEY_BASE_MODEL_VERSION, None)
 
         # merge
-        logger.info(f"merging...")
+        print(f"merging...")
         for key in tqdm(list(lora_sd.keys())):
             if "lora_down" not in key:
                 continue
@@ -75,7 +73,7 @@ def merge_lora_models(models, ratios, new_rank, new_conv_rank, device, merge_dty
             out_dim = up_weight.size()[0]
             conv2d = len(down_weight.size()) == 4
             kernel_size = None if not conv2d else down_weight.size()[2:4]
-            # logger.info(lora_module_name, network_dim, alpha, in_dim, out_dim, kernel_size)
+            # 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:
@@ -112,7 +110,7 @@ def merge_lora_models(models, ratios, new_rank, new_conv_rank, device, merge_dty
             merged_sd[lora_module_name] = weight
 
     # extract from merged weights
-    logger.info("extract new lora...")
+    print("extract new lora...")
     merged_lora_sd = {}
     with torch.no_grad():
         for lora_module_name, mat in tqdm(list(merged_sd.items())):
@@ -190,7 +188,7 @@ def merge(args):
         args.models, args.ratios, args.new_rank, new_conv_rank, args.device, merge_dtype
     )
 
-    logger.info(f"calculating hashes and creating metadata...")
+    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
@@ -205,12 +203,12 @@ def merge(args):
         )
         if v2:
             # TODO read sai modelspec
-            logger.warning(
+            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)
 
-    logger.info(f"saving model to: {args.save_to}")
+    print(f"saving model to: {args.save_to}")
     save_to_file(args.save_to, state_dict, save_dtype, metadata)
 
 

requirements.txt

@@ -1,20 +1,20 @@
-accelerate==0.25.0
-transformers==4.36.2
-diffusers[torch]==0.25.0
+accelerate==0.23.0
+transformers==4.30.2
+diffusers[torch]==0.21.2
 ftfy==6.1.1
 # albumentations==1.3.0
 opencv-python==4.7.0.68
-einops==0.7.0
+einops==0.6.0
 pytorch-lightning==1.9.0
 # bitsandbytes==0.39.1
 tensorboard==2.10.1
-safetensors==0.4.2
+safetensors==0.3.1
 # gradio==3.16.2
 altair==4.2.2
 easygui==0.98.3
 toml==0.10.2
 voluptuous==0.13.1
-huggingface-hub==0.20.1
+huggingface-hub==0.15.1
 # for BLIP captioning
 # requests==2.28.2
 # timm==0.6.12
@@ -29,7 +29,5 @@ huggingface-hub==0.20.1
 # protobuf==3.20.3
 # open clip for SDXL
 open-clip-torch==2.20.0
-# For logging
-rich==13.7.0
 # for kohya_ss library
 -e .

sdxl_minimal_inference.py

@@ -8,11 +8,14 @@ import os
 import random
 from einops import repeat
 import numpy as np
-
 import torch
-from library.device_utils import init_ipex, get_preferred_device
-init_ipex()
-
+try:
+    import intel_extension_for_pytorch as ipex
+    if torch.xpu.is_available():
+        from library.ipex import ipex_init
+        ipex_init()
+except Exception:
+    pass
 from tqdm import tqdm
 from transformers import CLIPTokenizer
 from diffusers import EulerDiscreteScheduler
@@ -22,10 +25,6 @@ from safetensors.torch import load_file
 
 from library import model_util, sdxl_model_util
 import networks.lora as lora
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
 
 # scheduler: このあたりの設定はSD1/2と同じでいいらしい
 # scheduler: The settings around here seem to be the same as SD1/2
@@ -88,7 +87,7 @@ if __name__ == "__main__":
     guidance_scale = 7
     seed = None  # 1
 
-    DEVICE = get_preferred_device()
+    DEVICE = "cuda"
     DTYPE = torch.float16  # bfloat16 may work
 
     parser = argparse.ArgumentParser()
@@ -143,7 +142,7 @@ if __name__ == "__main__":
 
     vae_dtype = DTYPE
     if DTYPE == torch.float16:
-        logger.info("use float32 for vae")
+        print("use float32 for vae")
         vae_dtype = torch.float32
     vae.to(DEVICE, dtype=vae_dtype)
     vae.eval()
@@ -190,7 +189,7 @@ if __name__ == "__main__":
             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)
-            # logger.info("emb1", emb1.shape)
+            # 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
 
@@ -220,7 +219,7 @@ if __name__ == "__main__":
 
                 enc_out = text_model2(tokens, output_hidden_states=True, return_dict=True)
                 text_embedding2_penu = enc_out["hidden_states"][-2]
-                # logger.info("hidden_states2", text_embedding2_penu.shape)
+                # print("hidden_states2", text_embedding2_penu.shape)
                 text_embedding2_pool = enc_out["text_embeds"]   # do not support Textual Inversion
 
             # 連結して終了 concat and finish
@@ -229,7 +228,7 @@ if __name__ == "__main__":
 
         # cond
         c_ctx, c_ctx_pool = call_text_encoder(prompt, prompt2)
-        # logger.info(c_ctx.shape, c_ctx_p.shape, c_vector.shape)
+        # print(c_ctx.shape, c_ctx_p.shape, c_vector.shape)
         c_vector = torch.cat([c_ctx_pool, c_vector], dim=1)
 
         # uncond
@@ -326,4 +325,4 @@ if __name__ == "__main__":
                 seed = int(seed)
             generate_image(prompt, prompt2, negative_prompt, seed)
 
-    logger.info("Done!")
+    print("Done!")

sdxl_train.py

@@ -1,6 +1,7 @@
 # training with captions
 
 import argparse
+import gc
 import math
 import os
 from multiprocessing import Value
@@ -8,24 +9,19 @@ from typing import List
 import toml
 
 from tqdm import tqdm
-
 import torch
-from library.device_utils import init_ipex, clean_memory_on_device
-init_ipex()
-
+try:
+    import intel_extension_for_pytorch as ipex
+    if torch.xpu.is_available():
+        from library.ipex import ipex_init
+        ipex_init()
+except Exception:
+    pass
 from accelerate.utils import set_seed
 from diffusers import DDPMScheduler
 from library import sdxl_model_util
 
 import library.train_util as train_util
-
-from library.utils import setup_logging, add_logging_arguments
-
-setup_logging()
-import logging
-
-logger = logging.getLogger(__name__)
-
 import library.config_util as config_util
 import library.sdxl_train_util as sdxl_train_util
 from library.config_util import (
@@ -38,7 +34,6 @@ from library.custom_train_functions import (
     prepare_scheduler_for_custom_training,
     scale_v_prediction_loss_like_noise_prediction,
     add_v_prediction_like_loss,
-    apply_debiased_estimation,
 )
 from library.sdxl_original_unet import SdxlUNet2DConditionModel
 
@@ -75,33 +70,41 @@ def get_block_params_to_optimize(unet: SdxlUNet2DConditionModel, block_lrs: List
 
 
 def append_block_lr_to_logs(block_lrs, logs, lr_scheduler, optimizer_type):
-    names = []
+    lrs = lr_scheduler.get_last_lr()
+
+    lr_index = 0
     block_index = 0
-    while block_index < UNET_NUM_BLOCKS_FOR_BLOCK_LR + 2:
+    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
-            names.append(f"block{block_index}")
         elif block_index == UNET_NUM_BLOCKS_FOR_BLOCK_LR:
-            names.append("text_encoder1")
+            name = "text_encoder1"
         elif block_index == UNET_NUM_BLOCKS_FOR_BLOCK_LR + 1:
-            names.append("text_encoder2")
+            name = "text_encoder2"
+        else:
+            raise ValueError(f"unexpected block_index: {block_index}")
 
         block_index += 1
 
-    train_util.append_lr_to_logs_with_names(logs, lr_scheduler, optimizer_type, names)
+        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)
-    setup_logging(args, reset=True)
 
-    assert (
-        not args.weighted_captions
-    ), "weighted_captions is not supported currently / weighted_captionsは現在サポートされていません"
+    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はサポートされていません"
@@ -126,18 +129,18 @@ def train(args):
     if args.dataset_class is None:
         blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, False, True))
         if args.dataset_config is not None:
-            logger.info(f"Load dataset config from {args.dataset_config}")
+            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):
-                logger.warning(
+                print(
                     "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
                         ", ".join(ignored)
                     )
                 )
         else:
             if use_dreambooth_method:
-                logger.info("Using DreamBooth method.")
+                print("Using DreamBooth method.")
                 user_config = {
                     "datasets": [
                         {
@@ -148,7 +151,7 @@ def train(args):
                     ]
                 }
             else:
-                logger.info("Training with captions.")
+                print("Training with captions.")
                 user_config = {
                     "datasets": [
                         {
@@ -178,7 +181,7 @@ def train(args):
         train_util.debug_dataset(train_dataset_group, True)
         return
     if len(train_dataset_group) == 0:
-        logger.error(
+        print(
             "No data found. Please verify the metadata file and train_data_dir option. / 画像がありません。メタデータおよびtrain_data_dirオプションを確認してください。"
         )
         return
@@ -194,7 +197,7 @@ def train(args):
         ), "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を準備する
-    logger.info("prepare accelerator")
+    print("prepare accelerator")
     accelerator = train_util.prepare_accelerator(args)
 
     # mixed precisionに対応した型を用意しておき適宜castする
@@ -261,16 +264,17 @@ def train(args):
         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")
-        clean_memory_on_device(accelerator.device)
+        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()
-    train_unet = args.learning_rate > 0
-    train_text_encoder1 = False
-    train_text_encoder2 = False
+    training_models.append(unet)
 
     if args.train_text_encoder:
         # TODO each option for two text encoders?
@@ -278,23 +282,10 @@ def train(args):
         if args.gradient_checkpointing:
             text_encoder1.gradient_checkpointing_enable()
             text_encoder2.gradient_checkpointing_enable()
-        lr_te1 = args.learning_rate_te1 if args.learning_rate_te1 is not None else args.learning_rate  # 0 means not train
-        lr_te2 = args.learning_rate_te2 if args.learning_rate_te2 is not None else args.learning_rate  # 0 means not train
-        train_text_encoder1 = lr_te1 > 0
-        train_text_encoder2 = lr_te2 > 0
-
-        # caching one text encoder output is not supported
-        if not train_text_encoder1:
-            text_encoder1.to(weight_dtype)
-        if not train_text_encoder2:
-            text_encoder2.to(weight_dtype)
-        text_encoder1.requires_grad_(train_text_encoder1)
-        text_encoder2.requires_grad_(train_text_encoder2)
-        text_encoder1.train(train_text_encoder1)
-        text_encoder2.train(train_text_encoder2)
+        training_models.append(text_encoder1)
+        training_models.append(text_encoder2)
+        # set require_grad=True later
     else:
-        text_encoder1.to(weight_dtype)
-        text_encoder2.to(weight_dtype)
         text_encoder1.requires_grad_(False)
         text_encoder2.requires_grad_(False)
         text_encoder1.eval()
@@ -303,7 +294,7 @@ def train(args):
         # TextEncoderの出力をキャッシュする
         if args.cache_text_encoder_outputs:
             # Text Encodes are eval and no grad
-            with torch.no_grad(), accelerator.autocast():
+            with torch.no_grad():
                 train_dataset_group.cache_text_encoder_outputs(
                     (tokenizer1, tokenizer2),
                     (text_encoder1, text_encoder2),
@@ -319,33 +310,30 @@ def train(args):
         vae.eval()
         vae.to(accelerator.device, dtype=vae_dtype)
 
-    unet.requires_grad_(train_unet)
-    if not train_unet:
-        unet.to(accelerator.device, dtype=weight_dtype)  # because of unet is not prepared
+    for m in training_models:
+        m.requires_grad_(True)
 
-    training_models = []
-    params_to_optimize = []
-    if train_unet:
-        training_models.append(unet)
-        if block_lrs is None:
-            params_to_optimize.append({"params": list(unet.parameters()), "lr": args.learning_rate})
-        else:
-            params_to_optimize.extend(get_block_params_to_optimize(unet, block_lrs))
+    if block_lrs is None:
+        params = []
+        for m in training_models:
+            params.extend(m.parameters())
+        params_to_optimize = params
 
-    if train_text_encoder1:
-        training_models.append(text_encoder1)
-        params_to_optimize.append({"params": list(text_encoder1.parameters()), "lr": args.learning_rate_te1 or args.learning_rate})
-    if train_text_encoder2:
-        training_models.append(text_encoder2)
-        params_to_optimize.append({"params": list(text_encoder2.parameters()), "lr": args.learning_rate_te2 or args.learning_rate})
-
-    # calculate number of trainable parameters
-    n_params = 0
-    for params in params_to_optimize:
-        for p in params["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"train unet: {train_unet}, text_encoder1: {train_text_encoder1}, text_encoder2: {train_text_encoder2}")
     accelerator.print(f"number of models: {len(training_models)}")
     accelerator.print(f"number of trainable parameters: {n_params}")
 
@@ -354,8 +342,8 @@ def train(args):
     _, _, optimizer = train_util.get_optimizer(args, trainable_params=params_to_optimize)
 
     # dataloaderを準備する
-    # DataLoaderのプロセス数：0 は persistent_workers が使えないので注意
-    n_workers = min(args.max_data_loader_n_workers, os.cpu_count())  # cpu_count or max_data_loader_n_workers
+    # 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,
@@ -370,9 +358,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}"
-        )
+        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)
@@ -399,24 +385,26 @@ def train(args):
         text_encoder2.to(weight_dtype)
 
     # acceleratorがなんかよろしくやってくれるらしい
-    if train_unet:
-        unet = accelerator.prepare(unet)
-    if train_text_encoder1:
-        # freeze last layer and final_layer_norm in te1 since we use the output of the penultimate layer
-        text_encoder1.text_model.encoder.layers[-1].requires_grad_(False)
-        text_encoder1.text_model.final_layer_norm.requires_grad_(False)
-        text_encoder1 = accelerator.prepare(text_encoder1)
-    if train_text_encoder2:
-        text_encoder2 = accelerator.prepare(text_encoder2)
+    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
+        )
 
-    optimizer, train_dataloader, lr_scheduler = accelerator.prepare(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)
-        clean_memory_on_device(accelerator.device)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
     else:
         # make sure Text Encoders are on GPU
         text_encoder1.to(accelerator.device)
@@ -441,9 +429,7 @@ def train(args):
     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"  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}"
     # )
@@ -462,18 +448,10 @@ def train(args):
 
     if accelerator.is_main_process:
         init_kwargs = {}
-        if args.wandb_run_name:
-            init_kwargs["wandb"] = {"name": args.wandb_run_name}
         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 --sample_at_first
-    sdxl_train_util.sample_images(
-        accelerator, args, 0, global_step, accelerator.device, vae, [tokenizer1, tokenizer2], [text_encoder1, text_encoder2], unet
-    )
-
-    loss_recorder = train_util.LossRecorder()
     for epoch in range(num_train_epochs):
         accelerator.print(f"\nepoch {epoch+1}/{num_train_epochs}")
         current_epoch.value = epoch + 1
@@ -481,9 +459,10 @@ def train(args):
         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):
+            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:
@@ -494,7 +473,7 @@ def train(args):
                         # NaNが含まれていれば警告を表示し0に置き換える
                         if torch.any(torch.isnan(latents)):
                             accelerator.print("NaN found in latents, replacing with zeros")
-                            latents = torch.nan_to_num(latents, 0, out=latents)
+                            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:
@@ -515,7 +494,6 @@ def train(args):
                         # else:
                         input_ids1 = input_ids1.to(accelerator.device)
                         input_ids2 = input_ids2.to(accelerator.device)
-                        # unwrap_model is fine for models not wrapped by accelerator
                         encoder_hidden_states1, encoder_hidden_states2, pool2 = train_util.get_hidden_states_sdxl(
                             args.max_token_length,
                             input_ids1,
@@ -525,7 +503,6 @@ def train(args):
                             text_encoder1,
                             text_encoder2,
                             None if not args.full_fp16 else weight_dtype,
-                            accelerator=accelerator,
                         )
                 else:
                     encoder_hidden_states1 = batch["text_encoder_outputs1_list"].to(accelerator.device).to(weight_dtype)
@@ -547,7 +524,7 @@ def train(args):
                     # 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
-                    # logger.info("text encoder outputs verified")
+                    # print("text encoder outputs verified")
 
                 # get size embeddings
                 orig_size = batch["original_sizes_hw"]
@@ -571,12 +548,7 @@ def train(args):
 
                 target = noise
 
-                if (
-                    args.min_snr_gamma
-                    or args.scale_v_pred_loss_like_noise_pred
-                    or args.v_pred_like_loss
-                    or args.debiased_estimation_loss
-                ):
+                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])
@@ -587,8 +559,6 @@ def train(args):
                         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)
-                    if args.debiased_estimation_loss:
-                        loss = apply_debiased_estimation(loss, timesteps, noise_scheduler)
 
                     loss = loss.mean()  # mean over batch dimension
                 else:
@@ -650,22 +620,29 @@ def train(args):
             if args.logging_dir is not None:
                 logs = {"loss": current_loss}
                 if block_lrs is None:
-                    train_util.append_lr_to_logs(logs, lr_scheduler, args.optimizer_type, including_unet=train_unet)
+                    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)  # U-Net is included in block_lrs
+                    append_block_lr_to_logs(block_lrs, logs, lr_scheduler, args.optimizer_type)
 
                 accelerator.log(logs, step=global_step)
 
-            loss_recorder.add(epoch=epoch, step=step, loss=current_loss)
-            avr_loss: float = loss_recorder.moving_average
-            logs = {"avr_loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
+            # 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_recorder.moving_average}
+            logs = {"loss/epoch": loss_total / len(train_dataloader)}
             accelerator.log(logs, step=epoch + 1)
 
         accelerator.wait_for_everyone()
@@ -734,13 +711,12 @@ def train(args):
             logit_scale,
             ckpt_info,
         )
-        logger.info("model saved.")
+        print("model saved.")
 
 
 def setup_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser()
 
-    add_logging_arguments(parser)
     train_util.add_sd_models_arguments(parser)
     train_util.add_dataset_arguments(parser, True, True, True)
     train_util.add_training_arguments(parser, False)
@@ -750,22 +726,7 @@ def setup_parser() -> argparse.ArgumentParser:
     custom_train_functions.add_custom_train_arguments(parser)
     sdxl_train_util.add_sdxl_training_arguments(parser)
 
-    parser.add_argument(
-        "--learning_rate_te1",
-        type=float,
-        default=None,
-        help="learning rate for text encoder 1 (ViT-L) / text encoder 1 (ViT-L)の学習率",
-    )
-    parser.add_argument(
-        "--learning_rate_te2",
-        type=float,
-        default=None,
-        help="learning rate for text encoder 2 (BiG-G) / text encoder 2 (BiG-G)の学習率",
-    )
-
-    parser.add_argument(
-        "--diffusers_xformers", action="store_true", help="use xformers by diffusers / Diffusersでxformersを使用する"
-    )
+    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",

sdxl_train_control_net_lllite.py

@@ -2,6 +2,7 @@
 # training code for ControlNet-LLLite with passing cond_image to U-Net's forward
 
 import argparse
+import gc
 import json
 import math
 import os
@@ -12,11 +13,14 @@ from types import SimpleNamespace
 import toml
 
 from tqdm import tqdm
-
 import torch
-from library.device_utils import init_ipex, clean_memory_on_device
-init_ipex()
-
+try:
+    import intel_extension_for_pytorch as ipex
+    if torch.xpu.is_available():
+        from library.ipex import ipex_init
+        ipex_init()
+except Exception:
+    pass
 from torch.nn.parallel import DistributedDataParallel as DDP
 from accelerate.utils import set_seed
 import accelerate
@@ -40,15 +44,8 @@ from library.custom_train_functions import (
     pyramid_noise_like,
     apply_noise_offset,
     scale_v_prediction_loss_like_noise_prediction,
-    apply_debiased_estimation,
 )
 import networks.control_net_lllite_for_train as control_net_lllite_for_train
-from library.utils import setup_logging, add_logging_arguments
-
-setup_logging()
-import logging
-
-logger = logging.getLogger(__name__)
 
 
 # TODO 他のスクリプトと共通化する
@@ -69,7 +66,6 @@ def train(args):
     train_util.verify_training_args(args)
     train_util.prepare_dataset_args(args, True)
     sdxl_train_util.verify_sdxl_training_args(args)
-    setup_logging(args, reset=True)
 
     cache_latents = args.cache_latents
     use_user_config = args.dataset_config is not None
@@ -83,11 +79,11 @@ def train(args):
     # データセットを準備する
     blueprint_generator = BlueprintGenerator(ConfigSanitizer(False, False, True, True))
     if use_user_config:
-        logger.info(f"Load dataset config from {args.dataset_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):
-            logger.warning(
+            print(
                 "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
                     ", ".join(ignored)
                 )
@@ -119,7 +115,7 @@ def train(args):
         train_util.debug_dataset(train_dataset_group)
         return
     if len(train_dataset_group) == 0:
-        logger.error(
+        print(
             "No data found. Please verify arguments (train_data_dir must be the parent of folders with images) / 画像がありません。引数指定を確認してください（train_data_dirには画像があるフォルダではなく、画像があるフォルダの親フォルダを指定する必要があります）"
         )
         return
@@ -129,9 +125,7 @@ def train(args):
             train_dataset_group.is_latent_cacheable()
         ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
     else:
-        logger.warning(
-            "WARNING: random_crop is not supported yet for ControlNet training / ControlNetの学習ではrandom_cropはまだサポートされていません"
-        )
+        print("WARNING: random_crop is not supported yet for ControlNet training / ControlNetの学習ではrandom_cropはまだサポートされていません")
 
     if args.cache_text_encoder_outputs:
         assert (
@@ -139,7 +133,7 @@ def train(args):
         ), "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を準備する
-    logger.info("prepare accelerator")
+    print("prepare accelerator")
     accelerator = train_util.prepare_accelerator(args)
     is_main_process = accelerator.is_main_process
 
@@ -171,7 +165,9 @@ def train(args):
                 accelerator.is_main_process,
             )
         vae.to("cpu")
-        clean_memory_on_device(accelerator.device)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
 
         accelerator.wait_for_everyone()
 
@@ -236,14 +232,14 @@ def train(args):
     accelerator.print("prepare optimizer, data loader etc.")
 
     trainable_params = list(unet.prepare_params())
-    logger.info(f"trainable params count: {len(trainable_params)}")
-    logger.info(f"number of trainable parameters: {sum(p.numel() for p in trainable_params if p.requires_grad)}")
+    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 は persistent_workers が使えないので注意
-    n_workers = min(args.max_data_loader_n_workers, os.cpu_count())  # cpu_count or max_data_loader_n_workers
+    # 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,
@@ -259,9 +255,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}"
-        )
+        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)
@@ -288,6 +282,9 @@ def train(args):
     # 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:
@@ -298,7 +295,8 @@ def train(args):
         # 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)
-        clean_memory_on_device(accelerator.device)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
     else:
         # make sure Text Encoders are on GPU
         text_encoder1.to(accelerator.device)
@@ -329,10 +327,8 @@ def train(args):
     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])}"
-    )
-    # logger.info(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
+    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}")
 
@@ -348,15 +344,14 @@ def train(args):
 
     if accelerator.is_main_process:
         init_kwargs = {}
-        if args.wandb_run_name:
-            init_kwargs["wandb"] = {"name": args.wandb_run_name}
         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_recorder = train_util.LossRecorder()
+    loss_list = []
+    loss_total = 0.0
     del train_dataset_group
 
     # function for saving/removing
@@ -402,7 +397,7 @@ def train(args):
                         # NaNが含まれていれば警告を表示し0に置き換える
                         if torch.any(torch.isnan(latents)):
                             accelerator.print("NaN found in latents, replacing with zeros")
-                            latents = torch.nan_to_num(latents, 0, out=latents)
+                            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:
@@ -465,13 +460,11 @@ def train(args):
                 loss = loss * loss_weights
 
                 if args.min_snr_gamma:
-                    loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization)
+                    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)
-                if args.debiased_estimation_loss:
-                    loss = apply_debiased_estimation(loss, timesteps, noise_scheduler)
 
                 loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし
 
@@ -507,9 +500,14 @@ def train(args):
                             remove_model(remove_ckpt_name)
 
             current_loss = loss.detach().item()
-            loss_recorder.add(epoch=epoch, step=step, loss=current_loss)
-            avr_loss: float = loss_recorder.moving_average
-            logs = {"avr_loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
+            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:
@@ -520,7 +518,7 @@ def train(args):
                 break
 
         if args.logging_dir is not None:
-            logs = {"loss/epoch": loss_recorder.moving_average}
+            logs = {"loss/epoch": loss_total / len(loss_list)}
             accelerator.log(logs, step=epoch + 1)
 
         accelerator.wait_for_everyone()
@@ -556,13 +554,12 @@ def train(args):
         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)
 
-        logger.info("model saved.")
+        print("model saved.")
 
 
 def setup_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser()
 
-    add_logging_arguments(parser)
     train_util.add_sd_models_arguments(parser)
     train_util.add_dataset_arguments(parser, False, True, True)
     train_util.add_training_arguments(parser, False)
@@ -578,12 +575,8 @@ def setup_parser() -> argparse.ArgumentParser:
         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("--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",

sdxl_train_control_net_lllite_old.py

@@ -1,4 +1,5 @@
 import argparse
+import gc
 import json
 import math
 import os
@@ -9,11 +10,14 @@ from types import SimpleNamespace
 import toml
 
 from tqdm import tqdm
-
 import torch
-from library.device_utils import init_ipex, clean_memory_on_device
-init_ipex()
-
+try:
+    import intel_extension_for_pytorch as ipex
+    if torch.xpu.is_available():
+        from library.ipex import ipex_init
+        ipex_init()
+except Exception:
+    pass
 from torch.nn.parallel import DistributedDataParallel as DDP
 from accelerate.utils import set_seed
 from diffusers import DDPMScheduler, ControlNetModel
@@ -36,15 +40,8 @@ from library.custom_train_functions import (
     pyramid_noise_like,
     apply_noise_offset,
     scale_v_prediction_loss_like_noise_prediction,
-    apply_debiased_estimation,
 )
 import networks.control_net_lllite as control_net_lllite
-from library.utils import setup_logging, add_logging_arguments
-
-setup_logging()
-import logging
-
-logger = logging.getLogger(__name__)
 
 
 # TODO 他のスクリプトと共通化する
@@ -65,7 +62,6 @@ def train(args):
     train_util.verify_training_args(args)
     train_util.prepare_dataset_args(args, True)
     sdxl_train_util.verify_sdxl_training_args(args)
-    setup_logging(args, reset=True)
 
     cache_latents = args.cache_latents
     use_user_config = args.dataset_config is not None
@@ -79,11 +75,11 @@ def train(args):
     # データセットを準備する
     blueprint_generator = BlueprintGenerator(ConfigSanitizer(False, False, True, True))
     if use_user_config:
-        logger.info(f"Load dataset config from {args.dataset_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):
-            logger.warning(
+            print(
                 "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
                     ", ".join(ignored)
                 )
@@ -115,7 +111,7 @@ def train(args):
         train_util.debug_dataset(train_dataset_group)
         return
     if len(train_dataset_group) == 0:
-        logger.error(
+        print(
             "No data found. Please verify arguments (train_data_dir must be the parent of folders with images) / 画像がありません。引数指定を確認してください（train_data_dirには画像があるフォルダではなく、画像があるフォルダの親フォルダを指定する必要があります）"
         )
         return
@@ -125,9 +121,7 @@ def train(args):
             train_dataset_group.is_latent_cacheable()
         ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
     else:
-        logger.warning(
-            "WARNING: random_crop is not supported yet for ControlNet training / ControlNetの学習ではrandom_cropはまだサポートされていません"
-        )
+        print("WARNING: random_crop is not supported yet for ControlNet training / ControlNetの学習ではrandom_cropはまだサポートされていません")
 
     if args.cache_text_encoder_outputs:
         assert (
@@ -135,7 +129,7 @@ def train(args):
         ), "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を準備する
-    logger.info("prepare accelerator")
+    print("prepare accelerator")
     accelerator = train_util.prepare_accelerator(args)
     is_main_process = accelerator.is_main_process
 
@@ -170,7 +164,9 @@ def train(args):
                 accelerator.is_main_process,
             )
         vae.to("cpu")
-        clean_memory_on_device(accelerator.device)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
 
         accelerator.wait_for_everyone()
 
@@ -204,14 +200,14 @@ def train(args):
     accelerator.print("prepare optimizer, data loader etc.")
 
     trainable_params = list(network.prepare_optimizer_params())
-    logger.info(f"trainable params count: {len(trainable_params)}")
-    logger.info(f"number of trainable parameters: {sum(p.numel() for p in trainable_params if p.requires_grad)}")
+    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 は persistent_workers が使えないので注意
-    n_workers = min(args.max_data_loader_n_workers, os.cpu_count())  # cpu_count or max_data_loader_n_workers
+    # 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,
@@ -227,9 +223,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}"
-        )
+        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)
@@ -259,6 +253,9 @@ def train(args):
     )
     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:
@@ -271,7 +268,8 @@ def train(args):
         # 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)
-        clean_memory_on_device(accelerator.device)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
     else:
         # make sure Text Encoders are on GPU
         text_encoder1.to(accelerator.device)
@@ -302,10 +300,8 @@ def train(args):
     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])}"
-    )
-    # logger.info(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
+    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}")
 
@@ -327,7 +323,8 @@ def train(args):
             "lllite_control_net_train" if args.log_tracker_name is None else args.log_tracker_name, init_kwargs=init_kwargs
         )
 
-    loss_recorder = train_util.LossRecorder()
+    loss_list = []
+    loss_total = 0.0
     del train_dataset_group
 
     # function for saving/removing
@@ -369,7 +366,7 @@ def train(args):
                         # NaNが含まれていれば警告を表示し0に置き換える
                         if torch.any(torch.isnan(latents)):
                             accelerator.print("NaN found in latents, replacing with zeros")
-                            latents = torch.nan_to_num(latents, 0, out=latents)
+                            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:
@@ -433,13 +430,11 @@ def train(args):
                 loss = loss * loss_weights
 
                 if args.min_snr_gamma:
-                    loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization)
+                    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)
-                if args.debiased_estimation_loss:
-                    loss = apply_debiased_estimation(loss, timesteps, noise_scheduler)
 
                 loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし
 
@@ -475,9 +470,14 @@ def train(args):
                             remove_model(remove_ckpt_name)
 
             current_loss = loss.detach().item()
-            loss_recorder.add(epoch=epoch, step=step, loss=current_loss)
-            avr_loss: float = loss_recorder.moving_average
-            logs = {"avr_loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
+            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:
@@ -488,7 +488,7 @@ def train(args):
                 break
 
         if args.logging_dir is not None:
-            logs = {"loss/epoch": loss_recorder.moving_average}
+            logs = {"loss/epoch": loss_total / len(loss_list)}
             accelerator.log(logs, step=epoch + 1)
 
         accelerator.wait_for_everyone()
@@ -524,13 +524,12 @@ def train(args):
         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)
 
-        logger.info("model saved.")
+        print("model saved.")
 
 
 def setup_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser()
 
-    add_logging_arguments(parser)
     train_util.add_sd_models_arguments(parser)
     train_util.add_dataset_arguments(parser, False, True, True)
     train_util.add_training_arguments(parser, False)
@@ -546,12 +545,8 @@ def setup_parser() -> argparse.ArgumentParser:
         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("--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",

sdxl_train_network.py

@@ -1,15 +1,15 @@
 import argparse
-
 import torch
-from library.device_utils import init_ipex, clean_memory_on_device
-init_ipex()
-
+try:
+    import intel_extension_for_pytorch as ipex
+    if torch.xpu.is_available():
+        from library.ipex import ipex_init
+        ipex_init()
+except Exception:
+    pass
 from library import sdxl_model_util, sdxl_train_util, train_util
 import train_network
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 class SdxlNetworkTrainer(train_network.NetworkTrainer):
     def __init__(self):
@@ -62,36 +62,36 @@ class SdxlNetworkTrainer(train_network.NetworkTrainer):
         if args.cache_text_encoder_outputs:
             if not args.lowram:
                 # メモリ消費を減らす
-                logger.info("move vae and unet to cpu to save memory")
+                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")
-                clean_memory_on_device(accelerator.device)
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
 
-            # When TE is not be trained, it will not be prepared so we need to use explicit autocast
-            with accelerator.autocast():
-                dataset.cache_text_encoder_outputs(
-                    tokenizers,
-                    text_encoders,
-                    accelerator.device,
-                    weight_dtype,
-                    args.cache_text_encoder_outputs_to_disk,
-                    accelerator.is_main_process,
-                )
+            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)
-            clean_memory_on_device(accelerator.device)
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
 
             if not args.lowram:
-                logger.info("move vae and unet back to original device")
+                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, dtype=weight_dtype)
-            text_encoders[1].to(accelerator.device, dtype=weight_dtype)
+            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:
@@ -121,7 +121,6 @@ class SdxlNetworkTrainer(train_network.NetworkTrainer):
                     text_encoders[0],
                     text_encoders[1],
                     None if not args.full_fp16 else weight_dtype,
-                    accelerator=accelerator,
                 )
         else:
             encoder_hidden_states1 = batch["text_encoder_outputs1_list"].to(accelerator.device).to(weight_dtype)
@@ -143,7 +142,7 @@ class SdxlNetworkTrainer(train_network.NetworkTrainer):
             # 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
-            # logger.info("text encoder outputs verified")
+            # print("text encoder outputs verified")
 
         return encoder_hidden_states1, encoder_hidden_states2, pool2
 

sdxl_train_textual_inversion.py

@@ -2,11 +2,14 @@ import argparse
 import os
 
 import regex
-
 import torch
-from library.device_utils import init_ipex
-init_ipex()
-
+try:
+    import intel_extension_for_pytorch as ipex
+    if torch.xpu.is_available():
+        from library.ipex import ipex_init
+        ipex_init()
+except Exception:
+    pass
 import open_clip
 from library import sdxl_model_util, sdxl_train_util, train_util
 
@@ -61,7 +64,6 @@ class SdxlTextualInversionTrainer(train_textual_inversion.TextualInversionTraine
                 text_encoders[0],
                 text_encoders[1],
                 None if not args.full_fp16 else weight_dtype,
-                accelerator=accelerator,
             )
         return encoder_hidden_states1, encoder_hidden_states2, pool2
 

tools/cache_latents.py

@@ -16,10 +16,7 @@ from library.config_util import (
     ConfigSanitizer,
     BlueprintGenerator,
 )
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 def cache_to_disk(args: argparse.Namespace) -> None:
     train_util.prepare_dataset_args(args, True)
@@ -44,18 +41,18 @@ def cache_to_disk(args: argparse.Namespace) -> None:
     if args.dataset_class is None:
         blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, False, True))
         if args.dataset_config is not None:
-            logger.info(f"Load dataset config from {args.dataset_config}")
+            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):
-                logger.warning(
+                print(
                     "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
                         ", ".join(ignored)
                     )
                 )
         else:
             if use_dreambooth_method:
-                logger.info("Using DreamBooth method.")
+                print("Using DreamBooth method.")
                 user_config = {
                     "datasets": [
                         {
@@ -66,7 +63,7 @@ def cache_to_disk(args: argparse.Namespace) -> None:
                     ]
                 }
             else:
-                logger.info("Training with captions.")
+                print("Training with captions.")
                 user_config = {
                     "datasets": [
                         {
@@ -93,7 +90,7 @@ def cache_to_disk(args: argparse.Namespace) -> None:
     collator = train_util.collator_class(current_epoch, current_step, ds_for_collator)
 
     # acceleratorを準備する
-    logger.info("prepare accelerator")
+    print("prepare accelerator")
     accelerator = train_util.prepare_accelerator(args)
 
     # mixed precisionに対応した型を用意しておき適宜castする
@@ -101,7 +98,7 @@ def cache_to_disk(args: argparse.Namespace) -> None:
     vae_dtype = torch.float32 if args.no_half_vae else weight_dtype
 
     # モデルを読み込む
-    logger.info("load model")
+    print("load model")
     if args.sdxl:
         (_, _, _, vae, _, _, _) = sdxl_train_util.load_target_model(args, accelerator, "sdxl", weight_dtype)
     else:
@@ -116,8 +113,8 @@ def cache_to_disk(args: argparse.Namespace) -> None:
     # dataloaderを準備する
     train_dataset_group.set_caching_mode("latents")
 
-    # DataLoaderのプロセス数：0 は persistent_workers が使えないので注意
-    n_workers = min(args.max_data_loader_n_workers, os.cpu_count())  # cpu_count or max_data_loader_n_workers
+    # 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,
@@ -155,7 +152,7 @@ def cache_to_disk(args: argparse.Namespace) -> None:
 
                 if args.skip_existing:
                     if train_util.is_disk_cached_latents_is_expected(image_info.bucket_reso, image_info.latents_npz, flip_aug):
-                        logger.warning(f"Skipping {image_info.latents_npz} because it already exists.")
+                        print(f"Skipping {image_info.latents_npz} because it already exists.")
                         continue
 
                 image_infos.append(image_info)

tools/cache_text_encoder_outputs.py

@@ -16,10 +16,7 @@ from library.config_util import (
     ConfigSanitizer,
     BlueprintGenerator,
 )
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 def cache_to_disk(args: argparse.Namespace) -> None:
     train_util.prepare_dataset_args(args, True)
@@ -51,18 +48,18 @@ def cache_to_disk(args: argparse.Namespace) -> None:
     if args.dataset_class is None:
         blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, False, True))
         if args.dataset_config is not None:
-            logger.info(f"Load dataset config from {args.dataset_config}")
+            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):
-                logger.warning(
+                print(
                     "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
                         ", ".join(ignored)
                     )
                 )
         else:
             if use_dreambooth_method:
-                logger.info("Using DreamBooth method.")
+                print("Using DreamBooth method.")
                 user_config = {
                     "datasets": [
                         {
@@ -73,7 +70,7 @@ def cache_to_disk(args: argparse.Namespace) -> None:
                     ]
                 }
             else:
-                logger.info("Training with captions.")
+                print("Training with captions.")
                 user_config = {
                     "datasets": [
                         {
@@ -98,14 +95,14 @@ def cache_to_disk(args: argparse.Namespace) -> None:
     collator = train_util.collator_class(current_epoch, current_step, ds_for_collator)
 
     # acceleratorを準備する
-    logger.info("prepare accelerator")
+    print("prepare accelerator")
     accelerator = train_util.prepare_accelerator(args)
 
     # mixed precisionに対応した型を用意しておき適宜castする
     weight_dtype, _ = train_util.prepare_dtype(args)
 
     # モデルを読み込む
-    logger.info("load model")
+    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]
@@ -121,8 +118,8 @@ def cache_to_disk(args: argparse.Namespace) -> None:
     # dataloaderを準備する
     train_dataset_group.set_caching_mode("text")
 
-    # DataLoaderのプロセス数：0 は persistent_workers が使えないので注意
-    n_workers = min(args.max_data_loader_n_workers, os.cpu_count())  # cpu_count or max_data_loader_n_workers
+    # 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,
@@ -150,7 +147,7 @@ def cache_to_disk(args: argparse.Namespace) -> None:
 
             if args.skip_existing:
                 if os.path.exists(image_info.text_encoder_outputs_npz):
-                    logger.warning(f"Skipping {image_info.text_encoder_outputs_npz} because it already exists.")
+                    print(f"Skipping {image_info.text_encoder_outputs_npz} because it already exists.")
                     continue
                 
             image_info.input_ids1 = input_ids1

tools/canny.py

@@ -1,10 +1,6 @@
 import argparse
 import cv2
 
-import logging
-from library.utils import setup_logging
-setup_logging()
-logger = logging.getLogger(__name__)
 
 def canny(args):
   img = cv2.imread(args.input)
@@ -14,7 +10,7 @@ def canny(args):
   # canny_img = 255 - canny_img
 
   cv2.imwrite(args.output, canny_img)
-  logger.info("done!")
+  print("done!")
 
 
 def setup_parser() -> argparse.ArgumentParser:

tools/convert_diffusers20_original_sd.py

@@ -6,10 +6,7 @@ import torch
 from diffusers import StableDiffusionPipeline
 
 import library.model_util as model_util
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 def convert(args):
     # 引数を確認する
@@ -26,23 +23,21 @@ def convert(args):
     is_load_ckpt = os.path.isfile(args.model_to_load)
     is_save_ckpt = len(os.path.splitext(args.model_to_save)[1]) > 0
 
-    assert not is_load_ckpt or args.v1 != args.v2, "v1 or v2 is required to load checkpoint / checkpointの読み込みにはv1/v2指定が必要です"
+    assert not is_load_ckpt or args.v1 != args.v2, f"v1 or v2 is required to load checkpoint / checkpointの読み込みにはv1/v2指定が必要です"
     # assert (
     #     is_save_ckpt or args.reference_model is not None
     # ), f"reference model is required to save as Diffusers / Diffusers形式での保存には参照モデルが必要です"
 
     # モデルを読み込む
     msg = "checkpoint" if is_load_ckpt else ("Diffusers" + (" as fp16" if args.fp16 else ""))
-    logger.info(f"loading {msg}: {args.model_to_load}")
+    print(f"loading {msg}: {args.model_to_load}")
 
     if is_load_ckpt:
         v2_model = args.v2
-        text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(
-            v2_model, args.model_to_load, unet_use_linear_projection_in_v2=args.unet_use_linear_projection
-        )
+        text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(v2_model, args.model_to_load, unet_use_linear_projection_in_v2=args.unet_use_linear_projection)
     else:
         pipe = StableDiffusionPipeline.from_pretrained(
-            args.model_to_load, torch_dtype=load_dtype, tokenizer=None, safety_checker=None, variant=args.variant
+            args.model_to_load, torch_dtype=load_dtype, tokenizer=None, safety_checker=None
         )
         text_encoder = pipe.text_encoder
         vae = pipe.vae
@@ -51,37 +46,26 @@ def convert(args):
         if args.v1 == args.v2:
             # 自動判定する
             v2_model = unet.config.cross_attention_dim == 1024
-            logger.info("checking model version: model is " + ("v2" if v2_model else "v1"))
+            print("checking model version: model is " + ("v2" if v2_model else "v1"))
         else:
             v2_model = not args.v1
 
     # 変換して保存する
     msg = ("checkpoint" + ("" if save_dtype is None else f" in {save_dtype}")) if is_save_ckpt else "Diffusers"
-    logger.info(f"converting and saving as {msg}: {args.model_to_save}")
+    print(f"converting and saving as {msg}: {args.model_to_save}")
 
     if is_save_ckpt:
         original_model = args.model_to_load if is_load_ckpt else None
         key_count = model_util.save_stable_diffusion_checkpoint(
-            v2_model,
-            args.model_to_save,
-            text_encoder,
-            unet,
-            original_model,
-            args.epoch,
-            args.global_step,
-            None if args.metadata is None else eval(args.metadata),
-            save_dtype=save_dtype,
-            vae=vae,
+            v2_model, args.model_to_save, text_encoder, unet, original_model, args.epoch, args.global_step, save_dtype, vae
         )
-        logger.info(f"model saved. total converted state_dict keys: {key_count}")
+        print(f"model saved. total converted state_dict keys: {key_count}")
     else:
-        logger.info(
-            f"copy scheduler/tokenizer config from: {args.reference_model if args.reference_model is not None else 'default model'}"
-        )
+        print(f"copy scheduler/tokenizer config from: {args.reference_model if args.reference_model is not None else 'default model'}")
         model_util.save_diffusers_checkpoint(
             v2_model, args.model_to_save, text_encoder, unet, args.reference_model, vae, args.use_safetensors
         )
-        logger.info("model saved.")
+        print(f"model saved.")
 
 
 def setup_parser() -> argparse.ArgumentParser:
@@ -93,9 +77,7 @@ def setup_parser() -> argparse.ArgumentParser:
         "--v2", action="store_true", help="load v2.0 model (v1 or v2 is required to load checkpoint) / 2.0のモデルを読み込む"
     )
     parser.add_argument(
-        "--unet_use_linear_projection",
-        action="store_true",
-        help="When saving v2 model as Diffusers, set U-Net config to `use_linear_projection=true` (to match stabilityai's model) / Diffusers形式でv2モデルを保存するときにU-Netの設定を`use_linear_projection=true`にする（stabilityaiのモデルと合わせる）",
+        "--unet_use_linear_projection", action="store_true", help="When saving v2 model as Diffusers, set U-Net config to `use_linear_projection=true` (to match stabilityai's model) / Diffusers形式でv2モデルを保存するときにU-Netの設定を`use_linear_projection=true`にする（stabilityaiのモデルと合わせる）"
     )
     parser.add_argument(
         "--fp16",
@@ -117,18 +99,6 @@ def setup_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--global_step", type=int, default=0, help="global_step to write to checkpoint / checkpointに記録するglobal_stepの値"
     )
-    parser.add_argument(
-        "--metadata",
-        type=str,
-        default=None,
-        help='モデルに保存されるメタデータ、Pythonの辞書形式で指定 / metadata: metadata written in to the model in Python Dictionary. Example metadata: \'{"name": "model_name", "resolution": "512x512"}\'',
-    )
-    parser.add_argument(
-        "--variant",
-        type=str,
-        default=None,
-        help="読む込むDiffusersのvariantを指定する、例: fp16 / variant: Diffusers variant to load. Example: fp16",
-    )
     parser.add_argument(
         "--reference_model",
         type=str,

tools/detect_face_rotate.py

@@ -15,10 +15,6 @@ import os
 from anime_face_detector import create_detector
 from tqdm import tqdm
 import numpy as np
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
 
 KP_REYE = 11
 KP_LEYE = 19
@@ -28,7 +24,7 @@ SCORE_THRES = 0.90
 
 def detect_faces(detector, image, min_size):
   preds = detector(image)                     # bgr
-  # logger.info(len(preds))
+  # print(len(preds))
 
   faces = []
   for pred in preds:
@@ -82,7 +78,7 @@ def process(args):
   assert args.crop_ratio is None or args.resize_face_size is None, f"crop_ratio指定時はresize_face_sizeは指定できません"
 
   # アニメ顔検出モデルを読み込む
-  logger.info("loading face detector.")
+  print("loading face detector.")
   detector = create_detector('yolov3')
 
   # cropの引数を解析する
@@ -101,7 +97,7 @@ def process(args):
     crop_h_ratio, crop_v_ratio = [float(t) for t in tokens]
 
   # 画像を処理する
-  logger.info("processing.")
+  print("processing.")
   output_extension = ".png"
 
   os.makedirs(args.dst_dir, exist_ok=True)
@@ -115,7 +111,7 @@ def process(args):
     if len(image.shape) == 2:
       image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
     if image.shape[2] == 4:
-      logger.warning(f"image has alpha. ignore / 画像の透明度が設定されているため無視します: {path}")
+      print(f"image has alpha. ignore / 画像の透明度が設定されているため無視します: {path}")
       image = image[:, :, :3].copy()                    # copyをしないと内部的に透明度情報が付いたままになるらしい
 
     h, w = image.shape[:2]
@@ -148,11 +144,11 @@ def process(args):
           # 顔サイズを基準にリサイズする
           scale = args.resize_face_size / face_size
           if scale < cur_crop_width / w:
-            logger.warning(
+            print(
                 f"image width too small in face size based resizing / 顔を基準にリサイズすると画像の幅がcrop sizeより小さい（顔が相対的に大きすぎる）ので顔サイズが変わります: {path}")
             scale = cur_crop_width / w
           if scale < cur_crop_height / h:
-            logger.warning(
+            print(
                 f"image height too small in face size based resizing / 顔を基準にリサイズすると画像の高さがcrop sizeより小さい（顔が相対的に大きすぎる）ので顔サイズが変わります: {path}")
             scale = cur_crop_height / h
         elif crop_h_ratio is not None:
@@ -161,10 +157,10 @@ def process(args):
         else:
           # 切り出しサイズ指定あり
           if w < cur_crop_width:
-            logger.warning(f"image width too small/ 画像の幅がcrop sizeより小さいので画質が劣化します: {path}")
+            print(f"image width too small/ 画像の幅がcrop sizeより小さいので画質が劣化します: {path}")
             scale = cur_crop_width / w
           if h < cur_crop_height:
-            logger.warning(f"image height too small/ 画像の高さがcrop sizeより小さいので画質が劣化します: {path}")
+            print(f"image height too small/ 画像の高さがcrop sizeより小さいので画質が劣化します: {path}")
             scale = cur_crop_height / h
           if args.resize_fit:
             scale = max(cur_crop_width / w, cur_crop_height / h)
@@ -202,7 +198,7 @@ def process(args):
         face_img = face_img[y:y + cur_crop_height]
 
       # # debug
-      # logger.info(path, cx, cy, angle)
+      # print(path, cx, cy, angle)
       # crp = cv2.resize(image, (image.shape[1]//8, image.shape[0]//8))
       # cv2.imshow("image", crp)
       # if cv2.waitKey() == 27:

tools/latent_upscaler.py

@@ -11,16 +11,10 @@ from typing import Dict, List
 import numpy as np
 
 import torch
-from library.device_utils import init_ipex, get_preferred_device
-init_ipex()
-
 from torch import nn
 from tqdm import tqdm
 from PIL import Image
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 class ResidualBlock(nn.Module):
     def __init__(self, in_channels, out_channels=None, kernel_size=3, stride=1, padding=1):
@@ -222,7 +216,7 @@ class Upscaler(nn.Module):
         upsampled_images = upsampled_images / 127.5 - 1.0
 
         # convert upsample images to latents with batch size
-        # logger.info("Encoding upsampled (LANCZOS4) images...")
+        # print("Encoding upsampled (LANCZOS4) images...")
         upsampled_latents = []
         for i in tqdm(range(0, upsampled_images.shape[0], vae_batch_size)):
             batch = upsampled_images[i : i + vae_batch_size].to(vae.device)
@@ -233,7 +227,7 @@ class Upscaler(nn.Module):
         upsampled_latents = torch.cat(upsampled_latents, dim=0)
 
         # upscale (refine) latents with this model with batch size
-        logger.info("Upscaling latents...")
+        print("Upscaling latents...")
         upscaled_latents = []
         for i in range(0, upsampled_latents.shape[0], batch_size):
             with torch.no_grad():
@@ -248,7 +242,7 @@ def create_upscaler(**kwargs):
     weights = kwargs["weights"]
     model = Upscaler()
 
-    logger.info(f"Loading weights from {weights}...")
+    print(f"Loading weights from {weights}...")
     if os.path.splitext(weights)[1] == ".safetensors":
         from safetensors.torch import load_file
 
@@ -261,20 +255,20 @@ def create_upscaler(**kwargs):
 
 # another interface: upscale images with a model for given images from command line
 def upscale_images(args: argparse.Namespace):
-    DEVICE = get_preferred_device()
+    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     us_dtype = torch.float16  # TODO: support fp32/bf16
     os.makedirs(args.output_dir, exist_ok=True)
 
     # load VAE with Diffusers
     assert args.vae_path is not None, "VAE path is required"
-    logger.info(f"Loading VAE from {args.vae_path}...")
+    print(f"Loading VAE from {args.vae_path}...")
     vae = AutoencoderKL.from_pretrained(args.vae_path, subfolder="vae")
     vae.to(DEVICE, dtype=us_dtype)
 
     # prepare model
-    logger.info("Preparing model...")
+    print("Preparing model...")
     upscaler: Upscaler = create_upscaler(weights=args.weights)
-    # logger.info("Loading weights from", args.weights)
+    # print("Loading weights from", args.weights)
     # upscaler.load_state_dict(torch.load(args.weights))
     upscaler.eval()
     upscaler.to(DEVICE, dtype=us_dtype)
@@ -309,14 +303,14 @@ def upscale_images(args: argparse.Namespace):
             image_debug.save(dest_file_name)
 
     # upscale
-    logger.info("Upscaling...")
+    print("Upscaling...")
     upscaled_latents = upscaler.upscale(
         vae, images, None, us_dtype, width * 2, height * 2, batch_size=args.batch_size, vae_batch_size=args.vae_batch_size
     )
     upscaled_latents /= 0.18215
 
     # decode with batch
-    logger.info("Decoding...")
+    print("Decoding...")
     upscaled_images = []
     for i in tqdm(range(0, upscaled_latents.shape[0], args.vae_batch_size)):
         with torch.no_grad():

tools/merge_models.py

@@ -5,10 +5,7 @@ import torch
 from safetensors import safe_open
 from safetensors.torch import load_file, save_file
 from tqdm import tqdm
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 def is_unet_key(key):
     # VAE or TextEncoder, the last one is for SDXL
@@ -48,10 +45,10 @@ def merge(args):
     # check if all models are safetensors
     for model in args.models:
         if not model.endswith("safetensors"):
-            logger.info(f"Model {model} is not a safetensors model")
+            print(f"Model {model} is not a safetensors model")
             exit()
         if not os.path.isfile(model):
-            logger.info(f"Model {model} does not exist")
+            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"
@@ -68,7 +65,7 @@ def merge(args):
 
         if merged_sd is None:
             # load first model
-            logger.info(f"Loading model {model}, ratio = {ratio}...")
+            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()):
@@ -84,11 +81,11 @@ def merge(args):
                     value = ratio * value.to(dtype)  # first model's value * ratio
                     merged_sd[key] = value
 
-            logger.info(f"Model has {len(merged_sd)} keys " + ("(UNet only)" if args.unet_only else ""))
+            print(f"Model has {len(merged_sd)} keys " + ("(UNet only)" if args.unet_only else ""))
             continue
 
         # load other models
-        logger.info(f"Loading model {model}, ratio = {ratio}...")
+        print(f"Loading model {model}, ratio = {ratio}...")
 
         with safe_open(model, framework="pt", device=args.device) as f:
             model_keys = f.keys()
@@ -96,7 +93,7 @@ def merge(args):
                 _, 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:
-                        logger.info(f"Skip: {new_key}")
+                        print(f"Skip: {new_key}")
                     continue
 
                 value = f.get_tensor(key)
@@ -107,7 +104,7 @@ def merge(args):
             for key in merged_sd.keys():
                 if key in model_keys:
                     continue
-                logger.warning(f"Key {key} not in model {model}, use first model's value")
+                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:
@@ -115,7 +112,7 @@ def merge(args):
 
     # add supplementary keys' value (including VAE and TextEncoder)
     if len(supplementary_key_ratios) > 0:
-        logger.info("add first model's value")
+        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)
@@ -123,7 +120,7 @@ def merge(args):
                     continue
 
                 if is_unet_key(new_key):  # not VAE or TextEncoder
-                    logger.warning(f"Key {new_key} not in all models, ratio = {supplementary_key_ratios[new_key]}")
+                    print(f"Key {new_key} not in all models, ratio = {supplementary_key_ratios[new_key]}")
 
                 value = f.get_tensor(key)  # original key
 
@@ -137,7 +134,7 @@ def merge(args):
     if not output_file.endswith(".safetensors"):
         output_file = output_file + ".safetensors"
 
-    logger.info(f"Saving to {output_file}...")
+    print(f"Saving to {output_file}...")
 
     # convert to save_dtype
     for k in merged_sd.keys():
@@ -145,7 +142,7 @@ def merge(args):
 
     save_file(merged_sd, output_file)
 
-    logger.info("Done!")
+    print("Done!")
 
 
 if __name__ == "__main__":

tools/original_control_net.py

@@ -7,10 +7,7 @@ from safetensors.torch import load_file
 from library.original_unet import UNet2DConditionModel, SampleOutput
 
 import library.model_util as model_util
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 class ControlNetInfo(NamedTuple):
     unet: Any
@@ -54,7 +51,7 @@ def load_control_net(v2, unet, model):
 
     # control sdからキー変換しつつU-Netに対応する部分のみ取り出し、DiffusersのU-Netに読み込む
     # state dictを読み込む
-    logger.info(f"ControlNet: loading control SD model : {model}")
+    print(f"ControlNet: loading control SD model : {model}")
 
     if model_util.is_safetensors(model):
         ctrl_sd_sd = load_file(model)
@@ -64,7 +61,7 @@ def load_control_net(v2, unet, model):
 
     # 重みをU-Netに読み込めるようにする。ControlNetはSD版のstate dictなので、それを読み込む
     is_difference = "difference" in ctrl_sd_sd
-    logger.info(f"ControlNet: loading difference: {is_difference}")
+    print("ControlNet: loading difference:", is_difference)
 
     # ControlNetには存在しないキーがあるので、まず現在のU-NetでSD版の全keyを作っておく
     # またTransfer Controlの元weightとなる
@@ -92,13 +89,13 @@ def load_control_net(v2, unet, model):
     # ControlNetのU-Netを作成する
     ctrl_unet = UNet2DConditionModel(**unet_config)
     info = ctrl_unet.load_state_dict(ctrl_unet_du_sd)
-    logger.info(f"ControlNet: loading Control U-Net: {info}")
+    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)
-    logger.info("ControlNet: loading ControlNet: {info}")
+    print("ControlNet: loading ControlNet:", info)
 
     ctrl_unet.to(unet.device, dtype=unet.dtype)
     ctrl_net.to(unet.device, dtype=unet.dtype)
@@ -120,7 +117,7 @@ def load_preprocess(prep_type: str):
 
         return canny
 
-    logger.info(f"Unsupported prep type: {prep_type}")
+    print("Unsupported prep type:", prep_type)
     return None
 
 
@@ -177,26 +174,13 @@ def call_unet_and_control_net(
     cnet_idx = step % cnet_cnt
     cnet_info = control_nets[cnet_idx]
 
-    # logger.info(current_ratio, cnet_info.prep, cnet_info.weight, cnet_info.ratio)
+    # 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]
-
-    # gradual latent support: match the size of guided_hint to the size of sample
-    if guided_hint.shape[-2:] != sample.shape[-2:]:
-        # print(f"guided_hint.shape={guided_hint.shape}, sample.shape={sample.shape}")
-        org_dtype = guided_hint.dtype
-        if org_dtype == torch.bfloat16:
-            guided_hint = guided_hint.to(torch.float32)
-        guided_hint = torch.nn.functional.interpolate(guided_hint, size=sample.shape[-2:], mode="bicubic")
-        if org_dtype == torch.bfloat16:
-            guided_hint = guided_hint.to(org_dtype)
-
     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 = 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]
 
     # U-Net
@@ -208,7 +192,7 @@ def call_unet_and_control_net(
   # ControlNet
   cnet_outs_list = []
   for i, cnet_info in enumerate(control_nets):
-    # logger.info(current_ratio, cnet_info.prep, cnet_info.weight, cnet_info.ratio)
+    # print(current_ratio, cnet_info.prep, cnet_info.weight, cnet_info.ratio)
     if cnet_info.ratio < current_ratio:
       continue
     guided_hint = guided_hints[i]
@@ -248,7 +232,7 @@ def unet_forward(
     upsample_size = None
 
     if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
-        logger.info("Forward upsample size to force interpolation output size.")
+        print("Forward upsample size to force interpolation output size.")
         forward_upsample_size = True
 
     # 1. time

tools/resize_images_to_resolution.py

@@ -6,10 +6,7 @@ import shutil
 import math
 from PIL import Image
 import numpy as np
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
+
 
 def resize_images(src_img_folder, dst_img_folder, max_resolution="512x512", divisible_by=2, interpolation=None, save_as_png=False, copy_associated_files=False):
   # Split the max_resolution string by "," and strip any whitespaces
@@ -86,7 +83,7 @@ def resize_images(src_img_folder, dst_img_folder, max_resolution="512x512", divi
       image.save(os.path.join(dst_img_folder, new_filename), quality=100)
 
       proc = "Resized" if current_pixels > max_pixels else "Saved"
-      logger.info(f"{proc} image: {filename} with size {img.shape[0]}x{img.shape[1]} as {new_filename}")
+      print(f"{proc} image: {filename} with size {img.shape[0]}x{img.shape[1]} as {new_filename}")
 
     # If other files with same basename, copy them with resolution suffix
     if copy_associated_files:
@@ -97,7 +94,7 @@ def resize_images(src_img_folder, dst_img_folder, max_resolution="512x512", divi
           continue
         for max_resolution in max_resolutions:
           new_asoc_file = base + '+' + max_resolution + ext
-          logger.info(f"Copy {asoc_file} as {new_asoc_file}")
+          print(f"Copy {asoc_file} as {new_asoc_file}")
           shutil.copy(os.path.join(src_img_folder, asoc_file), os.path.join(dst_img_folder, new_asoc_file))
 
 

tools/show_metadata.py

@@ -1,10 +1,6 @@
 import json
 import argparse
 from safetensors import safe_open
-from library.utils import setup_logging
-setup_logging()
-import logging
-logger = logging.getLogger(__name__)
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--model", type=str, required=True)
@@ -14,10 +10,10 @@ with safe_open(args.model, framework="pt") as f:
     metadata = f.metadata()
 
 if metadata is None:
-    logger.error("No metadata found")
+    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))
 
-    
+    

tools/split_ti_embeddings.py

@@ -0,0 +1,84 @@
+import argparse
+import os
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import load_file, save_file
+from tqdm import tqdm
+
+
+def split(args):
+    # load embedding
+    if args.embedding.endswith(".safetensors"):
+        embedding = load_file(args.embedding)
+        with safe_open(args.embedding, framework="pt") as f:
+            metadata = f.metadata()
+    else:
+        embedding = torch.load(args.embedding)
+        metadata = None
+
+    # check format
+    if "emb_params" in embedding:
+        # SD1/2
+        keys = ["emb_params"]
+    elif "clip_l" in embedding:
+        # SDXL
+        keys = ["clip_l", "clip_g"]
+    else:
+        print("Unknown embedding format")
+        exit()
+    num_vectors = embedding[keys[0]].shape[0]
+
+    # prepare output directory
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # prepare splits
+    if args.vectors_per_split is not None:
+        num_splits = (num_vectors + args.vectors_per_split - 1) // args.vectors_per_split
+        vectors_for_split = [args.vectors_per_split] * num_splits
+        if sum(vectors_for_split) > num_vectors:
+            vectors_for_split[-1] -= sum(vectors_for_split) - num_vectors
+            assert sum(vectors_for_split) == num_vectors
+    elif args.vectors is not None:
+        vectors_for_split = args.vectors
+        num_splits = len(vectors_for_split)
+    else:
+        print("Must specify either --vectors_per_split or --vectors / --vectors_per_split または --vectors のどちらかを指定する必要があります")
+        exit()
+
+    assert (
+        sum(vectors_for_split) == num_vectors
+    ), "Sum of vectors must be equal to the number of vectors in the embedding / 分割したベクトルの合計はembeddingのベクトル数と等しくなければなりません"
+
+    # split
+    basename = os.path.splitext(os.path.basename(args.embedding))[0]
+    done_vectors = 0
+    for i, num_vectors in enumerate(vectors_for_split):
+        print(f"Splitting {num_vectors} vectors...")
+
+        split_embedding = {}
+        for key in keys:
+            split_embedding[key] = embedding[key][done_vectors : done_vectors + num_vectors]
+
+        output_file = os.path.join(args.output_dir, f"{basename}_{i}.safetensors")
+        save_file(split_embedding, output_file, metadata)
+        print(f"Saved to {output_file}")
+
+        done_vectors += num_vectors
+
+    print("Done")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Merge models")
+    parser.add_argument("--embedding", type=str, help="Embedding to split")
+    parser.add_argument("--output_dir", type=str, help="Output directory")
+    parser.add_argument(
+        "--vectors_per_split",
+        type=int,
+        default=None,
+        help="Number of vectors per split. If num_vectors is 8 and vectors_per_split is 3, then 3, 3, 2 vectors will be split",
+    )
+    parser.add_argument("--vectors", type=int, default=None, nargs="*", help="number of vectors for each split. e.g. 3 3 2")
+    args = parser.parse_args()
+    split(args)

train_controlnet.py

@@ -1,4 +1,5 @@
 import argparse
+import gc
 import json
 import math
 import os
@@ -9,11 +10,14 @@ from types import SimpleNamespace
 import toml
 
 from tqdm import tqdm
-
 import torch
-from library.device_utils import init_ipex, clean_memory_on_device
-init_ipex()
-
+try:
+    import intel_extension_for_pytorch as ipex
+    if torch.xpu.is_available():
+        from library.ipex import ipex_init
+        ipex_init()
+except Exception:
+    pass
 from torch.nn.parallel import DistributedDataParallel as DDP
 from accelerate.utils import set_seed
 from diffusers import DDPMScheduler, ControlNetModel
@@ -33,12 +37,6 @@ from library.custom_train_functions import (
     pyramid_noise_like,
     apply_noise_offset,
 )
-from library.utils import setup_logging, add_logging_arguments
-
-setup_logging()
-import logging
-
-logger = logging.getLogger(__name__)
 
 
 # TODO 他のスクリプトと共通化する
@@ -60,7 +58,6 @@ def train(args):
     # training_started_at = time.time()
     train_util.verify_training_args(args)
     train_util.prepare_dataset_args(args, True)
-    setup_logging(args, reset=True)
 
     cache_latents = args.cache_latents
     use_user_config = args.dataset_config is not None
@@ -74,11 +71,11 @@ def train(args):
     # データセットを準備する
     blueprint_generator = BlueprintGenerator(ConfigSanitizer(False, False, True, True))
     if use_user_config:
-        logger.info(f"Load dataset config from {args.dataset_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):
-            logger.warning(
+            print(
                 "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
                     ", ".join(ignored)
                 )
@@ -108,7 +105,7 @@ def train(args):
         train_util.debug_dataset(train_dataset_group)
         return
     if len(train_dataset_group) == 0:
-        logger.error(
+        print(
             "No data found. Please verify arguments (train_data_dir must be the parent of folders with images) / 画像がありません。引数指定を確認してください（train_data_dirには画像があるフォルダではなく、画像があるフォルダの親フォルダを指定する必要があります）"
         )
         return
@@ -119,7 +116,7 @@ def train(args):
         ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
 
     # acceleratorを準備する
-    logger.info("prepare accelerator")
+    print("prepare accelerator")
     accelerator = train_util.prepare_accelerator(args)
     is_main_process = accelerator.is_main_process
 
@@ -224,8 +221,10 @@ def train(args):
                 accelerator.is_main_process,
             )
         vae.to("cpu")
-        clean_memory_on_device(accelerator.device)
-        
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
+
         accelerator.wait_for_everyone()
 
     if args.gradient_checkpointing:
@@ -239,8 +238,8 @@ def train(args):
     _, _, optimizer = train_util.get_optimizer(args, trainable_params)
 
     # dataloaderを準備する
-    # DataLoaderのプロセス数：0 は persistent_workers が使えないので注意
-    n_workers = min(args.max_data_loader_n_workers, os.cpu_count())  # cpu_count or max_data_loader_n_workers
+    # 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,
@@ -256,9 +255,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}"
-        )
+        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)
@@ -314,10 +311,8 @@ def train(args):
     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])}"
-    )
-    # logger.info(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
+    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}")
 
@@ -338,15 +333,12 @@ def train(args):
     )
     if accelerator.is_main_process:
         init_kwargs = {}
-        if args.wandb_run_name:
-            init_kwargs["wandb"] = {"name": args.wandb_run_name}
         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
-        )
+        accelerator.init_trackers("controlnet_train" if args.log_tracker_name is None else args.log_tracker_name, init_kwargs=init_kwargs)
 
-    loss_recorder = train_util.LossRecorder()
+    loss_list = []
+    loss_total = 0.0
     del train_dataset_group
 
     # function for saving/removing
@@ -380,11 +372,6 @@ def train(args):
             accelerator.print(f"removing old checkpoint: {old_ckpt_file}")
             os.remove(old_ckpt_file)
 
-    # For --sample_at_first
-    train_util.sample_images(
-        accelerator, args, 0, global_step, accelerator.device, vae, tokenizer, text_encoder, unet, controlnet=controlnet
-    )
-
     # training loop
     for epoch in range(num_train_epochs):
         if is_main_process:
@@ -463,7 +450,7 @@ def train(args):
                 loss = loss * loss_weights
 
                 if args.min_snr_gamma:
-                    loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization)
+                    loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma)
 
                 loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし
 
@@ -513,9 +500,14 @@ def train(args):
                             remove_model(remove_ckpt_name)
 
             current_loss = loss.detach().item()
-            loss_recorder.add(epoch=epoch, step=step, loss=current_loss)
-            avr_loss: float = loss_recorder.moving_average
-            logs = {"avr_loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
+            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:
@@ -526,7 +518,7 @@ def train(args):
                 break
 
         if args.logging_dir is not None:
-            logs = {"loss/epoch": loss_recorder.moving_average}
+            logs = {"loss/epoch": loss_total / len(loss_list)}
             accelerator.log(logs, step=epoch + 1)
 
         accelerator.wait_for_everyone()
@@ -574,13 +566,12 @@ def train(args):
         ckpt_name = train_util.get_last_ckpt_name(args, "." + args.save_model_as)
         save_model(ckpt_name, controlnet, force_sync_upload=True)
 
-        logger.info("model saved.")
+        print("model saved.")
 
 
 def setup_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser()
 
-    add_logging_arguments(parser)
     train_util.add_sd_models_arguments(parser)
     train_util.add_dataset_arguments(parser, False, True, True)
     train_util.add_training_arguments(parser, False)

train_db.py

@@ -1,6 +1,7 @@
 # DreamBooth training
 # XXX dropped option: fine_tune
 
+import gc
 import argparse
 import itertools
 import math
@@ -9,11 +10,14 @@ from multiprocessing import Value
 import toml
 
 from tqdm import tqdm
-
 import torch
-from library.device_utils import init_ipex, clean_memory_on_device
-init_ipex()
-
+try:
+    import intel_extension_for_pytorch as ipex
+    if torch.xpu.is_available():
+        from library.ipex import ipex_init
+        ipex_init()
+except Exception:
+    pass
 from accelerate.utils import set_seed
 from diffusers import DDPMScheduler
 
@@ -31,14 +35,7 @@ from library.custom_train_functions import (
     pyramid_noise_like,
     apply_noise_offset,
     scale_v_prediction_loss_like_noise_prediction,
-    apply_debiased_estimation,
 )
-from library.utils import setup_logging, add_logging_arguments
-
-setup_logging()
-import logging
-
-logger = logging.getLogger(__name__)
 
 # perlin_noise,
 
@@ -46,7 +43,6 @@ logger = logging.getLogger(__name__)
 def train(args):
     train_util.verify_training_args(args)
     train_util.prepare_dataset_args(args, False)
-    setup_logging(args, reset=True)
 
     cache_latents = args.cache_latents
 
@@ -59,11 +55,11 @@ def train(args):
     if args.dataset_class is None:
         blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, False, False, True))
         if args.dataset_config is not None:
-            logger.info(f"Load dataset config from {args.dataset_config}")
+            print(f"Load dataset config from {args.dataset_config}")
             user_config = config_util.load_user_config(args.dataset_config)
             ignored = ["train_data_dir", "reg_data_dir"]
             if any(getattr(args, attr) is not None for attr in ignored):
-                logger.warning(
+                print(
                     "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
                         ", ".join(ignored)
                     )
@@ -98,13 +94,13 @@ def train(args):
         ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
 
     # acceleratorを準備する
-    logger.info("prepare accelerator")
+    print("prepare accelerator")
 
     if args.gradient_accumulation_steps > 1:
-        logger.warning(
+        print(
             f"gradient_accumulation_steps is {args.gradient_accumulation_steps}. accelerate does not support gradient_accumulation_steps when training multiple models (U-Net and Text Encoder), so something might be wrong"
         )
-        logger.warning(
+        print(
             f"gradient_accumulation_stepsが{args.gradient_accumulation_steps}に設定されています。accelerateは複数モデル（U-NetおよびText Encoder）の学習時にgradient_accumulation_stepsをサポートしていないため結果は未知数です"
         )
 
@@ -112,7 +108,6 @@ def train(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
 
     # モデルを読み込む
     text_encoder, vae, unet, load_stable_diffusion_format = train_util.load_target_model(args, weight_dtype, accelerator)
@@ -137,13 +132,15 @@ def train(args):
 
     # 学習を準備する
     if cache_latents:
-        vae.to(accelerator.device, dtype=vae_dtype)
+        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")
-        clean_memory_on_device(accelerator.device)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
 
         accelerator.wait_for_everyone()
 
@@ -166,22 +163,16 @@ def train(args):
     # 学習に必要なクラスを準備する
     accelerator.print("prepare optimizer, data loader etc.")
     if train_text_encoder:
-        if args.learning_rate_te is None:
-            # wightout list, adamw8bit is crashed
-            trainable_params = list(itertools.chain(unet.parameters(), text_encoder.parameters()))
-        else:
-            trainable_params = [
-                {"params": list(unet.parameters()), "lr": args.learning_rate},
-                {"params": list(text_encoder.parameters()), "lr": args.learning_rate_te},
-            ]
+        # wightout list, adamw8bit is crashed
+        trainable_params = list(itertools.chain(unet.parameters(), text_encoder.parameters()))
     else:
         trainable_params = unet.parameters()
-
+    
     _, _, optimizer = train_util.get_optimizer(args, trainable_params)
 
     # dataloaderを準備する
-    # DataLoaderのプロセス数：0 は persistent_workers が使えないので注意
-    n_workers = min(args.max_data_loader_n_workers, os.cpu_count())  # cpu_count or max_data_loader_n_workers
+    # 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,
@@ -196,9 +187,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}"
-        )
+        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)
@@ -226,6 +215,9 @@ def train(args):
     else:
         unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(unet, optimizer, train_dataloader, lr_scheduler)
 
+    # transform DDP after prepare
+    text_encoder, unet = train_util.transform_if_model_is_DDP(text_encoder, unet)
+
     if not train_text_encoder:
         text_encoder.to(accelerator.device, dtype=weight_dtype)  # to avoid 'cpu' vs 'cuda' error
 
@@ -268,16 +260,12 @@ def train(args):
 
     if accelerator.is_main_process:
         init_kwargs = {}
-        if args.wandb_run_name:
-            init_kwargs["wandb"] = {"name": args.wandb_run_name}
         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)
 
-    # For --sample_at_first
-    train_util.sample_images(accelerator, args, 0, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
-
-    loss_recorder = train_util.LossRecorder()
+    loss_list = []
+    loss_total = 0.0
     for epoch in range(num_train_epochs):
         accelerator.print(f"\nepoch {epoch+1}/{num_train_epochs}")
         current_epoch.value = epoch + 1
@@ -345,11 +333,9 @@ def train(args):
                 loss = loss * loss_weights
 
                 if args.min_snr_gamma:
-                    loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization)
+                    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.debiased_estimation_loss:
-                    loss = apply_debiased_estimation(loss, timesteps, noise_scheduler)
 
                 loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし
 
@@ -397,20 +383,30 @@ def train(args):
 
             current_loss = loss.detach().item()
             if args.logging_dir is not None:
-                logs = {"loss": current_loss}
-                train_util.append_lr_to_logs(logs, lr_scheduler, args.optimizer_type, including_unet=True)
+                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
+                    logs["lr/d*lr"] = (
+                        lr_scheduler.optimizers[0].param_groups[0]["d"] * lr_scheduler.optimizers[0].param_groups[0]["lr"]
+                    )
                 accelerator.log(logs, step=global_step)
 
-            loss_recorder.add(epoch=epoch, step=step, loss=current_loss)
-            avr_loss: float = loss_recorder.moving_average
-            logs = {"avr_loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
+            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 global_step >= args.max_train_steps:
                 break
 
         if args.logging_dir is not None:
-            logs = {"loss/epoch": loss_recorder.moving_average}
+            logs = {"loss/epoch": loss_total / len(loss_list)}
             accelerator.log(logs, step=epoch + 1)
 
         accelerator.wait_for_everyone()
@@ -454,13 +450,12 @@ def train(args):
         train_util.save_sd_model_on_train_end(
             args, src_path, save_stable_diffusion_format, use_safetensors, save_dtype, epoch, global_step, text_encoder, unet, vae
         )
-        logger.info("model saved.")
+        print("model saved.")
 
 
 def setup_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser()
 
-    add_logging_arguments(parser)
     train_util.add_sd_models_arguments(parser)
     train_util.add_dataset_arguments(parser, True, False, True)
     train_util.add_training_arguments(parser, True)
@@ -469,12 +464,6 @@ def setup_parser() -> argparse.ArgumentParser:
     config_util.add_config_arguments(parser)
     custom_train_functions.add_custom_train_arguments(parser)
 
-    parser.add_argument(
-        "--learning_rate_te",
-        type=float,
-        default=None,
-        help="learning rate for text encoder, default is same as unet / Text Encoderの学習率、デフォルトはunetと同じ",
-    )
     parser.add_argument(
         "--no_token_padding",
         action="store_true",
@@ -486,11 +475,6 @@ def setup_parser() -> argparse.ArgumentParser:
         default=None,
         help="steps to stop text encoder training, -1 for no training / Text Encoderの学習を止めるステップ数、-1で最初から学習しない",
     )
-    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
 

train_network.py

@@ -1,5 +1,6 @@
 import importlib
 import argparse
+import gc
 import math
 import os
 import sys
@@ -10,13 +11,17 @@ from multiprocessing import Value
 import toml
 
 from tqdm import tqdm
-
 import torch
-from library.device_utils import init_ipex, clean_memory_on_device
-init_ipex()
 
-from torch.nn.parallel import DistributedDataParallel as DDP
+try:
+    import intel_extension_for_pytorch as ipex
 
+    if torch.xpu.is_available():
+        from library.ipex import ipex_init
+
+        ipex_init()
+except Exception:
+    pass
 from accelerate.utils import set_seed
 from diffusers import DDPMScheduler
 from library import model_util
@@ -38,14 +43,7 @@ from library.custom_train_functions import (
     prepare_scheduler_for_custom_training,
     scale_v_prediction_loss_like_noise_prediction,
     add_v_prediction_like_loss,
-    apply_debiased_estimation,
 )
-from library.utils import setup_logging, add_logging_arguments
-
-setup_logging()
-import logging
-
-logger = logging.getLogger(__name__)
 
 
 class NetworkTrainer:
@@ -110,14 +108,11 @@ class NetworkTrainer:
     def is_text_encoder_outputs_cached(self, args):
         return False
 
-    def is_train_text_encoder(self, args):
-        return not args.network_train_unet_only and not self.is_text_encoder_outputs_cached(args)
-
     def cache_text_encoder_outputs_if_needed(
         self, args, accelerator, unet, vae, tokenizers, text_encoders, data_loader, weight_dtype
     ):
         for t_enc in text_encoders:
-            t_enc.to(accelerator.device, dtype=weight_dtype)
+            t_enc.to(accelerator.device)
 
     def get_text_cond(self, args, accelerator, batch, tokenizers, text_encoders, weight_dtype):
         input_ids = batch["input_ids"].to(accelerator.device)
@@ -128,11 +123,6 @@ class NetworkTrainer:
         noise_pred = unet(noisy_latents, timesteps, text_conds).sample
         return noise_pred
 
-    def all_reduce_network(self, accelerator, network):
-        for param in network.parameters():
-            if param.grad is not None:
-                param.grad = accelerator.reduce(param.grad, reduction="mean")
-
     def sample_images(self, accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet):
         train_util.sample_images(accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet)
 
@@ -141,7 +131,6 @@ class NetworkTrainer:
         training_started_at = time.time()
         train_util.verify_training_args(args)
         train_util.prepare_dataset_args(args, True)
-        setup_logging(args, reset=True)
 
         cache_latents = args.cache_latents
         use_dreambooth_method = args.in_json is None
@@ -159,18 +148,18 @@ class NetworkTrainer:
         if args.dataset_class is None:
             blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, False, True))
             if use_user_config:
-                logger.info(f"Loading dataset config from {args.dataset_config}")
+                print(f"Loading dataset config from {args.dataset_config}")
                 user_config = config_util.load_user_config(args.dataset_config)
                 ignored = ["train_data_dir", "reg_data_dir", "in_json"]
                 if any(getattr(args, attr) is not None for attr in ignored):
-                    logger.warning(
+                    print(
                         "ignoring the following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
                             ", ".join(ignored)
                         )
                     )
             else:
                 if use_dreambooth_method:
-                    logger.info("Using DreamBooth method.")
+                    print("Using DreamBooth method.")
                     user_config = {
                         "datasets": [
                             {
@@ -181,7 +170,7 @@ class NetworkTrainer:
                         ]
                     }
                 else:
-                    logger.info("Training with captions.")
+                    print("Training with captions.")
                     user_config = {
                         "datasets": [
                             {
@@ -210,7 +199,7 @@ class NetworkTrainer:
             train_util.debug_dataset(train_dataset_group)
             return
         if len(train_dataset_group) == 0:
-            logger.error(
+            print(
                 "No data found. Please verify arguments (train_data_dir must be the parent of folders with images) / 画像がありません。引数指定を確認してください（train_data_dirには画像があるフォルダではなく、画像があるフォルダの親フォルダを指定する必要があります）"
             )
             return
@@ -223,7 +212,7 @@ class NetworkTrainer:
         self.assert_extra_args(args, train_dataset_group)
 
         # acceleratorを準備する
-        logger.info("preparing accelerator")
+        print("preparing accelerator")
         accelerator = train_util.prepare_accelerator(args)
         is_main_process = accelerator.is_main_process
 
@@ -272,12 +261,13 @@ class NetworkTrainer:
             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")
-            clean_memory_on_device(accelerator.device)
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            gc.collect()
 
             accelerator.wait_for_everyone()
 
         # 必要ならテキストエンコーダーの出力をキャッシュする: Text Encoderはcpuまたはgpuへ移される
-        # cache text encoder outputs if needed: Text Encoder is moved to cpu or gpu
         self.cache_text_encoder_outputs_if_needed(
             args, accelerator, unet, vae, tokenizers, text_encoders, train_dataset_group, weight_dtype
         )
@@ -309,18 +299,17 @@ class NetworkTrainer:
             )
         if network is None:
             return
-        network_has_multiplier = hasattr(network, "set_multiplier")
 
         if hasattr(network, "prepare_network"):
             network.prepare_network(args)
         if args.scale_weight_norms and not hasattr(network, "apply_max_norm_regularization"):
-            logger.warning(
+            print(
                 "warning: scale_weight_norms is specified but the network does not support it / scale_weight_normsが指定されていますが、ネットワークが対応していません"
             )
             args.scale_weight_norms = False
 
         train_unet = not args.network_train_text_encoder_only
-        train_text_encoder = self.is_train_text_encoder(args)
+        train_text_encoder = not args.network_train_unet_only and not self.is_text_encoder_outputs_cached(args)
         network.apply_to(text_encoder, unet, train_text_encoder, train_unet)
 
         if args.network_weights is not None:
@@ -349,8 +338,8 @@ class NetworkTrainer:
         optimizer_name, optimizer_args, optimizer = train_util.get_optimizer(args, trainable_params)
 
         # dataloaderを準備する
-        # DataLoaderのプロセス数：0 は persistent_workers が使えないので注意
-        n_workers = min(args.max_data_loader_n_workers, os.cpu_count())  # cpu_count or max_data_loader_n_workers
+        # 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,
@@ -390,43 +379,51 @@ class NetworkTrainer:
             accelerator.print("enable full bf16 training.")
             network.to(weight_dtype)
 
-        unet_weight_dtype = te_weight_dtype = weight_dtype
-        # Experimental Feature: Put base model into fp8 to save vram
-        if args.fp8_base:
-            assert torch.__version__ >= "2.1.0", "fp8_base requires torch>=2.1.0 / fp8を使う場合はtorch>=2.1.0が必要です。"
-            assert (
-                args.mixed_precision != "no"
-            ), "fp8_base requires mixed precision='fp16' or 'bf16' / fp8を使う場合はmixed_precision='fp16'または'bf16'が必要です。"
-            accelerator.print("enable fp8 training.")
-            unet_weight_dtype = torch.float8_e4m3fn
-            te_weight_dtype = torch.float8_e4m3fn
-
         unet.requires_grad_(False)
-        unet.to(dtype=unet_weight_dtype)
+        unet.to(dtype=weight_dtype)
         for t_enc in text_encoders:
             t_enc.requires_grad_(False)
 
-            # in case of cpu, dtype is already set to fp32 because cpu does not support fp8/fp16/bf16
-            if t_enc.device.type != "cpu":
-                t_enc.to(dtype=te_weight_dtype)
-                # nn.Embedding not support FP8
-                t_enc.text_model.embeddings.to(dtype=(weight_dtype if te_weight_dtype != weight_dtype else te_weight_dtype))
-
-        # acceleratorがなんかよろしくやってくれるらしい / accelerator will do something good
-        if train_unet:
-            unet = accelerator.prepare(unet)
-        else:
-            unet.to(accelerator.device, dtype=unet_weight_dtype)  # move to device because unet is not prepared by accelerator
-        if train_text_encoder:
+        # acceleratorがなんかよろしくやってくれるらしい
+        # TODO めちゃくちゃ冗長なのでコードを整理する
+        if train_unet and train_text_encoder:
             if len(text_encoders) > 1:
-                text_encoder = text_encoders = [accelerator.prepare(t_enc) for t_enc in text_encoders]
+                unet, t_enc1, t_enc2, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+                    unet, text_encoders[0], text_encoders[1], network, optimizer, train_dataloader, lr_scheduler
+                )
+                text_encoder = text_encoders = [t_enc1, t_enc2]
+                del t_enc1, t_enc2
             else:
-                text_encoder = accelerator.prepare(text_encoder)
+                unet, text_encoder, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+                    unet, text_encoder, network, optimizer, train_dataloader, lr_scheduler
+                )
+                text_encoders = [text_encoder]
+        elif train_unet:
+            unet, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+                unet, network, optimizer, train_dataloader, lr_scheduler
+            )
+        elif train_text_encoder:
+            if len(text_encoders) > 1:
+                t_enc1, t_enc2, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+                    text_encoders[0], text_encoders[1], network, optimizer, train_dataloader, lr_scheduler
+                )
+                text_encoder = text_encoders = [t_enc1, t_enc2]
+                del t_enc1, t_enc2
+            else:
+                text_encoder, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+                    text_encoder, network, optimizer, train_dataloader, lr_scheduler
+                )
                 text_encoders = [text_encoder]
-        else:
-            pass  # if text_encoder is not trained, no need to prepare. and device and dtype are already set
 
-        network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(network, optimizer, train_dataloader, lr_scheduler)
+            unet.to(accelerator.device, dtype=weight_dtype)  # move to device because unet is not prepared by accelerator
+        else:
+            network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+                network, optimizer, train_dataloader, lr_scheduler
+            )
+
+        # transform DDP after prepare (train_network here only)
+        text_encoders = train_util.transform_models_if_DDP(text_encoders)
+        unet, network = train_util.transform_models_if_DDP([unet, network])
 
         if args.gradient_checkpointing:
             # according to TI example in Diffusers, train is required
@@ -438,6 +435,9 @@ class NetworkTrainer:
                 if train_text_encoder:
                     t_enc.text_model.embeddings.requires_grad_(True)
 
+            # set top parameter requires_grad = True for gradient checkpointing works
+            if not train_text_encoder:  # train U-Net only
+                unet.parameters().__next__().requires_grad_(True)
         else:
             unet.eval()
             for t_enc in text_encoders:
@@ -445,7 +445,7 @@ class NetworkTrainer:
 
         del t_enc
 
-        accelerator.unwrap_model(network).prepare_grad_etc(text_encoder, unet)
+        network.prepare_grad_etc(text_encoder, unet)
 
         if not cache_latents:  # キャッシュしない場合はVAEを使うのでVAEを準備する
             vae.requires_grad_(False)
@@ -528,7 +528,6 @@ class NetworkTrainer:
             "ss_min_snr_gamma": args.min_snr_gamma,
             "ss_scale_weight_norms": args.scale_weight_norms,
             "ss_ip_noise_gamma": args.ip_noise_gamma,
-            "ss_debiased_estimation": bool(args.debiased_estimation_loss),
         }
 
         if use_user_config:
@@ -698,20 +697,19 @@ class NetworkTrainer:
 
         if accelerator.is_main_process:
             init_kwargs = {}
-            if args.wandb_run_name:
-                init_kwargs["wandb"] = {"name": args.wandb_run_name}
             if args.log_tracker_config is not None:
                 init_kwargs = toml.load(args.log_tracker_config)
             accelerator.init_trackers(
                 "network_train" if args.log_tracker_name is None else args.log_tracker_name, init_kwargs=init_kwargs
             )
 
-        loss_recorder = train_util.LossRecorder()
+        loss_list = []
+        loss_total = 0.0
         del train_dataset_group
 
         # callback for step start
-        if hasattr(accelerator.unwrap_model(network), "on_step_start"):
-            on_step_start = accelerator.unwrap_model(network).on_step_start
+        if hasattr(network, "on_step_start"):
+            on_step_start = network.on_step_start
         else:
             on_step_start = lambda *args, **kwargs: None
 
@@ -739,9 +737,6 @@ class NetworkTrainer:
                 accelerator.print(f"removing old checkpoint: {old_ckpt_file}")
                 os.remove(old_ckpt_file)
 
-        # For --sample_at_first
-        self.sample_images(accelerator, args, 0, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
-
         # training loop
         for epoch in range(num_train_epochs):
             accelerator.print(f"\nepoch {epoch+1}/{num_train_epochs}")
@@ -749,7 +744,7 @@ class NetworkTrainer:
 
             metadata["ss_epoch"] = str(epoch + 1)
 
-            accelerator.unwrap_model(network).on_epoch_start(text_encoder, unet)
+            network.on_epoch_start(text_encoder, unet)
 
             for step, batch in enumerate(train_dataloader):
                 current_step.value = global_step
@@ -766,21 +761,11 @@ class NetworkTrainer:
                             # NaNが含まれていれば警告を表示し0に置き換える
                             if torch.any(torch.isnan(latents)):
                                 accelerator.print("NaN found in latents, replacing with zeros")
-                                latents = torch.nan_to_num(latents, 0, out=latents)
+                                latents = torch.where(torch.isnan(latents), torch.zeros_like(latents), latents)
                         latents = latents * self.vae_scale_factor
+                    b_size = latents.shape[0]
 
-                    # get multiplier for each sample
-                    if network_has_multiplier:
-                        multipliers = batch["network_multipliers"]
-                        # if all multipliers are same, use single multiplier
-                        if torch.all(multipliers == multipliers[0]):
-                            multipliers = multipliers[0].item()
-                        else:
-                            raise NotImplementedError("multipliers for each sample is not supported yet")
-                        # print(f"set multiplier: {multipliers}")
-                        accelerator.unwrap_model(network).set_multiplier(multipliers)
-
-                    with torch.set_grad_enabled(train_text_encoder), accelerator.autocast():
+                    with torch.set_grad_enabled(train_text_encoder):
                         # Get the text embedding for conditioning
                         if args.weighted_captions:
                             text_encoder_conds = get_weighted_text_embeddings(
@@ -802,24 +787,10 @@ class NetworkTrainer:
                         args, noise_scheduler, latents
                     )
 
-                    # ensure the hidden state will require grad
-                    if args.gradient_checkpointing:
-                        for x in noisy_latents:
-                            x.requires_grad_(True)
-                        for t in text_encoder_conds:
-                            t.requires_grad_(True)
-
                     # Predict the noise residual
                     with accelerator.autocast():
                         noise_pred = self.call_unet(
-                            args,
-                            accelerator,
-                            unet,
-                            noisy_latents.requires_grad_(train_unet),
-                            timesteps,
-                            text_encoder_conds,
-                            batch,
-                            weight_dtype,
+                            args, accelerator, unet, noisy_latents, timesteps, text_encoder_conds, batch, weight_dtype
                         )
 
                     if args.v_parameterization:
@@ -835,29 +806,25 @@ class NetworkTrainer:
                     loss = loss * loss_weights
 
                     if args.min_snr_gamma:
-                        loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization)
+                        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)
-                    if args.debiased_estimation_loss:
-                        loss = apply_debiased_estimation(loss, timesteps, noise_scheduler)
 
                     loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし
 
                     accelerator.backward(loss)
-                    if accelerator.sync_gradients:
-                        self.all_reduce_network(accelerator, network)  # sync DDP grad manually
-                        if args.max_grad_norm != 0.0:
-                            params_to_clip = accelerator.unwrap_model(network).get_trainable_params()
-                            accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+                    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)
 
                 if args.scale_weight_norms:
-                    keys_scaled, mean_norm, maximum_norm = accelerator.unwrap_model(network).apply_max_norm_regularization(
+                    keys_scaled, mean_norm, maximum_norm = network.apply_max_norm_regularization(
                         args.scale_weight_norms, accelerator.device
                     )
                     max_mean_logs = {"Keys Scaled": keys_scaled, "Average key norm": mean_norm}
@@ -887,9 +854,14 @@ class NetworkTrainer:
                                 remove_model(remove_ckpt_name)
 
                 current_loss = loss.detach().item()
-                loss_recorder.add(epoch=epoch, step=step, loss=current_loss)
-                avr_loss: float = loss_recorder.moving_average
-                logs = {"avr_loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
+                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.scale_weight_norms:
@@ -903,7 +875,7 @@ class NetworkTrainer:
                     break
 
             if args.logging_dir is not None:
-                logs = {"loss/epoch": loss_recorder.moving_average}
+                logs = {"loss/epoch": loss_total / len(loss_list)}
                 accelerator.log(logs, step=epoch + 1)
 
             accelerator.wait_for_everyone()
@@ -942,13 +914,12 @@ class NetworkTrainer:
             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)
 
-            logger.info("model saved.")
+            print("model saved.")
 
 
 def setup_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser()
 
-    add_logging_arguments(parser)
     train_util.add_sd_models_arguments(parser)
     train_util.add_dataset_arguments(parser, True, True, True)
     train_util.add_training_arguments(parser, True)
@@ -956,9 +927,7 @@ def setup_parser() -> argparse.ArgumentParser:
     config_util.add_config_arguments(parser)
     custom_train_functions.add_custom_train_arguments(parser)
 
-    parser.add_argument(
-        "--no_metadata", action="store_true", help="do not save metadata in output model / メタデータを出力先モデルに保存しない"
-    )
+    parser.add_argument("--no_metadata", action="store_true", help="do not save metadata in output model / メタデータを出力先モデルに保存しない")
     parser.add_argument(
         "--save_model_as",
         type=str,
@@ -970,17 +939,10 @@ def setup_parser() -> argparse.ArgumentParser:
     parser.add_argument("--unet_lr", type=float, default=None, help="learning rate for U-Net / U-Netの学習率")
     parser.add_argument("--text_encoder_lr", type=float, default=None, help="learning rate for Text Encoder / Text Encoderの学習率")
 
+    parser.add_argument("--network_weights", type=str, default=None, help="pretrained weights for network / 学習するネットワークの初期重み")
+    parser.add_argument("--network_module", type=str, default=None, help="network module to train / 学習対象のネットワークのモジュール")
     parser.add_argument(
-        "--network_weights", type=str, default=None, help="pretrained weights for network / 学習するネットワークの初期重み"
-    )
-    parser.add_argument(
-        "--network_module", type=str, default=None, help="network module to train / 学習対象のネットワークのモジュール"
-    )
-    parser.add_argument(
-        "--network_dim",
-        type=int,
-        default=None,
-        help="network dimensions (depends on each network) / モジュールの次元数（ネットワークにより定義は異なります）",
+        "--network_dim", type=int, default=None, help="network dimensions (depends on each network) / モジュールの次元数（ネットワークにより定義は異なります）"
     )
     parser.add_argument(
         "--network_alpha",
@@ -995,25 +957,14 @@ def setup_parser() -> argparse.ArgumentParser:
         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(
-        "--network_args",
-        type=str,
-        default=None,
-        nargs="*",
-        help="additional arguments for network (key=value) / ネットワークへの追加の引数",
+        "--network_args", type=str, default=None, nargs="*", help="additional arguments for network (key=value) / ネットワークへの追加の引数"
+    )
+    parser.add_argument("--network_train_unet_only", action="store_true", help="only training U-Net part / U-Net関連部分のみ学習する")
+    parser.add_argument(
+        "--network_train_text_encoder_only", action="store_true", help="only training Text Encoder part / Text Encoder関連部分のみ学習する"
     )
     parser.add_argument(
-        "--network_train_unet_only", action="store_true", help="only training U-Net part / U-Net関連部分のみ学習する"
-    )
-    parser.add_argument(
-        "--network_train_text_encoder_only",
-        action="store_true",
-        help="only training Text Encoder part / Text Encoder関連部分のみ学習する",
-    )
-    parser.add_argument(
-        "--training_comment",
-        type=str,
-        default=None,
-        help="arbitrary comment string stored in metadata / メタデータに記録する任意のコメント文字列",
+        "--training_comment", type=str, default=None, help="arbitrary comment string stored in metadata / メタデータに記録する任意のコメント文字列"
     )
     parser.add_argument(
         "--dim_from_weights",

train_textual_inversion.py

@@ -1,15 +1,22 @@
 import argparse
+import gc
 import math
 import os
 from multiprocessing import Value
 import toml
 
 from tqdm import tqdm
-
 import torch
-from library.device_utils import init_ipex, clean_memory_on_device
-init_ipex()
 
+try:
+    import intel_extension_for_pytorch as ipex
+
+    if torch.xpu.is_available():
+        from library.ipex import ipex_init
+
+        ipex_init()
+except Exception:
+    pass
 from accelerate.utils import set_seed
 from diffusers import DDPMScheduler
 from transformers import CLIPTokenizer
@@ -28,14 +35,7 @@ from library.custom_train_functions import (
     prepare_scheduler_for_custom_training,
     scale_v_prediction_loss_like_noise_prediction,
     add_v_prediction_like_loss,
-    apply_debiased_estimation,
 )
-from library.utils import setup_logging, add_logging_arguments
-
-setup_logging()
-import logging
-
-logger = logging.getLogger(__name__)
 
 imagenet_templates_small = [
     "a photo of a {}",
@@ -170,9 +170,15 @@ class TextualInversionTrainer:
             args.output_name = args.token_string
         use_template = args.use_object_template or args.use_style_template
 
+        assert (
+            args.token_string is not None or args.token_strings is not None
+        ), "token_string or token_strings must be specified / token_stringまたはtoken_stringsを指定してください"
+        assert (
+            not use_template or args.token_strings is None
+        ), "token_strings cannot be used with template / token_stringsはテンプレートと一緒に使えません"
+
         train_util.verify_training_args(args)
         train_util.prepare_dataset_args(args, True)
-        setup_logging(args, reset=True)
 
         cache_latents = args.cache_latents
 
@@ -183,7 +189,7 @@ class TextualInversionTrainer:
         tokenizers = tokenizer_or_list if isinstance(tokenizer_or_list, list) else [tokenizer_or_list]
 
         # acceleratorを準備する
-        logger.info("prepare accelerator")
+        print("prepare accelerator")
         accelerator = train_util.prepare_accelerator(args)
 
         # mixed precisionに対応した型を用意しておき適宜castする
@@ -219,9 +225,17 @@ class TextualInversionTrainer:
         # add new word to tokenizer, count is num_vectors_per_token
         # if token_string is hoge, "hoge", "hoge1", "hoge2", ... are added
 
-        self.assert_token_string(args.token_string, tokenizers)
+        if args.token_strings is not None:
+            token_strings = args.token_strings
+            assert (
+                len(token_strings) == args.num_vectors_per_token
+            ), f"num_vectors_per_token is mismatch for token_strings / token_stringsの数がnum_vectors_per_tokenと合いません: {len(token_strings)}"
+            for token_string in token_strings:
+                self.assert_token_string(token_string, tokenizers)
+        else:
+            self.assert_token_string(args.token_string, tokenizers)
+            token_strings = [args.token_string] + [f"{args.token_string}{i+1}" for i in range(args.num_vectors_per_token - 1)]
 
-        token_strings = [args.token_string] + [f"{args.token_string}{i+1}" for i in range(args.num_vectors_per_token - 1)]
         token_ids_list = []
         token_embeds_list = []
         for i, (tokenizer, text_encoder, init_token_ids) in enumerate(zip(tokenizers, text_encoders, init_token_ids_list)):
@@ -293,7 +307,7 @@ class TextualInversionTrainer:
                         ]
                     }
                 else:
-                    logger.info("Train with captions.")
+                    print("Train with captions.")
                     user_config = {
                         "datasets": [
                             {
@@ -336,7 +350,7 @@ class TextualInversionTrainer:
                 prompt_replacement = None
         else:
             # サンプル生成用
-            if args.num_vectors_per_token > 1:
+            if args.num_vectors_per_token > 1 and args.token_strings is None:
                 replace_to = " ".join(token_strings)
                 train_dataset_group.add_replacement(args.token_string, replace_to)
                 prompt_replacement = (args.token_string, replace_to)
@@ -368,7 +382,9 @@ class TextualInversionTrainer:
             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")
-            clean_memory_on_device(accelerator.device)
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            gc.collect()
 
             accelerator.wait_for_everyone()
 
@@ -385,8 +401,8 @@ class TextualInversionTrainer:
         _, _, optimizer = train_util.get_optimizer(args, trainable_params)
 
         # dataloaderを準備する
-        # DataLoaderのプロセス数：0 は persistent_workers が使えないので注意
-        n_workers = min(args.max_data_loader_n_workers, os.cpu_count())  # cpu_count or max_data_loader_n_workers
+        # 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,
@@ -416,11 +432,15 @@ class TextualInversionTrainer:
             text_encoder_or_list, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
                 text_encoder_or_list, optimizer, train_dataloader, lr_scheduler
             )
+            # transform DDP after prepare
+            text_encoder_or_list, unet = train_util.transform_if_model_is_DDP(text_encoder_or_list, unet)
 
         elif len(text_encoders) == 2:
             text_encoder1, text_encoder2, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
                 text_encoders[0], text_encoders[1], optimizer, train_dataloader, lr_scheduler
             )
+            # transform DDP after prepare
+            text_encoder1, text_encoder2, unet = train_util.transform_if_model_is_DDP(text_encoder1, text_encoder2, unet)
 
             text_encoder_or_list = text_encoders = [text_encoder1, text_encoder2]
 
@@ -439,10 +459,9 @@ class TextualInversionTrainer:
 
             # Freeze all parameters except for the token embeddings in text encoder
             text_encoder.requires_grad_(True)
-            unwrapped_text_encoder = accelerator.unwrap_model(text_encoder)
-            unwrapped_text_encoder.text_model.encoder.requires_grad_(False)
-            unwrapped_text_encoder.text_model.final_layer_norm.requires_grad_(False)
-            unwrapped_text_encoder.text_model.embeddings.position_embedding.requires_grad_(False)
+            text_encoder.text_model.encoder.requires_grad_(False)
+            text_encoder.text_model.final_layer_norm.requires_grad_(False)
+            text_encoder.text_model.embeddings.position_embedding.requires_grad_(False)
             # text_encoder.text_model.embeddings.token_embedding.requires_grad_(True)
 
         unet.requires_grad_(False)
@@ -502,8 +521,6 @@ class TextualInversionTrainer:
 
         if accelerator.is_main_process:
             init_kwargs = {}
-            if args.wandb_run_name:
-                init_kwargs["wandb"] = {"name": args.wandb_run_name}
             if args.log_tracker_config is not None:
                 init_kwargs = toml.load(args.log_tracker_config)
             accelerator.init_trackers(
@@ -529,20 +546,6 @@ class TextualInversionTrainer:
                 accelerator.print(f"removing old checkpoint: {old_ckpt_file}")
                 os.remove(old_ckpt_file)
 
-        # For --sample_at_first
-        self.sample_images(
-            accelerator,
-            args,
-            0,
-            global_step,
-            accelerator.device,
-            vae,
-            tokenizer_or_list,
-            text_encoder_or_list,
-            unet,
-            prompt_replacement,
-        )
-
         # training loop
         for epoch in range(num_train_epochs):
             accelerator.print(f"\nepoch {epoch+1}/{num_train_epochs}")
@@ -592,19 +595,17 @@ class TextualInversionTrainer:
                     loss = loss * loss_weights
 
                     if args.min_snr_gamma:
-                        loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization)
+                        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)
-                    if args.debiased_estimation_loss:
-                        loss = apply_debiased_estimation(loss, timesteps, noise_scheduler)
 
                     loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし
 
                     accelerator.backward(loss)
                     if accelerator.sync_gradients and args.max_grad_norm != 0.0:
-                        params_to_clip = accelerator.unwrap_model(text_encoder).get_input_embeddings().parameters()
+                        params_to_clip = text_encoder.get_input_embeddings().parameters()
                         accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
 
                     optimizer.step()
@@ -616,11 +617,9 @@ class TextualInversionTrainer:
                         for text_encoder, orig_embeds_params, index_no_updates in zip(
                             text_encoders, orig_embeds_params_list, index_no_updates_list
                         ):
-                            # if full_fp16/bf16, input_embeddings_weight is fp16/bf16, orig_embeds_params is fp32
-                            input_embeddings_weight = accelerator.unwrap_model(text_encoder).get_input_embeddings().weight
-                            input_embeddings_weight[index_no_updates] = orig_embeds_params.to(input_embeddings_weight.dtype)[
+                            accelerator.unwrap_model(text_encoder).get_input_embeddings().weight[
                                 index_no_updates
-                            ]
+                            ] = orig_embeds_params[index_no_updates]
 
                 # Checks if the accelerator has performed an optimization step behind the scenes
                 if accelerator.sync_gradients:
@@ -728,24 +727,24 @@ class TextualInversionTrainer:
         is_main_process = accelerator.is_main_process
         if is_main_process:
             text_encoder = accelerator.unwrap_model(text_encoder)
-            updated_embs = text_encoder.get_input_embeddings().weight[token_ids].data.detach().clone()
 
         accelerator.end_training()
 
         if args.save_state and is_main_process:
             train_util.save_state_on_train_end(args, accelerator)
 
+        updated_embs = text_encoder.get_input_embeddings().weight[token_ids].data.detach().clone()
+
         if is_main_process:
             ckpt_name = train_util.get_last_ckpt_name(args, "." + args.save_model_as)
             save_model(ckpt_name, updated_embs_list, global_step, num_train_epochs, force_sync_upload=True)
 
-            logger.info("model saved.")
+            print("model saved.")
 
 
 def setup_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser()
 
-    add_logging_arguments(parser)
     train_util.add_sd_models_arguments(parser)
     train_util.add_dataset_arguments(parser, True, True, False)
     train_util.add_training_arguments(parser, True)
@@ -761,9 +760,7 @@ def setup_parser() -> argparse.ArgumentParser:
         help="format to save the model (default is .pt) / モデル保存時の形式（デフォルトはpt）",
     )
 
-    parser.add_argument(
-        "--weights", type=str, default=None, help="embedding weights to initialize / 学習するネットワークの初期重み"
-    )
+    parser.add_argument("--weights", type=str, default=None, help="embedding weights to initialize / 学習するネットワークの初期重み")
     parser.add_argument(
         "--num_vectors_per_token", type=int, default=1, help="number of vectors per token / トークンに割り当てるembeddingsの要素数"
     )
@@ -774,8 +771,13 @@ def setup_parser() -> argparse.ArgumentParser:
         help="token string used in training, must not exist in tokenizer / 学習時に使用されるトークン文字列、tokenizerに存在しない文字であること",
     )
     parser.add_argument(
-        "--init_word", type=str, default=None, help="words to initialize vector / ベクトルを初期化に使用する単語、複数可"
+        "--token_strings",
+        type=str,
+        default=None,
+        nargs="*",
+        help="token strings used in training for multiple embedding / 複数のembeddingsの個別学習時に使用されるトークン文字列",
     )
+    parser.add_argument("--init_word", type=str, default=None, help="words to initialize vector / ベクトルを初期化に使用する単語、複数可")
     parser.add_argument(
         "--use_object_template",
         action="store_true",

train_textual_inversion_XTI.py

@@ -1,16 +1,20 @@
 import importlib
 import argparse
+import gc
 import math
 import os
 import toml
 from multiprocessing import Value
 
 from tqdm import tqdm
-
 import torch
-from library.device_utils import init_ipex, clean_memory_on_device
-init_ipex()
-
+try:
+    import intel_extension_for_pytorch as ipex
+    if torch.xpu.is_available():
+        from library.ipex import ipex_init
+        ipex_init()
+except Exception:
+    pass
 from accelerate.utils import set_seed
 import diffusers
 from diffusers import DDPMScheduler
@@ -30,16 +34,9 @@ from library.custom_train_functions import (
     pyramid_noise_like,
     apply_noise_offset,
     scale_v_prediction_loss_like_noise_prediction,
-    apply_debiased_estimation,
 )
 import library.original_unet as original_unet
 from XTI_hijack import unet_forward_XTI, downblock_forward_XTI, upblock_forward_XTI
-from library.utils import setup_logging, add_logging_arguments
-
-setup_logging()
-import logging
-
-logger = logging.getLogger(__name__)
 
 imagenet_templates_small = [
     "a photo of a {}",
@@ -98,13 +95,12 @@ def train(args):
     if args.output_name is None:
         args.output_name = args.token_string
     use_template = args.use_object_template or args.use_style_template
-    setup_logging(args, reset=True)
 
     train_util.verify_training_args(args)
     train_util.prepare_dataset_args(args, True)
 
     if args.sample_every_n_steps is not None or args.sample_every_n_epochs is not None:
-        logger.warning(
+        print(
             "sample_every_n_steps and sample_every_n_epochs are not supported in this script currently / sample_every_n_stepsとsample_every_n_epochsは現在このスクリプトではサポートされていません"
         )
     assert (
@@ -119,7 +115,7 @@ def train(args):
     tokenizer = train_util.load_tokenizer(args)
 
     # acceleratorを準備する
-    logger.info("prepare accelerator")
+    print("prepare accelerator")
     accelerator = train_util.prepare_accelerator(args)
 
     # mixed precisionに対応した型を用意しておき適宜castする
@@ -132,7 +128,7 @@ def train(args):
     if args.init_word is not None:
         init_token_ids = tokenizer.encode(args.init_word, add_special_tokens=False)
         if len(init_token_ids) > 1 and len(init_token_ids) != args.num_vectors_per_token:
-            logger.warning(
+            print(
                 f"token length for init words is not same to num_vectors_per_token, init words is repeated or truncated / 初期化単語のトークン長がnum_vectors_per_tokenと合わないため、繰り返しまたは切り捨てが発生します: length {len(init_token_ids)}"
             )
     else:
@@ -146,7 +142,7 @@ def train(args):
     ), f"tokenizer has same word to token string. please use another one / 指定したargs.token_stringは既に存在します。別の単語を使ってください: {args.token_string}"
 
     token_ids = tokenizer.convert_tokens_to_ids(token_strings)
-    logger.info(f"tokens are added: {token_ids}")
+    print(f"tokens are added: {token_ids}")
     assert min(token_ids) == token_ids[0] and token_ids[-1] == token_ids[0] + len(token_ids) - 1, f"token ids is not ordered"
     assert len(tokenizer) - 1 == token_ids[-1], f"token ids is not end of tokenize: {len(tokenizer)}"
 
@@ -174,7 +170,7 @@ def train(args):
 
     tokenizer.add_tokens(token_strings_XTI)
     token_ids_XTI = tokenizer.convert_tokens_to_ids(token_strings_XTI)
-    logger.info(f"tokens are added (XTI): {token_ids_XTI}")
+    print(f"tokens are added (XTI): {token_ids_XTI}")
     # Resize the token embeddings as we are adding new special tokens to the tokenizer
     text_encoder.resize_token_embeddings(len(tokenizer))
 
@@ -183,7 +179,7 @@ def train(args):
     if init_token_ids is not None:
         for i, token_id in enumerate(token_ids_XTI):
             token_embeds[token_id] = token_embeds[init_token_ids[(i // 16) % len(init_token_ids)]]
-            # logger.info(token_id, token_embeds[token_id].mean(), token_embeds[token_id].min())
+            # print(token_id, token_embeds[token_id].mean(), token_embeds[token_id].min())
 
     # load weights
     if args.weights is not None:
@@ -191,22 +187,22 @@ def train(args):
         assert len(token_ids) == len(
             embeddings
         ), f"num_vectors_per_token is mismatch for weights / 指定した重みとnum_vectors_per_tokenの値が異なります: {len(embeddings)}"
-        # logger.info(token_ids, embeddings.size())
+        # print(token_ids, embeddings.size())
         for token_id, embedding in zip(token_ids_XTI, embeddings):
             token_embeds[token_id] = embedding
-            # logger.info(token_id, token_embeds[token_id].mean(), token_embeds[token_id].min())
-        logger.info(f"weighs loaded")
+            # print(token_id, token_embeds[token_id].mean(), token_embeds[token_id].min())
+        print(f"weighs loaded")
 
-    logger.info(f"create embeddings for {args.num_vectors_per_token} tokens, for {args.token_string}")
+    print(f"create embeddings for {args.num_vectors_per_token} tokens, for {args.token_string}")
 
     # データセットを準備する
     blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, False, False))
     if args.dataset_config is not None:
-        logger.info(f"Load dataset config from {args.dataset_config}")
+        print(f"Load dataset config from {args.dataset_config}")
         user_config = config_util.load_user_config(args.dataset_config)
         ignored = ["train_data_dir", "reg_data_dir", "in_json"]
         if any(getattr(args, attr) is not None for attr in ignored):
-            logger.info(
+            print(
                 "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
                     ", ".join(ignored)
                 )
@@ -214,14 +210,14 @@ def train(args):
     else:
         use_dreambooth_method = args.in_json is None
         if use_dreambooth_method:
-            logger.info("Use DreamBooth method.")
+            print("Use DreamBooth method.")
             user_config = {
                 "datasets": [
                     {"subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(args.train_data_dir, args.reg_data_dir)}
                 ]
             }
         else:
-            logger.info("Train with captions.")
+            print("Train with captions.")
             user_config = {
                 "datasets": [
                     {
@@ -245,7 +241,7 @@ def train(args):
 
     # make captions: tokenstring tokenstring1 tokenstring2 ...tokenstringn という文字列に書き換える超乱暴な実装
     if use_template:
-        logger.info(f"use template for training captions. is object: {args.use_object_template}")
+        print(f"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 = []
@@ -269,7 +265,7 @@ def train(args):
         train_util.debug_dataset(train_dataset_group, show_input_ids=True)
         return
     if len(train_dataset_group) == 0:
-        logger.error("No data found. Please verify arguments / 画像がありません。引数指定を確認してください")
+        print("No data found. Please verify arguments / 画像がありません。引数指定を確認してください")
         return
 
     if cache_latents:
@@ -291,7 +287,9 @@ def train(args):
         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")
-        clean_memory_on_device(accelerator.device)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
 
         accelerator.wait_for_everyone()
 
@@ -300,13 +298,13 @@ def train(args):
         text_encoder.gradient_checkpointing_enable()
 
     # 学習に必要なクラスを準備する
-    logger.info("prepare optimizer, data loader etc.")
+    print("prepare optimizer, data loader etc.")
     trainable_params = text_encoder.get_input_embeddings().parameters()
     _, _, optimizer = train_util.get_optimizer(args, trainable_params)
 
     # dataloaderを準備する
-    # DataLoaderのプロセス数：0 は persistent_workers が使えないので注意
-    n_workers = min(args.max_data_loader_n_workers, os.cpu_count())  # cpu_count or max_data_loader_n_workers
+    # 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,
@@ -321,9 +319,7 @@ def train(args):
         args.max_train_steps = args.max_train_epochs * math.ceil(
             len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps
         )
-        logger.info(
-            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)
@@ -336,8 +332,11 @@ def train(args):
         text_encoder, optimizer, train_dataloader, lr_scheduler
     )
 
+    # transform DDP after prepare
+    text_encoder, unet = train_util.transform_if_model_is_DDP(text_encoder, unet)
+
     index_no_updates = torch.arange(len(tokenizer)) < token_ids_XTI[0]
-    # logger.info(len(index_no_updates), torch.sum(index_no_updates))
+    # 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()
 
     # Freeze all parameters except for the token embeddings in text encoder
@@ -375,17 +374,15 @@ def train(args):
 
     # 学習する
     total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
-    logger.info("running training / 学習開始")
-    logger.info(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
-    logger.info(f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
-    logger.info(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
-    logger.info(f"  num epochs / epoch数: {num_train_epochs}")
-    logger.info(f"  batch size per device / バッチサイズ: {args.train_batch_size}")
-    logger.info(
-        f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}"
-    )
-    logger.info(f"  gradient ccumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
-    logger.info(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
@@ -399,21 +396,16 @@ def train(args):
 
     if accelerator.is_main_process:
         init_kwargs = {}
-        if args.wandb_run_name:
-            init_kwargs["wandb"] = {"name": args.wandb_run_name}
         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, init_kwargs=init_kwargs)
 
     # function for saving/removing
     def save_model(ckpt_name, embs, 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)
 
-        logger.info("")
-        logger.info(f"saving checkpoint: {ckpt_file}")
+        print(f"\nsaving checkpoint: {ckpt_file}")
         save_weights(ckpt_file, embs, save_dtype)
         if args.huggingface_repo_id is not None:
             huggingface_util.upload(args, ckpt_file, "/" + ckpt_name, force_sync_upload=force_sync_upload)
@@ -421,13 +413,12 @@ def train(args):
     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):
-            logger.info(f"removing old checkpoint: {old_ckpt_file}")
+            print(f"removing old checkpoint: {old_ckpt_file}")
             os.remove(old_ckpt_file)
 
     # training loop
     for epoch in range(num_train_epochs):
-        logger.info("")
-        logger.info(f"epoch {epoch+1}/{num_train_epochs}")
+        print(f"\nepoch {epoch+1}/{num_train_epochs}")
         current_epoch.value = epoch + 1
 
         text_encoder.train()
@@ -477,11 +468,9 @@ def train(args):
 
                 loss = loss * loss_weights
                 if args.min_snr_gamma:
-                    loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization)
+                    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.debiased_estimation_loss:
-                    loss = apply_debiased_estimation(loss, timesteps, noise_scheduler)
 
                 loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし
 
@@ -597,7 +586,7 @@ def train(args):
         ckpt_name = train_util.get_last_ckpt_name(args, "." + args.save_model_as)
         save_model(ckpt_name, updated_embs, global_step, num_train_epochs, force_sync_upload=True)
 
-        logger.info("model saved.")
+        print("model saved.")
 
 
 def save_weights(file, updated_embs, save_dtype):
@@ -658,7 +647,6 @@ def load_weights(file):
 def setup_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser()
 
-    add_logging_arguments(parser)
     train_util.add_sd_models_arguments(parser)
     train_util.add_dataset_arguments(parser, True, True, False)
     train_util.add_training_arguments(parser, True)
@@ -674,9 +662,7 @@ def setup_parser() -> argparse.ArgumentParser:
         help="format to save the model (default is .pt) / モデル保存時の形式（デフォルトはpt）",
     )
 
-    parser.add_argument(
-        "--weights", type=str, default=None, help="embedding weights to initialize / 学習するネットワークの初期重み"
-    )
+    parser.add_argument("--weights", type=str, default=None, help="embedding weights to initialize / 学習するネットワークの初期重み")
     parser.add_argument(
         "--num_vectors_per_token", type=int, default=1, help="number of vectors per token / トークンに割り当てるembeddingsの要素数"
     )
@@ -686,9 +672,7 @@ def setup_parser() -> argparse.ArgumentParser:
         default=None,
         help="token string used in training, must not exist in tokenizer / 学習時に使用されるトークン文字列、tokenizerに存在しない文字であること",
     )
-    parser.add_argument(
-        "--init_word", type=str, default=None, help="words to initialize vector / ベクトルを初期化に使用する単語、複数可"
-    )
+    parser.add_argument("--init_word", type=str, default=None, help="words to initialize vector / ベクトルを初期化に使用する単語、複数可")
     parser.add_argument(
         "--use_object_template",
         action="store_true",