diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index d35fe39250..95b3808df8 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -22,7 +22,7 @@ jobs: matrix: os: [ubuntu-latest] python-version: ["3.10"] # Python versions to test - pytorch-version: ["2.4.0", "2.6.0"] # PyTorch versions to test + pytorch-version: ["2.6.0", "2.8.0"] # PyTorch versions to test (diffusers >= 0.40 requires torch >= 2.6) steps: - uses: actions/checkout@v4 @@ -46,6 +46,12 @@ jobs: pip install dadaptation==3.2 torch==${{ matrix.pytorch-version }} torchvision pytest==8.3.4 pip install -r requirements.txt + - name: Check that the pinned PyTorch version is still installed + run: | + # If a requirement needs a newer torch, pip silently upgrades torch and leaves torchvision behind, + # which breaks every import ("operator torchvision::nms does not exist"). Fail loudly instead. + python -c "import torch, torchvision; print(torch.__version__, torchvision.__version__); assert torch.__version__.startswith('${{ matrix.pytorch-version }}'), 'torch was upgraded by requirements.txt'" + - name: Test with pytest run: pytest # See pytest.ini for configuration diff --git a/.gitignore b/.gitignore index 701cc20a10..5744a6b450 100644 --- a/.gitignore +++ b/.gitignore @@ -21,3 +21,7 @@ tests/test_data/ tests/test_output_*/ tests/downloaded_data/ tests/mask_viz/ + +# local-only regression tests (local model paths, recorded references) +tests/local/models.toml +tests/local/references/ diff --git a/README-ja.md b/README-ja.md index 1133f68abc..2128b98f75 100644 --- a/README-ja.md +++ b/README-ja.md @@ -50,7 +50,20 @@ Stable Diffusion等の画像生成モデルの学習、モデルによる画像 ### 更新履歴 -- **次のリリースに含まれる予定の更新:** 次のリリースに含まれる予定の主な変更点は以下の通りです。リリース前の変更点は予告なく変更される可能性があります。 +- **Version 0.12.0 (2026-09-24):** + - Windows on ARM64（NVIDIA RTX Spark PC など）に対応しました。[PR #2430](https://github.com/kohya-ss/sd-scripts/pull/2430)、[PR #2431](https://github.com/kohya-ss/sd-scripts/pull/2431)、[PR #2433](https://github.com/kohya-ss/sd-scripts/pull/2433) + - `opencv-python` がオプションになり（未インストール時は Pillow/NumPy による代替実装を使用）、`requirements.txt` が Windows ARM64 用 wheel のあるパッケージを自動的に選択するようになりました。詳細は[OpenCVなしでのインストール／Windows on ARM64について](#opencvなしでのインストールwindows-on-arm64について)をご覧ください。 + - `requirements.txt` の `transformers`、`schedulefree`、`safetensors` を Windows ARM64 用 wheel が提供されているバージョンに更新しました。 + - `requirements.txt` の依存関係を更新しました: `transformers` 4.57.6 → 5.5.4、`diffusers` 0.32.1 → 0.40.0、`accelerate` 1.6.0 → 1.15.0、`huggingface-hub` 0.34.3 → 1.32.0。[PR #2436](https://github.com/kohya-ss/sd-scripts/pull/2436) + - 主にセキュリティ上の保守のための更新です（`transformers` 4.x 系と `diffusers` 0.38 未満には修正が提供されなくなっています）。以前のバージョンのライブラリでもこのリリースは動作しますので、すぐに更新する必要はありませんが、なるべく早めに `pip install --upgrade -r requirements.txt` を実行することをお勧めします。 + - `diffusers` 0.40 は PyTorch 2.6 以降を必要とします（sd-scripts はすでに PyTorch 2.6.0 以降を必要としています）。CI は PyTorch 2.6.0 と 2.8.0 でテストするようになりました。 + - `transformers` 5.x の `CLIPTokenizer` は、オリジナルの CLIP トークナイザが行っていた `ftfy` によるテキスト正規化（曲がった引用符の直線化、全角文字の半角化など）を行わなくなりました。sd-scripts 側で同じ正規化を行うようにしたため、トークナイズ結果は従来と変わりません。 + - `tests/local` のローカル回帰テストにより、Text Encoder の出力、VAE の出力、ノイズスケジューラが従来のバージョンと同一であることを確認しています。 + - `diffusers` 0.40 は `.to(dtype)` のたびに "There are modules in AutoencoderKL that should be kept in float32: [] ..." という誤った警告を出します（リストが空でも警告する diffusers 側のバグ）。sd-scripts ではリストが空の場合にこの警告を抑制しています。 + - `transformers` 5.6 以降に対応し、`requirements.txt` の `transformers` を 5.17.0 に更新しました。[PR #2437](https://github.com/kohya-ss/sd-scripts/pull/2437) + - `transformers` 5.6 で `CLIPTextModel` の内部構造が変わりました（`text_model` サブモジュールが無くなりました）。sd-scripts ではモデルをラップすることで、チェックポイントのキー、Text Encoder の LoRA の重み名（`lora_te_text_model_...`）、`text_encoder.text_model.*` によるアクセスが従来と変わらないようにしています。`transformers` 5.6 未満ではラップは行われず、何も変わりません。 + - また `transformers` 5.6 では T5（FLUX.1 / SD3 の T5-XXL、HunyuanImage の byT5）の attention の実装が SDPA に切り替わり、高速化・省メモリ化が期待できます。T5 の bf16/fp16 出力は従来のバージョンからごくわずかに変わります（T5-XXL でコサイン類似度 ≈ 0.998、fp32 に対する精度は同等）。生成画像や学習結果の細部が変わる可能性がありますが、以前のバージョンでキャッシュした Text Encoder 出力もそのまま使えます。他の Text Encoder の出力は同一です。 + - 上記と同様に、以前のバージョンの `transformers` でも動作しますが、`pip install --upgrade -r requirements.txt` での更新をお勧めします。 - SD1.x / SD2.x / SDXL の学習向けに、OFTv2 と BOFT のネットワークモジュール（`networks.oft_v2`、`networks.boft`）を追加しました。[PR #2357](https://github.com/kohya-ss/sd-scripts/pull/2357) - PEFT の実装に準拠した直交変換系のアダプタです。PEFT 形式の重みも読み込めます。umisetokikaze 氏に感謝します。 - これらのモジュールでは `--network_dim` はブロックサイズを意味します。詳細は[ドキュメント](./docs/train_network_oft_boft.md)をご覧ください。 @@ -58,6 +71,9 @@ Stable Diffusion等の画像生成モデルの学習、モデルによる画像 - データセットのサブセットごとに、timestep のサンプリング分布を低ノイズ側または高ノイズ側へ偏らせることができます。詳細は[ドキュメント](./docs/timestep_sampling_offset.md)をご覧ください。 - `--show_timesteps` 使用時に offset を適用した timestep の分布を確認できる `--show_timesteps_offset` を追加しました。[PR #2410](https://github.com/kohya-ss/sd-scripts/pull/2410) - `shift` / `flux_shift` の timestep sampling における offset の挙動もドキュメントに記載しました。 + - SD1.x / SD2.x 用の古い画像生成スクリプト `gen_img_diffusers.py` を削除しました。しばらく前から動作しておらず（リファクタリングで削除された関数に依存していました）、実験的な CLIP / VGG16 guidance を除くすべての機能は `gen_img.py` でサポートされています。代わりに `gen_img.py` をお使いください（[gen_img_README-ja.md](./docs/gen_img_README-ja.md) をご覧ください）。ファイルは以前のリリースから取得できます。[PR #2439](https://github.com/kohya-ss/sd-scripts/pull/2439) + - 学習中のサンプル画像生成の `--sample_sampler`、および `gen_img.py` / `sdxl_gen_img.py` の `--sampler` で、`dpmsolver` と `dpmsingle` を指定すると最近のバージョンの `diffusers` でエラーになる問題を修正しました。[PR #2438](https://github.com/kohya-ss/sd-scripts/pull/2438) + - `lms` / `k_lms` サンプラーには `requirements.txt` に含まれない `scipy` パッケージが必要です。`scipy` が未インストールの場合、（最初のサンプル生成時ではなく）起動時に分かりやすいエラーを表示するようにしました。使用する場合は `pip install scipy` を実行してください。 - **Version 0.11.1 (2026-06-16):** - Anima LoRA／LLLite学習でtorch.compileサポートを追加しました。[PR #2379](https://github.com/kohya-ss/sd-scripts/pull/2379) @@ -255,6 +271,22 @@ PyTorchは環境によってバージョンが異なるため、requirements.txt RTX 50シリーズGPUの場合、PyTorch 2.8.0とCUDA 12.8/12.9を使用してください。`requirements.txt`はこのバージョンでも動作します。 +### OpenCVなしでのインストール／Windows on ARM64について + +`opencv-python` は `requirements.txt` に含まれていますが、学習・データセット処理パイプラインが利用している OpenCV 機能は限定的です（主に `cv2.resize`、`cv2.cvtColor`、およびデバッグ用の `cv2.imshow`）。`opencv-python` がインストールされていない場合、Pillow と NumPy による軽量な代替実装（`library/_cv2_stub`）が自動的に `cv2` として登録されるため、既存のスクリプトはそのまま動作します。OpenCV の大きなインストールを避けたい場合は、requirements のインストール後にアンインストールしてください： + +```bash +pip uninstall opencv-python +``` + +Windows on ARM64（たとえば NVIDIA RTX Spark PC など）では `opencv-python` のビルド済み wheel が提供されていないため、`requirements.txt` の環境マーカーにより自動的にスキップされます。通常どおり `pip install --upgrade -r requirements.txt` でインストールできます。同じ理由で、このプラットフォームでは `tensorboard` の代わりに `tensorboardX` がインストールされます（TensorBoard 2.x が依存する `grpcio` に Windows ARM64 用 wheel がないため）。`--log_with tensorboard` でのログ出力は `tensorboardX` 経由でそのまま動作します。ログの閲覧は別のマシンの TensorBoard で行ってください。 + +以下に注意してください： + +- OpenCV を含むデフォルトのインストールが推奨される経路です。代替実装は、データセット処理がデフォルトで使う `INTER_AREA` と `INTER_LINEAR` のリサイズを NumPy で OpenCV と同じ計算で再現しているため、学習結果は丸め誤差の範囲で一致しますが、OpenCV より低速です（2400 万画素の画像 1 枚あたり 0.1 秒程度）。`INTER_CUBIC` / `INTER_LANCZOS4` は Pillow を経由するため、わずかに結果が異なります。 +- 次のツールは実際の `opencv-python` を必要とし、未インストール時は明確なメッセージで終了します：`tools/canny.py`、`tools/detect_face_rotate.py`、および `gen_img.py` / `sdxl_gen_img.py` の ControlNet `canny` プリプロセッサ。 +- データセット確認時の `cv2.imshow` は、OpenCV が無い場合 Pillow 標準のビューア（`PIL.Image.show`）で表示され、`cv2.waitKey` はターミナルでの `input()` 待ちに置き換わります（1枚ずつ確認できます）。 + ### xformersのインストール（オプション） xformersをインストールするには、仮想環境を有効にした状態で以下のコマンドを実行してください。 diff --git a/README.md b/README.md index bc9e19e471..7b48be4e02 100644 --- a/README.md +++ b/README.md @@ -47,7 +47,20 @@ If you find this project helpful, please consider supporting its development via ### Change History -- **Changes planned for the next release:** The following are the main changes planned for the next release. Please note that these changes may be subject to change without notice before the release. +- **Version 0.12.0 (2026-09-24):** + - Added support for Windows on ARM64 (e.g. NVIDIA RTX Spark PCs). [PR #2430](https://github.com/kohya-ss/sd-scripts/pull/2430), [PR #2431](https://github.com/kohya-ss/sd-scripts/pull/2431), [PR #2433](https://github.com/kohya-ss/sd-scripts/pull/2433) + - `opencv-python` is now optional (a Pillow/NumPy fallback is used when it is missing), and `requirements.txt` selects the packages that have Windows ARM64 wheels automatically. For details, please refer to [Installing without OpenCV / Windows on ARM64](#installing-without-opencv--windows-on-arm64). + - `transformers`, `schedulefree` and `safetensors` in `requirements.txt` have been updated to versions that provide Windows ARM64 wheels. + - Updated the dependencies in `requirements.txt`: `transformers` 4.57.6 -> 5.5.4, `diffusers` 0.32.1 -> 0.40.0, `accelerate` 1.6.0 -> 1.15.0, `huggingface-hub` 0.34.3 -> 1.32.0. [PR #2436](https://github.com/kohya-ss/sd-scripts/pull/2436) + - This is mainly a security maintenance update (the 4.x line of `transformers` and `diffusers` < 0.38 no longer receive fixes). The previous versions of the libraries still work with this release, so you do not have to update them immediately, but it is recommended to run `pip install --upgrade -r requirements.txt` at your earliest convenience. + - `diffusers` 0.40 requires PyTorch 2.6 or later (PyTorch 2.6.0 or later has been the requirement of sd-scripts already). CI now tests with PyTorch 2.6.0 and 2.8.0. + - In `transformers` 5.x, `CLIPTokenizer` no longer applies the `ftfy` text normalization of the original CLIP tokenizer (straightening curly quotes, converting full-width characters, etc.). sd-scripts now applies it itself, so tokenization is unchanged from previous versions. + - Text encoder outputs, VAE outputs and the noise schedulers were verified to be identical to the previous versions with the local regression tests in `tests/local`. + - `diffusers` 0.40 prints a spurious warning "There are modules in AutoencoderKL that should be kept in float32: [] ..." on every `.to(dtype)` call (a bug in diffusers: the check fires even when the list is empty). sd-scripts suppresses this warning when the list is empty. + - Added support for `transformers` 5.6 and later, and updated `transformers` in `requirements.txt` to 5.17.0. [PR #2437](https://github.com/kohya-ss/sd-scripts/pull/2437) + - `transformers` 5.6 changed the internal structure of `CLIPTextModel` (the `text_model` submodule was removed). sd-scripts now wraps the model so that the checkpoint keys, the LoRA weight names of the text encoders (`lora_te_text_model_...`) and `text_encoder.text_model.*` access stay the same as before. Nothing changes for `transformers` < 5.6, where the wrapper is not applied. + - `transformers` 5.6 also switched the attention implementation of T5 (T5-XXL of FLUX.1 / SD3, byT5 of HunyuanImage) to SDPA, which should be faster and use less memory. The bf16/fp16 outputs of T5 differ very slightly from previous versions (cosine similarity ≈ 0.998 for T5-XXL, the accuracy against fp32 is the same). This may change generated images or trained weights in minor details, and cached text encoder outputs from previous versions are still usable. The outputs of the other text encoders are identical. + - As above, the previous versions of `transformers` still work, but updating with `pip install --upgrade -r requirements.txt` is recommended. - Added OFTv2 and BOFT network modules (`networks.oft_v2`, `networks.boft`) for SD1.x / SD2.x / SDXL training. [PR #2357](https://github.com/kohya-ss/sd-scripts/pull/2357) - Orthogonal fine-tuning adapters following the PEFT implementation. Weights in PEFT format can also be loaded. Thanks to umisetokikaze. - Note that `--network_dim` means the block size for these modules. For details, please refer to the [documentation](./docs/train_network_oft_boft.md). @@ -55,6 +68,9 @@ If you find this project helpful, please consider supporting its development via - Shifts the timestep sampling distribution of each dataset subset toward lower- or higher-noise timesteps. For details, please refer to the [documentation](./docs/timestep_sampling_offset.md). - Added `--show_timesteps_offset` to preview the timestep distribution with the offset applied when using `--show_timesteps`. [PR #2410](https://github.com/kohya-ss/sd-scripts/pull/2410) - The documentation also describes how the offset behaves with `shift` / `flux_shift` timestep sampling. + - Removed `gen_img_diffusers.py`, the old image generation script for SD1.x / SD2.x. It had not worked for a while (it depended on a function removed by the refactoring) and `gen_img.py` supports everything it did except the experimental CLIP / VGG16 guidance. Please use `gen_img.py` instead (see [gen_img_README.md](./docs/gen_img_README.md)). The file is still available in the previous releases. [PR #2439](https://github.com/kohya-ss/sd-scripts/pull/2439) + - Fixed the `dpmsolver` and `dpmsingle` samplers of `--sample_sampler` (sample image generation during training) and `--sampler` of `gen_img.py` / `sdxl_gen_img.py`, which failed with an error on recent versions of `diffusers`. [PR #2438](https://github.com/kohya-ss/sd-scripts/pull/2438) + - The `lms` / `k_lms` samplers require the `scipy` package, which is not included in `requirements.txt`. A clear error message is now shown at startup (instead of at the first sample generation) when `scipy` is missing. Please run `pip install scipy` to use them. - **Version 0.11.1 (2026-06-16):** - Added support for torch.compile in Anima LoRA/LLLite training. [PR #2379](https://github.com/kohya-ss/sd-scripts/pull/2379) @@ -224,6 +240,22 @@ The scripts are tested with PyTorch 2.6.0. PyTorch 2.6.0 or later is required. For RTX 50 series GPUs, PyTorch 2.8.0 with CUDA 12.8/12.9 should be used. `requirements.txt` will work with this version. +### Installing without OpenCV / Windows on ARM64 + +`opencv-python` is listed in `requirements.txt`, but the core training / dataset pipeline only uses a small subset of OpenCV (mainly `cv2.resize`, `cv2.cvtColor`, and a debug-only `cv2.imshow`). When `opencv-python` is not available, a lightweight Pillow/NumPy fallback under `library/_cv2_stub` is automatically registered as `cv2`, so existing scripts continue to work. If you would rather avoid the large OpenCV install, simply uninstall it after installing the requirements: + +```bash +pip uninstall opencv-python +``` + +On Windows on ARM64 (e.g. NVIDIA RTX Spark PCs), `opencv-python` has no prebuilt wheel, so `requirements.txt` skips it automatically through an environment marker and the usual `pip install --upgrade -r requirements.txt` works as is. For the same reason `tensorboardX` is installed instead of `tensorboard` on that platform (TensorBoard 2.x depends on `grpcio`, which has no Windows ARM64 wheel). Logging with `--log_with tensorboard` works unchanged through `tensorboardX`; view the logs with TensorBoard on another machine. + +Note that: + +- The default install with OpenCV remains the recommended path. The fallback reproduces OpenCV's `INTER_AREA` and `INTER_LINEAR` resizing (the modes the dataset pipeline uses by default) in NumPy, so training results match up to rounding, but it is slower than OpenCV (roughly 0.1 s per 24-megapixel image). `INTER_CUBIC` / `INTER_LANCZOS4` go through Pillow and differ slightly. +- The following tools still require real `opencv-python` and will exit with a clear message when it is missing: `tools/canny.py`, `tools/detect_face_rotate.py`, and the ControlNet `canny` preprocessor used by `gen_img.py` / `sdxl_gen_img.py`. +- Debug-only features such as `cv2.imshow` during dataset inspection fall back to Pillow's default image viewer (`PIL.Image.show`), and `cv2.waitKey` blocks on `input()` in the terminal so you can page through images one at a time. + ### xformers installation (optional) To install xformers, run the following command in your activated virtual environment: diff --git a/docs/train_network_README-ja.md b/docs/train_network_README-ja.md index 55c80c4b0a..32a56bd198 100644 --- a/docs/train_network_README-ja.md +++ b/docs/train_network_README-ja.md @@ -339,7 +339,7 @@ python networks\svd_merge_lora.py ## 当リポジトリ内の画像生成スクリプトで生成する -gen_img_diffusers.pyに、--network_module、--network_weightsの各オプションを追加してください。意味は学習時と同様です。 +gen_img.pyに、--network_module、--network_weightsの各オプションを追加してください。意味は学習時と同様です。詳細は[画像生成スクリプトのドキュメント](gen_img_README-ja.md)を参照してください。 --network_mulオプションで0~1.0の数値を指定すると、LoRAの適用率を変えられます。 diff --git a/docs/train_network_README-zh.md b/docs/train_network_README-zh.md index 830014f72d..db7cae84c4 100644 --- a/docs/train_network_README-zh.md +++ b/docs/train_network_README-zh.md @@ -317,7 +317,7 @@ python networks\svd_merge_lora.py ## 在此存储库中生成图像的脚本中 -请在`gen_img_diffusers.py`中添加`--network_module`和`--network_weights`选项。其含义与训练时相同。 +请在`gen_img.py`中添加`--network_module`和`--network_weights`选项。其含义与训练时相同。详情请参考[图像生成脚本的文档](gen_img_README.md)。 通过`--network_mul`选项，可以指定0~1.0的数字来改变LoRA的应用率。 diff --git a/docs/train_ti_README-ja.md b/docs/train_ti_README-ja.md index 86f45a5dcf..7558201108 100644 --- a/docs/train_ti_README-ja.md +++ b/docs/train_ti_README-ja.md @@ -101,5 +101,5 @@ tokenizerがすでに持っている単語（一般的な単語）は使用で ## 当リポジトリ内の画像生成スクリプトで生成する -gen_img_diffusers.pyに、``--textual_inversion_embeddings`` オプションで学習したembeddingsファイルを指定してください（複数可）。プロンプトでembeddingsファイルのファイル名（拡張子を除く）を使うと、そのembeddingsが適用されます。 +gen_img.pyに、``--textual_inversion_embeddings`` オプションで学習したembeddingsファイルを指定してください（複数可）。プロンプトでembeddingsファイルのファイル名（拡張子を除く）を使うと、そのembeddingsが適用されます。 diff --git a/finetune/prepare_buckets_latents.py b/finetune/prepare_buckets_latents.py index 9f559b6957..26799c42fd 100644 --- a/finetune/prepare_buckets_latents.py +++ b/finetune/prepare_buckets_latents.py @@ -8,7 +8,6 @@ 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 diff --git a/gen_img.py b/gen_img.py index 645c9085fc..f343cfed97 100644 --- a/gen_img.py +++ b/gen_img.py @@ -45,7 +45,8 @@ ) from einops import rearrange from tqdm import tqdm -from transformers import CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection, CLIPImageProcessor +from transformers import CLIPTextModel, CLIPVisionModelWithProjection, CLIPImageProcessor +from library.clip_tokenizer import CLIPTokenizer, as_legacy_clip_tokenizer # transformers.CLIPTokenizer with the legacy (ftfy) text normalization from accelerate import init_empty_weights import PIL from PIL import Image @@ -65,6 +66,7 @@ from library.custom_train_functions import pyramid_noise_like from networks.control_net_lllite import ControlNetLLLite from library.utils import GradualLatent, EulerAncestralDiscreteSchedulerGL +from library.sampling import check_sampler_requirements from library.utils import setup_logging, add_logging_arguments setup_logging() @@ -1610,6 +1612,8 @@ def __call__(self, *args, **kwargs): def main(args): + check_sampler_requirements(args.sampler) # fail before loading the models if the sampler needs a missing package + if args.fp16: dtype = torch.float16 elif args.bf16: @@ -1665,7 +1669,7 @@ def main(args): text_encoder = loading_pipe.text_encoder vae = loading_pipe.vae unet = loading_pipe.unet - tokenizer = loading_pipe.tokenizer + tokenizer = as_legacy_clip_tokenizer(loading_pipe.tokenizer) del loading_pipe # Diffusers U-Net to original U-Net @@ -1734,6 +1738,9 @@ def main(args): elif args.sampler == "dpmsolver" or args.sampler == "dpmsolver++": scheduler_cls = DPMSolverMultistepScheduler sched_init_args["algorithm_type"] = args.sampler + if args.sampler == "dpmsolver": + # diffusers rejects the default final_sigmas_type="zero" for the (non-++) dpmsolver algorithm + sched_init_args["final_sigmas_type"] = "sigma_min" scheduler_module = diffusers.schedulers.scheduling_dpmsolver_multistep has_clip_sample = False elif args.sampler == "dpmsingle": diff --git a/gen_img_diffusers.py b/gen_img_diffusers.py deleted file mode 100644 index adf77c8c56..0000000000 --- a/gen_img_diffusers.py +++ /dev/null @@ -1,4027 +0,0 @@ -""" -VGG( - (features): Sequential( - (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (1): ReLU(inplace=True) - (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (3): ReLU(inplace=True) - (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) - (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (6): ReLU(inplace=True) - (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (8): ReLU(inplace=True) - (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) - (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (11): ReLU(inplace=True) - (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (13): ReLU(inplace=True) - (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (15): ReLU(inplace=True) - (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) - (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (18): ReLU(inplace=True) - (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (20): ReLU(inplace=True) - (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (22): ReLU(inplace=True) - (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) - (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (25): ReLU(inplace=True) - (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (27): ReLU(inplace=True) - (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (29): ReLU(inplace=True) - (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) - ) - (avgpool): AdaptiveAvgPool2d(output_size=(7, 7)) - (classifier): Sequential( - (0): Linear(in_features=25088, out_features=4096, bias=True) - (1): ReLU(inplace=True) - (2): Dropout(p=0.5, inplace=False) - (3): Linear(in_features=4096, out_features=4096, bias=True) - (4): ReLU(inplace=True) - (5): Dropout(p=0.5, inplace=False) - (6): Linear(in_features=4096, out_features=1000, bias=True) - ) -) -""" - -import itertools -import json -from typing import Any, Dict, List, NamedTuple, Optional, Tuple, Union, Callable -import glob -import importlib -import inspect -import time -import zipfile -from diffusers.utils import deprecate -from diffusers.configuration_utils import FrozenDict -import argparse -import math -import os -import random -import re - -import diffusers -import numpy as np - -import torch -from library.device_utils import init_ipex, clean_memory, get_preferred_device -init_ipex() - -import torchvision -from diffusers import ( - AutoencoderKL, - DDPMScheduler, - EulerAncestralDiscreteScheduler, - DPMSolverMultistepScheduler, - DPMSolverSinglestepScheduler, - LMSDiscreteScheduler, - PNDMScheduler, - DDIMScheduler, - EulerDiscreteScheduler, - HeunDiscreteScheduler, - KDPM2DiscreteScheduler, - KDPM2AncestralDiscreteScheduler, - # UNet2DConditionModel, - StableDiffusionPipeline, -) -from einops import rearrange -from tqdm import tqdm -from torchvision import transforms -from transformers import CLIPTextModel, CLIPTokenizer, CLIPModel, CLIPTextConfig -import PIL -from PIL import Image -from PIL.PngImagePlugin import PngInfo - -import library.model_util as model_util -import library.train_util as train_util -from networks.lora import LoRANetwork -import tools.original_control_net as original_control_net -from tools.original_control_net import ControlNetInfo -from library.original_unet import UNet2DConditionModel, InferUNet2DConditionModel, SampleOutput -from library.original_unet import FlashAttentionFunction -from library.utils import GradualLatent, EulerAncestralDiscreteSchedulerGL -from library.utils import setup_logging, add_logging_arguments - -setup_logging() -import logging - -logger = logging.getLogger(__name__) - - -# XTI hijack: monkey-patches for UNet2DConditionModel / CrossAttnDownBlock2D / CrossAttnUpBlock2D -# used when --XTI_embeddings is specified. -def unet_forward_XTI( - self, - sample: torch.FloatTensor, - timestep: Union[torch.Tensor, float, int], - encoder_hidden_states: torch.Tensor, - class_labels: Optional[torch.Tensor] = None, - return_dict: bool = True, -) -> Union[Dict, Tuple]: - default_overall_up_factor = 2**self.num_upsamplers - - forward_upsample_size = False - upsample_size = None - - if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]): - forward_upsample_size = True - - # 1. time - timesteps = timestep - timesteps = self.handle_unusual_timesteps(sample, timesteps) - - t_emb = self.time_proj(timesteps) - t_emb = t_emb.to(dtype=self.dtype) - emb = self.time_embedding(t_emb) - - # 2. pre-process - sample = self.conv_in(sample) - - # 3. down - down_block_res_samples = (sample,) - down_i = 0 - for downsample_block in self.down_blocks: - if downsample_block.has_cross_attention: - sample, res_samples = downsample_block( - hidden_states=sample, - temb=emb, - encoder_hidden_states=encoder_hidden_states[down_i : down_i + 2], - ) - down_i += 2 - else: - sample, res_samples = downsample_block(hidden_states=sample, temb=emb) - - down_block_res_samples += res_samples - - # 4. mid - sample = self.mid_block(sample, emb, encoder_hidden_states=encoder_hidden_states[6]) - - # 5. up - up_i = 7 - 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)] - - 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[up_i : up_i + 3], - upsample_size=upsample_size, - ) - up_i += 3 - else: - sample = upsample_block( - hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size - ) - - # 6. post-process - sample = self.conv_norm_out(sample) - sample = self.conv_act(sample) - sample = self.conv_out(sample) - - if not return_dict: - return (sample,) - - return SampleOutput(sample=sample) - - -def downblock_forward_XTI( - self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None, cross_attention_kwargs=None -): - output_states = () - i = 0 - - for resnet, attn in zip(self.resnets, self.attentions): - if self.training and self.gradient_checkpointing: - - def create_custom_forward(module, return_dict=None): - def custom_forward(*inputs): - if return_dict is not None: - return module(*inputs, return_dict=return_dict) - else: - return module(*inputs) - - return custom_forward - - hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb) - hidden_states = torch.utils.checkpoint.checkpoint( - create_custom_forward(attn, return_dict=False), hidden_states, encoder_hidden_states[i] - )[0] - else: - hidden_states = resnet(hidden_states, temb) - hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states[i]).sample - - output_states += (hidden_states,) - i += 1 - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - output_states += (hidden_states,) - - return hidden_states, output_states - - -def upblock_forward_XTI( - self, - hidden_states, - res_hidden_states_tuple, - temb=None, - encoder_hidden_states=None, - upsample_size=None, -): - i = 0 - for resnet, attn in zip(self.resnets, self.attentions): - 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: - - def create_custom_forward(module, return_dict=None): - def custom_forward(*inputs): - if return_dict is not None: - return module(*inputs, return_dict=return_dict) - else: - return module(*inputs) - - return custom_forward - - hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb) - hidden_states = torch.utils.checkpoint.checkpoint( - create_custom_forward(attn, return_dict=False), hidden_states, encoder_hidden_states[i] - )[0] - else: - hidden_states = resnet(hidden_states, temb) - hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states[i]).sample - - i += 1 - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, upsample_size) - - return hidden_states - - -# scheduler: -SCHEDULER_LINEAR_START = 0.00085 -SCHEDULER_LINEAR_END = 0.0120 -SCHEDULER_TIMESTEPS = 1000 -SCHEDLER_SCHEDULE = "scaled_linear" - -# その他の設定 -LATENT_CHANNELS = 4 -DOWNSAMPLING_FACTOR = 8 - -# CLIP_ID_L14_336 = "openai/clip-vit-large-patch14-336" - -# CLIP guided SD関連 -CLIP_MODEL_PATH = "laion/CLIP-ViT-B-32-laion2B-s34B-b79K" -FEATURE_EXTRACTOR_SIZE = (224, 224) -FEATURE_EXTRACTOR_IMAGE_MEAN = [0.48145466, 0.4578275, 0.40821073] -FEATURE_EXTRACTOR_IMAGE_STD = [0.26862954, 0.26130258, 0.27577711] - -VGG16_IMAGE_MEAN = [0.485, 0.456, 0.406] -VGG16_IMAGE_STD = [0.229, 0.224, 0.225] -VGG16_INPUT_RESIZE_DIV = 4 - -# CLIP特徴量の取得時にcutoutを使うか：使う場合にはソースを書き換えてください -NUM_CUTOUTS = 4 -USE_CUTOUTS = False - -# region モジュール入れ替え部 -""" -高速化のためのモジュール入れ替え -""" - - -def replace_unet_modules(unet: diffusers.models.unet_2d_condition.UNet2DConditionModel, mem_eff_attn, xformers, sdpa): - if mem_eff_attn: - logger.info("Enable memory efficient attention for U-Net") - - # これはDiffusersのU-Netではなく自前のU-Netなので置き換えなくても良い - unet.set_use_memory_efficient_attention(False, True) - elif xformers: - logger.info("Enable xformers for U-Net") - try: - import xformers.ops - except ImportError: - raise ImportError("No xformers / xformersがインストールされていないようです") - - unet.set_use_memory_efficient_attention(True, False) - elif sdpa: - logger.info("Enable SDPA for U-Net") - unet.set_use_memory_efficient_attention(False, False) - unet.set_use_sdpa(True) - - -# TODO common train_util.py -def replace_vae_modules(vae: diffusers.models.AutoencoderKL, mem_eff_attn, xformers, sdpa): - if mem_eff_attn: - replace_vae_attn_to_memory_efficient() - elif xformers: - replace_vae_attn_to_xformers() - elif sdpa: - replace_vae_attn_to_sdpa() - - -def replace_vae_attn_to_memory_efficient(): - logger.info("VAE Attention.forward has been replaced to FlashAttention (not xformers)") - flash_func = FlashAttentionFunction - - def forward_flash_attn(self, hidden_states, **kwargs): - q_bucket_size = 512 - k_bucket_size = 1024 - - residual = hidden_states - batch, channel, height, width = hidden_states.shape - - # norm - hidden_states = self.group_norm(hidden_states) - - hidden_states = hidden_states.view(batch, channel, height * width).transpose(1, 2) - - # proj to q, k, v - query_proj = self.to_q(hidden_states) - key_proj = self.to_k(hidden_states) - value_proj = self.to_v(hidden_states) - - query_proj, key_proj, value_proj = map( - lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.heads), (query_proj, key_proj, value_proj) - ) - - out = flash_func.apply(query_proj, key_proj, value_proj, None, False, q_bucket_size, k_bucket_size) - - out = rearrange(out, "b h n d -> b n (h d)") - - # compute next hidden_states - # linear proj - hidden_states = self.to_out[0](hidden_states) - # dropout - hidden_states = self.to_out[1](hidden_states) - - hidden_states = hidden_states.transpose(-1, -2).reshape(batch, channel, height, width) - - # res connect and rescale - hidden_states = (hidden_states + residual) / self.rescale_output_factor - return hidden_states - - def forward_flash_attn_0_14(self, hidden_states, **kwargs): - if not hasattr(self, "to_q"): - self.to_q = self.query - self.to_k = self.key - self.to_v = self.value - self.to_out = [self.proj_attn, torch.nn.Identity()] - self.heads = self.num_heads - return forward_flash_attn(self, hidden_states, **kwargs) - - if diffusers.__version__ < "0.15.0": - diffusers.models.attention.AttentionBlock.forward = forward_flash_attn_0_14 - else: - diffusers.models.attention_processor.Attention.forward = forward_flash_attn - - -def replace_vae_attn_to_xformers(): - logger.info("VAE: Attention.forward has been replaced to xformers") - import xformers.ops - - def forward_xformers(self, hidden_states, **kwargs): - residual = hidden_states - batch, channel, height, width = hidden_states.shape - - # norm - hidden_states = self.group_norm(hidden_states) - - hidden_states = hidden_states.view(batch, channel, height * width).transpose(1, 2) - - # proj to q, k, v - query_proj = self.to_q(hidden_states) - key_proj = self.to_k(hidden_states) - value_proj = self.to_v(hidden_states) - - query_proj, key_proj, value_proj = map( - lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.heads), (query_proj, key_proj, value_proj) - ) - - query_proj = query_proj.contiguous() - key_proj = key_proj.contiguous() - value_proj = value_proj.contiguous() - out = xformers.ops.memory_efficient_attention(query_proj, key_proj, value_proj, attn_bias=None) - - out = rearrange(out, "b h n d -> b n (h d)") - - # compute next hidden_states - # linear proj - hidden_states = self.to_out[0](hidden_states) - # dropout - hidden_states = self.to_out[1](hidden_states) - - hidden_states = hidden_states.transpose(-1, -2).reshape(batch, channel, height, width) - - # res connect and rescale - hidden_states = (hidden_states + residual) / self.rescale_output_factor - return hidden_states - - def forward_xformers_0_14(self, hidden_states, **kwargs): - if not hasattr(self, "to_q"): - self.to_q = self.query - self.to_k = self.key - self.to_v = self.value - self.to_out = [self.proj_attn, torch.nn.Identity()] - self.heads = self.num_heads - return forward_xformers(self, hidden_states, **kwargs) - - if diffusers.__version__ < "0.15.0": - diffusers.models.attention.AttentionBlock.forward = forward_xformers_0_14 - else: - diffusers.models.attention_processor.Attention.forward = forward_xformers - - -def replace_vae_attn_to_sdpa(): - logger.info("VAE: Attention.forward has been replaced to sdpa") - - def forward_sdpa(self, hidden_states, **kwargs): - residual = hidden_states - batch, channel, height, width = hidden_states.shape - - # norm - hidden_states = self.group_norm(hidden_states) - - hidden_states = hidden_states.view(batch, channel, height * width).transpose(1, 2) - - # proj to q, k, v - query_proj = self.to_q(hidden_states) - key_proj = self.to_k(hidden_states) - value_proj = self.to_v(hidden_states) - - query_proj, key_proj, value_proj = map( - lambda t: rearrange(t, "b n (h d) -> b n h d", h=self.heads), (query_proj, key_proj, value_proj) - ) - - out = torch.nn.functional.scaled_dot_product_attention( - query_proj, key_proj, value_proj, attn_mask=None, dropout_p=0.0, is_causal=False - ) - - out = rearrange(out, "b n h d -> b n (h d)") - - # compute next hidden_states - # linear proj - hidden_states = self.to_out[0](hidden_states) - # dropout - hidden_states = self.to_out[1](hidden_states) - - hidden_states = hidden_states.transpose(-1, -2).reshape(batch, channel, height, width) - - # res connect and rescale - hidden_states = (hidden_states + residual) / self.rescale_output_factor - return hidden_states - - def forward_sdpa_0_14(self, hidden_states, **kwargs): - if not hasattr(self, "to_q"): - self.to_q = self.query - self.to_k = self.key - self.to_v = self.value - self.to_out = [self.proj_attn, torch.nn.Identity()] - self.heads = self.num_heads - return forward_sdpa(self, hidden_states, **kwargs) - - if diffusers.__version__ < "0.15.0": - diffusers.models.attention.AttentionBlock.forward = forward_sdpa_0_14 - else: - diffusers.models.attention_processor.Attention.forward = forward_sdpa - - -# endregion - -# region 画像生成の本体：lpw_stable_diffusion.py （ASL）からコピーして修正 -# https://github.com/huggingface/diffusers/blob/main/examples/community/lpw_stable_diffusion.py -# Pipelineだけ独立して使えないのと機能追加するのとでコピーして修正 - - -class PipelineLike: - r""" - Pipeline for text-to-image generation using Stable Diffusion without tokens length limit, and support parsing - weighting in prompt. - This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the - library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`CLIPTextModel`]): - Frozen text-encoder. Stable Diffusion uses the text portion of - [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically - the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant. - tokenizer (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of - [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. - safety_checker ([`StableDiffusionSafetyChecker`]): - Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/CompVis/stable-diffusion-v1-4) for details. - feature_extractor ([`CLIPFeatureExtractor`]): - Model that extracts features from generated images to be used as inputs for the `safety_checker`. - """ - - def __init__( - self, - device, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: InferUNet2DConditionModel, - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], - clip_skip: int, - clip_model: CLIPModel, - clip_guidance_scale: float, - clip_image_guidance_scale: float, - vgg16_model: torchvision.models.VGG, - vgg16_guidance_scale: float, - vgg16_layer_no: int, - # safety_checker: StableDiffusionSafetyChecker, - # feature_extractor: CLIPFeatureExtractor, - ): - super().__init__() - self.device = device - self.clip_skip = clip_skip - - if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." - " `clip_sample` should be set to False in the configuration file. Please make sure to update the" - " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" - " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" - " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" - ) - deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["clip_sample"] = False - scheduler._internal_dict = FrozenDict(new_config) - - self.vae = vae - self.text_encoder = text_encoder - self.tokenizer = tokenizer - self.unet = unet - self.scheduler = scheduler - self.safety_checker = None - - # Textual Inversion - self.token_replacements = {} - - # XTI - self.token_replacements_XTI = {} - - # CLIP guidance - self.clip_guidance_scale = clip_guidance_scale - self.clip_image_guidance_scale = clip_image_guidance_scale - self.clip_model = clip_model - self.normalize = transforms.Normalize(mean=FEATURE_EXTRACTOR_IMAGE_MEAN, std=FEATURE_EXTRACTOR_IMAGE_STD) - self.make_cutouts = MakeCutouts(FEATURE_EXTRACTOR_SIZE) - - # VGG16 guidance - self.vgg16_guidance_scale = vgg16_guidance_scale - if self.vgg16_guidance_scale > 0.0: - return_layers = {f"{vgg16_layer_no}": "feat"} - self.vgg16_feat_model = torchvision.models._utils.IntermediateLayerGetter( - vgg16_model.features, return_layers=return_layers - ) - self.vgg16_normalize = transforms.Normalize(mean=VGG16_IMAGE_MEAN, std=VGG16_IMAGE_STD) - - # ControlNet - self.control_nets: List[ControlNetInfo] = [] - self.control_net_enabled = True # control_netsが空ならTrueでもFalseでもControlNetは動作しない - - self.gradual_latent: GradualLatent = None - - # Textual Inversion - def add_token_replacement(self, target_token_id, rep_token_ids): - self.token_replacements[target_token_id] = rep_token_ids - - def set_enable_control_net(self, en: bool): - self.control_net_enabled = en - - def replace_token(self, tokens, layer=None): - new_tokens = [] - for token in tokens: - if token in self.token_replacements: - replacer_ = self.token_replacements[token] - if layer: - replacer = [] - for r in replacer_: - if r in self.token_replacements_XTI: - replacer.append(self.token_replacements_XTI[r][layer]) - else: - replacer = replacer_ - new_tokens.extend(replacer) - else: - new_tokens.append(token) - return new_tokens - - def add_token_replacement_XTI(self, target_token_id, rep_token_ids): - self.token_replacements_XTI[target_token_id] = rep_token_ids - - def set_control_nets(self, ctrl_nets): - self.control_nets = ctrl_nets - - def set_gradual_latent(self, gradual_latent): - if gradual_latent is None: - logger.info("gradual_latent is disabled") - self.gradual_latent = None - else: - logger.info(f"gradual_latent is enabled: {gradual_latent}") - self.gradual_latent = gradual_latent # (ds_ratio, start_timesteps, every_n_steps, ratio_step) - - # region xformersとか使う部分：独自に書き換えるので関係なし - - def enable_xformers_memory_efficient_attention(self): - r""" - Enable memory efficient attention as implemented in xformers. - When this option is enabled, you should observe lower GPU memory usage and a potential speed up at inference - time. Speed up at training time is not guaranteed. - Warning: When Memory Efficient Attention and Sliced attention are both enabled, the Memory Efficient Attention - is used. - """ - self.unet.set_use_memory_efficient_attention_xformers(True) - - def disable_xformers_memory_efficient_attention(self): - r""" - Disable memory efficient attention as implemented in xformers. - """ - self.unet.set_use_memory_efficient_attention_xformers(False) - - def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"): - r""" - Enable sliced attention computation. - When this option is enabled, the attention module will split the input tensor in slices, to compute attention - in several steps. This is useful to save some memory in exchange for a small speed decrease. - Args: - slice_size (`str` or `int`, *optional*, defaults to `"auto"`): - When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If - a number is provided, uses as many slices as `attention_head_dim // slice_size`. In this case, - `attention_head_dim` must be a multiple of `slice_size`. - """ - if slice_size == "auto": - # half the attention head size is usually a good trade-off between - # speed and memory - slice_size = self.unet.config.attention_head_dim // 2 - self.unet.set_attention_slice(slice_size) - - def disable_attention_slicing(self): - r""" - Disable sliced attention computation. If `enable_attention_slicing` was previously invoked, this method will go - back to computing attention in one step. - """ - # set slice_size = `None` to disable `attention slicing` - self.enable_attention_slicing(None) - - def enable_sequential_cpu_offload(self): - r""" - Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet, - text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a - `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called. - """ - # accelerateが必要になるのでとりあえず省略 - raise NotImplementedError("cpu_offload is omitted.") - # if is_accelerate_available(): - # from accelerate import cpu_offload - # else: - # raise ImportError("Please install accelerate via `pip install accelerate`") - - # device = self.device - - # for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae, self.safety_checker]: - # if cpu_offloaded_model is not None: - # cpu_offload(cpu_offloaded_model, device) - - # endregion - - @torch.no_grad() - def __call__( - self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, - init_image: Union[torch.FloatTensor, PIL.Image.Image, List[PIL.Image.Image]] = None, - mask_image: Union[torch.FloatTensor, PIL.Image.Image, List[PIL.Image.Image]] = None, - height: int = 512, - width: int = 512, - num_inference_steps: int = 50, - guidance_scale: float = 7.5, - negative_scale: float = None, - strength: float = 0.8, - # num_images_per_prompt: Optional[int] = 1, - eta: float = 0.0, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.FloatTensor] = None, - max_embeddings_multiples: Optional[int] = 3, - output_type: Optional[str] = "pil", - vae_batch_size: float = None, - return_latents: bool = False, - # return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, - is_cancelled_callback: Optional[Callable[[], bool]] = None, - callback_steps: Optional[int] = 1, - img2img_noise=None, - clip_prompts=None, - clip_guide_images=None, - networks: Optional[List[LoRANetwork]] = None, - **kwargs, - ): - r""" - Function invoked when calling the pipeline for generation. - Args: - prompt (`str` or `List[str]`): - The prompt or prompts to guide the image generation. - negative_prompt (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored - if `guidance_scale` is less than `1`). - init_image (`torch.FloatTensor` or `PIL.Image.Image`): - `Image`, or tensor representing an image batch, that will be used as the starting point for the - process. - mask_image (`torch.FloatTensor` or `PIL.Image.Image`): - `Image`, or tensor representing an image batch, to mask `init_image`. White pixels in the mask will be - replaced by noise and therefore repainted, while black pixels will be preserved. If `mask_image` is a - PIL image, it will be converted to a single channel (luminance) before use. If it's a tensor, it should - contain one color channel (L) instead of 3, so the expected shape would be `(B, H, W, 1)`. - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - guidance_scale (`float`, *optional*, defaults to 7.5): - Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). - `guidance_scale` is defined as `w` of equation 2. of [Imagen - Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > - 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, - usually at the expense of lower image quality. - strength (`float`, *optional*, defaults to 0.8): - Conceptually, indicates how much to transform the reference `init_image`. Must be between 0 and 1. - `init_image` will be used as a starting point, adding more noise to it the larger the `strength`. The - number of denoising steps depends on the amount of noise initially added. When `strength` is 1, added - noise will be maximum and the denoising process will run for the full number of iterations specified in - `num_inference_steps`. A value of 1, therefore, essentially ignores `init_image`. - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - eta (`float`, *optional*, defaults to 0.0): - Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to - [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator`, *optional*): - A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation - deterministic. - latents (`torch.FloatTensor`, *optional*): - Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - tensor will ge generated by sampling using the supplied random `generator`. - max_embeddings_multiples (`int`, *optional*, defaults to `3`): - The max multiple length of prompt embeddings compared to the max output length of text encoder. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between - [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that will be called every `callback_steps` steps during inference. The function will be - called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. - is_cancelled_callback (`Callable`, *optional*): - A function that will be called every `callback_steps` steps during inference. If the function returns - `True`, the inference will be cancelled. - callback_steps (`int`, *optional*, defaults to 1): - The frequency at which the `callback` function will be called. If not specified, the callback will be - called at every step. - Returns: - `None` if cancelled by `is_cancelled_callback`, - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple. - When returning a tuple, the first element is a list with the generated images, and the second element is a - list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work" - (nsfw) content, according to the `safety_checker`. - """ - num_images_per_prompt = 1 # fixed - - if isinstance(prompt, str): - batch_size = 1 - prompt = [prompt] - elif isinstance(prompt, list): - batch_size = len(prompt) - else: - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - reginonal_network = " AND " in prompt[0] - - vae_batch_size = ( - batch_size - if vae_batch_size is None - else (int(vae_batch_size) if vae_batch_size >= 1 else max(1, int(batch_size * vae_batch_size))) - ) - - if strength < 0 or strength > 1: - raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") - - if height % 8 != 0 or width % 8 != 0: - raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") - - if (callback_steps is None) or ( - callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) - ): - raise ValueError( - f"`callback_steps` has to be a positive integer but is {callback_steps} of type" f" {type(callback_steps)}." - ) - - # get prompt text embeddings - - # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) - # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` - # corresponds to doing no classifier free guidance. - do_classifier_free_guidance = guidance_scale > 1.0 - - if not do_classifier_free_guidance and negative_scale is not None: - logger.warning(f"negative_scale is ignored if guidance scalle <= 1.0") - negative_scale = None - - # get unconditional embeddings for classifier free guidance - if negative_prompt is None: - negative_prompt = [""] * batch_size - elif isinstance(negative_prompt, str): - negative_prompt = [negative_prompt] * batch_size - if batch_size != len(negative_prompt): - raise ValueError( - f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" - f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" - " the batch size of `prompt`." - ) - - if not self.token_replacements_XTI: - text_embeddings, uncond_embeddings, prompt_tokens = get_weighted_text_embeddings( - pipe=self, - prompt=prompt, - uncond_prompt=negative_prompt if do_classifier_free_guidance else None, - max_embeddings_multiples=max_embeddings_multiples, - clip_skip=self.clip_skip, - **kwargs, - ) - - if negative_scale is not None: - _, real_uncond_embeddings, _ = get_weighted_text_embeddings( - pipe=self, - prompt=prompt, # こちらのトークン長に合わせてuncondを作るので75トークン超で必須 - uncond_prompt=[""] * batch_size, - max_embeddings_multiples=max_embeddings_multiples, - clip_skip=self.clip_skip, - **kwargs, - ) - - if self.token_replacements_XTI: - text_embeddings_concat = [] - for layer in [ - "IN01", - "IN02", - "IN04", - "IN05", - "IN07", - "IN08", - "MID", - "OUT03", - "OUT04", - "OUT05", - "OUT06", - "OUT07", - "OUT08", - "OUT09", - "OUT10", - "OUT11", - ]: - text_embeddings, uncond_embeddings, prompt_tokens = get_weighted_text_embeddings( - pipe=self, - prompt=prompt, - uncond_prompt=negative_prompt if do_classifier_free_guidance else None, - max_embeddings_multiples=max_embeddings_multiples, - clip_skip=self.clip_skip, - layer=layer, - **kwargs, - ) - if do_classifier_free_guidance: - if negative_scale is None: - text_embeddings_concat.append(torch.cat([uncond_embeddings, text_embeddings])) - else: - text_embeddings_concat.append(torch.cat([uncond_embeddings, text_embeddings, real_uncond_embeddings])) - text_embeddings = torch.stack(text_embeddings_concat) - else: - if do_classifier_free_guidance: - if negative_scale is None: - text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) - else: - text_embeddings = torch.cat([uncond_embeddings, text_embeddings, real_uncond_embeddings]) - - # CLIP guidanceで使用するembeddingsを取得する - if self.clip_guidance_scale > 0: - clip_text_input = prompt_tokens - if clip_text_input.shape[1] > self.tokenizer.model_max_length: - # TODO 75文字を超えたら警告を出す？ - logger.info(f"trim text input {clip_text_input.shape}") - clip_text_input = torch.cat( - [clip_text_input[:, : self.tokenizer.model_max_length - 1], clip_text_input[:, -1].unsqueeze(1)], dim=1 - ) - logger.info(f"trimmed {clip_text_input.shape}") - - for i, clip_prompt in enumerate(clip_prompts): - if clip_prompt is not None: # clip_promptがあれば上書きする - clip_text_input[i] = self.tokenizer( - clip_prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ).input_ids.to(self.device) - - text_embeddings_clip = self.clip_model.get_text_features(clip_text_input) - text_embeddings_clip = text_embeddings_clip / text_embeddings_clip.norm(p=2, dim=-1, keepdim=True) # prompt複数件でもOK - - if ( - self.clip_image_guidance_scale > 0 - or self.vgg16_guidance_scale > 0 - and clip_guide_images is not None - or self.control_nets - ): - if isinstance(clip_guide_images, PIL.Image.Image): - clip_guide_images = [clip_guide_images] - - if self.clip_image_guidance_scale > 0: - clip_guide_images = [preprocess_guide_image(im) for im in clip_guide_images] - clip_guide_images = torch.cat(clip_guide_images, dim=0) - - clip_guide_images = self.normalize(clip_guide_images).to(self.device).to(text_embeddings.dtype) - image_embeddings_clip = self.clip_model.get_image_features(clip_guide_images) - image_embeddings_clip = image_embeddings_clip / image_embeddings_clip.norm(p=2, dim=-1, keepdim=True) - if len(image_embeddings_clip) == 1: - image_embeddings_clip = image_embeddings_clip.repeat((batch_size, 1, 1, 1)) - elif self.vgg16_guidance_scale > 0: - size = (width // VGG16_INPUT_RESIZE_DIV, height // VGG16_INPUT_RESIZE_DIV) # とりあえず1/4に（小さいか?） - clip_guide_images = [preprocess_vgg16_guide_image(im, size) for im in clip_guide_images] - clip_guide_images = torch.cat(clip_guide_images, dim=0) - - clip_guide_images = self.vgg16_normalize(clip_guide_images).to(self.device).to(text_embeddings.dtype) - image_embeddings_vgg16 = self.vgg16_feat_model(clip_guide_images)["feat"] - if len(image_embeddings_vgg16) == 1: - image_embeddings_vgg16 = image_embeddings_vgg16.repeat((batch_size, 1, 1, 1)) - else: - # ControlNetのhintにguide imageを流用する - # 前処理はControlNet側で行う - pass - - # set timesteps - self.scheduler.set_timesteps(num_inference_steps, self.device) - - latents_dtype = text_embeddings.dtype - init_latents_orig = None - mask = None - - if init_image is None: - # get the initial random noise unless the user supplied it - - # Unlike in other pipelines, latents need to be generated in the target device - # for 1-to-1 results reproducibility with the CompVis implementation. - # However this currently doesn't work in `mps`. - latents_shape = ( - batch_size * num_images_per_prompt, - self.unet.in_channels, - height // 8, - width // 8, - ) - - if latents is None: - if self.device.type == "mps": - # randn does not exist on mps - latents = torch.randn( - latents_shape, - generator=generator, - device="cpu", - dtype=latents_dtype, - ).to(self.device) - else: - latents = torch.randn( - latents_shape, - generator=generator, - device=self.device, - dtype=latents_dtype, - ) - else: - if latents.shape != latents_shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") - latents = latents.to(self.device) - - timesteps = self.scheduler.timesteps.to(self.device) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - else: - # image to tensor - if isinstance(init_image, PIL.Image.Image): - init_image = [init_image] - if isinstance(init_image[0], PIL.Image.Image): - init_image = [preprocess_image(im) for im in init_image] - init_image = torch.cat(init_image) - if isinstance(init_image, list): - init_image = torch.stack(init_image) - - # mask image to tensor - if mask_image is not None: - if isinstance(mask_image, PIL.Image.Image): - mask_image = [mask_image] - if isinstance(mask_image[0], PIL.Image.Image): - mask_image = torch.cat([preprocess_mask(im) for im in mask_image]) # H*W, 0 for repaint - - # encode the init image into latents and scale the latents - init_image = init_image.to(device=self.device, dtype=latents_dtype) - if init_image.size()[-2:] == (height // 8, width // 8): - init_latents = init_image - else: - if vae_batch_size >= batch_size: - init_latent_dist = self.vae.encode(init_image).latent_dist - init_latents = init_latent_dist.sample(generator=generator) - else: - clean_memory() - init_latents = [] - for i in tqdm(range(0, min(batch_size, len(init_image)), vae_batch_size)): - init_latent_dist = self.vae.encode( - init_image[i : i + vae_batch_size] if vae_batch_size > 1 else init_image[i].unsqueeze(0) - ).latent_dist - init_latents.append(init_latent_dist.sample(generator=generator)) - init_latents = torch.cat(init_latents) - - init_latents = 0.18215 * init_latents - - if len(init_latents) == 1: - init_latents = init_latents.repeat((batch_size, 1, 1, 1)) - init_latents_orig = init_latents - - # preprocess mask - if mask_image is not None: - mask = mask_image.to(device=self.device, dtype=latents_dtype) - if len(mask) == 1: - mask = mask.repeat((batch_size, 1, 1, 1)) - - # check sizes - if not mask.shape == init_latents.shape: - raise ValueError("The mask and init_image should be the same size!") - - # get the original timestep using init_timestep - offset = self.scheduler.config.get("steps_offset", 0) - init_timestep = int(num_inference_steps * strength) + offset - init_timestep = min(init_timestep, num_inference_steps) - - timesteps = self.scheduler.timesteps[-init_timestep] - timesteps = torch.tensor([timesteps] * batch_size * num_images_per_prompt, device=self.device) - - # add noise to latents using the timesteps - latents = self.scheduler.add_noise(init_latents, img2img_noise, timesteps) - - t_start = max(num_inference_steps - init_timestep + offset, 0) - timesteps = self.scheduler.timesteps[t_start:].to(self.device) - - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 - # and should be between [0, 1] - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - num_latent_input = (3 if negative_scale is not None else 2) if do_classifier_free_guidance else 1 - - if self.control_nets: - guided_hints = original_control_net.get_guided_hints(self.control_nets, num_latent_input, batch_size, clip_guide_images) - - if reginonal_network: - num_sub_and_neg_prompts = len(text_embeddings) // batch_size - # last subprompt and negative prompt - text_emb_last = [] - for j in range(batch_size): - text_emb_last.append(text_embeddings[(j + 1) * num_sub_and_neg_prompts - 2]) - text_emb_last.append(text_embeddings[(j + 1) * num_sub_and_neg_prompts - 1]) - text_emb_last = torch.stack(text_emb_last) - else: - text_emb_last = text_embeddings - - enable_gradual_latent = False - if self.gradual_latent: - if not hasattr(self.scheduler, "set_gradual_latent_params"): - logger.info("gradual_latent is not supported for this scheduler. Ignoring.") - logger.info(f'{self.scheduler.__class__.__name__}') - else: - enable_gradual_latent = True - step_elapsed = 1000 - current_ratio = self.gradual_latent.ratio - - # first, we downscale the latents to the specified ratio / 最初に指定された比率にlatentsをダウンスケールする - height, width = latents.shape[-2:] - org_dtype = latents.dtype - if org_dtype == torch.bfloat16: - latents = latents.float() - latents = torch.nn.functional.interpolate( - latents, scale_factor=current_ratio, mode="bicubic", align_corners=False - ).to(org_dtype) - - # apply unsharp mask / アンシャープマスクを適用する - if self.gradual_latent.gaussian_blur_ksize: - latents = self.gradual_latent.apply_unshark_mask(latents) - - for i, t in enumerate(tqdm(timesteps)): - resized_size = None - if enable_gradual_latent: - # gradually upscale the latents / latentsを徐々にアップスケールする - if ( - t < self.gradual_latent.start_timesteps - and current_ratio < 1.0 - and step_elapsed >= self.gradual_latent.every_n_steps - ): - current_ratio = min(current_ratio + self.gradual_latent.ratio_step, 1.0) - # make divisible by 8 because size of latents must be divisible at bottom of UNet - h = int(height * current_ratio) // 8 * 8 - w = int(width * current_ratio) // 8 * 8 - resized_size = (h, w) - self.scheduler.set_gradual_latent_params(resized_size, self.gradual_latent) - step_elapsed = 0 - else: - self.scheduler.set_gradual_latent_params(None, None) - step_elapsed += 1 - - # expand the latents if we are doing classifier free guidance - latent_model_input = latents.repeat((num_latent_input, 1, 1, 1)) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - if self.control_nets and self.control_net_enabled: - noise_pred = original_control_net.call_unet_and_control_net( - i, - num_latent_input, - self.unet, - self.control_nets, - guided_hints, - i / len(timesteps), - latent_model_input, - t, - text_embeddings, - text_emb_last, - ).sample - else: - noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample - - # perform guidance - if do_classifier_free_guidance: - if negative_scale is None: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(num_latent_input) # uncond by negative prompt - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - else: - noise_pred_negative, noise_pred_text, noise_pred_uncond = noise_pred.chunk( - num_latent_input - ) # uncond is real uncond - noise_pred = ( - noise_pred_uncond - + guidance_scale * (noise_pred_text - noise_pred_uncond) - - negative_scale * (noise_pred_negative - noise_pred_uncond) - ) - - # perform clip guidance - if self.clip_guidance_scale > 0 or self.clip_image_guidance_scale > 0 or self.vgg16_guidance_scale > 0: - text_embeddings_for_guidance = ( - text_embeddings.chunk(num_latent_input)[1] if do_classifier_free_guidance else text_embeddings - ) - - if self.clip_guidance_scale > 0: - noise_pred, latents = self.cond_fn( - latents, - t, - i, - text_embeddings_for_guidance, - noise_pred, - text_embeddings_clip, - self.clip_guidance_scale, - NUM_CUTOUTS, - USE_CUTOUTS, - ) - if self.clip_image_guidance_scale > 0 and clip_guide_images is not None: - noise_pred, latents = self.cond_fn( - latents, - t, - i, - text_embeddings_for_guidance, - noise_pred, - image_embeddings_clip, - self.clip_image_guidance_scale, - NUM_CUTOUTS, - USE_CUTOUTS, - ) - if self.vgg16_guidance_scale > 0 and clip_guide_images is not None: - noise_pred, latents = self.cond_fn_vgg16( - latents, t, i, text_embeddings_for_guidance, noise_pred, image_embeddings_vgg16, self.vgg16_guidance_scale - ) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - if mask is not None: - # masking - init_latents_proper = self.scheduler.add_noise(init_latents_orig, img2img_noise, torch.tensor([t])) - latents = (init_latents_proper * mask) + (latents * (1 - mask)) - - # call the callback, if provided - if i % callback_steps == 0: - if callback is not None: - callback(i, t, latents) - if is_cancelled_callback is not None and is_cancelled_callback(): - return None - - if return_latents: - return (latents, False) - - latents = 1 / 0.18215 * latents - if vae_batch_size >= batch_size: - image = self.vae.decode(latents).sample - else: - clean_memory() - images = [] - for i in tqdm(range(0, batch_size, vae_batch_size)): - images.append( - self.vae.decode(latents[i : i + vae_batch_size] if vae_batch_size > 1 else latents[i].unsqueeze(0)).sample - ) - image = torch.cat(images) - - image = (image / 2 + 0.5).clamp(0, 1) - - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16 - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - if self.safety_checker is not None: - safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(self.device) - image, has_nsfw_concept = self.safety_checker( - images=image, - clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype), - ) - else: - has_nsfw_concept = None - - if output_type == "pil": - # image = self.numpy_to_pil(image) - image = (image * 255).round().astype("uint8") - image = [Image.fromarray(im) for im in image] - - # if not return_dict: - return (image, has_nsfw_concept) - - # return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) - - def text2img( - self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, - height: int = 512, - width: int = 512, - num_inference_steps: int = 50, - guidance_scale: float = 7.5, - num_images_per_prompt: Optional[int] = 1, - eta: float = 0.0, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.FloatTensor] = None, - max_embeddings_multiples: Optional[int] = 3, - output_type: Optional[str] = "pil", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, - callback_steps: Optional[int] = 1, - **kwargs, - ): - r""" - Function for text-to-image generation. - Args: - prompt (`str` or `List[str]`): - The prompt or prompts to guide the image generation. - negative_prompt (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored - if `guidance_scale` is less than `1`). - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - guidance_scale (`float`, *optional*, defaults to 7.5): - Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). - `guidance_scale` is defined as `w` of equation 2. of [Imagen - Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > - 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, - usually at the expense of lower image quality. - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - eta (`float`, *optional*, defaults to 0.0): - Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to - [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator`, *optional*): - A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation - deterministic. - latents (`torch.FloatTensor`, *optional*): - Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - tensor will ge generated by sampling using the supplied random `generator`. - max_embeddings_multiples (`int`, *optional*, defaults to `3`): - The max multiple length of prompt embeddings compared to the max output length of text encoder. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between - [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that will be called every `callback_steps` steps during inference. The function will be - called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. - callback_steps (`int`, *optional*, defaults to 1): - The frequency at which the `callback` function will be called. If not specified, the callback will be - called at every step. - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple. - When returning a tuple, the first element is a list with the generated images, and the second element is a - list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work" - (nsfw) content, according to the `safety_checker`. - """ - return self.__call__( - prompt=prompt, - negative_prompt=negative_prompt, - height=height, - width=width, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - num_images_per_prompt=num_images_per_prompt, - eta=eta, - generator=generator, - latents=latents, - max_embeddings_multiples=max_embeddings_multiples, - output_type=output_type, - return_dict=return_dict, - callback=callback, - callback_steps=callback_steps, - **kwargs, - ) - - def img2img( - self, - init_image: Union[torch.FloatTensor, PIL.Image.Image], - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, - strength: float = 0.8, - num_inference_steps: Optional[int] = 50, - guidance_scale: Optional[float] = 7.5, - num_images_per_prompt: Optional[int] = 1, - eta: Optional[float] = 0.0, - generator: Optional[torch.Generator] = None, - max_embeddings_multiples: Optional[int] = 3, - output_type: Optional[str] = "pil", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, - callback_steps: Optional[int] = 1, - **kwargs, - ): - r""" - Function for image-to-image generation. - Args: - init_image (`torch.FloatTensor` or `PIL.Image.Image`): - `Image`, or tensor representing an image batch, that will be used as the starting point for the - process. - prompt (`str` or `List[str]`): - The prompt or prompts to guide the image generation. - negative_prompt (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored - if `guidance_scale` is less than `1`). - strength (`float`, *optional*, defaults to 0.8): - Conceptually, indicates how much to transform the reference `init_image`. Must be between 0 and 1. - `init_image` will be used as a starting point, adding more noise to it the larger the `strength`. The - number of denoising steps depends on the amount of noise initially added. When `strength` is 1, added - noise will be maximum and the denoising process will run for the full number of iterations specified in - `num_inference_steps`. A value of 1, therefore, essentially ignores `init_image`. - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. This parameter will be modulated by `strength`. - guidance_scale (`float`, *optional*, defaults to 7.5): - Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). - `guidance_scale` is defined as `w` of equation 2. of [Imagen - Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > - 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, - usually at the expense of lower image quality. - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - eta (`float`, *optional*, defaults to 0.0): - Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to - [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator`, *optional*): - A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation - deterministic. - max_embeddings_multiples (`int`, *optional*, defaults to `3`): - The max multiple length of prompt embeddings compared to the max output length of text encoder. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between - [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that will be called every `callback_steps` steps during inference. The function will be - called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. - callback_steps (`int`, *optional*, defaults to 1): - The frequency at which the `callback` function will be called. If not specified, the callback will be - called at every step. - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple. - When returning a tuple, the first element is a list with the generated images, and the second element is a - list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work" - (nsfw) content, according to the `safety_checker`. - """ - return self.__call__( - prompt=prompt, - negative_prompt=negative_prompt, - init_image=init_image, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - strength=strength, - num_images_per_prompt=num_images_per_prompt, - eta=eta, - generator=generator, - max_embeddings_multiples=max_embeddings_multiples, - output_type=output_type, - return_dict=return_dict, - callback=callback, - callback_steps=callback_steps, - **kwargs, - ) - - def inpaint( - self, - init_image: Union[torch.FloatTensor, PIL.Image.Image], - mask_image: Union[torch.FloatTensor, PIL.Image.Image], - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, - strength: float = 0.8, - num_inference_steps: Optional[int] = 50, - guidance_scale: Optional[float] = 7.5, - num_images_per_prompt: Optional[int] = 1, - eta: Optional[float] = 0.0, - generator: Optional[torch.Generator] = None, - max_embeddings_multiples: Optional[int] = 3, - output_type: Optional[str] = "pil", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, - callback_steps: Optional[int] = 1, - **kwargs, - ): - r""" - Function for inpaint. - Args: - init_image (`torch.FloatTensor` or `PIL.Image.Image`): - `Image`, or tensor representing an image batch, that will be used as the starting point for the - process. This is the image whose masked region will be inpainted. - mask_image (`torch.FloatTensor` or `PIL.Image.Image`): - `Image`, or tensor representing an image batch, to mask `init_image`. White pixels in the mask will be - replaced by noise and therefore repainted, while black pixels will be preserved. If `mask_image` is a - PIL image, it will be converted to a single channel (luminance) before use. If it's a tensor, it should - contain one color channel (L) instead of 3, so the expected shape would be `(B, H, W, 1)`. - prompt (`str` or `List[str]`): - The prompt or prompts to guide the image generation. - negative_prompt (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored - if `guidance_scale` is less than `1`). - strength (`float`, *optional*, defaults to 0.8): - Conceptually, indicates how much to inpaint the masked area. Must be between 0 and 1. When `strength` - is 1, the denoising process will be run on the masked area for the full number of iterations specified - in `num_inference_steps`. `init_image` will be used as a reference for the masked area, adding more - noise to that region the larger the `strength`. If `strength` is 0, no inpainting will occur. - num_inference_steps (`int`, *optional*, defaults to 50): - The reference number of denoising steps. More denoising steps usually lead to a higher quality image at - the expense of slower inference. This parameter will be modulated by `strength`, as explained above. - guidance_scale (`float`, *optional*, defaults to 7.5): - Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). - `guidance_scale` is defined as `w` of equation 2. of [Imagen - Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > - 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, - usually at the expense of lower image quality. - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - eta (`float`, *optional*, defaults to 0.0): - Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to - [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator`, *optional*): - A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation - deterministic. - max_embeddings_multiples (`int`, *optional*, defaults to `3`): - The max multiple length of prompt embeddings compared to the max output length of text encoder. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between - [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that will be called every `callback_steps` steps during inference. The function will be - called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. - callback_steps (`int`, *optional*, defaults to 1): - The frequency at which the `callback` function will be called. If not specified, the callback will be - called at every step. - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple. - When returning a tuple, the first element is a list with the generated images, and the second element is a - list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work" - (nsfw) content, according to the `safety_checker`. - """ - return self.__call__( - prompt=prompt, - negative_prompt=negative_prompt, - init_image=init_image, - mask_image=mask_image, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - strength=strength, - num_images_per_prompt=num_images_per_prompt, - eta=eta, - generator=generator, - max_embeddings_multiples=max_embeddings_multiples, - output_type=output_type, - return_dict=return_dict, - callback=callback, - callback_steps=callback_steps, - **kwargs, - ) - - # CLIP guidance StableDiffusion - # copy from https://github.com/huggingface/diffusers/blob/main/examples/community/clip_guided_stable_diffusion.py - - # バッチを分解して1件ずつ処理する - def cond_fn( - self, - latents, - timestep, - index, - text_embeddings, - noise_pred_original, - guide_embeddings_clip, - clip_guidance_scale, - num_cutouts, - use_cutouts=True, - ): - if len(latents) == 1: - return self.cond_fn1( - latents, - timestep, - index, - text_embeddings, - noise_pred_original, - guide_embeddings_clip, - clip_guidance_scale, - num_cutouts, - use_cutouts, - ) - - noise_pred = [] - cond_latents = [] - for i in range(len(latents)): - lat1 = latents[i].unsqueeze(0) - tem1 = text_embeddings[i].unsqueeze(0) - npo1 = noise_pred_original[i].unsqueeze(0) - gem1 = guide_embeddings_clip[i].unsqueeze(0) - npr1, cla1 = self.cond_fn1(lat1, timestep, index, tem1, npo1, gem1, clip_guidance_scale, num_cutouts, use_cutouts) - noise_pred.append(npr1) - cond_latents.append(cla1) - - noise_pred = torch.cat(noise_pred) - cond_latents = torch.cat(cond_latents) - return noise_pred, cond_latents - - @torch.enable_grad() - def cond_fn1( - self, - latents, - timestep, - index, - text_embeddings, - noise_pred_original, - guide_embeddings_clip, - clip_guidance_scale, - num_cutouts, - use_cutouts=True, - ): - latents = latents.detach().requires_grad_() - - if isinstance(self.scheduler, LMSDiscreteScheduler): - sigma = self.scheduler.sigmas[index] - # the model input needs to be scaled to match the continuous ODE formulation in K-LMS - latent_model_input = latents / ((sigma**2 + 1) ** 0.5) - else: - latent_model_input = latents - - # predict the noise residual - noise_pred = self.unet(latent_model_input, timestep, encoder_hidden_states=text_embeddings).sample - - if isinstance(self.scheduler, (PNDMScheduler, DDIMScheduler)): - alpha_prod_t = self.scheduler.alphas_cumprod[timestep] - beta_prod_t = 1 - alpha_prod_t - # compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf - pred_original_sample = (latents - beta_prod_t ** (0.5) * noise_pred) / alpha_prod_t ** (0.5) - - fac = torch.sqrt(beta_prod_t) - sample = pred_original_sample * (fac) + latents * (1 - fac) - elif isinstance(self.scheduler, LMSDiscreteScheduler): - sigma = self.scheduler.sigmas[index] - sample = latents - sigma * noise_pred - else: - raise ValueError(f"scheduler type {type(self.scheduler)} not supported") - - sample = 1 / 0.18215 * sample - image = self.vae.decode(sample).sample - image = (image / 2 + 0.5).clamp(0, 1) - - if use_cutouts: - image = self.make_cutouts(image, num_cutouts) - else: - image = transforms.Resize(FEATURE_EXTRACTOR_SIZE)(image) - image = self.normalize(image).to(latents.dtype) - - image_embeddings_clip = self.clip_model.get_image_features(image) - image_embeddings_clip = image_embeddings_clip / image_embeddings_clip.norm(p=2, dim=-1, keepdim=True) - - if use_cutouts: - dists = spherical_dist_loss(image_embeddings_clip, guide_embeddings_clip) - dists = dists.view([num_cutouts, sample.shape[0], -1]) - loss = dists.sum(2).mean(0).sum() * clip_guidance_scale - else: - # バッチサイズが複数だと正しく動くかわからない - loss = spherical_dist_loss(image_embeddings_clip, guide_embeddings_clip).mean() * clip_guidance_scale - - grads = -torch.autograd.grad(loss, latents)[0] - - if isinstance(self.scheduler, LMSDiscreteScheduler): - latents = latents.detach() + grads * (sigma**2) - noise_pred = noise_pred_original - else: - noise_pred = noise_pred_original - torch.sqrt(beta_prod_t) * grads - return noise_pred, latents - - # バッチを分解して一件ずつ処理する - def cond_fn_vgg16(self, latents, timestep, index, text_embeddings, noise_pred_original, guide_embeddings, guidance_scale): - if len(latents) == 1: - return self.cond_fn_vgg16_b1( - latents, timestep, index, text_embeddings, noise_pred_original, guide_embeddings, guidance_scale - ) - - noise_pred = [] - cond_latents = [] - for i in range(len(latents)): - lat1 = latents[i].unsqueeze(0) - tem1 = text_embeddings[i].unsqueeze(0) - npo1 = noise_pred_original[i].unsqueeze(0) - gem1 = guide_embeddings[i].unsqueeze(0) - npr1, cla1 = self.cond_fn_vgg16_b1(lat1, timestep, index, tem1, npo1, gem1, guidance_scale) - noise_pred.append(npr1) - cond_latents.append(cla1) - - noise_pred = torch.cat(noise_pred) - cond_latents = torch.cat(cond_latents) - return noise_pred, cond_latents - - # 1件だけ処理する - @torch.enable_grad() - def cond_fn_vgg16_b1(self, latents, timestep, index, text_embeddings, noise_pred_original, guide_embeddings, guidance_scale): - latents = latents.detach().requires_grad_() - - if isinstance(self.scheduler, LMSDiscreteScheduler): - sigma = self.scheduler.sigmas[index] - # the model input needs to be scaled to match the continuous ODE formulation in K-LMS - latent_model_input = latents / ((sigma**2 + 1) ** 0.5) - else: - latent_model_input = latents - - # predict the noise residual - noise_pred = self.unet(latent_model_input, timestep, encoder_hidden_states=text_embeddings).sample - - if isinstance(self.scheduler, (PNDMScheduler, DDIMScheduler)): - alpha_prod_t = self.scheduler.alphas_cumprod[timestep] - beta_prod_t = 1 - alpha_prod_t - # compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf - pred_original_sample = (latents - beta_prod_t ** (0.5) * noise_pred) / alpha_prod_t ** (0.5) - - fac = torch.sqrt(beta_prod_t) - sample = pred_original_sample * (fac) + latents * (1 - fac) - elif isinstance(self.scheduler, LMSDiscreteScheduler): - sigma = self.scheduler.sigmas[index] - sample = latents - sigma * noise_pred - else: - raise ValueError(f"scheduler type {type(self.scheduler)} not supported") - - sample = 1 / 0.18215 * sample - image = self.vae.decode(sample).sample - image = (image / 2 + 0.5).clamp(0, 1) - image = transforms.Resize((image.shape[-2] // VGG16_INPUT_RESIZE_DIV, image.shape[-1] // VGG16_INPUT_RESIZE_DIV))(image) - image = self.vgg16_normalize(image).to(latents.dtype) - - image_embeddings = self.vgg16_feat_model(image)["feat"] - - # バッチサイズが複数だと正しく動くかわからない - loss = ( - (image_embeddings - guide_embeddings) ** 2 - ).mean() * guidance_scale # MSE style transferでコンテンツの損失はMSEなので - - grads = -torch.autograd.grad(loss, latents)[0] - if isinstance(self.scheduler, LMSDiscreteScheduler): - latents = latents.detach() + grads * (sigma**2) - noise_pred = noise_pred_original - else: - noise_pred = noise_pred_original - torch.sqrt(beta_prod_t) * grads - return noise_pred, latents - - -class MakeCutouts(torch.nn.Module): - def __init__(self, cut_size, cut_power=1.0): - super().__init__() - - self.cut_size = cut_size - self.cut_power = cut_power - - def forward(self, pixel_values, num_cutouts): - sideY, sideX = pixel_values.shape[2:4] - max_size = min(sideX, sideY) - min_size = min(sideX, sideY, self.cut_size) - cutouts = [] - for _ in range(num_cutouts): - size = int(torch.rand([]) ** self.cut_power * (max_size - min_size) + min_size) - offsetx = torch.randint(0, sideX - size + 1, ()) - offsety = torch.randint(0, sideY - size + 1, ()) - cutout = pixel_values[:, :, offsety : offsety + size, offsetx : offsetx + size] - cutouts.append(torch.nn.functional.adaptive_avg_pool2d(cutout, self.cut_size)) - return torch.cat(cutouts) - - -def spherical_dist_loss(x, y): - x = torch.nn.functional.normalize(x, dim=-1) - y = torch.nn.functional.normalize(y, dim=-1) - return (x - y).norm(dim=-1).div(2).arcsin().pow(2).mul(2) - - -re_attention = re.compile( - r""" -\\\(| -\\\)| -\\\[| -\\]| -\\\\| -\\| -\(| -\[| -:([+-]?[.\d]+)\)| -\)| -]| -[^\\()\[\]:]+| -: -""", - re.X, -) - - -def parse_prompt_attention(text): - """ - Parses a string with attention tokens and returns a list of pairs: text and its associated weight. - Accepted tokens are: - (abc) - increases attention to abc by a multiplier of 1.1 - (abc:3.12) - increases attention to abc by a multiplier of 3.12 - [abc] - decreases attention to abc by a multiplier of 1.1 - \( - literal character '(' - \[ - literal character '[' - \) - literal character ')' - \] - literal character ']' - \\ - literal character '\' - anything else - just text - >>> parse_prompt_attention('normal text') - [['normal text', 1.0]] - >>> parse_prompt_attention('an (important) word') - [['an ', 1.0], ['important', 1.1], [' word', 1.0]] - >>> parse_prompt_attention('(unbalanced') - [['unbalanced', 1.1]] - >>> parse_prompt_attention('\(literal\]') - [['(literal]', 1.0]] - >>> parse_prompt_attention('(unnecessary)(parens)') - [['unnecessaryparens', 1.1]] - >>> parse_prompt_attention('a (((house:1.3)) [on] a (hill:0.5), sun, (((sky))).') - [['a ', 1.0], - ['house', 1.5730000000000004], - [' ', 1.1], - ['on', 1.0], - [' a ', 1.1], - ['hill', 0.55], - [', sun, ', 1.1], - ['sky', 1.4641000000000006], - ['.', 1.1]] - """ - - res = [] - round_brackets = [] - square_brackets = [] - - round_bracket_multiplier = 1.1 - square_bracket_multiplier = 1 / 1.1 - - def multiply_range(start_position, multiplier): - for p in range(start_position, len(res)): - res[p][1] *= multiplier - - # keep break as separate token - text = text.replace("BREAK", "\\BREAK\\") - - for m in re_attention.finditer(text): - text = m.group(0) - weight = m.group(1) - - if text.startswith("\\"): - res.append([text[1:], 1.0]) - elif text == "(": - round_brackets.append(len(res)) - elif text == "[": - square_brackets.append(len(res)) - elif weight is not None and len(round_brackets) > 0: - multiply_range(round_brackets.pop(), float(weight)) - elif text == ")" and len(round_brackets) > 0: - multiply_range(round_brackets.pop(), round_bracket_multiplier) - elif text == "]" and len(square_brackets) > 0: - multiply_range(square_brackets.pop(), square_bracket_multiplier) - else: - res.append([text, 1.0]) - - for pos in round_brackets: - multiply_range(pos, round_bracket_multiplier) - - for pos in square_brackets: - multiply_range(pos, square_bracket_multiplier) - - if len(res) == 0: - res = [["", 1.0]] - - # merge runs of identical weights - i = 0 - while i + 1 < len(res): - if res[i][1] == res[i + 1][1] and res[i][0].strip() != "BREAK" and res[i + 1][0].strip() != "BREAK": - res[i][0] += res[i + 1][0] - res.pop(i + 1) - else: - i += 1 - - return res - - -def get_prompts_with_weights(pipe: PipelineLike, prompt: List[str], max_length: int, layer=None): - r""" - Tokenize a list of prompts and return its tokens with weights of each token. - No padding, starting or ending token is included. - """ - tokens = [] - weights = [] - truncated = False - - for text in prompt: - texts_and_weights = parse_prompt_attention(text) - text_token = [] - text_weight = [] - for word, weight in texts_and_weights: - if word.strip() == "BREAK": - # pad until next multiple of tokenizer's max token length - pad_len = pipe.tokenizer.model_max_length - (len(text_token) % pipe.tokenizer.model_max_length) - logger.info(f"BREAK pad_len: {pad_len}") - for i in range(pad_len): - # v2のときEOSをつけるべきかどうかわからないぜ - # if i == 0: - # text_token.append(pipe.tokenizer.eos_token_id) - # else: - text_token.append(pipe.tokenizer.pad_token_id) - text_weight.append(1.0) - continue - - # tokenize and discard the starting and the ending token - token = pipe.tokenizer(word).input_ids[1:-1] - - token = pipe.replace_token(token, layer=layer) - - text_token += token - # copy the weight by length of token - text_weight += [weight] * len(token) - # stop if the text is too long (longer than truncation limit) - if len(text_token) > max_length: - truncated = True - break - # truncate - if len(text_token) > max_length: - truncated = True - text_token = text_token[:max_length] - text_weight = text_weight[:max_length] - 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") - return tokens, weights - - -def pad_tokens_and_weights(tokens, weights, max_length, bos, eos, pad, no_boseos_middle=True, chunk_length=77): - r""" - Pad the tokens (with starting and ending tokens) and weights (with 1.0) to max_length. - """ - max_embeddings_multiples = (max_length - 2) // (chunk_length - 2) - weights_length = max_length if no_boseos_middle else max_embeddings_multiples * chunk_length - for i in range(len(tokens)): - tokens[i] = [bos] + tokens[i] + [eos] + [pad] * (max_length - 2 - len(tokens[i])) - if no_boseos_middle: - weights[i] = [1.0] + weights[i] + [1.0] * (max_length - 1 - len(weights[i])) - else: - w = [] - if len(weights[i]) == 0: - w = [1.0] * weights_length - else: - for j in range(max_embeddings_multiples): - w.append(1.0) # weight for starting token in this chunk - w += weights[i][j * (chunk_length - 2) : min(len(weights[i]), (j + 1) * (chunk_length - 2))] - w.append(1.0) # weight for ending token in this chunk - w += [1.0] * (weights_length - len(w)) - weights[i] = w[:] - - return tokens, weights - - -def get_unweighted_text_embeddings( - pipe: PipelineLike, - text_input: torch.Tensor, - chunk_length: int, - clip_skip: int, - eos: int, - pad: int, - no_boseos_middle: Optional[bool] = True, -): - """ - When the length of tokens is a multiple of the capacity of the text encoder, - it should be split into chunks and sent to the text encoder individually. - """ - max_embeddings_multiples = (text_input.shape[1] - 2) // (chunk_length - 2) - if max_embeddings_multiples > 1: - text_embeddings = [] - for i in range(max_embeddings_multiples): - # extract the i-th chunk - text_input_chunk = text_input[:, i * (chunk_length - 2) : (i + 1) * (chunk_length - 2) + 2].clone() - - # cover the head and the tail by the starting and the ending tokens - text_input_chunk[:, 0] = text_input[0, 0] - if pad == eos: # v1 - text_input_chunk[:, -1] = text_input[0, -1] - else: # v2 - for j in range(len(text_input_chunk)): - if text_input_chunk[j, -1] != eos and text_input_chunk[j, -1] != pad: # 最後に普通の文字がある - text_input_chunk[j, -1] = eos - if text_input_chunk[j, 1] == pad: # BOSだけであとはPAD - text_input_chunk[j, 1] = eos - - if clip_skip is None or clip_skip == 1: - text_embedding = pipe.text_encoder(text_input_chunk)[0] - else: - enc_out = pipe.text_encoder(text_input_chunk, output_hidden_states=True, return_dict=True) - text_embedding = enc_out["hidden_states"][-clip_skip] - text_embedding = pipe.text_encoder.text_model.final_layer_norm(text_embedding) - - if no_boseos_middle: - if i == 0: - # discard the ending token - text_embedding = text_embedding[:, :-1] - elif i == max_embeddings_multiples - 1: - # discard the starting token - text_embedding = text_embedding[:, 1:] - else: - # discard both starting and ending tokens - text_embedding = text_embedding[:, 1:-1] - - text_embeddings.append(text_embedding) - text_embeddings = torch.concat(text_embeddings, axis=1) - else: - if clip_skip is None or clip_skip == 1: - text_embeddings = pipe.text_encoder(text_input)[0] - else: - enc_out = pipe.text_encoder(text_input, output_hidden_states=True, return_dict=True) - text_embeddings = enc_out["hidden_states"][-clip_skip] - text_embeddings = pipe.text_encoder.text_model.final_layer_norm(text_embeddings) - return text_embeddings - - -def get_weighted_text_embeddings( - pipe: PipelineLike, - prompt: Union[str, List[str]], - uncond_prompt: Optional[Union[str, List[str]]] = None, - max_embeddings_multiples: Optional[int] = 1, - no_boseos_middle: Optional[bool] = False, - skip_parsing: Optional[bool] = False, - skip_weighting: Optional[bool] = False, - clip_skip=None, - layer=None, - **kwargs, -): - r""" - Prompts can be assigned with local weights using brackets. For example, - prompt 'A (very beautiful) masterpiece' highlights the words 'very beautiful', - and the embedding tokens corresponding to the words get multiplied by a constant, 1.1. - Also, to regularize of the embedding, the weighted embedding would be scaled to preserve the original mean. - Args: - pipe (`DiffusionPipeline`): - Pipe to provide access to the tokenizer and the text encoder. - prompt (`str` or `List[str]`): - The prompt or prompts to guide the image generation. - uncond_prompt (`str` or `List[str]`): - The unconditional prompt or prompts for guide the image generation. If unconditional prompt - is provided, the embeddings of prompt and uncond_prompt are concatenated. - max_embeddings_multiples (`int`, *optional*, defaults to `1`): - The max multiple length of prompt embeddings compared to the max output length of text encoder. - no_boseos_middle (`bool`, *optional*, defaults to `False`): - If the length of text token is multiples of the capacity of text encoder, whether reserve the starting and - ending token in each of the chunk in the middle. - skip_parsing (`bool`, *optional*, defaults to `False`): - Skip the parsing of brackets. - skip_weighting (`bool`, *optional*, defaults to `False`): - Skip the weighting. When the parsing is skipped, it is forced True. - """ - max_length = (pipe.tokenizer.model_max_length - 2) * max_embeddings_multiples + 2 - if isinstance(prompt, str): - prompt = [prompt] - - # split the prompts with "AND". each prompt must have the same number of splits - new_prompts = [] - for p in prompt: - new_prompts.extend(p.split(" AND ")) - prompt = new_prompts - - if not skip_parsing: - prompt_tokens, prompt_weights = get_prompts_with_weights(pipe, prompt, max_length - 2, layer=layer) - if uncond_prompt is not None: - if isinstance(uncond_prompt, str): - uncond_prompt = [uncond_prompt] - uncond_tokens, uncond_weights = get_prompts_with_weights(pipe, uncond_prompt, max_length - 2, layer=layer) - else: - prompt_tokens = [token[1:-1] for token in pipe.tokenizer(prompt, max_length=max_length, truncation=True).input_ids] - prompt_weights = [[1.0] * len(token) for token in prompt_tokens] - if uncond_prompt is not None: - if isinstance(uncond_prompt, str): - uncond_prompt = [uncond_prompt] - uncond_tokens = [ - token[1:-1] for token in pipe.tokenizer(uncond_prompt, max_length=max_length, truncation=True).input_ids - ] - uncond_weights = [[1.0] * len(token) for token in uncond_tokens] - - # round up the longest length of tokens to a multiple of (model_max_length - 2) - max_length = max([len(token) for token in prompt_tokens]) - if uncond_prompt is not None: - max_length = max(max_length, max([len(token) for token in uncond_tokens])) - - max_embeddings_multiples = min( - max_embeddings_multiples, - (max_length - 1) // (pipe.tokenizer.model_max_length - 2) + 1, - ) - max_embeddings_multiples = max(1, max_embeddings_multiples) - max_length = (pipe.tokenizer.model_max_length - 2) * max_embeddings_multiples + 2 - - # pad the length of tokens and weights - bos = pipe.tokenizer.bos_token_id - eos = pipe.tokenizer.eos_token_id - pad = pipe.tokenizer.pad_token_id - prompt_tokens, prompt_weights = pad_tokens_and_weights( - prompt_tokens, - prompt_weights, - max_length, - bos, - eos, - pad, - no_boseos_middle=no_boseos_middle, - chunk_length=pipe.tokenizer.model_max_length, - ) - prompt_tokens = torch.tensor(prompt_tokens, dtype=torch.long, device=pipe.device) - if uncond_prompt is not None: - uncond_tokens, uncond_weights = pad_tokens_and_weights( - uncond_tokens, - uncond_weights, - max_length, - bos, - eos, - pad, - no_boseos_middle=no_boseos_middle, - chunk_length=pipe.tokenizer.model_max_length, - ) - uncond_tokens = torch.tensor(uncond_tokens, dtype=torch.long, device=pipe.device) - - # get the embeddings - text_embeddings = get_unweighted_text_embeddings( - pipe, - prompt_tokens, - pipe.tokenizer.model_max_length, - clip_skip, - eos, - pad, - no_boseos_middle=no_boseos_middle, - ) - prompt_weights = torch.tensor(prompt_weights, dtype=text_embeddings.dtype, device=pipe.device) - if uncond_prompt is not None: - uncond_embeddings = get_unweighted_text_embeddings( - pipe, - uncond_tokens, - pipe.tokenizer.model_max_length, - clip_skip, - eos, - pad, - no_boseos_middle=no_boseos_middle, - ) - uncond_weights = torch.tensor(uncond_weights, dtype=uncond_embeddings.dtype, device=pipe.device) - - # assign weights to the prompts and normalize in the sense of mean - # TODO: should we normalize by chunk or in a whole (current implementation)? - # →全体でいいんじゃないかな - if (not skip_parsing) and (not skip_weighting): - previous_mean = text_embeddings.float().mean(axis=[-2, -1]).to(text_embeddings.dtype) - text_embeddings *= prompt_weights.unsqueeze(-1) - current_mean = text_embeddings.float().mean(axis=[-2, -1]).to(text_embeddings.dtype) - text_embeddings *= (previous_mean / current_mean).unsqueeze(-1).unsqueeze(-1) - if uncond_prompt is not None: - previous_mean = uncond_embeddings.float().mean(axis=[-2, -1]).to(uncond_embeddings.dtype) - uncond_embeddings *= uncond_weights.unsqueeze(-1) - current_mean = uncond_embeddings.float().mean(axis=[-2, -1]).to(uncond_embeddings.dtype) - uncond_embeddings *= (previous_mean / current_mean).unsqueeze(-1).unsqueeze(-1) - - if uncond_prompt is not None: - return text_embeddings, uncond_embeddings, prompt_tokens - return text_embeddings, None, prompt_tokens - - -def preprocess_guide_image(image): - image = image.resize(FEATURE_EXTRACTOR_SIZE, resample=Image.NEAREST) # cond_fnと合わせる - image = np.array(image).astype(np.float32) / 255.0 - image = image[None].transpose(0, 3, 1, 2) # nchw - image = torch.from_numpy(image) - return image # 0 to 1 - - -# VGG16の入力は任意サイズでよいので入力画像を適宜リサイズする -def preprocess_vgg16_guide_image(image, size): - image = image.resize(size, resample=Image.NEAREST) # cond_fnと合わせる - image = np.array(image).astype(np.float32) / 255.0 - image = image[None].transpose(0, 3, 1, 2) # nchw - image = torch.from_numpy(image) - return image # 0 to 1 - - -def preprocess_image(image): - w, h = image.size - w, h = map(lambda x: x - x % 32, (w, h)) # resize to integer multiple of 32 - image = image.resize((w, h), resample=PIL.Image.LANCZOS) - image = np.array(image).astype(np.float32) / 255.0 - image = image[None].transpose(0, 3, 1, 2) - image = torch.from_numpy(image) - return 2.0 * image - 1.0 - - -def preprocess_mask(mask): - mask = mask.convert("L") - w, h = mask.size - w, h = map(lambda x: x - x % 32, (w, h)) # resize to integer multiple of 32 - mask = mask.resize((w // 8, h // 8), resample=PIL.Image.BILINEAR) # LANCZOS) - mask = np.array(mask).astype(np.float32) / 255.0 - mask = np.tile(mask, (4, 1, 1)) - mask = mask[None].transpose(0, 1, 2, 3) # what does this step do? - mask = 1 - mask # repaint white, keep black - mask = torch.from_numpy(mask) - return mask - - -# regular expression for dynamic prompt: -# starts and ends with "{" and "}" -# contains at least one variant divided by "|" -# optional framgments divided by "$$" at start -# if the first fragment is "E" or "e", enumerate all variants -# if the second fragment is a number or two numbers, repeat the variants in the range -# if the third fragment is a string, use it as a separator - -RE_DYNAMIC_PROMPT = re.compile(r"\{((e|E)\$\$)?(([\d\-]+)\$\$)?(([^\|\}]+?)\$\$)?(.+?((\|).+?)*?)\}") - - -def handle_dynamic_prompt_variants(prompt, repeat_count): - founds = list(RE_DYNAMIC_PROMPT.finditer(prompt)) - if not founds: - return [prompt] - - # make each replacement for each variant - enumerating = False - replacers = [] - for found in founds: - # if "e$$" is found, enumerate all variants - found_enumerating = found.group(2) is not None - enumerating = enumerating or found_enumerating - - separator = ", " if found.group(6) is None else found.group(6) - variants = found.group(7).split("|") - - # parse count range - count_range = found.group(4) - if count_range is None: - count_range = [1, 1] - else: - count_range = count_range.split("-") - if len(count_range) == 1: - count_range = [int(count_range[0]), int(count_range[0])] - elif len(count_range) == 2: - count_range = [int(count_range[0]), int(count_range[1])] - else: - logger.warning(f"invalid count range: {count_range}") - count_range = [1, 1] - if count_range[0] > count_range[1]: - count_range = [count_range[1], count_range[0]] - if count_range[0] < 0: - count_range[0] = 0 - if count_range[1] > len(variants): - count_range[1] = len(variants) - - if found_enumerating: - # make function to enumerate all combinations - def make_replacer_enum(vari, cr, sep): - def replacer(): - values = [] - for count in range(cr[0], cr[1] + 1): - for comb in itertools.combinations(vari, count): - values.append(sep.join(comb)) - return values - - return replacer - - replacers.append(make_replacer_enum(variants, count_range, separator)) - else: - # make function to choose random combinations - def make_replacer_single(vari, cr, sep): - def replacer(): - count = random.randint(cr[0], cr[1]) - comb = random.sample(vari, count) - return [sep.join(comb)] - - return replacer - - replacers.append(make_replacer_single(variants, count_range, separator)) - - # make each prompt - if not enumerating: - # if not enumerating, repeat the prompt, replace each variant randomly - prompts = [] - for _ in range(repeat_count): - current = prompt - for found, replacer in zip(founds, replacers): - current = current.replace(found.group(0), replacer()[0], 1) - prompts.append(current) - else: - # if enumerating, iterate all combinations for previous prompts - prompts = [prompt] - - for found, replacer in zip(founds, replacers): - if found.group(2) is not None: - # make all combinations for existing prompts - new_prompts = [] - for current in prompts: - replecements = replacer() - for replecement in replecements: - new_prompts.append(current.replace(found.group(0), replecement, 1)) - prompts = new_prompts - - for found, replacer in zip(founds, replacers): - # make random selection for existing prompts - if found.group(2) is None: - for i in range(len(prompts)): - prompts[i] = prompts[i].replace(found.group(0), replacer()[0], 1) - - return prompts - - -# endregion - - -# def load_clip_l14_336(dtype): -# logger.info(f"loading CLIP: {CLIP_ID_L14_336}") -# text_encoder = CLIPTextModel.from_pretrained(CLIP_ID_L14_336, torch_dtype=dtype) -# return text_encoder - - -class BatchDataBase(NamedTuple): - # バッチ分割が必要ないデータ - step: int - prompt: str - negative_prompt: str - seed: int - init_image: Any - mask_image: Any - clip_prompt: str - guide_image: Any - raw_prompt: str - - -class BatchDataExt(NamedTuple): - # バッチ分割が必要なデータ - width: int - height: int - steps: int - scale: float - negative_scale: float - strength: float - network_muls: Tuple[float] - num_sub_prompts: int - - -class BatchData(NamedTuple): - return_latents: bool - base: BatchDataBase - ext: BatchDataExt - - -def main(args): - if args.fp16: - dtype = torch.float16 - elif args.bf16: - dtype = torch.bfloat16 - else: - dtype = torch.float32 - - highres_fix = args.highres_fix_scale is not None - # assert not highres_fix or args.image_path is None, f"highres_fix doesn't work with img2img / highres_fixはimg2imgと同時に使えません" - - if args.v2 and args.clip_skip is not None: - logger.warning("v2 with clip_skip will be unexpected / v2でclip_skipを使用することは想定されていません") - - # モデルを読み込む - if not os.path.isfile(args.ckpt): # ファイルがないならパターンで探し、一つだけ該当すればそれを使う - files = glob.glob(args.ckpt) - if len(files) == 1: - args.ckpt = files[0] - - use_stable_diffusion_format = os.path.isfile(args.ckpt) - if use_stable_diffusion_format: - logger.info("load StableDiffusion checkpoint") - text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.ckpt) - else: - logger.info("load Diffusers pretrained models") - loading_pipe = StableDiffusionPipeline.from_pretrained(args.ckpt, safety_checker=None, torch_dtype=dtype) - text_encoder = loading_pipe.text_encoder - vae = loading_pipe.vae - unet = loading_pipe.unet - tokenizer = loading_pipe.tokenizer - del loading_pipe - - # Diffusers U-Net to original U-Net - original_unet = UNet2DConditionModel( - unet.config.sample_size, - unet.config.attention_head_dim, - unet.config.cross_attention_dim, - unet.config.use_linear_projection, - unet.config.upcast_attention, - ) - original_unet.load_state_dict(unet.state_dict()) - unet = original_unet - unet: InferUNet2DConditionModel = InferUNet2DConditionModel(unet) - - # VAEを読み込む - if args.vae is not None: - vae = model_util.load_vae(args.vae, dtype) - logger.info("additional VAE loaded") - - # # 置換するCLIPを読み込む - # if args.replace_clip_l14_336: - # text_encoder = load_clip_l14_336(dtype) - # logger.info(f"large clip {CLIP_ID_L14_336} is loaded") - - if args.clip_guidance_scale > 0.0 or args.clip_image_guidance_scale: - logger.info("prepare clip model") - clip_model = CLIPModel.from_pretrained(CLIP_MODEL_PATH, torch_dtype=dtype) - else: - clip_model = None - - if args.vgg16_guidance_scale > 0.0: - logger.info("prepare resnet model") - vgg16_model = torchvision.models.vgg16(torchvision.models.VGG16_Weights.IMAGENET1K_V1) - else: - vgg16_model = None - - # xformers、Hypernetwork対応 - if not args.diffusers_xformers: - mem_eff = not (args.xformers or args.sdpa) - replace_unet_modules(unet, mem_eff, args.xformers, args.sdpa) - replace_vae_modules(vae, mem_eff, args.xformers, args.sdpa) - - # tokenizerを読み込む - logger.info("loading tokenizer") - if use_stable_diffusion_format: - tokenizer = train_util.load_tokenizer(args) - - # schedulerを用意する - sched_init_args = {} - scheduler_num_noises_per_step = 1 - if args.sampler == "ddim": - scheduler_cls = DDIMScheduler - scheduler_module = diffusers.schedulers.scheduling_ddim - elif args.sampler == "ddpm": # ddpmはおかしくなるのでoptionから外してある - scheduler_cls = DDPMScheduler - scheduler_module = diffusers.schedulers.scheduling_ddpm - elif args.sampler == "pndm": - scheduler_cls = PNDMScheduler - scheduler_module = diffusers.schedulers.scheduling_pndm - elif args.sampler == "lms" or args.sampler == "k_lms": - scheduler_cls = LMSDiscreteScheduler - scheduler_module = diffusers.schedulers.scheduling_lms_discrete - elif args.sampler == "euler" or args.sampler == "k_euler": - scheduler_cls = EulerDiscreteScheduler - scheduler_module = diffusers.schedulers.scheduling_euler_discrete - elif args.sampler == "euler_a" or args.sampler == "k_euler_a": - scheduler_cls = EulerAncestralDiscreteSchedulerGL - scheduler_module = diffusers.schedulers.scheduling_euler_ancestral_discrete - elif args.sampler == "dpmsolver" or args.sampler == "dpmsolver++": - scheduler_cls = DPMSolverMultistepScheduler - sched_init_args["algorithm_type"] = args.sampler - scheduler_module = diffusers.schedulers.scheduling_dpmsolver_multistep - elif args.sampler == "dpmsingle": - scheduler_cls = DPMSolverSinglestepScheduler - scheduler_module = diffusers.schedulers.scheduling_dpmsolver_singlestep - elif args.sampler == "heun": - scheduler_cls = HeunDiscreteScheduler - scheduler_module = diffusers.schedulers.scheduling_heun_discrete - elif args.sampler == "dpm_2" or args.sampler == "k_dpm_2": - scheduler_cls = KDPM2DiscreteScheduler - scheduler_module = diffusers.schedulers.scheduling_k_dpm_2_discrete - elif args.sampler == "dpm_2_a" or args.sampler == "k_dpm_2_a": - scheduler_cls = KDPM2AncestralDiscreteScheduler - scheduler_module = diffusers.schedulers.scheduling_k_dpm_2_ancestral_discrete - scheduler_num_noises_per_step = 2 - - if args.v_parameterization: - sched_init_args["prediction_type"] = "v_prediction" - - # samplerの乱数をあらかじめ指定するための処理 - - # replace randn - class NoiseManager: - def __init__(self): - self.sampler_noises = None - self.sampler_noise_index = 0 - - def reset_sampler_noises(self, noises): - self.sampler_noise_index = 0 - self.sampler_noises = noises - - def randn(self, shape, device=None, dtype=None, layout=None, generator=None): - # logger.info(f"replacing {shape} {len(self.sampler_noises)} {self.sampler_noise_index}") - if self.sampler_noises is not None and self.sampler_noise_index < len(self.sampler_noises): - noise = self.sampler_noises[self.sampler_noise_index] - if shape != noise.shape: - noise = None - else: - noise = None - - if noise == None: - logger.warning(f"unexpected noise request: {self.sampler_noise_index}, {shape}") - noise = torch.randn(shape, dtype=dtype, device=device, generator=generator) - - self.sampler_noise_index += 1 - return noise - - class TorchRandReplacer: - def __init__(self, noise_manager): - self.noise_manager = noise_manager - - def __getattr__(self, item): - if item == "randn": - return self.noise_manager.randn - if hasattr(torch, item): - return getattr(torch, item) - raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, item)) - - noise_manager = NoiseManager() - if scheduler_module is not None: - scheduler_module.torch = TorchRandReplacer(noise_manager) - - scheduler = scheduler_cls( - num_train_timesteps=SCHEDULER_TIMESTEPS, - beta_start=SCHEDULER_LINEAR_START, - beta_end=SCHEDULER_LINEAR_END, - beta_schedule=SCHEDLER_SCHEDULE, - **sched_init_args, - ) - - # clip_sample=Trueにする - if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is False: - logger.info("set clip_sample to True") - scheduler.config.clip_sample = True - - # deviceを決定する - device = get_preferred_device() - - # custom pipelineをコピったやつを生成する - if args.vae_slices: - from library.slicing_vae import SlicingAutoencoderKL - - sli_vae = SlicingAutoencoderKL( - act_fn="silu", - block_out_channels=(128, 256, 512, 512), - down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"], - in_channels=3, - latent_channels=4, - layers_per_block=2, - norm_num_groups=32, - out_channels=3, - sample_size=512, - up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"], - num_slices=args.vae_slices, - ) - sli_vae.load_state_dict(vae.state_dict()) # vaeのパラメータをコピーする - vae = sli_vae - del sli_vae - vae.to(dtype).to(device) - vae.eval() - - text_encoder.to(dtype).to(device) - unet.to(dtype).to(device) - - text_encoder.eval() - unet.eval() - - if clip_model is not None: - clip_model.to(dtype).to(device) - clip_model.eval() - if vgg16_model is not None: - vgg16_model.to(dtype).to(device) - vgg16_model.eval() - - # networkを組み込む - if args.network_module: - networks = [] - network_default_muls = [] - network_pre_calc = args.network_pre_calc - - # merge関連の引数を統合する - if args.network_merge: - network_merge = len(args.network_module) # all networks are merged - elif args.network_merge_n_models: - network_merge = args.network_merge_n_models - else: - network_merge = 0 - - for i, network_module in enumerate(args.network_module): - logger.info(f"import network module: {network_module}") - imported_module = importlib.import_module(network_module) - - network_mul = 1.0 if args.network_mul is None or len(args.network_mul) <= i else args.network_mul[i] - - net_kwargs = {} - if args.network_args and i < len(args.network_args): - network_args = args.network_args[i] - # TODO escape special chars - network_args = network_args.split(";") - for net_arg in network_args: - key, value = net_arg.split("=") - net_kwargs[key] = value - - if args.network_weights is None or len(args.network_weights) <= i: - raise ValueError("No weight. Weight is required.") - - network_weight = args.network_weights[i] - logger.info(f"load network weights from: {network_weight}") - - if model_util.is_safetensors(network_weight) and args.network_show_meta: - from safetensors.torch import safe_open - - with safe_open(network_weight, framework="pt") as f: - metadata = f.metadata() - if metadata is not None: - logger.info(f"metadata for: {network_weight}: {metadata}") - - network, weights_sd = imported_module.create_network_from_weights( - network_mul, network_weight, vae, text_encoder, unet, for_inference=True, **net_kwargs - ) - if network is None: - return - - mergeable = network.is_mergeable() - if network_merge and not mergeable: - logger.warning("network is not mergiable. ignore merge option.") - - if not mergeable or i >= network_merge: - # not merging - network.apply_to(text_encoder, unet) - info = network.load_state_dict(weights_sd, False) # network.load_weightsを使うようにするとよい - logger.info(f"weights are loaded: {info}") - - if args.opt_channels_last: - network.to(memory_format=torch.channels_last) - network.to(dtype).to(device) - - if network_pre_calc: - logger.info("backup original weights") - network.backup_weights() - - networks.append(network) - network_default_muls.append(network_mul) - else: - network.merge_to(text_encoder, unet, weights_sd, dtype, device) - - else: - networks = [] - - # upscalerの指定があれば取得する - upscaler = None - if args.highres_fix_upscaler: - logger.info(f"import upscaler module {args.highres_fix_upscaler}") - imported_module = importlib.import_module(args.highres_fix_upscaler) - - us_kwargs = {} - if args.highres_fix_upscaler_args: - for net_arg in args.highres_fix_upscaler_args.split(";"): - key, value = net_arg.split("=") - us_kwargs[key] = value - - logger.info("create upscaler") - upscaler = imported_module.create_upscaler(**us_kwargs) - upscaler.to(dtype).to(device) - - # ControlNetの処理 - control_nets: List[ControlNetInfo] = [] - if args.control_net_models: - for i, model in enumerate(args.control_net_models): - prep_type = None if not args.control_net_preps or len(args.control_net_preps) <= i else args.control_net_preps[i] - weight = 1.0 if not args.control_net_weights or len(args.control_net_weights) <= i else args.control_net_weights[i] - ratio = 1.0 if not args.control_net_ratios or len(args.control_net_ratios) <= i else args.control_net_ratios[i] - - ctrl_unet, ctrl_net = original_control_net.load_control_net(args.v2, unet, model) - prep = original_control_net.load_preprocess(prep_type) - control_nets.append(ControlNetInfo(ctrl_unet, ctrl_net, prep, weight, ratio)) - - if args.opt_channels_last: - logger.info(f"set optimizing: channels last") - text_encoder.to(memory_format=torch.channels_last) - vae.to(memory_format=torch.channels_last) - unet.to(memory_format=torch.channels_last) - if clip_model is not None: - clip_model.to(memory_format=torch.channels_last) - if networks: - for network in networks: - network.to(memory_format=torch.channels_last) - if vgg16_model is not None: - vgg16_model.to(memory_format=torch.channels_last) - - for cn in control_nets: - cn.unet.to(memory_format=torch.channels_last) - cn.net.to(memory_format=torch.channels_last) - - pipe = PipelineLike( - device, - vae, - text_encoder, - tokenizer, - unet, - scheduler, - args.clip_skip, - clip_model, - args.clip_guidance_scale, - args.clip_image_guidance_scale, - vgg16_model, - args.vgg16_guidance_scale, - args.vgg16_guidance_layer, - ) - pipe.set_control_nets(control_nets) - logger.info("pipeline is ready.") - - if args.diffusers_xformers: - pipe.enable_xformers_memory_efficient_attention() - - # Deep Shrink - if args.ds_depth_1 is not None: - unet.set_deep_shrink(args.ds_depth_1, args.ds_timesteps_1, args.ds_depth_2, args.ds_timesteps_2, args.ds_ratio) - - # Gradual Latent - if args.gradual_latent_timesteps is not None: - if args.gradual_latent_unsharp_params: - us_params = args.gradual_latent_unsharp_params.split(",") - us_ksize, us_sigma, us_strength = [float(v) for v in us_params[:3]] - us_target_x = True if len(us_params) <= 3 else bool(int(us_params[3])) - us_ksize = int(us_ksize) - else: - us_ksize, us_sigma, us_strength, us_target_x = None, None, None, None - - gradual_latent = GradualLatent( - args.gradual_latent_ratio, - args.gradual_latent_timesteps, - args.gradual_latent_every_n_steps, - args.gradual_latent_ratio_step, - args.gradual_latent_s_noise, - us_ksize, - us_sigma, - us_strength, - us_target_x, - ) - pipe.set_gradual_latent(gradual_latent) - - # Extended Textual Inversion および Textual Inversionを処理する - if args.XTI_embeddings: - diffusers.models.UNet2DConditionModel.forward = unet_forward_XTI - diffusers.models.unet_2d_blocks.CrossAttnDownBlock2D.forward = downblock_forward_XTI - diffusers.models.unet_2d_blocks.CrossAttnUpBlock2D.forward = upblock_forward_XTI - - if args.textual_inversion_embeddings: - token_ids_embeds = [] - for embeds_file in args.textual_inversion_embeddings: - if model_util.is_safetensors(embeds_file): - from safetensors.torch import load_file - - data = load_file(embeds_file) - else: - data = torch.load(embeds_file, map_location="cpu") - - if "string_to_param" in data: - data = data["string_to_param"] - embeds = next(iter(data.values())) - - if type(embeds) != torch.Tensor: - raise ValueError( - f"weight file does not contains Tensor / 重みファイルのデータがTensorではありません: {embeds_file}" - ) - - num_vectors_per_token = embeds.size()[0] - token_string = os.path.splitext(os.path.basename(embeds_file))[0] - token_strings = [token_string] + [f"{token_string}{i+1}" for i in range(num_vectors_per_token - 1)] - - # add new word to tokenizer, count is num_vectors_per_token - num_added_tokens = tokenizer.add_tokens(token_strings) - assert ( - num_added_tokens == num_vectors_per_token - ), f"tokenizer has same word to token string (filename). please rename the file / 指定した名前（ファイル名）のトークンが既に存在します。ファイルをリネームしてください: {embeds_file}" - - token_ids = tokenizer.convert_tokens_to_ids(token_strings) - logger.info(f"Textual Inversion embeddings `{token_string}` loaded. 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)}" - - if num_vectors_per_token > 1: - pipe.add_token_replacement(token_ids[0], token_ids) - - token_ids_embeds.append((token_ids, embeds)) - - text_encoder.resize_token_embeddings(len(tokenizer)) - token_embeds = text_encoder.get_input_embeddings().weight.data - for token_ids, embeds in token_ids_embeds: - for token_id, embed in zip(token_ids, embeds): - token_embeds[token_id] = embed - - if args.XTI_embeddings: - XTI_layers = [ - "IN01", - "IN02", - "IN04", - "IN05", - "IN07", - "IN08", - "MID", - "OUT03", - "OUT04", - "OUT05", - "OUT06", - "OUT07", - "OUT08", - "OUT09", - "OUT10", - "OUT11", - ] - token_ids_embeds_XTI = [] - for embeds_file in args.XTI_embeddings: - if model_util.is_safetensors(embeds_file): - from safetensors.torch import load_file - - data = load_file(embeds_file) - else: - data = torch.load(embeds_file, map_location="cpu") - if set(data.keys()) != set(XTI_layers): - raise ValueError("NOT XTI") - embeds = torch.concat(list(data.values())) - num_vectors_per_token = data["MID"].size()[0] - - token_string = os.path.splitext(os.path.basename(embeds_file))[0] - token_strings = [token_string] + [f"{token_string}{i+1}" for i in range(num_vectors_per_token - 1)] - - # add new word to tokenizer, count is num_vectors_per_token - num_added_tokens = tokenizer.add_tokens(token_strings) - assert ( - num_added_tokens == num_vectors_per_token - ), f"tokenizer has same word to token string (filename). please rename the file / 指定した名前（ファイル名）のトークンが既に存在します。ファイルをリネームしてください: {embeds_file}" - - token_ids = tokenizer.convert_tokens_to_ids(token_strings) - logger.info(f"XTI embeddings `{token_string}` loaded. Tokens are added: {token_ids}") - - # if num_vectors_per_token > 1: - pipe.add_token_replacement(token_ids[0], token_ids) - - token_strings_XTI = [] - for layer_name in XTI_layers: - token_strings_XTI += [f"{t}_{layer_name}" for t in token_strings] - tokenizer.add_tokens(token_strings_XTI) - token_ids_XTI = tokenizer.convert_tokens_to_ids(token_strings_XTI) - token_ids_embeds_XTI.append((token_ids_XTI, embeds)) - for t in token_ids: - t_XTI_dic = {} - for i, layer_name in enumerate(XTI_layers): - t_XTI_dic[layer_name] = t + (i + 1) * num_added_tokens - pipe.add_token_replacement_XTI(t, t_XTI_dic) - - text_encoder.resize_token_embeddings(len(tokenizer)) - token_embeds = text_encoder.get_input_embeddings().weight.data - for token_ids, embeds in token_ids_embeds_XTI: - for token_id, embed in zip(token_ids, embeds): - token_embeds[token_id] = embed - - # promptを取得する - if args.from_file is not None: - logger.info(f"reading prompts from {args.from_file}") - with open(args.from_file, "r", encoding="utf-8") as f: - prompt_list = f.read().splitlines() - prompt_list = [d for d in prompt_list if len(d.strip()) > 0 and d[0] != "#"] - elif args.prompt is not None: - prompt_list = [args.prompt] - else: - prompt_list = [] - - if args.interactive: - args.n_iter = 1 - - # img2imgの前処理、画像の読み込みなど - def load_images(path): - if os.path.isfile(path): - paths = [path] - else: - paths = ( - glob.glob(os.path.join(path, "*.png")) - + glob.glob(os.path.join(path, "*.jpg")) - + glob.glob(os.path.join(path, "*.jpeg")) - + glob.glob(os.path.join(path, "*.webp")) - ) - paths.sort() - - images = [] - for p in paths: - image = Image.open(p) - if image.mode != "RGB": - logger.info(f"convert image to RGB from {image.mode}: {p}") - image = image.convert("RGB") - images.append(image) - - return images - - def resize_images(imgs, size): - resized = [] - for img in imgs: - r_img = img.resize(size, Image.Resampling.LANCZOS) - if hasattr(img, "filename"): # filename属性がない場合があるらしい - r_img.filename = img.filename - resized.append(r_img) - return resized - - if args.image_path is not None: - logger.info(f"load image for img2img: {args.image_path}") - init_images = load_images(args.image_path) - assert len(init_images) > 0, f"No image / 画像がありません: {args.image_path}" - logger.info(f"loaded {len(init_images)} images for img2img") - else: - init_images = None - - if args.mask_path is not None: - logger.info(f"load mask for inpainting: {args.mask_path}") - mask_images = load_images(args.mask_path) - assert len(mask_images) > 0, f"No mask image / マスク画像がありません: {args.image_path}" - logger.info(f"loaded {len(mask_images)} mask images for inpainting") - else: - mask_images = None - - # promptがないとき、画像のPngInfoから取得する - if init_images is not None and len(prompt_list) == 0 and not args.interactive: - logger.info("get prompts from images' meta data") - for img in init_images: - if "prompt" in img.text: - prompt = img.text["prompt"] - if "negative-prompt" in img.text: - prompt += " --n " + img.text["negative-prompt"] - prompt_list.append(prompt) - - # プロンプトと画像を一致させるため指定回数だけ繰り返す（画像を増幅する） - l = [] - for im in init_images: - l.extend([im] * args.images_per_prompt) - init_images = l - - if mask_images is not None: - l = [] - for im in mask_images: - l.extend([im] * args.images_per_prompt) - mask_images = l - - # 画像サイズにオプション指定があるときはリサイズする - if args.W is not None and args.H is not None: - # highres fix を考慮に入れる - w, h = args.W, args.H - if highres_fix: - w = int(w * args.highres_fix_scale + 0.5) - h = int(h * args.highres_fix_scale + 0.5) - - if init_images is not None: - logger.info(f"resize img2img source images to {w}*{h}") - init_images = resize_images(init_images, (w, h)) - if mask_images is not None: - logger.info(f"resize img2img mask images to {w}*{h}") - mask_images = resize_images(mask_images, (w, h)) - - regional_network = False - if networks and mask_images: - # mask を領域情報として流用する、現在は一回のコマンド呼び出しで1枚だけ対応 - regional_network = True - logger.info("use mask as region") - - size = None - for i, network in enumerate(networks): - if (i < 3 and args.network_regional_mask_max_color_codes is None) or i < args.network_regional_mask_max_color_codes: - np_mask = np.array(mask_images[0]) - - if args.network_regional_mask_max_color_codes: - # カラーコードでマスクを指定する - ch0 = (i + 1) & 1 - ch1 = ((i + 1) >> 1) & 1 - ch2 = ((i + 1) >> 2) & 1 - np_mask = np.all(np_mask == np.array([ch0, ch1, ch2]) * 255, axis=2) - np_mask = np_mask.astype(np.uint8) * 255 - else: - np_mask = np_mask[:, :, i] - size = np_mask.shape - else: - np_mask = np.full(size, 255, dtype=np.uint8) - mask = torch.from_numpy(np_mask.astype(np.float32) / 255.0) - network.set_region(i, i == len(networks) - 1, mask) - mask_images = None - - prev_image = None # for VGG16 guided - if args.guide_image_path is not None: - logger.info(f"load image for CLIP/VGG16/ControlNet guidance: {args.guide_image_path}") - guide_images = [] - for p in args.guide_image_path: - guide_images.extend(load_images(p)) - - logger.info(f"loaded {len(guide_images)} guide images for guidance") - if len(guide_images) == 0: - logger.info( - f"No guide image, use previous generated image. / ガイド画像がありません。直前に生成した画像を使います: {args.image_path}" - ) - guide_images = None - else: - guide_images = None - - # seed指定時はseedを決めておく - if args.seed is not None: - # dynamic promptを使うと足りなくなる→images_per_promptを適当に大きくしておいてもらう - random.seed(args.seed) - predefined_seeds = [random.randint(0, 0x7FFFFFFF) for _ in range(args.n_iter * len(prompt_list) * args.images_per_prompt)] - if len(predefined_seeds) == 1: - predefined_seeds[0] = args.seed - else: - predefined_seeds = None - - # デフォルト画像サイズを設定する：img2imgではこれらの値は無視される（またはW*Hにリサイズ済み） - if args.W is None: - args.W = 512 - if args.H is None: - args.H = 512 - - # 画像生成のループ - os.makedirs(args.outdir, exist_ok=True) - max_embeddings_multiples = 1 if args.max_embeddings_multiples is None else args.max_embeddings_multiples - - for gen_iter in range(args.n_iter): - logger.info(f"iteration {gen_iter+1}/{args.n_iter}") - iter_seed = random.randint(0, 0x7FFFFFFF) - - # shuffle prompt list - if args.shuffle_prompts: - random.shuffle(prompt_list) - - # バッチ処理の関数 - def process_batch(batch: List[BatchData], highres_fix, highres_1st=False): - batch_size = len(batch) - - # highres_fixの処理 - if highres_fix and not highres_1st: - # 1st stageのバッチを作成して呼び出す：サイズを小さくして呼び出す - is_1st_latent = upscaler.support_latents() if upscaler else args.highres_fix_latents_upscaling - - logger.info("process 1st stage") - batch_1st = [] - for _, base, ext in batch: - width_1st = int(ext.width * args.highres_fix_scale + 0.5) - height_1st = int(ext.height * args.highres_fix_scale + 0.5) - width_1st = width_1st - width_1st % 32 - height_1st = height_1st - height_1st % 32 - - strength_1st = ext.strength if args.highres_fix_strength is None else args.highres_fix_strength - - ext_1st = BatchDataExt( - width_1st, - height_1st, - args.highres_fix_steps, - ext.scale, - ext.negative_scale, - strength_1st, - ext.network_muls, - ext.num_sub_prompts, - ) - batch_1st.append(BatchData(is_1st_latent, base, ext_1st)) - - pipe.set_enable_control_net(True) # 1st stageではControlNetを有効にする - images_1st = process_batch(batch_1st, True, True) - - # 2nd stageのバッチを作成して以下処理する - logger.info("process 2nd stage") - width_2nd, height_2nd = batch[0].ext.width, batch[0].ext.height - - if upscaler: - # upscalerを使って画像を拡大する - lowreso_imgs = None if is_1st_latent else images_1st - lowreso_latents = None if not is_1st_latent else images_1st - - # 戻り値はPIL.Image.Imageかtorch.Tensorのlatents - batch_size = len(images_1st) - vae_batch_size = ( - batch_size - if args.vae_batch_size is None - else (max(1, int(batch_size * args.vae_batch_size)) if args.vae_batch_size < 1 else args.vae_batch_size) - ) - vae_batch_size = int(vae_batch_size) - images_1st = upscaler.upscale( - vae, lowreso_imgs, lowreso_latents, dtype, width_2nd, height_2nd, batch_size, vae_batch_size - ) - - elif args.highres_fix_latents_upscaling: - # latentを拡大する - org_dtype = images_1st.dtype - if images_1st.dtype == torch.bfloat16: - images_1st = images_1st.to(torch.float) # interpolateがbf16をサポートしていない - images_1st = torch.nn.functional.interpolate( - images_1st, (batch[0].ext.height // 8, batch[0].ext.width // 8), mode="bilinear" - ) # , antialias=True) - images_1st = images_1st.to(org_dtype) - - else: - # 画像をLANCZOSで拡大する - images_1st = [image.resize((width_2nd, height_2nd), resample=PIL.Image.LANCZOS) for image in images_1st] - - batch_2nd = [] - for i, (bd, image) in enumerate(zip(batch, images_1st)): - bd_2nd = BatchData(False, BatchDataBase(*bd.base[0:3], bd.base.seed + 1, image, None, *bd.base[6:]), bd.ext) - batch_2nd.append(bd_2nd) - batch = batch_2nd - - if args.highres_fix_disable_control_net: - pipe.set_enable_control_net(False) # オプション指定時、2nd stageではControlNetを無効にする - - # このバッチの情報を取り出す - ( - return_latents, - (step_first, _, _, _, init_image, mask_image, _, guide_image, _), - (width, height, steps, scale, negative_scale, strength, network_muls, num_sub_prompts), - ) = batch[0] - noise_shape = (LATENT_CHANNELS, height // DOWNSAMPLING_FACTOR, width // DOWNSAMPLING_FACTOR) - - prompts = [] - negative_prompts = [] - raw_prompts = [] - start_code = torch.zeros((batch_size, *noise_shape), device=device, dtype=dtype) - noises = [ - torch.zeros((batch_size, *noise_shape), device=device, dtype=dtype) - for _ in range(steps * scheduler_num_noises_per_step) - ] - seeds = [] - clip_prompts = [] - - if init_image is not None: # img2img? - i2i_noises = torch.zeros((batch_size, *noise_shape), device=device, dtype=dtype) - init_images = [] - - if mask_image is not None: - mask_images = [] - else: - mask_images = None - else: - i2i_noises = None - init_images = None - mask_images = None - - if guide_image is not None: # CLIP image guided? - guide_images = [] - else: - guide_images = None - - # バッチ内の位置に関わらず同じ乱数を使うためにここで乱数を生成しておく。あわせてimage/maskがbatch内で同一かチェックする - all_images_are_same = True - all_masks_are_same = True - all_guide_images_are_same = True - for i, ( - _, - (_, prompt, negative_prompt, seed, init_image, mask_image, clip_prompt, guide_image, raw_prompt), - _, - ) in enumerate(batch): - prompts.append(prompt) - negative_prompts.append(negative_prompt) - seeds.append(seed) - clip_prompts.append(clip_prompt) - raw_prompts.append(raw_prompt) - - if init_image is not None: - init_images.append(init_image) - if i > 0 and all_images_are_same: - all_images_are_same = init_images[-2] is init_image - - if mask_image is not None: - mask_images.append(mask_image) - if i > 0 and all_masks_are_same: - all_masks_are_same = mask_images[-2] is mask_image - - if guide_image is not None: - if type(guide_image) is list: - guide_images.extend(guide_image) - all_guide_images_are_same = False - else: - guide_images.append(guide_image) - if i > 0 and all_guide_images_are_same: - all_guide_images_are_same = guide_images[-2] is guide_image - - # make start code - torch.manual_seed(seed) - start_code[i] = torch.randn(noise_shape, device=device, dtype=dtype) - - # make each noises - for j in range(steps * scheduler_num_noises_per_step): - noises[j][i] = torch.randn(noise_shape, device=device, dtype=dtype) - - if i2i_noises is not None: # img2img noise - i2i_noises[i] = torch.randn(noise_shape, device=device, dtype=dtype) - - noise_manager.reset_sampler_noises(noises) - - # すべての画像が同じなら1枚だけpipeに渡すことでpipe側で処理を高速化する - if init_images is not None and all_images_are_same: - init_images = init_images[0] - if mask_images is not None and all_masks_are_same: - mask_images = mask_images[0] - if guide_images is not None and all_guide_images_are_same: - guide_images = guide_images[0] - - # ControlNet使用時はguide imageをリサイズする - if control_nets: - # TODO resampleのメソッド - guide_images = guide_images if type(guide_images) == list else [guide_images] - guide_images = [i.resize((width, height), resample=PIL.Image.LANCZOS) for i in guide_images] - if len(guide_images) == 1: - guide_images = guide_images[0] - - # generate - if networks: - # 追加ネットワークの処理 - shared = {} - for n, m in zip(networks, network_muls if network_muls else network_default_muls): - n.set_multiplier(m) - if regional_network: - n.set_current_generation(batch_size, num_sub_prompts, width, height, shared) - - if not regional_network and network_pre_calc: - for n in networks: - n.restore_weights() - for n in networks: - n.pre_calculation() - logger.info("pre-calculation... done") - - images = pipe( - prompts, - negative_prompts, - init_images, - mask_images, - height, - width, - steps, - scale, - negative_scale, - strength, - latents=start_code, - output_type="pil", - max_embeddings_multiples=max_embeddings_multiples, - img2img_noise=i2i_noises, - vae_batch_size=args.vae_batch_size, - return_latents=return_latents, - clip_prompts=clip_prompts, - clip_guide_images=guide_images, - )[0] - if highres_1st and not args.highres_fix_save_1st: # return images or latents - return images - - # save image - highres_prefix = ("0" if highres_1st else "1") if highres_fix else "" - ts_str = time.strftime("%Y%m%d%H%M%S", time.localtime()) - for i, (image, prompt, negative_prompts, seed, clip_prompt, raw_prompt) in enumerate( - zip(images, prompts, negative_prompts, seeds, clip_prompts, raw_prompts) - ): - if highres_fix: - seed -= 1 # record original seed - metadata = PngInfo() - metadata.add_text("prompt", prompt) - metadata.add_text("seed", str(seed)) - metadata.add_text("sampler", args.sampler) - metadata.add_text("steps", str(steps)) - metadata.add_text("scale", str(scale)) - if negative_prompt is not None: - metadata.add_text("negative-prompt", negative_prompt) - if negative_scale is not None: - metadata.add_text("negative-scale", str(negative_scale)) - if clip_prompt is not None: - metadata.add_text("clip-prompt", clip_prompt) - if raw_prompt is not None: - metadata.add_text("raw-prompt", raw_prompt) - - if args.use_original_file_name and init_images is not None: - if type(init_images) is list: - fln = os.path.splitext(os.path.basename(init_images[i % len(init_images)].filename))[0] + ".png" - else: - fln = os.path.splitext(os.path.basename(init_images.filename))[0] + ".png" - elif args.sequential_file_name: - fln = f"im_{highres_prefix}{step_first + i + 1:06d}.png" - else: - fln = f"im_{ts_str}_{highres_prefix}{i:03d}_{seed}.png" - - image.save(os.path.join(args.outdir, fln), pnginfo=metadata) - - if not args.no_preview and not highres_1st and args.interactive: - try: - import cv2 - - for prompt, image in zip(prompts, images): - cv2.imshow(prompt[:128], np.array(image)[:, :, ::-1]) # プロンプトが長いと死ぬ - cv2.waitKey() - cv2.destroyAllWindows() - except ImportError: - logger.info( - "opencv-python is not installed, cannot preview / opencv-pythonがインストールされていないためプレビューできません" - ) - - return images - - # 画像生成のプロンプトが一周するまでのループ - prompt_index = 0 - global_step = 0 - batch_data = [] - while args.interactive or prompt_index < len(prompt_list): - if len(prompt_list) == 0: - # interactive - valid = False - while not valid: - logger.info("") - logger.info("Type prompt:") - try: - raw_prompt = input() - except EOFError: - break - - valid = len(raw_prompt.strip().split(" --")[0].strip()) > 0 - if not valid: # EOF, end app - break - else: - raw_prompt = prompt_list[prompt_index] - - # sd-dynamic-prompts like variants: - # count is 1 (not dynamic) or images_per_prompt (no enumeration) or arbitrary (enumeration) - raw_prompts = handle_dynamic_prompt_variants(raw_prompt, args.images_per_prompt) - - # repeat prompt - for pi in range(args.images_per_prompt if len(raw_prompts) == 1 else len(raw_prompts)): - raw_prompt = raw_prompts[pi] if len(raw_prompts) > 1 else raw_prompts[0] - - if pi == 0 or len(raw_prompts) > 1: - # parse prompt: if prompt is not changed, skip parsing - width = args.W - height = args.H - scale = args.scale - negative_scale = args.negative_scale - steps = args.steps - seed = None - seeds = None - strength = 0.8 if args.strength is None else args.strength - negative_prompt = "" - clip_prompt = None - network_muls = None - - # Deep Shrink - ds_depth_1 = None # means no override - ds_timesteps_1 = args.ds_timesteps_1 - ds_depth_2 = args.ds_depth_2 - ds_timesteps_2 = args.ds_timesteps_2 - ds_ratio = args.ds_ratio - - # Gradual Latent - gl_timesteps = None # means no override - gl_ratio = args.gradual_latent_ratio - gl_every_n_steps = args.gradual_latent_every_n_steps - gl_ratio_step = args.gradual_latent_ratio_step - gl_s_noise = args.gradual_latent_s_noise - gl_unsharp_params = args.gradual_latent_unsharp_params - - prompt_args = raw_prompt.strip().split(" --") - prompt = prompt_args[0] - logger.info(f"prompt {prompt_index+1}/{len(prompt_list)}: {prompt}") - - for parg in prompt_args[1:]: - try: - m = re.match(r"w (\d+)", parg, re.IGNORECASE) - if m: - width = int(m.group(1)) - logger.info(f"width: {width}") - continue - - m = re.match(r"h (\d+)", parg, re.IGNORECASE) - if m: - height = int(m.group(1)) - logger.info(f"height: {height}") - continue - - m = re.match(r"s (\d+)", parg, re.IGNORECASE) - if m: # steps - steps = max(1, min(1000, int(m.group(1)))) - logger.info(f"steps: {steps}") - continue - - m = re.match(r"d ([\d,]+)", parg, re.IGNORECASE) - if m: # seed - seeds = [int(d) for d in m.group(1).split(",")] - logger.info(f"seeds: {seeds}") - continue - - m = re.match(r"l ([\d\.]+)", parg, re.IGNORECASE) - if m: # scale - scale = float(m.group(1)) - logger.info(f"scale: {scale}") - continue - - m = re.match(r"nl ([\d\.]+|none|None)", parg, re.IGNORECASE) - if m: # negative scale - if m.group(1).lower() == "none": - negative_scale = None - else: - negative_scale = float(m.group(1)) - logger.info(f"negative scale: {negative_scale}") - continue - - m = re.match(r"t ([\d\.]+)", parg, re.IGNORECASE) - if m: # strength - strength = float(m.group(1)) - logger.info(f"strength: {strength}") - continue - - m = re.match(r"n (.+)", parg, re.IGNORECASE) - if m: # negative prompt - negative_prompt = m.group(1) - logger.info(f"negative prompt: {negative_prompt}") - continue - - m = re.match(r"c (.+)", parg, re.IGNORECASE) - if m: # clip prompt - clip_prompt = m.group(1) - logger.info(f"clip prompt: {clip_prompt}") - continue - - m = re.match(r"am ([\d\.\-,]+)", parg, re.IGNORECASE) - if m: # network multiplies - network_muls = [float(v) for v in m.group(1).split(",")] - while len(network_muls) < len(networks): - network_muls.append(network_muls[-1]) - logger.info(f"network mul: {network_muls}") - continue - - # Deep Shrink - m = re.match(r"dsd1 ([\d\.]+)", parg, re.IGNORECASE) - if m: # deep shrink depth 1 - ds_depth_1 = int(m.group(1)) - logger.info(f"deep shrink depth 1: {ds_depth_1}") - continue - - m = re.match(r"dst1 ([\d\.]+)", parg, re.IGNORECASE) - if m: # deep shrink timesteps 1 - ds_timesteps_1 = int(m.group(1)) - ds_depth_1 = ds_depth_1 if ds_depth_1 is not None else -1 # -1 means override - logger.info(f"deep shrink timesteps 1: {ds_timesteps_1}") - continue - - m = re.match(r"dsd2 ([\d\.]+)", parg, re.IGNORECASE) - if m: # deep shrink depth 2 - ds_depth_2 = int(m.group(1)) - ds_depth_1 = ds_depth_1 if ds_depth_1 is not None else -1 # -1 means override - logger.info(f"deep shrink depth 2: {ds_depth_2}") - continue - - m = re.match(r"dst2 ([\d\.]+)", parg, re.IGNORECASE) - if m: # deep shrink timesteps 2 - ds_timesteps_2 = int(m.group(1)) - ds_depth_1 = ds_depth_1 if ds_depth_1 is not None else -1 # -1 means override - logger.info(f"deep shrink timesteps 2: {ds_timesteps_2}") - continue - - m = re.match(r"dsr ([\d\.]+)", parg, re.IGNORECASE) - if m: # deep shrink ratio - ds_ratio = float(m.group(1)) - ds_depth_1 = ds_depth_1 if ds_depth_1 is not None else -1 # -1 means override - logger.info(f"deep shrink ratio: {ds_ratio}") - continue - - # Gradual Latent - m = re.match(r"glt ([\d\.]+)", parg, re.IGNORECASE) - if m: # gradual latent timesteps - gl_timesteps = int(m.group(1)) - logger.info(f"gradual latent timesteps: {gl_timesteps}") - continue - - m = re.match(r"glr ([\d\.]+)", parg, re.IGNORECASE) - if m: # gradual latent ratio - gl_ratio = float(m.group(1)) - gl_timesteps = gl_timesteps if gl_timesteps is not None else -1 # -1 means override - logger.info(f"gradual latent ratio: {ds_ratio}") - continue - - m = re.match(r"gle ([\d\.]+)", parg, re.IGNORECASE) - if m: # gradual latent every n steps - gl_every_n_steps = int(m.group(1)) - gl_timesteps = gl_timesteps if gl_timesteps is not None else -1 # -1 means override - logger.info(f"gradual latent every n steps: {gl_every_n_steps}") - continue - - m = re.match(r"gls ([\d\.]+)", parg, re.IGNORECASE) - if m: # gradual latent ratio step - gl_ratio_step = float(m.group(1)) - gl_timesteps = gl_timesteps if gl_timesteps is not None else -1 # -1 means override - logger.info(f"gradual latent ratio step: {gl_ratio_step}") - continue - - m = re.match(r"glsn ([\d\.]+)", parg, re.IGNORECASE) - if m: # gradual latent s noise - gl_s_noise = float(m.group(1)) - gl_timesteps = gl_timesteps if gl_timesteps is not None else -1 # -1 means override - logger.info(f"gradual latent s noise: {gl_s_noise}") - continue - - m = re.match(r"glus ([\d\.\-,]+)", parg, re.IGNORECASE) - if m: # gradual latent unsharp params - gl_unsharp_params = m.group(1) - gl_timesteps = gl_timesteps if gl_timesteps is not None else -1 # -1 means override - logger.info(f"gradual latent unsharp params: {gl_unsharp_params}") - continue - - except ValueError as ex: - logger.info(f"Exception in parsing / 解析エラー: {parg}") - logger.info(ex) - - # override Deep Shrink - if ds_depth_1 is not None: - if ds_depth_1 < 0: - ds_depth_1 = args.ds_depth_1 or 3 - unet.set_deep_shrink(ds_depth_1, ds_timesteps_1, ds_depth_2, ds_timesteps_2, ds_ratio) - - # override Gradual Latent - if gl_timesteps is not None: - if gl_timesteps < 0: - gl_timesteps = args.gradual_latent_timesteps or 650 - if gl_unsharp_params is not None: - unsharp_params = gl_unsharp_params.split(",") - us_ksize, us_sigma, us_strength = [float(v) for v in unsharp_params[:3]] - logger.info(f'{unsharp_params}') - us_target_x = True if len(unsharp_params) < 4 else bool(int(unsharp_params[3])) - us_ksize = int(us_ksize) - else: - us_ksize, us_sigma, us_strength, us_target_x = None, None, None, None - gradual_latent = GradualLatent( - gl_ratio, - gl_timesteps, - gl_every_n_steps, - gl_ratio_step, - gl_s_noise, - us_ksize, - us_sigma, - us_strength, - us_target_x, - ) - pipe.set_gradual_latent(gradual_latent) - - # prepare seed - if seeds is not None: # given in prompt - # 数が足りないなら前のをそのまま使う - if len(seeds) > 0: - seed = seeds.pop(0) - else: - if predefined_seeds is not None: - if len(predefined_seeds) > 0: - seed = predefined_seeds.pop(0) - else: - logger.info("predefined seeds are exhausted") - seed = None - elif args.iter_same_seed: - seed = iter_seed - else: - seed = None # 前のを消す - - if seed is None: - seed = random.randint(0, 0x7FFFFFFF) - if args.interactive: - logger.info(f"seed: {seed}") - - # prepare init image, guide image and mask - init_image = mask_image = guide_image = None - - # 同一イメージを使うとき、本当はlatentに変換しておくと無駄がないが面倒なのでとりあえず毎回処理する - if init_images is not None: - init_image = init_images[global_step % len(init_images)] - - # img2imgの場合は、基本的に元画像のサイズで生成する。highres fixの場合はargs.W, args.Hとscaleに従いリサイズ済みなので無視する - # 32単位に丸めたやつにresizeされるので踏襲する - if not highres_fix: - width, height = init_image.size - width = width - width % 32 - height = height - height % 32 - if width != init_image.size[0] or height != init_image.size[1]: - logger.info( - f"img2img image size is not divisible by 32 so aspect ratio is changed / img2imgの画像サイズが32で割り切れないためリサイズされます。画像が歪みます" - ) - - if mask_images is not None: - mask_image = mask_images[global_step % len(mask_images)] - - if guide_images is not None: - if control_nets: # 複数件の場合あり - c = len(control_nets) - p = global_step % (len(guide_images) // c) - guide_image = guide_images[p * c : p * c + c] - else: - guide_image = guide_images[global_step % len(guide_images)] - elif args.clip_image_guidance_scale > 0 or args.vgg16_guidance_scale > 0: - if prev_image is None: - logger.info("Generate 1st image without guide image.") - else: - logger.info("Use previous image as guide image.") - guide_image = prev_image - - if regional_network: - num_sub_prompts = len(prompt.split(" AND ")) - assert ( - len(networks) <= num_sub_prompts - ), "Number of networks must be less than or equal to number of sub prompts." - else: - num_sub_prompts = None - - b1 = BatchData( - False, - BatchDataBase( - global_step, prompt, negative_prompt, seed, init_image, mask_image, clip_prompt, guide_image, raw_prompt - ), - BatchDataExt( - width, - height, - steps, - scale, - negative_scale, - strength, - tuple(network_muls) if network_muls else None, - num_sub_prompts, - ), - ) - if len(batch_data) > 0 and batch_data[-1].ext != b1.ext: # バッチ分割必要？ - process_batch(batch_data, highres_fix) - batch_data.clear() - - batch_data.append(b1) - if len(batch_data) == args.batch_size: - prev_image = process_batch(batch_data, highres_fix)[0] - batch_data.clear() - - global_step += 1 - - prompt_index += 1 - - if len(batch_data) > 0: - process_batch(batch_data, highres_fix) - batch_data.clear() - - logger.info("done!") - - -def setup_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser() - - add_logging_arguments(parser) - - parser.add_argument( - "--v2", action="store_true", help="load Stable Diffusion v2.0 model / Stable Diffusion 2.0のモデルを読み込む" - ) - parser.add_argument( - "--v_parameterization", action="store_true", help="enable v-parameterization training / v-parameterization学習を有効にする" - ) - parser.add_argument("--prompt", type=str, default=None, help="prompt / プロンプト") - parser.add_argument( - "--from_file", - type=str, - default=None, - help="if specified, load prompts from this file / 指定時はプロンプトをファイルから読み込む", - ) - parser.add_argument( - "--interactive", - action="store_true", - help="interactive mode (generates one image) / 対話モード（生成される画像は1枚になります）", - ) - parser.add_argument( - "--no_preview", action="store_true", help="do not show generated image in interactive mode / 対話モードで画像を表示しない" - ) - parser.add_argument( - "--image_path", type=str, default=None, help="image to inpaint or to generate from / img2imgまたはinpaintを行う元画像" - ) - parser.add_argument("--mask_path", type=str, default=None, help="mask in inpainting / inpaint時のマスク") - parser.add_argument("--strength", type=float, default=None, help="img2img strength / img2img時のstrength") - parser.add_argument("--images_per_prompt", type=int, default=1, help="number of images per prompt / プロンプトあたりの出力枚数") - parser.add_argument("--outdir", type=str, default="outputs", help="dir to write results to / 生成画像の出力先") - parser.add_argument( - "--sequential_file_name", action="store_true", help="sequential output file name / 生成画像のファイル名を連番にする" - ) - parser.add_argument( - "--use_original_file_name", - action="store_true", - help="prepend original file name in img2img / img2imgで元画像のファイル名を生成画像のファイル名の先頭に付ける", - ) - # parser.add_argument("--ddim_eta", type=float, default=0.0, help="ddim eta (eta=0.0 corresponds to deterministic sampling", ) - parser.add_argument("--n_iter", type=int, default=1, help="sample this often / 繰り返し回数") - parser.add_argument("--H", type=int, default=None, help="image height, in pixel space / 生成画像高さ") - parser.add_argument("--W", type=int, default=None, help="image width, in pixel space / 生成画像幅") - parser.add_argument("--batch_size", type=int, default=1, help="batch size / バッチサイズ") - parser.add_argument( - "--vae_batch_size", - type=float, - default=None, - help="batch size for VAE, < 1.0 for ratio / VAE処理時のバッチサイズ、1未満の値の場合は通常バッチサイズの比率", - ) - parser.add_argument( - "--vae_slices", - type=int, - default=None, - help="number of slices to split image into for VAE to reduce VRAM usage, None for no splitting (default), slower if specified. 16 or 32 recommended / VAE処理時にVRAM使用量削減のため画像を分割するスライス数、Noneの場合は分割しない（デフォルト）、指定すると遅くなる。16か32程度を推奨", - ) - parser.add_argument("--steps", type=int, default=50, help="number of ddim sampling steps / サンプリングステップ数") - parser.add_argument( - "--sampler", - type=str, - default="ddim", - choices=[ - "ddim", - "pndm", - "lms", - "euler", - "euler_a", - "heun", - "dpm_2", - "dpm_2_a", - "dpmsolver", - "dpmsolver++", - "dpmsingle", - "k_lms", - "k_euler", - "k_euler_a", - "k_dpm_2", - "k_dpm_2_a", - ], - help=f"sampler (scheduler) type / サンプラー（スケジューラ）の種類", - ) - parser.add_argument( - "--scale", - type=float, - default=7.5, - help="unconditional guidance scale: eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty)) / guidance scale", - ) - parser.add_argument( - "--ckpt", type=str, default=None, help="path to checkpoint of model / モデルのcheckpointファイルまたはディレクトリ" - ) - parser.add_argument( - "--vae", - type=str, - default=None, - help="path to checkpoint of vae to replace / VAEを入れ替える場合、VAEのcheckpointファイルまたはディレクトリ", - ) - parser.add_argument( - "--tokenizer_cache_dir", - type=str, - default=None, - help="directory for caching Tokenizer (for offline training) / Tokenizerをキャッシュするディレクトリ（ネット接続なしでの学習のため）", - ) - # parser.add_argument("--replace_clip_l14_336", action='store_true', - # help="Replace CLIP (Text Encoder) to l/14@336 / CLIP(Text Encoder)をl/14@336に入れ替える") - parser.add_argument( - "--seed", - type=int, - default=None, - help="seed, or seed of seeds in multiple generation / 1枚生成時のseed、または複数枚生成時の乱数seedを決めるためのseed", - ) - parser.add_argument( - "--iter_same_seed", - action="store_true", - help="use same seed for all prompts in iteration if no seed specified / 乱数seedの指定がないとき繰り返し内はすべて同じseedを使う（プロンプト間の差異の比較用）", - ) - parser.add_argument( - "--shuffle_prompts", - action="store_true", - help="shuffle prompts in iteration / 繰り返し内のプロンプトをシャッフルする", - ) - parser.add_argument("--fp16", action="store_true", help="use fp16 / fp16を指定し省メモリ化する") - parser.add_argument("--bf16", action="store_true", help="use bfloat16 / bfloat16を指定し省メモリ化する") - parser.add_argument("--xformers", action="store_true", help="use xformers / xformersを使用し高速化する") - parser.add_argument("--sdpa", action="store_true", help="use sdpa in PyTorch 2 / sdpa") - parser.add_argument( - "--diffusers_xformers", - action="store_true", - help="use xformers by diffusers (Hypernetworks doesn't work) / Diffusersでxformersを使用する（Hypernetwork利用不可）", - ) - parser.add_argument( - "--opt_channels_last", - action="store_true", - help="set channels last option to model / モデルにchannels lastを指定し最適化する", - ) - parser.add_argument( - "--network_module", - type=str, - default=None, - nargs="*", - help="additional network module to use / 追加ネットワークを使う時そのモジュール名", - ) - parser.add_argument( - "--network_weights", type=str, default=None, nargs="*", help="additional network weights to load / 追加ネットワークの重み" - ) - parser.add_argument( - "--network_mul", type=float, default=None, nargs="*", help="additional network multiplier / 追加ネットワークの効果の倍率" - ) - parser.add_argument( - "--network_args", - type=str, - default=None, - nargs="*", - help="additional arguments for network (key=value) / ネットワークへの追加の引数", - ) - parser.add_argument( - "--network_show_meta", action="store_true", help="show metadata of network model / ネットワークモデルのメタデータを表示する" - ) - parser.add_argument( - "--network_merge_n_models", - type=int, - default=None, - help="merge this number of networks / この数だけネットワークをマージする", - ) - parser.add_argument( - "--network_merge", action="store_true", help="merge network weights to original model / ネットワークの重みをマージする" - ) - parser.add_argument( - "--network_pre_calc", - action="store_true", - help="pre-calculate network for generation / ネットワークのあらかじめ計算して生成する", - ) - parser.add_argument( - "--network_regional_mask_max_color_codes", - type=int, - default=None, - help="max color codes for regional mask (default is None, mask by channel) / regional maskの最大色数（デフォルトはNoneでチャンネルごとのマスク）", - ) - parser.add_argument( - "--textual_inversion_embeddings", - type=str, - default=None, - nargs="*", - help="Embeddings files of Textual Inversion / Textual Inversionのembeddings", - ) - parser.add_argument( - "--XTI_embeddings", - type=str, - default=None, - nargs="*", - help="Embeddings files of Extended Textual Inversion / Extended Textual Inversionのembeddings", - ) - parser.add_argument( - "--clip_skip", type=int, default=None, help="layer number from bottom to use in CLIP / CLIPの後ろからn層目の出力を使う" - ) - parser.add_argument( - "--max_embeddings_multiples", - type=int, - default=None, - help="max embedding multiples, max token length is 75 * multiples / トークン長をデフォルトの何倍とするか 75*この値 がトークン長となる", - ) - parser.add_argument( - "--clip_guidance_scale", - type=float, - default=0.0, - help="enable CLIP guided SD, scale for guidance (DDIM, PNDM, LMS samplers only) / CLIP guided SDを有効にしてこのscaleを適用する（サンプラーはDDIM、PNDM、LMSのみ）", - ) - parser.add_argument( - "--clip_image_guidance_scale", - type=float, - default=0.0, - help="enable CLIP guided SD by image, scale for guidance / 画像によるCLIP guided SDを有効にしてこのscaleを適用する", - ) - parser.add_argument( - "--vgg16_guidance_scale", - type=float, - default=0.0, - help="enable VGG16 guided SD by image, scale for guidance / 画像によるVGG16 guided SDを有効にしてこのscaleを適用する", - ) - parser.add_argument( - "--vgg16_guidance_layer", - type=int, - default=20, - help="layer of VGG16 to calculate contents guide (1~30, 20 for conv4_2) / VGG16のcontents guideに使うレイヤー番号 (1~30、20はconv4_2)", - ) - parser.add_argument( - "--guide_image_path", type=str, default=None, nargs="*", help="image to CLIP guidance / CLIP guided SDでガイドに使う画像" - ) - parser.add_argument( - "--highres_fix_scale", - type=float, - default=None, - help="enable highres fix, reso scale for 1st stage / highres fixを有効にして最初の解像度をこのscaleにする", - ) - parser.add_argument( - "--highres_fix_steps", - type=int, - default=28, - help="1st stage steps for highres fix / highres fixの最初のステージのステップ数", - ) - parser.add_argument( - "--highres_fix_strength", - type=float, - default=None, - help="1st stage img2img strength for highres fix / highres fixの最初のステージのimg2img時のstrength、省略時はstrengthと同じ", - ) - parser.add_argument( - "--highres_fix_save_1st", - action="store_true", - help="save 1st stage images for highres fix / highres fixの最初のステージの画像を保存する", - ) - parser.add_argument( - "--highres_fix_latents_upscaling", - action="store_true", - help="use latents upscaling for highres fix / highres fixでlatentで拡大する", - ) - parser.add_argument( - "--highres_fix_upscaler", - type=str, - default=None, - help="upscaler module for highres fix / highres fixで使うupscalerのモジュール名", - ) - parser.add_argument( - "--highres_fix_upscaler_args", - type=str, - default=None, - help="additional arguments for upscaler (key=value) / upscalerへの追加の引数", - ) - parser.add_argument( - "--highres_fix_disable_control_net", - action="store_true", - help="disable ControlNet for highres fix / highres fixでControlNetを使わない", - ) - - parser.add_argument( - "--negative_scale", - type=float, - default=None, - help="set another guidance scale for negative prompt / ネガティブプロンプトのscaleを指定する", - ) - - parser.add_argument( - "--control_net_models", type=str, default=None, nargs="*", help="ControlNet models to use / 使用するControlNetのモデル名" - ) - parser.add_argument( - "--control_net_preps", - type=str, - default=None, - nargs="*", - help="ControlNet preprocess to use / 使用するControlNetのプリプロセス名", - ) - parser.add_argument("--control_net_weights", type=float, default=None, nargs="*", help="ControlNet weights / ControlNetの重み") - parser.add_argument( - "--control_net_ratios", - type=float, - default=None, - nargs="*", - help="ControlNet guidance ratio for steps / ControlNetでガイドするステップ比率", - ) - # parser.add_argument( - # "--control_net_image_path", type=str, default=None, nargs="*", help="image for ControlNet guidance / ControlNetでガイドに使う画像" - # ) - - # Deep Shrink - parser.add_argument( - "--ds_depth_1", - type=int, - default=None, - help="Enable Deep Shrink with this depth 1, valid values are 0 to 3 / Deep Shrinkをこのdepthで有効にする", - ) - parser.add_argument( - "--ds_timesteps_1", - type=int, - default=650, - help="Apply Deep Shrink depth 1 until this timesteps / Deep Shrink depth 1を適用するtimesteps", - ) - parser.add_argument("--ds_depth_2", type=int, default=None, help="Deep Shrink depth 2 / Deep Shrinkのdepth 2") - parser.add_argument( - "--ds_timesteps_2", - type=int, - default=650, - help="Apply Deep Shrink depth 2 until this timesteps / Deep Shrink depth 2を適用するtimesteps", - ) - parser.add_argument( - "--ds_ratio", type=float, default=0.5, help="Deep Shrink ratio for downsampling / Deep Shrinkのdownsampling比率" - ) - - # gradual latent - parser.add_argument( - "--gradual_latent_timesteps", - type=int, - default=None, - help="enable Gradual Latent hires fix and apply upscaling from this timesteps / Gradual Latent hires fixをこのtimestepsで有効にし、このtimestepsからアップスケーリングを適用する", - ) - parser.add_argument( - "--gradual_latent_ratio", - type=float, - default=0.5, - help=" this size ratio, 0.5 means 1/2 / Gradual Latent hires fixをこのサイズ比率で有効にする、0.5は1/2を意味する", - ) - parser.add_argument( - "--gradual_latent_ratio_step", - type=float, - default=0.125, - help="step to increase ratio for Gradual Latent / Gradual Latentのratioをどのくらいずつ上げるか", - ) - parser.add_argument( - "--gradual_latent_every_n_steps", - type=int, - default=3, - help="steps to increase size of latents every this steps for Gradual Latent / Gradual Latentでlatentsのサイズをこのステップごとに上げる", - ) - parser.add_argument( - "--gradual_latent_s_noise", - type=float, - default=1.0, - help="s_noise for Gradual Latent / Gradual Latentのs_noise", - ) - parser.add_argument( - "--gradual_latent_unsharp_params", - type=str, - default=None, - help="unsharp mask parameters for Gradual Latent: ksize, sigma, strength, target-x (1 means True). `3,0.5,0.5,1` or `3,1.0,1.0,0` is recommended /" - + " Gradual Latentのunsharp maskのパラメータ: ksize, sigma, strength, target-x. `3,0.5,0.5,1` または `3,1.0,1.0,0` が推奨", - ) - - return parser - - -if __name__ == "__main__": - parser = setup_parser() - - args = parser.parse_args() - setup_logging(args, reset=True) - main(args) diff --git a/inpainting_minimal_inference.py b/inpainting_minimal_inference.py index 421554b214..fb5e9fbd9d 100644 --- a/inpainting_minimal_inference.py +++ b/inpainting_minimal_inference.py @@ -38,7 +38,7 @@ from diffusers import EulerDiscreteScheduler from PIL import Image from tqdm import tqdm -from transformers import CLIPTokenizer +from library.clip_tokenizer import CLIPTokenizer # transformers.CLIPTokenizer with the legacy (ftfy) text normalization from library.device_utils import get_preferred_device, init_ipex diff --git a/library/_cv2_stub/__init__.py b/library/_cv2_stub/__init__.py new file mode 100644 index 0000000000..e5f398c955 --- /dev/null +++ b/library/_cv2_stub/__init__.py @@ -0,0 +1,523 @@ +""" +Pure-Python / NumPy / Pillow fallback for a subset of the OpenCV (cv2) API. + +This stub is automatically registered as the ``cv2`` module in ``sys.modules`` +by :mod:`library.cv2_compat` when the real ``opencv-python`` package is not +installed, so that the rest of the codebase can keep using ``import cv2`` +without modification. + +Only the API surface actually used by the training / core-library code paths +is implemented. In particular: + +- ``cvtColor`` supports BGR/RGB/BGRA/RGBA/HSV/GRAY conversions used in the + codebase (the HSV variant follows OpenCV's H=[0,180) convention). +- ``resize`` reproduces OpenCV's ``INTER_AREA`` and ``INTER_LINEAR`` sampling + in NumPy (these are the modes the dataset pipeline uses by default, so + results match a real OpenCV install up to rounding). ``INTER_CUBIC``, + ``INTER_LANCZOS4`` and ``INTER_NEAREST*`` go through Pillow's closest + ``PIL.Image.Resampling`` filter and differ slightly from OpenCV. +- ``imshow`` displays the image through Pillow's default image viewer + (``PIL.Image.show``). ``waitKey`` blocks on ``input()`` so that the caller + can page through images one at a time; it returns ``ord(s[0])`` of the + first character typed (or ``13`` for an empty Enter) so existing checks + like ``k == ord("e")`` continue to work. ``destroyAllWindows`` is a no-op. +- ``imwrite`` is implemented via Pillow (BGR -> RGB). + +Tools that rely on richer OpenCV functionality (Canny, warpAffine, face +detection, imdecode/imencode, etc.) check :data:`library.cv2_compat.HAS_OPENCV` +and bail out with a clear error message when real OpenCV is missing. +""" + +from __future__ import annotations + +import logging +from typing import Optional, Tuple + +import numpy as np +from PIL import Image + +logger = logging.getLogger(__name__) + +__version__ = "0.0.0-stub" + +# --------------------------------------------------------------------------- +# Constants (unique integers; values do not need to match real OpenCV) +# --------------------------------------------------------------------------- + +# Color conversion codes +COLOR_BGR2RGB = 4 +COLOR_RGB2BGR = 5 +COLOR_BGRA2RGBA = 6 +COLOR_RGBA2BGRA = 7 +COLOR_BGR2GRAY = 10 +COLOR_RGB2GRAY = 11 +COLOR_GRAY2BGR = 12 +COLOR_GRAY2RGB = 13 +COLOR_BGR2HSV = 40 +COLOR_HSV2BGR = 41 +COLOR_RGB2HSV = 42 +COLOR_HSV2RGB = 43 + +# Interpolation flags +INTER_NEAREST = 100 +INTER_LINEAR = 101 +INTER_CUBIC = 102 +INTER_AREA = 103 +INTER_LANCZOS4 = 104 +INTER_NEAREST_EXACT = 105 + +# Image read flags +IMREAD_UNCHANGED = -1 +IMREAD_GRAYSCALE = 0 +IMREAD_COLOR = 1 + +# Border types (only values actually referenced are needed) +BORDER_CONSTANT = 200 +BORDER_REFLECT = 201 +BORDER_REPLICATE = 202 + +# imwrite params +IMWRITE_JPEG_QUALITY = 300 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +# Interpolation modes that go through Pillow. INTER_AREA and INTER_LINEAR are +# implemented in NumPy instead (see _area_resample_rows / _linear_resample_rows) +# because Pillow's filters are antialiased and do not reproduce OpenCV's output. +_INTER_TO_PIL = { + INTER_NEAREST: Image.Resampling.NEAREST, + INTER_NEAREST_EXACT: Image.Resampling.NEAREST, + INTER_CUBIC: Image.Resampling.BICUBIC, + INTER_LANCZOS4: Image.Resampling.LANCZOS, +} + + +def _resize_single_channel(ch: np.ndarray, dsize: Tuple[int, int], pil_interp) -> np.ndarray: + return np.asarray(Image.fromarray(ch).resize(dsize, resample=pil_interp)) + + +def _area_resample_rows(x: np.ndarray, new_len: int) -> np.ndarray: + """Exact pixel-area resampling along axis 0 (OpenCV INTER_AREA when shrinking). + + Output row j is the average of the input over its footprint + ``[j * scale, (j + 1) * scale)``, which is exactly what OpenCV computes + when both axes shrink. Only used for ``scale > 1``. + """ + n = x.shape[0] + if new_len == n: + return x + scale = 1.0 / (new_len / n) # same rounding as OpenCV's scale = 1 / inv_scale + dt = _work_dtype(x, n) + + # The footprint of output row j is: a fractional first row idx_j, between + # k-1 and k whole rows (k = floor(scale)), and a fractional last row + # idx_{j+1}. The whole rows are summed with a fixed-width gather (rows + # idx_j+1 .. idx_j+k) followed by a plain sum, which is much faster in + # NumPy than cumsum or reduceat. Whenever the gather window runs past the + # whole rows it lands on row idx_{j+1} (`extra` times), which is taken + # out again through the last-row weight. + bounds = np.arange(new_len + 1, dtype=np.float64) * scale + bounds[-1] = n + idx = np.minimum(np.floor(bounds).astype(np.int64), n - 1) + frac = bounds - idx + k = int(np.floor(scale)) + + rows = np.minimum(idx[:-1, None] + 1 + np.arange(k, dtype=np.int64)[None, :], n - 1) + gathered = np.take(x, rows.ravel(), axis=0) + sums = gathered.reshape((new_len, k) + x.shape[1:]).sum(axis=1, dtype=dt) + + whole = idx[1:] - idx[:-1] - 1 # whole rows actually inside the footprint (k-1 or k) + extra = k - whole # how many gathered rows were row idx_{j+1} instead + w_first = _column((1.0 - frac[:-1]).astype(dt), x.ndim) + w_last = _column((frac[1:] - extra).astype(dt), x.ndim) + sums += w_first * np.take(x, idx[:-1], axis=0) + sums += w_last * np.take(x, idx[1:], axis=0) + sums /= scale + return sums + + +def _area_2tap_resample_rows(x: np.ndarray, new_len: int) -> np.ndarray: + """OpenCV's INTER_AREA fallback used when the image is not shrunk on both axes. + + OpenCV only performs true area averaging when ``scale >= 1`` for both + axes. Otherwise it runs its generic two-tap resampler with the + ``area_mode`` sample positions (``sx = floor(dx * scale)``, weight taken + from the overlap of the destination pixel with the next source pixel). + For an enlarging axis this equals exact area averaging; for a shrinking + axis it is a coarse two-tap approximation. Reproduced here so that mixed + up/down scaling matches OpenCV up to rounding. + """ + n = x.shape[0] + if new_len == n: + return x + # Mirror OpenCV's arithmetic exactly (inv_scale = dst/src as double, + # scale = 1/inv_scale, weight computed in double then truncated to float) + # so that floating-point ties at exact pixel boundaries resolve the same way. + inv_scale = new_len / n + scale = 1.0 / inv_scale + d = np.arange(new_len, dtype=np.float64) + lo = np.floor(d * scale).astype(np.int64) + w = ((d + 1.0) - (lo + 1) * inv_scale).astype(np.float32) + w = np.where(w <= 0, np.float32(0), w - np.floor(w)) + w[lo >= n - 1] = 0.0 + lo[lo >= n - 1] = n - 1 + return _two_tap_rows(x, lo, w) + + +def _linear_resample_rows(x: np.ndarray, new_len: int) -> np.ndarray: + """Two-tap bilinear resampling along axis 0 (OpenCV INTER_LINEAR semantics). + + Sample positions use half-pixel centres and are clamped at the borders the + same way OpenCV does; there is no antialiasing when downscaling. + """ + n = x.shape[0] + if new_len == n: + return x + # Same arithmetic as OpenCV: position in double, truncated to float + # before the floor, so boundary ties resolve identically. + scale = 1.0 / (new_len / n) + pos = ((np.arange(new_len, dtype=np.float64) + 0.5) * scale - 0.5).astype(np.float32) + lo = np.floor(pos).astype(np.int64) + w = pos - lo + w[lo < 0] = 0.0 + lo[lo < 0] = 0 + w[lo >= n - 1] = 0.0 + lo[lo >= n - 1] = n - 1 + return _two_tap_rows(x, lo, w) + + +def _column(v: np.ndarray, ndim: int) -> np.ndarray: + """Reshape a 1-D per-row weight vector so it broadcasts along axis 0.""" + return v.reshape((-1,) + (1,) * (ndim - 1)) + + +def _work_dtype(x: np.ndarray, n: int) -> type: + # float32 keeps sums of uint8 data exact as long as the total stays below + # 2**24 (i.e. fewer than ~65k pixels along the axis); otherwise float64. + if x.dtype == np.uint8 and n * 255 < (1 << 24): + return np.float32 + return np.float64 + + +def _two_tap_rows(x: np.ndarray, lo: np.ndarray, w: np.ndarray) -> np.ndarray: + """out = x[lo] * (1 - w) + x[lo + 1] * w along axis 0 (lo + 1 clamped).""" + n = x.shape[0] + hi = np.minimum(lo + 1, n - 1) + dt = np.float64 if x.dtype == np.float64 else np.float32 + w = _column(w.astype(dt), x.ndim) + return np.take(x, lo, axis=0) * (1 - w) + np.take(x, hi, axis=0) * w + + +def _resample_2d(src: np.ndarray, new_h: int, new_w: int, fn) -> np.ndarray: + """Apply a row resampler to both axes: rows first on the contiguous + input, then on a contiguous transposed copy of the (already smaller) + intermediate so the column pass also reduces over a contiguous block.""" + out = fn(src, new_h) + out = np.ascontiguousarray(np.swapaxes(out, 0, 1)) + out = fn(out, new_w) + return np.ascontiguousarray(np.swapaxes(out, 0, 1)) + + +def _cast_like(out: np.ndarray, dtype: np.dtype) -> np.ndarray: + if np.issubdtype(dtype, np.integer): + info = np.iinfo(dtype) + out += 0.5 + np.floor(out, out=out) + np.clip(out, info.min, info.max, out=out) + return out.astype(dtype) + return out.astype(dtype, copy=False) + + +def _swap_rgb_bgr_3ch(img: np.ndarray) -> np.ndarray: + return np.ascontiguousarray(img[..., ::-1]) + + +def _swap_rgb_bgr_4ch(img: np.ndarray) -> np.ndarray: + # Reverse the first three channels, keep the alpha. + out = np.empty_like(img) + out[..., 0] = img[..., 2] + out[..., 1] = img[..., 1] + out[..., 2] = img[..., 0] + out[..., 3] = img[..., 3] + return out + + +def _to_gray(img: np.ndarray, source_order: str) -> np.ndarray: + # ITU-R BT.601 luma coefficients, matching OpenCV. + if source_order == "BGR": + b, g, r = img[..., 0], img[..., 1], img[..., 2] + else: + r, g, b = img[..., 0], img[..., 1], img[..., 2] + gray = 0.114 * b.astype(np.float32) + 0.587 * g.astype(np.float32) + 0.299 * r.astype(np.float32) + return np.clip(gray + 0.5, 0, 255).astype(img.dtype if img.dtype == np.uint8 else np.uint8) + + +def _bgr_to_hsv_uint8(img: np.ndarray) -> np.ndarray: + """OpenCV-compatible BGR -> HSV (uint8, H in [0,180)).""" + if img.dtype != np.uint8: + raise NotImplementedError("HSV conversion stub only supports uint8 images") + + f = img.astype(np.float32) + b, g, r = f[..., 0], f[..., 1], f[..., 2] + v = np.maximum(np.maximum(b, g), r) + m = np.minimum(np.minimum(b, g), r) + diff = v - m + + s = np.where(v > 0, (diff / np.maximum(v, 1e-8)) * 255.0, 0.0) + + h = np.zeros_like(v) + # Avoid division by zero; where diff == 0, hue stays 0. + safe = diff > 0 + # v == r + mask_r = safe & (v == r) + mask_g = safe & (v == g) & ~mask_r + mask_b = safe & (v == b) & ~mask_r & ~mask_g + + h[mask_r] = (60.0 * (g[mask_r] - b[mask_r]) / diff[mask_r]) + h[mask_g] = (60.0 * (2.0 + (b[mask_g] - r[mask_g]) / diff[mask_g])) + h[mask_b] = (60.0 * (4.0 + (r[mask_b] - g[mask_b]) / diff[mask_b])) + h = np.where(h < 0, h + 360.0, h) + # Map [0,360) -> [0,180) to match OpenCV uint8 convention. Hues that + # round up to 180 wrap to 0 (they are the same colour), as in OpenCV. + h = np.floor(h / 2.0 + 0.5) % 180.0 + + out = np.stack( + [ + h.astype(np.uint8), + np.clip(s + 0.5, 0, 255).astype(np.uint8), + np.clip(v + 0.5, 0, 255).astype(np.uint8), + ], + axis=-1, + ) + return out + + +def _hsv_to_bgr_uint8(img: np.ndarray) -> np.ndarray: + """OpenCV-compatible HSV -> BGR (uint8, H in [0,180)).""" + if img.dtype != np.uint8: + raise NotImplementedError("HSV conversion stub only supports uint8 images") + + f = img.astype(np.float32) + h = f[..., 0] * 2.0 # back to [0,360) + s = f[..., 1] / 255.0 + v = f[..., 2] + + c = v * s + hh = h / 60.0 + x = c * (1.0 - np.abs((hh % 2.0) - 1.0)) + m = v - c + + zeros = np.zeros_like(h) + r = np.select( + [hh < 1, hh < 2, hh < 3, hh < 4, hh < 5, hh < 6], + [c, x, zeros, zeros, x, c], + default=0, + ) + g = np.select( + [hh < 1, hh < 2, hh < 3, hh < 4, hh < 5, hh < 6], + [x, c, c, x, zeros, zeros], + default=0, + ) + b = np.select( + [hh < 1, hh < 2, hh < 3, hh < 4, hh < 5, hh < 6], + [zeros, zeros, x, c, c, x], + default=0, + ) + + out = np.stack( + [ + np.clip(b + m + 0.5, 0, 255).astype(np.uint8), + np.clip(g + m + 0.5, 0, 255).astype(np.uint8), + np.clip(r + m + 0.5, 0, 255).astype(np.uint8), + ], + axis=-1, + ) + return out + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def cvtColor(src: np.ndarray, code: int) -> np.ndarray: + """Subset of OpenCV's ``cv2.cvtColor`` implemented with NumPy.""" + if src is None: + raise ValueError("cvtColor: src is None") + + if code in (COLOR_BGR2RGB, COLOR_RGB2BGR): + if src.ndim != 3 or src.shape[2] != 3: + raise ValueError(f"cvtColor: expected 3-channel image, got shape {src.shape}") + return _swap_rgb_bgr_3ch(src) + + if code in (COLOR_BGRA2RGBA, COLOR_RGBA2BGRA): + if src.ndim != 3 or src.shape[2] != 4: + raise ValueError(f"cvtColor: expected 4-channel image, got shape {src.shape}") + return _swap_rgb_bgr_4ch(src) + + if code == COLOR_BGR2GRAY: + return _to_gray(src, "BGR") + if code == COLOR_RGB2GRAY: + return _to_gray(src, "RGB") + if code in (COLOR_GRAY2BGR, COLOR_GRAY2RGB): + if src.ndim == 2: + return np.stack([src, src, src], axis=-1) + if src.ndim == 3 and src.shape[2] == 1: + return np.concatenate([src, src, src], axis=-1) + raise ValueError(f"cvtColor: GRAY2BGR expected single-channel image, got shape {src.shape}") + + if code == COLOR_BGR2HSV: + return _bgr_to_hsv_uint8(src) + if code == COLOR_HSV2BGR: + return _hsv_to_bgr_uint8(src) + if code == COLOR_RGB2HSV: + return _bgr_to_hsv_uint8(_swap_rgb_bgr_3ch(src)) + if code == COLOR_HSV2RGB: + return _swap_rgb_bgr_3ch(_hsv_to_bgr_uint8(src)) + + raise NotImplementedError( + f"cv2 stub: cvtColor code {code} is not implemented. " + "Install opencv-python to use this conversion." + ) + + +def resize( + src: np.ndarray, + dsize: Tuple[int, int], + dst: Optional[np.ndarray] = None, + fx: float = 0, + fy: float = 0, + interpolation: int = INTER_LINEAR, +) -> np.ndarray: + """Subset of OpenCV's ``cv2.resize`` (NumPy for AREA/LINEAR, Pillow otherwise).""" + if src is None: + raise ValueError("resize: src is None") + if src.ndim == 3 and src.shape[2] == 1: + src = src[..., 0] # OpenCV returns a 2-D array for single-channel input + + if dsize is None or dsize == (0, 0): + if fx <= 0 or fy <= 0: + raise ValueError("resize: either dsize or fx/fy must be provided") + dsize = (int(round(src.shape[1] * fx)), int(round(src.shape[0] * fy))) + + new_w, new_h = int(dsize[0]), int(dsize[1]) + + if interpolation in (INTER_AREA, INTER_LINEAR): + # NumPy implementations that reproduce OpenCV's sampling (see helpers). + if (new_h, new_w) == src.shape[:2]: + return src.copy() + if interpolation == INTER_LINEAR: + fn = _linear_resample_rows + elif new_h <= src.shape[0] and new_w <= src.shape[1]: + fn = _area_resample_rows + else: + fn = _area_2tap_resample_rows + out = _resample_2d(src, new_h, new_w, fn) + return _cast_like(out, src.dtype) + + pil_interp = _INTER_TO_PIL.get(interpolation, Image.Resampling.BICUBIC) + dsize = (new_w, new_h) + + if src.ndim == 2: + return _resize_single_channel(src, dsize, pil_interp) + + if src.ndim == 3 and src.shape[2] == 3 and src.dtype == np.uint8: + pil = Image.fromarray(src, mode="RGB") + return np.asarray(pil.resize(dsize, resample=pil_interp)) + + # Everything else (single channel with trailing dim, RGBA, non-uint8) is + # resized one channel at a time. In particular RGBA must NOT go through + # Pillow's "RGBA" mode: Pillow premultiplies by alpha while resampling, + # which alters the colour channels wherever alpha < 255, whereas OpenCV + # treats the four channels independently. + channels = [_resize_single_channel(src[..., c], dsize, pil_interp) for c in range(src.shape[2])] + return np.stack(channels, axis=-1) + + +def imshow(winname: str, mat: np.ndarray) -> None: + """Display ``mat`` via Pillow's default image viewer. + + OpenCV's ``imshow`` takes BGR(A) input, so the channel order is flipped + before handing the array to Pillow. Callers reaching this stub have + explicitly opted in to image display (e.g. ``--debug_dataset`` or the + interactive viewer in ``gen_img.py``), so surfacing the image is the + expected behavior; the accompanying ``waitKey`` call paces the display. + """ + if mat is None: + return + + if mat.ndim == 2: + pil = Image.fromarray(mat) + elif mat.ndim == 3 and mat.shape[2] == 1: + pil = Image.fromarray(mat[..., 0]) + elif mat.ndim == 3 and mat.shape[2] == 3: + pil = Image.fromarray(_swap_rgb_bgr_3ch(mat)) + elif mat.ndim == 3 and mat.shape[2] == 4: + pil = Image.fromarray(_swap_rgb_bgr_4ch(mat), mode="RGBA") + else: + logger.warning("cv2 stub: imshow got unsupported shape %s; skipping.", mat.shape) + return + + try: + pil.show(title=winname) + except Exception as e: # noqa: BLE001 + logger.error("cv2 stub: imshow failed to display '%s': %s", winname, e) + + +def waitKey(delay: int = 0) -> int: + """Block on ``input()`` in place of OpenCV's key-wait. + + The ``delay`` argument is ignored: in this codebase ``waitKey`` is only + used from interactive / debug code paths that want to pause between + images. Returns ``ord(s[0])`` of the first typed character so that + existing checks such as ``k == ord("e")`` / ``k == ord("s")`` keep + working; an empty Enter returns ``13`` (carriage return). ``q`` is + remapped to ``27`` (ESC) since ESC cannot be entered through + ``input()``, preserving ``k == 27`` branches in callers. + """ + try: + s = input("cv2 stub: press Enter to continue (or type 'q' to quit / 'e' / 's'): ") + except EOFError: + return -1 + if not s: + return 13 + # 'q' stands in for the ESC key (keycode 27) since ESC cannot be entered + # through input(); callers that branch on ``k == 27`` (e.g. the outer + # epoch loop in debug_dataset) keep working. + if s[0] == "q": + return 27 + return ord(s[0]) + + +def destroyAllWindows() -> None: # pragma: no cover - no-op + return None + + +def imwrite(filename: str, img: np.ndarray, params=None) -> bool: + """Minimal ``cv2.imwrite`` replacement. Assumes BGR / BGRA input.""" + if img is None: + return False + + if img.ndim == 2: + pil = Image.fromarray(img) + elif img.ndim == 3 and img.shape[2] == 3: + pil = Image.fromarray(_swap_rgb_bgr_3ch(img)) + elif img.ndim == 3 and img.shape[2] == 4: + pil = Image.fromarray(_swap_rgb_bgr_4ch(img), mode="RGBA") + else: + raise ValueError(f"imwrite: unsupported image shape {img.shape}") + + try: + pil.save(filename) + return True + except Exception as e: # noqa: BLE001 + logger.error("cv2 stub: imwrite failed for %s: %s", filename, e) + return False + + +# Sentinel that :mod:`library.cv2_compat` reads to decide whether the caller +# is running against the real OpenCV or this stub. +_IS_CV2_STUB = True diff --git a/library/args.py b/library/args.py index 832445df93..28c1f8f01d 100644 --- a/library/args.py +++ b/library/args.py @@ -839,6 +839,12 @@ def verify_training_args(args: argparse.Namespace): "cache_latents_to_disk is enabled, so cache_latents is also enabled / cache_latents_to_diskが有効なため、cache_latentsを有効にします" ) + if getattr(args, "sample_prompts", None) and getattr(args, "sample_sampler", None): + # e.g. lms needs scipy: fail here rather than at the first sample generation after loading the models + from library.sampling import check_sampler_requirements + + check_sampler_requirements(args.sample_sampler) + if getattr(args, "train_inpainting", False) and getattr(args, "cache_latents", False): raise ValueError( "train_inpainting and cache_latents cannot be used together. " diff --git a/library/clip_text_model.py b/library/clip_text_model.py new file mode 100644 index 0000000000..1f20b44229 --- /dev/null +++ b/library/clip_text_model.py @@ -0,0 +1,120 @@ +"""Compatibility wrapper for `transformers.CLIPTextModel`. + +transformers 5.6.0 flattened `CLIPTextModel`: the `text_model` submodule (`CLIPTextTransformer`) +was removed and `embeddings` / `encoder` / `final_layer_norm` became direct children. +`from_pretrained` / `save_pretrained` convert the keys transparently, but `load_state_dict`, +`state_dict()` and `named_modules()` expose the flattened layout. This changes + +- the state dict keys that our checkpoint loaders / savers produce (`text_model.` prefix), and +- the module names that LoRA / OFT etc. derive their weight names from + (`lora_te_text_model_encoder_layers_*`), which must stay stable for compatibility with + existing LoRA files and other tools. + +`CLIPTextModelWrapper` puts the flattened model back under a `text_model` attribute so that +state dict keys, module names and `text_encoder.text_model.*` attribute access keep the +pre-5.6 layout. Only the model construction sites need to call `wrap_clip_text_model`; +everything else keeps working unchanged. With transformers < 5.6 the wrapper is not applied. + +`CLIPTextModelWithProjection` is not affected (it still has `text_model`). +""" + +from typing import Any, Optional, Union + +import torch +from torch import nn + + +def is_flattened_clip_text_model(model: nn.Module) -> bool: + """True if `model` is a `CLIPTextModel` with the flattened (transformers >= 5.6) layout.""" + return ( + not isinstance(model, CLIPTextModelWrapper) + and hasattr(model, "encoder") + and hasattr(model, "final_layer_norm") + and not hasattr(model, "text_model") + ) + + +class CLIPTextModelWrapper(nn.Module): + """Holds a flattened `CLIPTextModel` as `self.text_model` and delegates everything to it. + + `state_dict()` / `load_state_dict()` / `named_modules()` see the `text_model.` prefix, so the + wrapped model is indistinguishable from a pre-5.6 `CLIPTextModel` for loaders, savers and + network modules. `forward` and the commonly used `PreTrainedModel` helpers are delegated; + any other attribute not found on the wrapper falls back to the inner model. + """ + + def __init__(self, text_model: nn.Module): + super().__init__() + self.text_model = text_model + + # --- forward ----------------------------------------------------------------------------- + + def forward(self, *args, **kwargs): + return self.text_model(*args, **kwargs) + + # --- PreTrainedModel-like helpers ---------------------------------------------------------- + + @property + def config(self): + return self.text_model.config + + @property + def dtype(self) -> torch.dtype: + return self.text_model.dtype + + @property + def device(self) -> torch.device: + return self.text_model.device + + def get_input_embeddings(self) -> nn.Module: + return self.text_model.embeddings.token_embedding + + def set_input_embeddings(self, value: nn.Module): + self.text_model.embeddings.token_embedding = value + + def resize_token_embeddings(self, *args, **kwargs): + return self.text_model.resize_token_embeddings(*args, **kwargs) + + def gradient_checkpointing_enable(self, *args, **kwargs): + return self.text_model.gradient_checkpointing_enable(*args, **kwargs) + + def gradient_checkpointing_disable(self, *args, **kwargs): + return self.text_model.gradient_checkpointing_disable(*args, **kwargs) + + def save_pretrained(self, *args, **kwargs): + # the inner model converts the keys back to the layout of the installed transformers + return self.text_model.save_pretrained(*args, **kwargs) + + # --- fallback ------------------------------------------------------------------------------ + + def __getattr__(self, name: str) -> Any: + # nn.Module.__getattr__ only resolves parameters / buffers / submodules; anything else + # (e.g. `eos_token_id`, `main_input_name`, `_no_split_modules`) comes from the inner model. + try: + return super().__getattr__(name) + except AttributeError: + if name == "text_model": # not yet registered (during __init__) or a real error + raise + return getattr(self.text_model, name) + + +def wrap_clip_text_model(model: nn.Module) -> nn.Module: + """Wrap `model` with `CLIPTextModelWrapper` if it has the flattened layout, else return as is. + + Call this right after constructing a `CLIPTextModel` (via `_from_config` / `from_pretrained` / + a Diffusers pipeline) and before `load_state_dict`. Idempotent. + """ + if is_flattened_clip_text_model(model): + return CLIPTextModelWrapper(model) + return model + + +def unwrap_clip_text_model(model: nn.Module) -> nn.Module: + """Return the inner `CLIPTextModel` if `model` is wrapped, else `model` itself. + + Use this when handing the model to code that needs the real transformers class, e.g. a + Diffusers pipeline for `save_pretrained` (which records the class name in model_index.json). + """ + if isinstance(model, CLIPTextModelWrapper): + return model.text_model + return model diff --git a/library/clip_tokenizer.py b/library/clip_tokenizer.py new file mode 100644 index 0000000000..9b74ea4243 --- /dev/null +++ b/library/clip_tokenizer.py @@ -0,0 +1,58 @@ +"""CLIPTokenizer with the text normalization of the legacy (slow) CLIP tokenizer. + +Up to transformers 4.x, ``CLIPTokenizer`` was a Python (slow) tokenizer that ran +``ftfy.fix_text`` and whitespace cleanup on the text before BPE, exactly like the +original OpenAI CLIP tokenizer the models were trained with. ``ftfy.fix_text`` +straightens curly quotes, converts full-width characters to ASCII, unescapes +HTML entities and repairs mojibake. + +In transformers 5.x ``CLIPTokenizer`` is the fast (Rust) tokenizer, whose +normalizer only applies NFC, lowercasing and whitespace collapsing, so such +strings are tokenized differently. This subclass restores the ftfy step for the +fast tokenizer, which makes the token ids identical across transformers +versions. With a slow tokenizer it is a no-op (the cleanup is already applied +inside the tokenizer). +""" + +import re +from typing import Any, List, Tuple, Union + +import ftfy +from transformers import CLIPTokenizer as _CLIPTokenizer + +_WHITESPACE_RE = re.compile(r"\s+") + + +def clean_clip_text(text: str) -> str: + """Same normalization as the legacy slow ``CLIPTokenizer`` (``whitespace_clean(ftfy.fix_text(text))``).""" + return _WHITESPACE_RE.sub(" ", ftfy.fix_text(text)).strip() + + +def _clean(value: Any) -> Any: + # str, pairs of str (text, text_pair), pre-tokenized lists and batches of those + if isinstance(value, str): + return clean_clip_text(value) + if isinstance(value, (list, tuple)): + return type(value)(_clean(v) for v in value) + return value + + +def as_legacy_clip_tokenizer(tokenizer: _CLIPTokenizer) -> _CLIPTokenizer: + """Give an existing ``transformers.CLIPTokenizer`` instance (e.g. one loaded by a diffusers pipeline) the legacy normalization.""" + if type(tokenizer) is _CLIPTokenizer: + tokenizer.__class__ = CLIPTokenizer + return tokenizer + + +class CLIPTokenizer(_CLIPTokenizer): + """Drop-in replacement for ``transformers.CLIPTokenizer`` that keeps the legacy text normalization.""" + + def _encode_plus(self, text, text_pair=None, *args, **kwargs): + if self.is_fast: + text, text_pair = _clean(text), _clean(text_pair) + return super()._encode_plus(text, text_pair, *args, **kwargs) + + def _batch_encode_plus(self, batch_text_or_text_pairs, *args, **kwargs): + if self.is_fast: + batch_text_or_text_pairs = _clean(batch_text_or_text_pairs) + return super()._batch_encode_plus(batch_text_or_text_pairs, *args, **kwargs) diff --git a/library/cv2_compat.py b/library/cv2_compat.py new file mode 100644 index 0000000000..0a21e054e5 --- /dev/null +++ b/library/cv2_compat.py @@ -0,0 +1,58 @@ +""" +OpenCV (``cv2``) optional-dependency bridge. + +Importing this module has one of two effects: + +1. If ``opencv-python`` is installed, nothing changes -- the real ``cv2`` + module is used as-is, and :data:`HAS_OPENCV` is set to ``True``. + +2. If it is not installed, a lightweight stub implemented in + :mod:`library._cv2_stub` is registered as ``cv2`` in + :data:`sys.modules`, so any subsequent ``import cv2`` call transparently + picks it up. :data:`HAS_OPENCV` is set to ``False``. + +Import this module BEFORE ``import cv2`` in any file that wants to tolerate +a missing OpenCV installation:: + + from library import cv2_compat # noqa: F401 - must come before `import cv2` + import cv2 + +Tools that genuinely require OpenCV should bail out early with a clear +message when :data:`HAS_OPENCV` is ``False``. +""" + +from __future__ import annotations + +import logging +import sys + +logger = logging.getLogger(__name__) + +try: + import cv2 as _real_cv2 # noqa: F401 # real OpenCV is importable + HAS_OPENCV = True +except ImportError: + HAS_OPENCV = False + if "cv2" not in sys.modules: + from library import _cv2_stub + + sys.modules["cv2"] = _cv2_stub + logger.info( + "opencv-python is not installed; using the built-in cv2 stub " + "(Pillow/NumPy fallback). Some tools that rely on the real " + "OpenCV will not work." + ) + + +def require_opencv(tool_name: str = "this tool") -> None: + """Abort execution with a clear message when real OpenCV is missing.""" + if HAS_OPENCV: + return + raise SystemExit( + f"{tool_name} requires opencv-python, but it is not installed.\n" + f"{tool_name} は opencv-python が必要ですが、インストールされていません。\n" + "Install it with: pip install opencv-python" + ) + + +__all__ = ["HAS_OPENCV", "require_opencv"] diff --git a/library/dataset.py b/library/dataset.py index cfbe65c12e..c6cb296e51 100644 --- a/library/dataset.py +++ b/library/dataset.py @@ -34,6 +34,7 @@ from concurrent.futures import Future, ThreadPoolExecutor from typing import Any, Callable, Dict, List, NamedTuple, Optional, Sequence, Tuple, Union +from library import cv2_compat # noqa: F401 - must be imported before `import cv2` import cv2 import imagesize import numpy as np diff --git a/library/flux_utils.py b/library/flux_utils.py index 410b34ce23..4eee4f5b03 100644 --- a/library/flux_utils.py +++ b/library/flux_utils.py @@ -19,6 +19,7 @@ from library import flux_models from library.safetensors_utils import load_safetensors +from library.clip_text_model import wrap_clip_text_model MODEL_VERSION_FLUX_V1 = "flux1" MODEL_NAME_DEV = "dev" @@ -342,7 +343,7 @@ def load_clip_l( } config = CLIPConfig(**CLIPL_CONFIG) with init_empty_weights(): - clip = CLIPTextModel._from_config(config) + clip = wrap_clip_text_model(CLIPTextModel._from_config(config)) if state_dict is not None: sd = state_dict diff --git a/library/mask_generator.py b/library/mask_generator.py index 1f92f56391..4d74240e62 100644 --- a/library/mask_generator.py +++ b/library/mask_generator.py @@ -30,6 +30,7 @@ import random from typing import Optional, Tuple +from library import cv2_compat # noqa: F401 - must be imported before `import cv2` import cv2 import numpy as np from PIL import Image, ImageDraw diff --git a/library/model_io.py b/library/model_io.py index 5466b4f8de..b56663c9aa 100644 --- a/library/model_io.py +++ b/library/model_io.py @@ -28,6 +28,7 @@ import library.sai_model_spec as sai_model_spec from library.device_utils import clean_memory_on_device +from library.clip_text_model import wrap_clip_text_model from library.utils import setup_logging if TYPE_CHECKING: @@ -372,7 +373,7 @@ def _load_target_model(args: argparse.Namespace, weight_dtype, device="cpu", une f"model is not found as a file or in Hugging Face, perhaps file name is wrong? / 指定したモデル名のファイル、またはHugging Faceのモデルが見つかりません。ファイル名が誤っているかもしれません: {name_or_path}" ) raise ex - text_encoder = pipe.text_encoder + text_encoder = wrap_clip_text_model(pipe.text_encoder) vae = pipe.vae unet = pipe.unet del pipe diff --git a/library/model_util.py b/library/model_util.py index 385f29b9fa..96e3f11c50 100644 --- a/library/model_util.py +++ b/library/model_util.py @@ -11,6 +11,7 @@ import diffusers from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig, logging +from library.clip_text_model import wrap_clip_text_model, unwrap_clip_text_model from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline # , UNet2DConditionModel from safetensors.torch import load_file, save_file from library.original_unet import UNet2DConditionModel @@ -52,7 +53,8 @@ # Diffusersの設定を読み込むための参照モデル DIFFUSERS_REF_MODEL_ID_V1 = "runwayml/stable-diffusion-v1-5" -DIFFUSERS_REF_MODEL_ID_V2 = "stabilityai/stable-diffusion-2-1" +# stabilityai/stable-diffusion-2-1 is no longer accessible on the Hub; sd2-community/stable-diffusion-2-1 is a mirror +DIFFUSERS_REF_MODEL_ID_V2 = "sd2-community/stable-diffusion-2-1" # region StableDiffusion->Diffusersの変換コード @@ -1048,7 +1050,7 @@ def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, device="cpu", dt torch_dtype="float32", transformers_version="4.25.0.dev0", ) - text_model = CLIPTextModel._from_config(cfg) + text_model = wrap_clip_text_model(CLIPTextModel._from_config(cfg)) info = text_model.load_state_dict(converted_text_encoder_checkpoint) else: converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v1(state_dict) @@ -1077,7 +1079,7 @@ def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, device="cpu", dt projection_dim=768, torch_dtype="float32", ) - text_model = CLIPTextModel._from_config(cfg) + text_model = wrap_clip_text_model(CLIPTextModel._from_config(cfg)) info = text_model.load_state_dict(converted_text_encoder_checkpoint) logger.info(f"loading text encoder: {info}") @@ -1258,7 +1260,7 @@ def save_diffusers_checkpoint(v2, output_dir, text_encoder, unet, pretrained_mod pipeline = StableDiffusionPipeline( unet=diffusers_unet, - text_encoder=text_encoder, + text_encoder=unwrap_clip_text_model(text_encoder), vae=vae, scheduler=scheduler, tokenizer=tokenizer, diff --git a/library/sampling.py b/library/sampling.py index 8c2ad41766..50de030c08 100644 --- a/library/sampling.py +++ b/library/sampling.py @@ -21,6 +21,7 @@ """ import argparse +import importlib.util import json import os import re @@ -69,12 +70,31 @@ SCHEDLER_SCHEDULE = "scaled_linear" +def check_sampler_requirements(sample_sampler: str) -> None: + """Raise a clear ``ImportError`` if an optional dependency of the sampler is missing. + + ``lms`` / ``k_lms`` (``LMSDiscreteScheduler``) need ``scipy``, which is not in ``requirements.txt``. + Without this check the failure surfaces only when the scheduler is built, i.e. at the first sample + generation after the models are loaded. + """ + if sample_sampler in ("lms", "k_lms") and importlib.util.find_spec("scipy") is None: + raise ImportError( + f"The sampler '{sample_sampler}' (LMSDiscreteScheduler) requires the `scipy` package, which is not included in" + " requirements.txt. Please run `pip install scipy`, or choose another sampler (e.g. euler, dpmsolver++)." + f" / サンプラー '{sample_sampler}' (LMSDiscreteScheduler) には requirements.txt に含まれない `scipy` パッケージが必要です。" + "`pip install scipy` を実行するか、他のサンプラー（euler、dpmsolver++ など）を指定してください。" + ) + + def get_my_scheduler( *, sample_sampler: str, v_parameterization: bool, ): + check_sampler_requirements(sample_sampler) + sched_init_args = {} + has_steps_offset = True if sample_sampler == "ddim": scheduler_cls = DDIMScheduler elif sample_sampler == "ddpm": # ddpmはおかしくなるのでoptionから外してある @@ -90,8 +110,12 @@ def get_my_scheduler( elif sample_sampler == "dpmsolver" or sample_sampler == "dpmsolver++": scheduler_cls = DPMSolverMultistepScheduler sched_init_args["algorithm_type"] = sample_sampler + if sample_sampler == "dpmsolver": + # diffusers rejects the default final_sigmas_type="zero" for the (non-++) dpmsolver algorithm + sched_init_args["final_sigmas_type"] = "sigma_min" elif sample_sampler == "dpmsingle": scheduler_cls = DPMSolverSinglestepScheduler + has_steps_offset = False # DPMSolverSinglestepScheduler has no steps_offset argument elif sample_sampler == "heun": scheduler_cls = HeunDiscreteScheduler elif sample_sampler == "dpm_2" or sample_sampler == "k_dpm_2": @@ -103,13 +127,14 @@ def get_my_scheduler( if v_parameterization: sched_init_args["prediction_type"] = "v_prediction" + if has_steps_offset: + sched_init_args["steps_offset"] = 1 scheduler = scheduler_cls( num_train_timesteps=SCHEDULER_TIMESTEPS, beta_start=SCHEDULER_LINEAR_START, beta_end=SCHEDULER_LINEAR_END, beta_schedule=SCHEDLER_SCHEDULE, - steps_offset=1, **sched_init_args, ) @@ -126,7 +151,7 @@ def sample_images(*args, **kwargs): def line_to_prompt_dict(line: str) -> dict: - # subset of gen_img_diffusers + # subset of gen_img.py prompt_args = line.split(" --") prompt_dict = {} prompt_dict["prompt"] = prompt_args[0] diff --git a/library/sdxl_model_util.py b/library/sdxl_model_util.py index 0613153029..dee7f53bb1 100644 --- a/library/sdxl_model_util.py +++ b/library/sdxl_model_util.py @@ -7,6 +7,7 @@ from typing import List from diffusers import AutoencoderKL, EulerDiscreteScheduler, UNet2DConditionModel from library import model_util +from library.clip_text_model import wrap_clip_text_model, unwrap_clip_text_model from library import sdxl_original_unet from library.utils import setup_logging @@ -236,7 +237,7 @@ def load_models_from_sdxl_checkpoint(model_version, ckpt_path, map_location, dty # transformers_version="4.25.0.dev0", ) with init_empty_weights(): - text_model1 = CLIPTextModel._from_config(text_model1_cfg) + text_model1 = wrap_clip_text_model(CLIPTextModel._from_config(text_model1_cfg)) # Text Encoder 2 is different from Stability AI's SDXL. SDXL uses open clip, but we use the model from HuggingFace. # Note: Tokenizer from HuggingFace is different from SDXL. We must use open clip's tokenizer. @@ -564,6 +565,7 @@ def remove_name_or_path(model): model.config._name_or_path = None model.config._name_or_path = None + text_encoder1 = unwrap_clip_text_model(text_encoder1) remove_name_or_path(diffusers_unet) remove_name_or_path(text_encoder1) remove_name_or_path(text_encoder2) diff --git a/library/sdxl_train_util.py b/library/sdxl_train_util.py index 9fec085e51..fb93309834 100644 --- a/library/sdxl_train_util.py +++ b/library/sdxl_train_util.py @@ -10,9 +10,10 @@ from accelerate import init_empty_weights from tqdm import tqdm -from transformers import CLIPTokenizer +from library.clip_tokenizer import CLIPTokenizer # transformers.CLIPTokenizer with the legacy (ftfy) text normalization from library import model_util, sdxl_model_util, checkpoint_io, sampling, sdxl_original_unet import library.model_io as model_io +from library.clip_text_model import wrap_clip_text_model from .utils import setup_logging setup_logging() @@ -112,7 +113,7 @@ def _load_target_model( ) raise ex - text_encoder1 = pipe.text_encoder + text_encoder1 = wrap_clip_text_model(pipe.text_encoder) text_encoder2 = pipe.text_encoder_2 # convert to fp32 for cache text_encoders outputs diff --git a/library/slicing_vae.py b/library/slicing_vae.py index ea76534296..19f11e1f01 100644 --- a/library/slicing_vae.py +++ b/library/slicing_vae.py @@ -23,9 +23,9 @@ from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.models.modeling_utils import ModelMixin -from diffusers.models.unet_2d_blocks import UNetMidBlock2D, get_down_block, get_up_block -from diffusers.models.vae import DecoderOutput, DiagonalGaussianDistribution -from diffusers.models.autoencoder_kl import AutoencoderKLOutput +from diffusers.models.unets.unet_2d_blocks import UNetMidBlock2D, get_down_block, get_up_block +from diffusers.models.autoencoders.vae import DecoderOutput, DiagonalGaussianDistribution +from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKLOutput from .utils import setup_logging setup_logging() import logging diff --git a/library/strategy_base.py b/library/strategy_base.py index 7da7f84d95..4469a2924f 100644 --- a/library/strategy_base.py +++ b/library/strategy_base.py @@ -50,17 +50,30 @@ def get_strategy(cls) -> Optional["TokenizeStrategy"]: return cls._strategy def _load_tokenizer( - self, model_class: Any, model_id: str, subfolder: Optional[str] = None, tokenizer_cache_dir: Optional[str] = None + self, + model_class: Any, + model_id: str, + subfolder: Optional[str] = None, + tokenizer_cache_dir: Optional[str] = None, + cache_name: Optional[str] = None, + **from_pretrained_kwargs, ) -> Any: + """ + Load a tokenizer from the Hub (or from `tokenizer_cache_dir` if it was saved there before). + + `cache_name` overrides the directory name under `tokenizer_cache_dir` (default: `model_id` with "/" replaced). + `from_pretrained_kwargs` are passed to `model_class.from_pretrained` when loading from the Hub; they are + persisted by `save_pretrained`, so the cached tokenizer does not need them. + """ tokenizer = None if tokenizer_cache_dir: - local_tokenizer_path = os.path.join(tokenizer_cache_dir, model_id.replace("/", "_")) + local_tokenizer_path = os.path.join(tokenizer_cache_dir, (cache_name or model_id).replace("/", "_")) if os.path.exists(local_tokenizer_path): logger.info(f"load tokenizer from cache: {local_tokenizer_path}") tokenizer = model_class.from_pretrained(local_tokenizer_path) # same for v1 and v2 if tokenizer is None: - tokenizer = model_class.from_pretrained(model_id, subfolder=subfolder) + tokenizer = model_class.from_pretrained(model_id, subfolder=subfolder, **from_pretrained_kwargs) if tokenizer_cache_dir and not os.path.exists(local_tokenizer_path): logger.info(f"save Tokenizer to cache: {local_tokenizer_path}") diff --git a/library/strategy_flux.py b/library/strategy_flux.py index 7d63fee3c8..2c265cf5da 100644 --- a/library/strategy_flux.py +++ b/library/strategy_flux.py @@ -3,7 +3,8 @@ from typing import Any, List, Optional, Tuple, Union import torch import numpy as np -from transformers import CLIPTokenizer, T5TokenizerFast +from transformers import T5TokenizerFast +from library.clip_tokenizer import CLIPTokenizer # transformers.CLIPTokenizer with the legacy (ftfy) text normalization from library import flux_utils import library.accelerator_setup as accelerator_setup diff --git a/library/strategy_lumina.py b/library/strategy_lumina.py index 0535a35fbd..caee1760b0 100644 --- a/library/strategy_lumina.py +++ b/library/strategy_lumina.py @@ -22,7 +22,11 @@ logger = logging.getLogger(__name__) -GEMMA_ID = "google/gemma-2-2b" +# google/gemma-2-2b is gated on the Hub. The tokenizer files in the official Lumina-Image-2.0 repository are identical +# to those of google/gemma-2-2b (same sha256 for tokenizer.json / tokenizer.model / tokenizer_config.json / +# special_tokens_map.json), so the tokenizer is loaded from there. +GEMMA_TOKENIZER_ID = "Alpha-VLLM/Lumina-Image-2.0" +GEMMA_TOKENIZER_SUBFOLDER = "tokenizer" class LuminaTokenizeStrategy(TokenizeStrategy): @@ -30,7 +34,7 @@ def __init__( self, system_prompt:str, max_length: Optional[int], tokenizer_cache_dir: Optional[str] = None ) -> None: self.tokenizer: GemmaTokenizerFast = AutoTokenizer.from_pretrained( - GEMMA_ID, cache_dir=tokenizer_cache_dir + GEMMA_TOKENIZER_ID, subfolder=GEMMA_TOKENIZER_SUBFOLDER, cache_dir=tokenizer_cache_dir ) self.tokenizer.padding_side = "right" diff --git a/library/strategy_sd.py b/library/strategy_sd.py index e9858c839c..54fbd3f3c6 100644 --- a/library/strategy_sd.py +++ b/library/strategy_sd.py @@ -4,7 +4,7 @@ import numpy as np import torch -from transformers import CLIPTokenizer +from library.clip_tokenizer import CLIPTokenizer # transformers.CLIPTokenizer with the legacy (ftfy) text normalization from library import accelerator_setup import library.device_utils as device_utils from library.strategy_base import LatentsCachingStrategy, TokenizeStrategy, TextEncodingStrategy @@ -17,7 +17,13 @@ TOKENIZER_ID = "openai/clip-vit-large-patch14" -V2_STABLE_DIFFUSION_ID = "stabilityai/stable-diffusion-2" # ここからtokenizerだけ使う v2とv2.1はtokenizer仕様は同じ + +# The SD2.x tokenizer (previously loaded from stabilityai/stable-diffusion-2, which is no longer accessible on the Hub) +# is the v1 (OpenAI CLIP) tokenizer with "!" (id 0) as the pad token instead of <|endoftext|>: same vocabulary, +# same merges, same token ids. So it is built from the v1 tokenizer. v2 and v2.1 use the same tokenizer. +V2_TOKENIZER_PAD_TOKEN = "!" +# directory name under --tokenizer_cache_dir, kept so that existing caches of the v2 tokenizer are still used +V2_TOKENIZER_CACHE_NAME = "stabilityai_stable-diffusion-2" class SdTokenizeStrategy(TokenizeStrategy): @@ -28,7 +34,11 @@ def __init__(self, v2: bool, max_length: Optional[int], tokenizer_cache_dir: Opt logger.info(f"Using {'v2' if v2 else 'v1'} tokenizer") if v2: self.tokenizer = self._load_tokenizer( - CLIPTokenizer, V2_STABLE_DIFFUSION_ID, subfolder="tokenizer", tokenizer_cache_dir=tokenizer_cache_dir + CLIPTokenizer, + TOKENIZER_ID, + tokenizer_cache_dir=tokenizer_cache_dir, + cache_name=V2_TOKENIZER_CACHE_NAME, + pad_token=V2_TOKENIZER_PAD_TOKEN, ) else: self.tokenizer = self._load_tokenizer(CLIPTokenizer, TOKENIZER_ID, tokenizer_cache_dir=tokenizer_cache_dir) diff --git a/library/strategy_sd3.py b/library/strategy_sd3.py index 8ba3d0bb38..dc0317d6c7 100644 --- a/library/strategy_sd3.py +++ b/library/strategy_sd3.py @@ -4,7 +4,8 @@ from typing import Any, List, Optional, Tuple, Union import torch import numpy as np -from transformers import CLIPTokenizer, T5TokenizerFast, CLIPTextModel, CLIPTextModelWithProjection, T5EncoderModel +from transformers import T5TokenizerFast, CLIPTextModel, CLIPTextModelWithProjection, T5EncoderModel +from library.clip_tokenizer import CLIPTokenizer # transformers.CLIPTokenizer with the legacy (ftfy) text normalization from library import sd3_utils import library.accelerator_setup as accelerator_setup diff --git a/library/strategy_sdxl.py b/library/strategy_sdxl.py index 6b3e2afa65..263e18cbec 100644 --- a/library/strategy_sdxl.py +++ b/library/strategy_sdxl.py @@ -3,7 +3,8 @@ import numpy as np import torch -from transformers import CLIPTokenizer, CLIPTextModel, CLIPTextModelWithProjection +from transformers import CLIPTextModel, CLIPTextModelWithProjection +from library.clip_tokenizer import CLIPTokenizer # transformers.CLIPTokenizer with the legacy (ftfy) text normalization from library.strategy_base import TokenizeStrategy, TextEncodingStrategy, TextEncoderOutputsCachingStrategy diff --git a/library/utils.py b/library/utils.py index 2cf1883603..ba93bfcfa2 100644 --- a/library/utils.py +++ b/library/utils.py @@ -9,6 +9,7 @@ from diffusers import EulerAncestralDiscreteScheduler import diffusers.schedulers.scheduling_euler_ancestral_discrete from diffusers.schedulers.scheduling_euler_ancestral_discrete import EulerAncestralDiscreteSchedulerOutput +from library import cv2_compat # noqa: F401 - must be imported before `import cv2` import cv2 from PIL import Image import numpy as np @@ -87,6 +88,23 @@ def setup_logging(args=None, log_level=None, reset=False): logger.info(msg_init) +class SpuriousDiffusersFp32WarningFilter(logging.Filter): + """Drops a spurious warning of diffusers >= 0.40. + + `ModelMixin.to()` warns "There are modules in that should be kept in float32: [] ..." on every + `.to(dtype)` call, even when the model has no module to keep in float32 (the check is + `fp32_modules = self._keep_in_fp32_modules or []` followed by `if ... and fp32_modules is not None`, + which is always true). sd-scripts casts the VAE / U-Net with `.to(dtype)` on purpose, so the warning + is noise. It is only dropped when the list is empty; a real "keep in float32" warning still shows. + """ + + def filter(self, record: logging.LogRecord) -> bool: + return "should be kept in float32: []" not in record.getMessage() + + +logging.getLogger("diffusers.models.modeling_utils").addFilter(SpuriousDiffusersFp32WarningFilter()) + + setup_logging() logger = logging.getLogger(__name__) diff --git a/networks/lora_interrogator.py b/networks/lora_interrogator.py index 2657995979..76deb4af4b 100644 --- a/networks/lora_interrogator.py +++ b/networks/lora_interrogator.py @@ -4,7 +4,7 @@ from library import model_util import library.model_io as model_io import argparse -from transformers import CLIPTokenizer +from library.clip_tokenizer import CLIPTokenizer # transformers.CLIPTokenizer with the legacy (ftfy) text normalization import torch from library.device_utils import init_ipex, get_preferred_device @@ -18,7 +18,7 @@ logger = logging.getLogger(__name__) TOKENIZER_PATH = "openai/clip-vit-large-patch14" -V2_STABLE_DIFFUSION_PATH = "stabilityai/stable-diffusion-2" # ここからtokenizerだけ使う +V2_TOKENIZER_PAD_TOKEN = "!" # the SD2 tokenizer is the v1 tokenizer with "!" as the pad token (see library/strategy_sd.py) DEVICE = get_preferred_device() @@ -48,7 +48,7 @@ def interrogate(args): logger.info("loading tokenizer") if args.v2: - tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(V2_STABLE_DIFFUSION_PATH, subfolder="tokenizer") + tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(TOKENIZER_PATH, pad_token=V2_TOKENIZER_PAD_TOKEN) else: tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(TOKENIZER_PATH) # , model_max_length=max_token_length + 2) diff --git a/pytest.ini b/pytest.ini index 34b7e9c1fb..7b24147a03 100644 --- a/pytest.ini +++ b/pytest.ini @@ -2,6 +2,8 @@ minversion = 6.0 testpaths = tests +# tests/local needs local model weights and is run explicitly: pytest tests/local +norecursedirs = local filterwarnings = ignore::DeprecationWarning ignore::UserWarning diff --git a/requirements.txt b/requirements.txt index bb5d2cd08a..17d6dc7c16 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,25 +1,34 @@ -accelerate==1.6.0 -transformers==4.54.1 -diffusers[torch]==0.32.1 +accelerate==1.15.0 +transformers==5.17.0 +diffusers[torch]==0.40.0 ftfy==6.3.1 # albumentations==1.3.0 -opencv-python==4.10.0.84 +# opencv-python has no wheel for Windows on ARM64 (e.g. NVIDIA RTX Spark PCs), so it is +# skipped there. It is optional: without it a Pillow/NumPy-based stub in library/_cv2_stub +# takes over automatically. Some tools under tools/ (canny.py, detect_face_rotate.py, and +# the ControlNet canny preprocessor) still require the real opencv-python. +opencv-python==4.10.0.84; platform_machine != "ARM64" einops==0.7.0 # pytorch-lightning==1.9.0 bitsandbytes lion-pytorch==0.2.3 -schedulefree==1.4 +schedulefree==1.4.1 pytorch-optimizer==3.10.0 prodigy-plus-schedule-free==1.9.2 prodigyopt==1.1.2 -tensorboard -safetensors==0.4.5 +# tensorboard 2.x depends on grpcio, which has no wheel for Windows on ARM64 (pip would +# silently fall back to the ancient tensorboard 1.10). tensorboardX is a pure-Python +# drop-in that accelerate uses automatically for --log_with tensorboard; view the logs +# with TensorBoard on another machine. +tensorboard; platform_machine != "ARM64" +tensorboardX; platform_machine == "ARM64" +safetensors==0.8.0 # gradio==3.16.2 # altair==4.2.2 # easygui==0.98.3 toml==0.10.2 voluptuous==0.15.2 -huggingface-hub==0.34.3 +huggingface-hub==1.32.0 # for Image utils imagesize==1.4.1 numpy @@ -37,7 +46,7 @@ numpy # for cuda 12.1(default 11.8) # onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ -# this is for onnx: +# this is for onnx: # protobuf==3.20.3 # open clip for SDXL # open-clip-torch==2.20.0 diff --git a/sdxl_gen_img.py b/sdxl_gen_img.py index f5ae50a4d9..9f7a47f3b9 100755 --- a/sdxl_gen_img.py +++ b/sdxl_gen_img.py @@ -47,7 +47,8 @@ from einops import rearrange from tqdm import tqdm from torchvision import transforms -from transformers import CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection, CLIPImageProcessor +from transformers import CLIPTextModel, CLIPVisionModelWithProjection, CLIPImageProcessor +from library.clip_tokenizer import CLIPTokenizer # transformers.CLIPTokenizer with the legacy (ftfy) text normalization import PIL from PIL import Image from PIL.PngImagePlugin import PngInfo @@ -61,6 +62,7 @@ from library.original_unet import FlashAttentionFunction from networks.control_net_lllite import ControlNetLLLite from library.utils import GradualLatent, EulerAncestralDiscreteSchedulerGL +from library.sampling import check_sampler_requirements from library.utils import setup_logging, add_logging_arguments setup_logging() @@ -1480,6 +1482,8 @@ class BatchData(NamedTuple): def main(args): + check_sampler_requirements(args.sampler) # fail before loading the models if the sampler needs a missing package + if args.fp16: dtype = torch.float16 elif args.bf16: @@ -1542,6 +1546,9 @@ def main(args): elif args.sampler == "dpmsolver" or args.sampler == "dpmsolver++": scheduler_cls = DPMSolverMultistepScheduler sched_init_args["algorithm_type"] = args.sampler + if args.sampler == "dpmsolver": + # diffusers rejects the default final_sigmas_type="zero" for the (non-++) dpmsolver algorithm + sched_init_args["final_sigmas_type"] = "sigma_min" scheduler_module = diffusers.schedulers.scheduling_dpmsolver_multistep has_clip_sample = False elif args.sampler == "dpmsingle": diff --git a/sdxl_minimal_inference.py b/sdxl_minimal_inference.py index a1e93b7f02..dfaa85e714 100644 --- a/sdxl_minimal_inference.py +++ b/sdxl_minimal_inference.py @@ -15,7 +15,7 @@ init_ipex() from tqdm import tqdm -from transformers import CLIPTokenizer +from library.clip_tokenizer import CLIPTokenizer # transformers.CLIPTokenizer with the legacy (ftfy) text normalization from diffusers import EulerDiscreteScheduler from PIL import Image diff --git a/tests/local/README.md b/tests/local/README.md new file mode 100644 index 0000000000..be45d65fef --- /dev/null +++ b/tests/local/README.md @@ -0,0 +1,39 @@ +# Local-only regression tests + +The tests in this directory need model weights on the local disk and are **not run by CI** (`pytest` from the repository root does not descend into `tests/local`). They exist to check dependency upgrades (transformers, diffusers, torch, ...) against the outputs of the versions currently in use. + +## Setup + +```bash +cp tests/local/models.example.toml tests/local/models.toml +# edit tests/local/models.toml: fill in the paths of the models you have, delete the rest +``` + +`models.toml` and the recorded references (`tests/local/references/`) are git-ignored. + +The tokenizers are fetched from the Hugging Face Hub on the first run and cached; afterwards `HF_HUB_OFFLINE=1` works. + +## Usage + +```bash +# 1. record references with the current dependency versions +python tests/local/regression_te.py record # all entries + schedulers +python tests/local/regression_te.py record --models sdxl,flux + +# 2. upgrade the dependencies, then compare +python tests/local/regression_te.py compare +pytest tests/local # same thing as pytest +pytest tests/local -k sdxl + +python tests/local/regression_te.py list # show the configured entries +``` + +## What is compared + +- **Text encoders** (per configured model): the trainers' own `TokenizeStrategy.tokenize` -> `TextEncodingStrategy.encode_tokens` path with the text encoders built by the family's loader, for a fixed prompt set (empty, short, tags, Japanese + emoji, > 77 tokens, > 225 tokens), one prompt at a time and once as a batch. Token ids / attention masks must match exactly; hidden states, pooled outputs etc. are compared with `numpy.allclose` using per-model tolerances. +- **VAE** (sd / sdxl): encode / decode of a fixed image through the diffusers `AutoencoderKL` built by the checkpoint loader (fp32). +- **Schedulers** (always, no weights needed): the diffusers noise schedulers from `library.sampling.get_my_scheduler` for every sampler name, with and without v-prediction: timesteps, `alphas_cumprod`, `sigmas`, and the sample after a few `step()` calls with fixed noise. + +For every array the maximum absolute / relative difference and the cosine similarity are printed, so drifts below the tolerance remain visible. The reference `.json` next to each `.npz` records the library versions the reference was made with. + +The prompt list in `regression_te.py` must not change once references are recorded (the comparison refuses to run if it did; re-record in that case). diff --git a/tests/local/models.example.toml b/tests/local/models.example.toml new file mode 100644 index 0000000000..5ae03eb7a8 --- /dev/null +++ b/tests/local/models.example.toml @@ -0,0 +1,102 @@ +# Template for tests/local/models.toml (git-ignored: it holds local paths). +# +# cp tests/local/models.example.toml tests/local/models.toml +# # edit the paths, delete the [[models]] entries you cannot run +# python tests/local/regression_te.py record # with the current dependency versions +# ... upgrade transformers / diffusers ... +# python tests/local/regression_te.py compare # or: pytest tests/local +# +# Every [[models]] entry needs `name` (used for the reference file name) and +# `type` (sd / sdxl / sd3 / flux / lumina / hunyuan_image / anima). +# Optional keys for every entry: +# enabled = false skip this entry in `record` / `compare` / pytest runs without deleting it +# (it still runs when named explicitly with --models) +# device = "cuda" | "cpu" overrides [common].device +# dtype = "float32" | "bfloat16" | "float16" +# dtype of the text encoders. Default: float32 for +# sd / sdxl, bfloat16 for the others (same as a +# typical training run). +# atol / rtol tolerances for numpy.allclose on float outputs. +# Defaults: fp32 1e-4 / 1e-3, bf16 5e-2 / 3e-2, fp16 2e-2 / 2e-2 +# tokenizer_cache_dir same meaning as --tokenizer_cache_dir of the trainers +# Token ids and attention masks must always match exactly. +# The diffusers noise schedulers ("schedulers") are always checked; no weights needed. + +[common] +device = "cuda" +# reference_dir = "tests/local/references" # default; git-ignored + +# ---- Stable Diffusion 1.x ------------------------------------------------- +[[models]] +name = "sd15" +type = "sd" +ckpt = "D:/models/sd/v1-5-pruned-emaonly.safetensors" +v2 = false +clip_skip = 2 # optional; omit for the last layer +max_token_length = 225 # optional: 75 (default) / 150 / 225; >75 exercises chunking +check_vae = true # encode/decode a fixed image with the diffusers AutoencoderKL + +# ---- Stable Diffusion 2.x ------------------------------------------------- +[[models]] +name = "sd21" +type = "sd" +ckpt = "D:/models/sd/v2-1_768-ema-pruned.safetensors" +v2 = true +max_token_length = 225 + +# ---- SDXL ----------------------------------------------------------------- +[[models]] +name = "sdxl" +type = "sdxl" +ckpt = "D:/models/sdxl/sd_xl_base_1.0.safetensors" +max_token_length = 225 +check_vae = true + +# ---- SD3 / SD3.5 ---------------------------------------------------------- +[[models]] +name = "sd3" +type = "sd3" +clip_l = "D:/models/sd3/clip_l.safetensors" +clip_g = "D:/models/sd3/clip_g.safetensors" +t5xxl = "D:/models/sd3/t5xxl_fp16.safetensors" +t5xxl_max_token_length = 256 +apply_lg_attn_mask = false +apply_t5_attn_mask = false +dtype = "bfloat16" + +# ---- FLUX.1 --------------------------------------------------------------- +[[models]] +name = "flux" +type = "flux" +clip_l = "D:/models/flux/clip_l.safetensors" +t5xxl = "D:/models/flux/t5xxl_fp16.safetensors" +t5xxl_max_token_length = 512 +apply_t5_attn_mask = false +dtype = "bfloat16" + +# ---- Lumina Image 2.0 (Gemma 2 2B) ----------------------------------------- +[[models]] +name = "lumina" +type = "lumina" +gemma2 = "D:/models/lumina/gemma_2_2b_fp16.safetensors" +system_prompt = "" # same as --system_prompt +max_token_length = 256 +dtype = "bfloat16" + +# ---- HunyuanImage-2.1 (Qwen2.5-VL 7B + byT5) -------------------------------- +[[models]] +name = "hunyuan_image" +type = "hunyuan_image" +qwen2_5_vl = "D:/models/hunyuan_image/qwen_2.5_vl_7b.safetensors" +byt5 = "D:/models/hunyuan_image/byt5_small_glyphxl_fp16.safetensors" +dtype = "bfloat16" + +# ---- Anima (Qwen3 0.6B + T5 tokenizer) -------------------------------------- +[[models]] +name = "anima" +type = "anima" +qwen3 = "D:/models/anima/qwen_3_06b_base.safetensors" # file (uses configs/qwen3_06b) or a model directory +# t5_tokenizer_path = "D:/models/anima/t5_tokenizer" # optional; bundled config is used when omitted +qwen3_max_token_length = 512 +t5_max_token_length = 512 +dtype = "bfloat16" diff --git a/tests/local/regression_te.py b/tests/local/regression_te.py new file mode 100644 index 0000000000..b0d78f70f2 --- /dev/null +++ b/tests/local/regression_te.py @@ -0,0 +1,662 @@ +"""Local-only regression harness for transformers / diffusers upgrades. + +This script is NOT run by CI. It needs locally available model weights (paths +are given in ``tests/local/models.toml``, which is git-ignored) and, on the +first run, the tokenizers from the Hugging Face Hub (cached afterwards, so +later runs work with ``HF_HUB_OFFLINE=1``). + +Workflow: + +1. With the *current* dependency versions, record reference outputs:: + + python tests/local/regression_te.py record + +2. Upgrade transformers / diffusers, then compare against the references:: + + python tests/local/regression_te.py compare + # or, as pytest: pytest tests/local + +What is recorded, per model listed in the config: + +* text encoders: the exact objects the training scripts use + (``TokenizeStrategy.tokenize`` -> ``TextEncodingStrategy.encode_tokens``, + text encoders built by the family's own loader), for a fixed set of prompts, + one prompt at a time and once as a batch (padding + attention masks); +* for SD / SDXL: encode / decode of a fixed image through the diffusers + ``AutoencoderKL`` that the checkpoint loader builds; +* always: the diffusers noise schedulers used for sample generation + (timesteps, alphas_cumprod, sigmas, and a few ``step()`` calls with fixed + noise). + +Token ids and masks must match exactly. Floating point outputs are compared +with ``numpy.allclose`` using per-model tolerances (see ``models.example.toml``) +and the maximum absolute / relative difference and cosine similarity are +printed for every tensor so that small drifts stay visible. +""" + +from __future__ import annotations + +import argparse +import datetime +import gc +import inspect +import json +import os +import platform +import sys +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional + +import numpy as np +import toml +import torch + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +LOCAL_DIR = Path(__file__).resolve().parent +DEFAULT_CONFIG = LOCAL_DIR / "models.toml" +DEFAULT_REFERENCE_DIR = LOCAL_DIR / "references" + +SCHEDULERS_NAME = "schedulers" + +# Fixed prompts. Keep this list stable: changing it invalidates recorded references. +PROMPTS: List[str] = [ + "", # empty (unconditional) + "a photo of a cat", + "1girl, solo, long hair, looking at viewer, smile, blue eyes, white shirt, outdoors, cherry blossoms, masterpiece, best quality", + # Japanese + emoji: exercises byte-level / unicode handling of the tokenizers + "桜の木の下で本を読む少女、夕暮れ、水彩画風🌸📖", + # quoted text: HunyuanImage extracts it for the byT5 glyph encoder (ASCII and CJK quotes) + 'A neon sign on a brick wall that says "OPEN 24 HOURS", with a small paper poster below reading “本日開店”', + # ~100 CLIP tokens: crosses the 77-token boundary (SD / SDXL chunking) + ( + "A highly detailed oil painting of an ancient library carved into the side of a mountain, " + "with towering wooden shelves, thousands of leather-bound books, warm candle light flickering " + "on the stone walls, dust particles floating in shafts of golden sunlight coming through tall " + "arched windows, a spiral staircase in the center, intricate carvings of owls and vines on the " + "railings, a single scholar in a dark green robe reading at a long oak table" + ), + # ~350 CLIP tokens: exceeds max_token_length=225 as well, so truncation is exercised too + ( + "Wide establishing shot of a futuristic coastal city at dawn: sleek white towers with hanging " + "gardens, transparent skybridges connecting them, small delivery drones drifting between the " + "buildings, a monorail curving along the shoreline, fishing boats with colorful sails leaving " + "the harbor, mist rising from the water, warm orange light on the east-facing facades and cool " + "blue shadows on the west, seagulls, a lighthouse on a rocky promontory, cargo ships on the " + "horizon, reflections in the wet streets, a few early commuters with umbrellas, street vendors " + "setting up stalls with fruit and flowers, tram wires overhead, ivy on old brick warehouses " + "converted into cafes, bicycles leaning against lamp posts, a large digital billboard showing " + "the weather forecast, a park with cherry trees in full bloom, children playing near a fountain, " + "an old man feeding pigeons, a dog chasing a ball, detailed textures, photorealistic, 8k, " + "cinematic lighting, shallow depth of field, shot on 35mm film, subtle grain" + ), +] + +SAMPLERS = [ + "ddim", + "ddpm", + "pndm", + "lms", + "euler", + "euler_a", + "dpmsolver", + "dpmsolver++", + "dpmsingle", + "heun", + "dpm_2", + "dpm_2_a", +] + +DTYPES = { + "float32": torch.float32, + "fp32": torch.float32, + "float16": torch.float16, + "fp16": torch.float16, + "bfloat16": torch.bfloat16, + "bf16": torch.bfloat16, +} + +# default tolerances (atol, rtol) by dtype of the text encoder computation +DEFAULT_TOLERANCES = { + torch.float32: (1e-4, 1e-3), + torch.float16: (2e-2, 2e-2), + torch.bfloat16: (5e-2, 3e-2), +} + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _to_numpy(x: Any) -> Optional[np.ndarray]: + if x is None: + return None + if isinstance(x, torch.Tensor): + x = x.detach().cpu() + if x.dtype in (torch.bfloat16, torch.float16): + x = x.float() + return x.numpy() + if isinstance(x, (int, float, bool)): + return np.asarray(x) + if isinstance(x, (list, tuple)) and all(isinstance(t, torch.Tensor) for t in x): + return np.stack([_to_numpy(t) for t in x]) + raise TypeError(f"unsupported output type: {type(x)}") + + +def _store(out: Dict[str, np.ndarray], prefix: str, values: Any) -> None: + """Flatten a list/tuple of tensors (or a single tensor) into ``out``.""" + if isinstance(values, (list, tuple)): + for i, v in enumerate(values): + if v is None: + continue + out[f"{prefix}.{i}"] = _to_numpy(v) + elif values is not None: + out[prefix] = _to_numpy(values) + + +def _versions() -> Dict[str, str]: + import diffusers + import transformers + + try: + import huggingface_hub + + hub = huggingface_hub.__version__ + except Exception: # pragma: no cover + hub = "?" + try: + import accelerate + + acc = accelerate.__version__ + except Exception: # pragma: no cover + acc = "?" + return { + "python": platform.python_version(), + "torch": torch.__version__, + "transformers": transformers.__version__, + "diffusers": diffusers.__version__, + "accelerate": acc, + "huggingface_hub": hub, + "platform": f"{platform.system()} {platform.machine()}", + "recorded_at": datetime.datetime.now().isoformat(timespec="seconds"), + } + + +def _fixed_image(size: int = 512, device: str = "cpu") -> torch.Tensor: + """Deterministic smooth-ish RGB image in [-1, 1], shape [1, 3, size, size].""" + rng = np.random.RandomState(0) + low = torch.from_numpy(rng.rand(1, 3, 32, 32).astype(np.float32)) + img = torch.nn.functional.interpolate(low, size=(size, size), mode="bicubic", align_corners=False) + # add a little high-frequency detail so the encoder sees more than a blur + noise = torch.from_numpy(rng.rand(1, 3, size, size).astype(np.float32)) * 0.05 + img = (img + noise).clamp(0, 1) * 2 - 1 + return img.to(device) + + +def _free(*objs) -> None: + for o in objs: + del o + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +# --------------------------------------------------------------------------- +# generic text encoder recording +# --------------------------------------------------------------------------- + + +@torch.no_grad() +def encode_prompts(tokenize_strategy, encoding_strategy, models: List[Any], encode_kwargs: Optional[dict] = None) -> Dict[str, np.ndarray]: + """Run tokenize -> encode_tokens for every prompt (batch size 1) and once for all prompts as a batch.""" + encode_kwargs = encode_kwargs or {} + out: Dict[str, np.ndarray] = {} + for i, prompt in enumerate(PROMPTS): + tokens = tokenize_strategy.tokenize(prompt) + _store(out, f"p{i}.tok", tokens) + encoded = encoding_strategy.encode_tokens(tokenize_strategy, models, tokens, **encode_kwargs) + _store(out, f"p{i}.out", encoded) + + tokens = tokenize_strategy.tokenize(list(PROMPTS)) + _store(out, "batch.tok", tokens) + encoded = encoding_strategy.encode_tokens(tokenize_strategy, models, tokens, **encode_kwargs) + _store(out, "batch.out", encoded) + return out + + +@torch.no_grad() +def encode_vae(vae, device: str) -> Dict[str, np.ndarray]: + """Encode / decode a fixed image with a diffusers AutoencoderKL (fp32).""" + vae = vae.to(device=device, dtype=torch.float32) + img = _fixed_image(512, device) + dist = vae.encode(img).latent_dist + decoded = vae.decode(dist.mean).sample + return { + "vae.latent_mean": _to_numpy(dist.mean), + "vae.latent_std": _to_numpy(dist.std), + "vae.decoded": _to_numpy(decoded), + } + + +# --------------------------------------------------------------------------- +# model families +# --------------------------------------------------------------------------- + + +@dataclass +class ModelSpec: + name: str + type: str + cfg: Dict[str, Any] + device: str + dtype: torch.dtype + atol: float + rtol: float + + @staticmethod + def from_config(entry: Dict[str, Any], common: Dict[str, Any]) -> "ModelSpec": + name = entry["name"] + mtype = entry["type"] + device = entry.get("device", common.get("device", "cuda" if torch.cuda.is_available() else "cpu")) + default_dtype = "float32" if mtype in ("sd", "sdxl") else "bfloat16" + dtype = DTYPES[str(entry.get("dtype", default_dtype)).lower()] + d_atol, d_rtol = DEFAULT_TOLERANCES[dtype] + return ModelSpec( + name=name, + type=mtype, + cfg=entry, + device=device, + dtype=dtype, + atol=float(entry.get("atol", d_atol)), + rtol=float(entry.get("rtol", d_rtol)), + ) + + +def run_sd(spec: ModelSpec) -> Dict[str, np.ndarray]: + from library import model_util, strategy_sd + + cfg = spec.cfg + v2 = bool(cfg.get("v2", False)) + text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(v2, cfg["ckpt"], spec.device, spec.dtype) + text_encoder.eval() + tok = strategy_sd.SdTokenizeStrategy(v2, cfg.get("max_token_length"), cfg.get("tokenizer_cache_dir")) + enc = strategy_sd.SdTextEncodingStrategy(cfg.get("clip_skip")) + out = encode_prompts(tok, enc, [text_encoder]) + if cfg.get("check_vae", True): + out.update(encode_vae(vae, spec.device)) + _free(text_encoder, vae, unet) + return out + + +def run_sdxl(spec: ModelSpec) -> Dict[str, np.ndarray]: + from library import sdxl_model_util, strategy_sdxl + + cfg = spec.cfg + text_model1, text_model2, vae, unet, _logit_scale, _ckpt_info = sdxl_model_util.load_models_from_sdxl_checkpoint( + sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0, cfg["ckpt"], spec.device, spec.dtype + ) + text_model1.to(spec.device).eval() + text_model2.to(spec.device).eval() + tok = strategy_sdxl.SdxlTokenizeStrategy(cfg.get("max_token_length"), cfg.get("tokenizer_cache_dir")) + enc = strategy_sdxl.SdxlTextEncodingStrategy() + out = encode_prompts(tok, enc, [text_model1, text_model2]) + if cfg.get("check_vae", True): + out.update(encode_vae(vae, spec.device)) + _free(text_model1, text_model2, vae, unet) + return out + + +def run_sd3(spec: ModelSpec) -> Dict[str, np.ndarray]: + from library import sd3_utils, strategy_sd3 + + cfg = spec.cfg + clip_l = sd3_utils.load_clip_l(cfg["clip_l"], spec.dtype, spec.device) + clip_g = sd3_utils.load_clip_g(cfg["clip_g"], spec.dtype, spec.device) + t5xxl = sd3_utils.load_t5xxl(cfg["t5xxl"], spec.dtype, spec.device) + for m in (clip_l, clip_g, t5xxl): + # the loaders build the models with init_empty_weights + assign=True; the trainers move them + # explicitly afterwards, so do the same here + m.to(spec.device).eval() + tok = strategy_sd3.Sd3TokenizeStrategy(int(cfg.get("t5xxl_max_token_length", 256)), cfg.get("tokenizer_cache_dir")) + enc = strategy_sd3.Sd3TextEncodingStrategy( + apply_lg_attn_mask=bool(cfg.get("apply_lg_attn_mask", False)), + apply_t5_attn_mask=bool(cfg.get("apply_t5_attn_mask", False)), + ) + out = encode_prompts(tok, enc, [clip_l, clip_g, t5xxl], {"enable_dropout": False}) + _free(clip_l, clip_g, t5xxl) + return out + + +def run_flux(spec: ModelSpec) -> Dict[str, np.ndarray]: + from library import flux_utils, strategy_flux + + cfg = spec.cfg + clip_l = flux_utils.load_clip_l(cfg["clip_l"], spec.dtype, spec.device) + t5xxl = flux_utils.load_t5xxl(cfg["t5xxl"], spec.dtype, spec.device) + clip_l.to(spec.device).eval() + t5xxl.to(spec.device).eval() + tok = strategy_flux.FluxTokenizeStrategy(int(cfg.get("t5xxl_max_token_length", 512)), cfg.get("tokenizer_cache_dir")) + enc = strategy_flux.FluxTextEncodingStrategy(apply_t5_attn_mask=bool(cfg.get("apply_t5_attn_mask", False))) + out = encode_prompts(tok, enc, [clip_l, t5xxl]) + _free(clip_l, t5xxl) + return out + + +def run_lumina(spec: ModelSpec) -> Dict[str, np.ndarray]: + from library import lumina_util, strategy_lumina + + cfg = spec.cfg + gemma2 = lumina_util.load_gemma2(cfg["gemma2"], spec.dtype, spec.device) + gemma2.to(spec.device).eval() + tok = strategy_lumina.LuminaTokenizeStrategy(cfg.get("system_prompt"), cfg.get("max_token_length"), cfg.get("tokenizer_cache_dir")) + enc = strategy_lumina.LuminaTextEncodingStrategy() + out = encode_prompts(tok, enc, [gemma2]) + _free(gemma2) + return out + + +def run_hunyuan_image(spec: ModelSpec) -> Dict[str, np.ndarray]: + from library import hunyuan_image_text_encoder, strategy_hunyuan_image + + cfg = spec.cfg + _vlm_tokenizer, qwen2vl = hunyuan_image_text_encoder.load_qwen2_5_vl(cfg["qwen2_5_vl"], spec.dtype, spec.device) + _byt5_tokenizer, byt5 = hunyuan_image_text_encoder.load_byt5(cfg["byt5"], spec.dtype, spec.device) # (tokenizer, model) + qwen2vl.to(spec.device).eval() + byt5.to(spec.device).eval() + tok = strategy_hunyuan_image.HunyuanImageTokenizeStrategy(cfg.get("tokenizer_cache_dir")) + enc = strategy_hunyuan_image.HunyuanImageTextEncodingStrategy() + out = encode_prompts(tok, enc, [qwen2vl, byt5]) + # The VLM hidden states at padded positions are garbage that the model masks out downstream + # (encoder_attention_mask); they also depend on the transformers version. Zero them so that + # only the positions that matter are compared. out.0 = embeddings, out.1 = attention mask. + for prefix in [f"p{i}" for i in range(len(PROMPTS))] + ["batch"]: + emb, mask = out.get(f"{prefix}.out.0"), out.get(f"{prefix}.out.1") + if emb is not None and mask is not None: + out[f"{prefix}.out.0"] = emb * mask[..., None].astype(emb.dtype) + _free(qwen2vl, byt5) + return out + + +def run_anima(spec: ModelSpec) -> Dict[str, np.ndarray]: + from library import anima_utils, strategy_anima + + cfg = spec.cfg + qwen3, _tokenizer = anima_utils.load_qwen3_text_encoder(cfg["qwen3"], spec.dtype, spec.device) + qwen3.eval() + tok = strategy_anima.AnimaTokenizeStrategy( + qwen3_path=cfg["qwen3"], + t5_tokenizer_path=cfg.get("t5_tokenizer_path"), + qwen3_max_length=int(cfg.get("qwen3_max_token_length", 512)), + t5_max_length=int(cfg.get("t5_max_token_length", 512)), + ) + enc = strategy_anima.AnimaTextEncodingStrategy() + out = encode_prompts(tok, enc, [qwen3]) + _free(qwen3) + return out + + +@torch.no_grad() +def run_schedulers(spec: Optional[ModelSpec] = None) -> Dict[str, np.ndarray]: + """Exercise the diffusers schedulers used for sample image generation (no weights needed).""" + from library import sampling + + out: Dict[str, np.ndarray] = {} + for sampler in SAMPLERS: + for v_param in (False, True): + key = f"{sampler}.v{int(v_param)}" + try: + sched = sampling.get_my_scheduler(sample_sampler=sampler, v_parameterization=v_param) + except Exception as e: + # e.g. LMSDiscreteScheduler needs scipy (not in requirements.txt); a scheduler that cannot be + # built with the installed diffusers is reported and skipped. If it starts working after an + # upgrade, the new keys show up as "not in reference" -- re-record in that case. + print(f" skipping sampler {sampler!r}: {type(e).__name__}: {str(e).strip().splitlines()[0]}") + break + sched.set_timesteps(20) + out[f"{key}.timesteps"] = _to_numpy(torch.as_tensor(sched.timesteps)) + out[f"{key}.alphas_cumprod"] = _to_numpy(sched.alphas_cumprod) + if getattr(sched, "sigmas", None) is not None: + out[f"{key}.sigmas"] = _to_numpy(torch.as_tensor(sched.sigmas)) + + gen = torch.Generator("cpu").manual_seed(1234) + sample = torch.randn(1, 4, 8, 8, generator=gen) + if hasattr(sched, "init_noise_sigma"): + sample = sample * sched.init_noise_sigma + step_params = inspect.signature(sched.step).parameters + for k in range(4): + t = sched.timesteps[k] + model_input = sched.scale_model_input(sample, t) if hasattr(sched, "scale_model_input") else sample + model_output = torch.randn(model_input.shape, generator=gen) * 0.5 + kwargs = {"generator": gen} if "generator" in step_params else {} + sample = sched.step(model_output, t, sample, **kwargs).prev_sample + out[f"{key}.sample_after_4_steps"] = _to_numpy(sample) + return out + + +RUNNERS: Dict[str, Callable[[ModelSpec], Dict[str, np.ndarray]]] = { + "sd": run_sd, + "sdxl": run_sdxl, + "sd3": run_sd3, + "flux": run_flux, + "lumina": run_lumina, + "hunyuan_image": run_hunyuan_image, + "anima": run_anima, +} + + +# --------------------------------------------------------------------------- +# config / references +# --------------------------------------------------------------------------- + + +def load_config(path: Path = DEFAULT_CONFIG) -> Dict[str, Any]: + if not path.exists(): + raise FileNotFoundError(f"{path} not found. Copy models.example.toml to models.toml and fill in your local model paths.") + return toml.load(path) + + +def model_specs(config: Dict[str, Any], include_disabled: bool = False) -> List[ModelSpec]: + common = config.get("common", {}) + specs = [] + for entry in config.get("models", []): + if not entry.get("enabled", True) and not include_disabled: + continue + if entry["type"] not in RUNNERS: + raise ValueError(f"unknown model type {entry['type']!r} for {entry['name']!r}; expected one of {sorted(RUNNERS)}") + specs.append(ModelSpec.from_config(entry, common)) + return specs + + +def reference_dir(config: Dict[str, Any]) -> Path: + d = config.get("common", {}).get("reference_dir") + return Path(d) if d else DEFAULT_REFERENCE_DIR + + +def schedulers_spec(config: Dict[str, Any]) -> ModelSpec: + atol, rtol = DEFAULT_TOLERANCES[torch.float32] + return ModelSpec(name=SCHEDULERS_NAME, type=SCHEDULERS_NAME, cfg={}, device="cpu", dtype=torch.float32, atol=atol, rtol=rtol) + + +def all_specs(config: Dict[str, Any], names: Optional[List[str]] = None, include_disabled: bool = False) -> List[ModelSpec]: + """All runnable specs. Entries with ``enabled = false`` are skipped unless they are named explicitly in ``names``.""" + if names: + # an explicit name overrides `enabled = false` + specs = [schedulers_spec(config)] + model_specs(config, include_disabled=True) + missing = set(names) - {s.name for s in specs} + if missing: + raise ValueError(f"unknown model name(s): {sorted(missing)}; available: {[s.name for s in specs]}") + return [s for s in specs if s.name in names] + return [schedulers_spec(config)] + model_specs(config, include_disabled=include_disabled) + + +def run_spec(spec: ModelSpec) -> Dict[str, np.ndarray]: + if spec.type == SCHEDULERS_NAME: + return run_schedulers(spec) + return RUNNERS[spec.type](spec) + + +def record(config: Dict[str, Any], names: Optional[List[str]] = None) -> None: + ref_dir = reference_dir(config) + ref_dir.mkdir(parents=True, exist_ok=True) + for spec in all_specs(config, names): + print(f"=== record {spec.name} ({spec.type}, device={spec.device}, dtype={spec.dtype})") + outputs = run_spec(spec) + np.savez_compressed(ref_dir / f"{spec.name}.npz", **outputs) + meta = { + "name": spec.name, + "type": spec.type, + "device": spec.device, + "dtype": str(spec.dtype).replace("torch.", ""), + "prompts": PROMPTS, + "keys": sorted(outputs.keys()), + "versions": _versions(), + } + (ref_dir / f"{spec.name}.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8") + print(f" saved {len(outputs)} arrays -> {ref_dir / (spec.name + '.npz')}") + + +# --------------------------------------------------------------------------- +# comparison +# --------------------------------------------------------------------------- + + +@dataclass +class CompareResult: + name: str + ok: bool + lines: List[str] = field(default_factory=list) + failures: List[str] = field(default_factory=list) + + @property + def summary(self) -> str: + return "\n".join(self.lines + ([""] + self.failures if self.failures else [])) + + +def _cosine(a: np.ndarray, b: np.ndarray) -> float: + a = a.astype(np.float64).ravel() + b = b.astype(np.float64).ravel() + na, nb = np.linalg.norm(a), np.linalg.norm(b) + if na == 0 or nb == 0: + return 1.0 if na == nb else 0.0 + return float(np.dot(a, b) / (na * nb)) + + +def compare_arrays(name: str, ref: Dict[str, np.ndarray], cur: Dict[str, np.ndarray], atol: float, rtol: float) -> CompareResult: + result = CompareResult(name=name, ok=True) + keys = sorted(set(ref) | set(cur)) + width = max(len(k) for k in keys) if keys else 10 + result.lines.append(f"--- {name}: {len(keys)} arrays, atol={atol:g}, rtol={rtol:g}") + for k in keys: + if k not in ref: + result.failures.append(f"{k}: not in reference (new output)") + result.ok = False + continue + if k not in cur: + result.failures.append(f"{k}: missing in current outputs") + result.ok = False + continue + r, c = ref[k], cur[k] + if r.shape != c.shape: + result.failures.append(f"{k}: shape mismatch reference={r.shape} current={c.shape}") + result.ok = False + continue + if r.dtype.kind in "iub" or c.dtype.kind in "iub": + n_diff = int(np.count_nonzero(r != c)) + status = "OK" if n_diff == 0 else "MISMATCH" + result.lines.append(f"{k:<{width}} {str(r.shape):<22} exact diff_elems={n_diff:<8} {status}") + if n_diff: + result.failures.append(f"{k}: {n_diff} differing integer elements") + result.ok = False + continue + r64, c64 = r.astype(np.float64), c.astype(np.float64) + finite = np.isfinite(r64) & np.isfinite(c64) + nan_mismatch = int(np.count_nonzero(np.isfinite(r64) != np.isfinite(c64))) + diff = np.abs(r64 - c64)[finite] + max_abs = float(diff.max()) if diff.size else 0.0 + denom = np.abs(r64)[finite] + max_rel = float((diff / np.maximum(denom, 1e-12)).max()) if diff.size else 0.0 + cos = _cosine(r64[finite], c64[finite]) if diff.size else 1.0 + close = np.allclose(r64[finite], c64[finite], atol=atol, rtol=rtol) and nan_mismatch == 0 + n_bad = int(np.count_nonzero(diff > atol + rtol * denom)) + status = "OK" if close else "MISMATCH" + result.lines.append( + f"{k:<{width}} {str(r.shape):<22} max_abs={max_abs:<10.3e} max_rel={max_rel:<10.3e} cos={cos:.6f} bad={n_bad:<8} {status}" + ) + if not close: + result.failures.append(f"{k}: max_abs={max_abs:.3e} max_rel={max_rel:.3e} cos={cos:.6f} bad_elems={n_bad} nan_mismatch={nan_mismatch}") + result.ok = False + return result + + +def compare_spec(spec: ModelSpec, ref_dir: Path) -> CompareResult: + ref_path = ref_dir / f"{spec.name}.npz" + if not ref_path.exists(): + raise FileNotFoundError(f"reference {ref_path} not found; run `python tests/local/regression_te.py record --models {spec.name}` first") + meta_path = ref_dir / f"{spec.name}.json" + if meta_path.exists(): + meta = json.loads(meta_path.read_text(encoding="utf-8")) + if meta.get("prompts") != PROMPTS: + raise RuntimeError(f"{spec.name}: PROMPTS changed since the reference was recorded; re-record") + ref_versions = meta.get("versions", {}) + else: + ref_versions = {} + with np.load(ref_path) as f: + ref = {k: f[k] for k in f.files} + print(f"=== compare {spec.name} ({spec.type}, device={spec.device}, dtype={spec.dtype})") + if ref_versions: + cur_versions = _versions() + for lib in ("transformers", "diffusers", "torch", "accelerate", "huggingface_hub"): + print(f" {lib:<16} reference={ref_versions.get(lib, '?'):<14} current={cur_versions.get(lib, '?')}") + cur = run_spec(spec) + return compare_arrays(spec.name, ref, cur, spec.atol, spec.rtol) + + +def compare(config: Dict[str, Any], names: Optional[List[str]] = None, verbose: bool = True) -> bool: + ref_dir = reference_dir(config) + all_ok = True + for spec in all_specs(config, names): + result = compare_spec(spec, ref_dir) + if verbose: + print(result.summary) + print(f" {spec.name}: {'OK' if result.ok else 'FAILED'}") + all_ok &= result.ok + print("ALL OK" if all_ok else "SOME COMPARISONS FAILED") + return all_ok + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def main(argv: Optional[List[str]] = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("mode", choices=["record", "compare", "list"]) + parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG, help="path to models.toml") + parser.add_argument("--models", type=str, default=None, help="comma separated model names to process (default: all)") + parser.add_argument("--quiet", action="store_true", help="compare: print only the per-model verdicts") + args = parser.parse_args(argv) + + config = load_config(args.config) + names = [n.strip() for n in args.models.split(",")] if args.models else None + + if args.mode == "list": + for spec in all_specs(config, include_disabled=True): + enabled = spec.cfg.get("enabled", True) + print( + f"{spec.name:<20} type={spec.type:<14} device={spec.device:<6} dtype={str(spec.dtype).replace('torch.', ''):<9} " + f"atol={spec.atol:g} rtol={spec.rtol:g}{'' if enabled else ' (enabled = false; runs only when named with --models)'}" + ) + return 0 + if args.mode == "record": + record(config, names) + return 0 + return 0 if compare(config, names, verbose=not args.quiet) else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/local/test_regression_te.py b/tests/local/test_regression_te.py new file mode 100644 index 0000000000..b99763e307 --- /dev/null +++ b/tests/local/test_regression_te.py @@ -0,0 +1,34 @@ +"""pytest wrapper for the local regression harness (see regression_te.py). + +Skipped entirely unless ``tests/local/models.toml`` exists. Each model listed +there (plus the weight-free ``schedulers`` check) becomes one test that compares +the current outputs against the recorded references. Tests whose reference has +not been recorded yet are skipped. + + pytest tests/local # run all + pytest tests/local -k sdxl # one model +""" + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import regression_te as rt # noqa: E402 + +if not rt.DEFAULT_CONFIG.exists(): + pytest.skip(f"{rt.DEFAULT_CONFIG} not found (local-only regression tests)", allow_module_level=True) + +_CONFIG = rt.load_config() +_SPECS = rt.all_specs(_CONFIG) +_REF_DIR = rt.reference_dir(_CONFIG) + + +@pytest.mark.parametrize("spec", _SPECS, ids=[s.name for s in _SPECS]) +def test_regression(spec: rt.ModelSpec): + if not (_REF_DIR / f"{spec.name}.npz").exists(): + pytest.skip(f"no reference recorded for {spec.name}; run `python tests/local/regression_te.py record --models {spec.name}`") + result = rt.compare_spec(spec, _REF_DIR) + print(result.summary) + assert result.ok, "\n".join(result.failures) diff --git a/tests/test_clip_text_model.py b/tests/test_clip_text_model.py new file mode 100644 index 0000000000..8e9e1c88e2 --- /dev/null +++ b/tests/test_clip_text_model.py @@ -0,0 +1,123 @@ +"""Tests for library/clip_text_model.py (CLIPTextModel layout compatibility wrapper). + +The tests run on any transformers version: with < 5.6 the wrapper is not applied by +`wrap_clip_text_model` (the model already has `text_model`), so the wrapper class itself is +exercised by wrapping the inner transformer directly. +""" + +import torch +from transformers import CLIPTextConfig, CLIPTextModel + +from library.clip_text_model import ( + CLIPTextModelWrapper, + is_flattened_clip_text_model, + unwrap_clip_text_model, + wrap_clip_text_model, +) + +TINY_CONFIG = dict( + vocab_size=64, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=2, + max_position_embeddings=8, + hidden_act="quick_gelu", + layer_norm_eps=1e-05, + pad_token_id=1, + bos_token_id=0, + eos_token_id=2, + model_type="clip_text_model", + projection_dim=16, +) + +EXPECTED_KEYS = [ + "text_model.embeddings.token_embedding.weight", + "text_model.embeddings.position_embedding.weight", + "text_model.encoder.layers.0.self_attn.q_proj.weight", + "text_model.encoder.layers.1.mlp.fc2.bias", + "text_model.final_layer_norm.weight", +] + + +def _make_model() -> torch.nn.Module: + torch.manual_seed(0) + return CLIPTextModel._from_config(CLIPTextConfig(**TINY_CONFIG)) + + +def _force_wrap(model: CLIPTextModel) -> CLIPTextModelWrapper: + # transformers >= 5.6: the model itself is flat; < 5.6: use the inner CLIPTextTransformer + inner = model if is_flattened_clip_text_model(model) else model.text_model + return CLIPTextModelWrapper(inner) + + +def test_wrap_keeps_legacy_state_dict_keys_and_module_names(): + model = wrap_clip_text_model(_make_model()) + keys = set(model.state_dict().keys()) + for k in EXPECTED_KEYS: + assert k in keys, k + assert all(k.startswith("text_model.") for k in keys) + + names = dict(model.named_modules()) + assert "text_model.encoder.layers.0.self_attn" in names + assert "text_model.encoder.layers.0.mlp" in names + assert names["text_model.encoder.layers.0.self_attn"].__class__.__name__ in ("CLIPAttention", "CLIPSdpaAttention") + assert names["text_model.encoder.layers.0.mlp"].__class__.__name__ == "CLIPMLP" + + +def test_wrap_is_idempotent_and_unwrap_returns_transformers_model(): + model = wrap_clip_text_model(_make_model()) + assert wrap_clip_text_model(model) is model + assert isinstance(unwrap_clip_text_model(model), CLIPTextModel) + assert unwrap_clip_text_model(unwrap_clip_text_model(model)) is unwrap_clip_text_model(model) + + +def test_wrapper_forward_and_attributes_match_inner(): + model = _force_wrap(_make_model()) + inner = model.text_model + input_ids = torch.tensor([[0, 5, 6, 7, 2, 1, 1, 1]]) + + model.eval() + with torch.no_grad(): + out = model(input_ids, output_hidden_states=True, return_dict=True) + ref = inner(input_ids, output_hidden_states=True, return_dict=True) + assert torch.equal(out.last_hidden_state, ref.last_hidden_state) + assert torch.equal(out.pooler_output, ref.pooler_output) + assert len(out.hidden_states) == TINY_CONFIG["num_hidden_layers"] + 1 + # the legacy attribute path used all over the training scripts + assert torch.equal(model.text_model.final_layer_norm(out.hidden_states[-1]), out.last_hidden_state) + + assert model.config.hidden_size == TINY_CONFIG["hidden_size"] + assert model.dtype == torch.float32 + assert model.device == next(inner.parameters()).device + assert model.get_input_embeddings() is inner.embeddings.token_embedding + assert model.state_dict().keys() == {"text_model." + k for k in inner.state_dict().keys()} + + +def test_wrapper_load_state_dict_with_legacy_keys(): + src = wrap_clip_text_model(_make_model()) + dst = _force_wrap(_make_model()) + legacy_sd = {k: v.clone() for k, v in src.state_dict().items()} + info = dst.load_state_dict(legacy_sd, strict=True) + assert not info.missing_keys and not info.unexpected_keys + for k, v in legacy_sd.items(): + assert torch.equal(dst.state_dict()[k], v) + + +def test_wrapper_dtype_device_follow_module_to(): + model = _force_wrap(_make_model()) + model.to(dtype=torch.bfloat16) + assert model.dtype == torch.bfloat16 + assert model.text_model.embeddings.token_embedding.weight.dtype == torch.bfloat16 + + +def test_is_flattened_detection(): + raw = _make_model() + # on any version: a wrapped model is never reported as flattened + assert not is_flattened_clip_text_model(wrap_clip_text_model(raw)) + # < 5.6 has text_model, >= 5.6 does not: exactly one of raw / inner looks flat + if hasattr(raw, "text_model"): + assert not is_flattened_clip_text_model(raw) + assert is_flattened_clip_text_model(raw.text_model) + else: + assert is_flattened_clip_text_model(raw) diff --git a/tests/test_clip_tokenizer.py b/tests/test_clip_tokenizer.py new file mode 100644 index 0000000000..f071b5971c --- /dev/null +++ b/tests/test_clip_tokenizer.py @@ -0,0 +1,28 @@ +"""Tests for library.clip_tokenizer (legacy CLIP text normalization on top of the fast tokenizer). + +The tokenizer itself needs files from the Hugging Face Hub, so only the pure normalization is tested here. +The end-to-end token ids are covered by the local regression harness (tests/local). +""" + +from library.clip_tokenizer import _clean, clean_clip_text + + +def test_clean_clip_text_matches_legacy_normalization(): + # curly quotes are straightened, full-width characters become ASCII, HTML entities are unescaped, + # mojibake is repaired, whitespace is collapsed and stripped (ftfy.fix_text + whitespace_clean) + assert clean_clip_text("“quoted”") == '"quoted"' + assert clean_clip_text("ＡＢＣ　１２３") == "ABC 123" + assert clean_clip_text("a & b") == "a & b" + assert clean_clip_text("résumé") == "résumé" # mojibake of "résumé" + assert clean_clip_text(" many spaces\n\tand lines ") == "many spaces and lines" + # plain text is unchanged + assert clean_clip_text("a photo of a cat, masterpiece") == "a photo of a cat, masterpiece" + assert clean_clip_text("") == "" + + +def test_clean_handles_batches_pairs_and_pretokenized_input(): + assert _clean(None) is None + assert _clean(["“a”", "b"]) == ['"a"', "b"] + assert _clean(("“a”", "“b”")) == ('"a"', '"b"') + assert _clean([["“a”", "b"], ["c"]]) == [['"a"', "b"], ["c"]] + assert _clean(3) == 3 diff --git a/tests/test_cv2_stub.py b/tests/test_cv2_stub.py new file mode 100644 index 0000000000..aac0d4ccbc --- /dev/null +++ b/tests/test_cv2_stub.py @@ -0,0 +1,171 @@ +"""Tests for the optional-OpenCV fallback (library/cv2_compat.py + library/_cv2_stub). + +The equivalence tests need the real opencv-python and are skipped without it. +The registration test runs in a subprocess with ``cv2`` blocked so it works +either way. +""" + +import subprocess +import sys +import textwrap + +import numpy as np +import pytest + +from library import _cv2_stub as stub + +real_cv2 = pytest.importorskip("cv2") +if getattr(real_cv2, "_IS_CV2_STUB", False): + real_cv2 = None + +needs_opencv = pytest.mark.skipif(real_cv2 is None, reason="real opencv-python not installed") + + +def _maxdiff(a: np.ndarray, b: np.ndarray) -> float: + assert a.shape == b.shape, (a.shape, b.shape) + assert a.dtype == b.dtype, (a.dtype, b.dtype) + return float(np.abs(a.astype(np.float64) - b.astype(np.float64)).max()) + + +def test_stub_is_registered_when_opencv_is_missing(): + code = textwrap.dedent( + """ + import importlib.abc, sys + class Block(importlib.abc.MetaPathFinder): + def find_spec(self, name, path=None, target=None): + if name == "cv2" or name.startswith("cv2."): + raise ImportError("blocked") + sys.meta_path.insert(0, Block()) + import library.utils # imports cv2_compat before cv2 + import cv2 + from library import cv2_compat + assert cv2_compat.HAS_OPENCV is False + assert getattr(cv2, "_IS_CV2_STUB", False) is True + import numpy as np + img = np.zeros((4, 6, 3), dtype=np.uint8) + assert cv2.resize(img, (3, 2), interpolation=cv2.INTER_AREA).shape == (2, 3, 3) + print("stub-ok") + """ + ) + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=300) + assert result.returncode == 0, result.stderr + assert "stub-ok" in result.stdout + + +def test_require_opencv_raises_without_opencv(monkeypatch): + from library import cv2_compat + + monkeypatch.setattr(cv2_compat, "HAS_OPENCV", False) + with pytest.raises(SystemExit): + cv2_compat.require_opencv("tools/canny.py") + monkeypatch.setattr(cv2_compat, "HAS_OPENCV", True) + cv2_compat.require_opencv("tools/canny.py") # no-op + + +@needs_opencv +@pytest.mark.parametrize( + "code_name, channels, tol", + [ + ("COLOR_BGR2RGB", 3, 0), + ("COLOR_RGB2BGR", 3, 0), + ("COLOR_BGRA2RGBA", 4, 0), + ("COLOR_RGBA2BGRA", 4, 0), + ("COLOR_BGR2GRAY", 3, 1), + ("COLOR_RGB2GRAY", 3, 1), + ("COLOR_BGR2HSV", 3, 1), + ("COLOR_RGB2HSV", 3, 1), + ], +) +def test_cvtcolor_matches_opencv(code_name, channels, tol): + rng = np.random.default_rng(0) + img = rng.integers(0, 256, size=(37, 53, channels), dtype=np.uint8) + # include pure hues and greys, where hue wrap-around / zero saturation matter + img[0, :3] = [[255, 0, 0], [0, 255, 0], [0, 0, 255]] if channels == 3 else [[255, 0, 0, 7], [0, 255, 0, 7], [0, 0, 255, 7]] + img[1, :2, :3] = [[0, 0, 0], [255, 255, 255]] + expected = real_cv2.cvtColor(img, getattr(real_cv2, code_name)) + got = stub.cvtColor(img, getattr(stub, code_name)) + assert _maxdiff(expected, got) <= tol + + +@needs_opencv +def test_hsv_to_bgr_matches_opencv(): + rng = np.random.default_rng(1) + hsv = np.stack( + [ + rng.integers(0, 180, (37, 53), dtype=np.uint8), + rng.integers(0, 256, (37, 53), dtype=np.uint8), + rng.integers(0, 256, (37, 53), dtype=np.uint8), + ], + axis=-1, + ) + assert _maxdiff(real_cv2.cvtColor(hsv, real_cv2.COLOR_HSV2BGR), stub.cvtColor(hsv, stub.COLOR_HSV2BGR)) <= 2 + assert _maxdiff(real_cv2.cvtColor(hsv, real_cv2.COLOR_HSV2RGB), stub.cvtColor(hsv, stub.COLOR_HSV2RGB)) <= 2 + + +@needs_opencv +def test_gray2bgr_matches_opencv(): + gray = np.random.default_rng(2).integers(0, 256, (11, 13), dtype=np.uint8) + assert _maxdiff(real_cv2.cvtColor(gray, real_cv2.COLOR_GRAY2BGR), stub.cvtColor(gray, stub.COLOR_GRAY2BGR)) == 0 + + +@needs_opencv +@pytest.mark.parametrize("interp", ["INTER_AREA", "INTER_LINEAR"]) +def test_resize_area_linear_match_opencv_random_sizes(interp): + """AREA and LINEAR are re-implemented in NumPy and must match OpenCV up to rounding, + for shrinking, enlarging and mixed scaling, and for 1/3/4-channel images.""" + rng = np.random.default_rng(3) + for _ in range(150): + h, w = int(rng.integers(1, 80)), int(rng.integers(1, 80)) + nh, nw = int(rng.integers(1, 90)), int(rng.integers(1, 90)) + c = int(rng.choice([0, 1, 3, 4])) + img = rng.integers(0, 256, size=(h, w) if c == 0 else (h, w, c), dtype=np.uint8) + expected = real_cv2.resize(img, (nw, nh), interpolation=getattr(real_cv2, interp)) + got = stub.resize(img, (nw, nh), interpolation=getattr(stub, interp)) + assert _maxdiff(expected, got) <= 1, (img.shape, (nh, nw)) + + +@needs_opencv +def test_resize_area_large_image_exact(): + rng = np.random.default_rng(4) + img = rng.integers(0, 256, size=(1536, 2048, 3), dtype=np.uint8) + expected = real_cv2.resize(img, (1024, 768), interpolation=real_cv2.INTER_AREA) + got = stub.resize(img, (1024, 768), interpolation=stub.INTER_AREA) + assert _maxdiff(expected, got) <= 1 + + +def test_resize_rgba_channels_are_independent(): + """Pillow premultiplies alpha when resizing RGBA; the stub must not (OpenCV does not).""" + rng = np.random.default_rng(5) + rgb = rng.integers(0, 256, size=(40, 60, 3), dtype=np.uint8) + alpha_a = np.full((40, 60, 1), 255, dtype=np.uint8) + alpha_b = rng.integers(0, 256, size=(40, 60, 1), dtype=np.uint8) + for interp in (stub.INTER_AREA, stub.INTER_LINEAR, stub.INTER_CUBIC, stub.INTER_LANCZOS4): + out_a = stub.resize(np.concatenate([rgb, alpha_a], axis=-1), (30, 20), interpolation=interp) + out_b = stub.resize(np.concatenate([rgb, alpha_b], axis=-1), (30, 20), interpolation=interp) + assert np.array_equal(out_a[..., :3], out_b[..., :3]) + + +def test_resize_shapes_and_dtypes(): + img = np.zeros((10, 20, 3), dtype=np.uint8) + assert stub.resize(img, (5, 4), interpolation=stub.INTER_AREA).shape == (4, 5, 3) + assert stub.resize(np.zeros((10, 20), np.uint8), (5, 4), interpolation=stub.INTER_LINEAR).shape == (4, 5) + # single-channel with trailing dim comes back 2-D, like OpenCV + assert stub.resize(np.zeros((10, 20, 1), np.uint8), (5, 4), interpolation=stub.INTER_AREA).shape == (4, 5) + f32 = np.random.default_rng(6).random((16, 16), dtype=np.float32) + out = stub.resize(f32, (64, 64), interpolation=stub.INTER_CUBIC) + assert out.shape == (64, 64) and out.dtype == np.float32 + # same size returns a copy, never the input itself + same = stub.resize(img, (20, 10), interpolation=stub.INTER_AREA) + assert same is not img and same.shape == img.shape + # fx / fy form + assert stub.resize(img, None, fx=0.5, fy=0.5, interpolation=stub.INTER_AREA).shape == (5, 10, 3) + + +def test_imwrite_roundtrip(tmp_path): + from PIL import Image + + bgr = np.random.default_rng(8).integers(0, 256, size=(8, 9, 3), dtype=np.uint8) + path = tmp_path / "x.png" + assert stub.imwrite(str(path), bgr) + rgb = np.asarray(Image.open(path).convert("RGB")) + assert np.array_equal(rgb, bgr[..., ::-1]) diff --git a/tests/test_diffusers_warning_filter.py b/tests/test_diffusers_warning_filter.py new file mode 100644 index 0000000000..1880451843 --- /dev/null +++ b/tests/test_diffusers_warning_filter.py @@ -0,0 +1,45 @@ +"""library.utils installs a filter that drops the spurious "should be kept in float32: []" warning of diffusers >= 0.40.""" + +import logging + +import torch +from diffusers import AutoencoderKL + +import library.utils # noqa: F401 - installs the filter at import time + + +def _tiny_vae() -> AutoencoderKL: + return AutoencoderKL( + in_channels=3, out_channels=3, down_block_types=("DownEncoderBlock2D",), up_block_types=("UpDecoderBlock2D",), + block_out_channels=(8,), layers_per_block=1, latent_channels=4, norm_num_groups=8, sample_size=16, + ) + + +def _messages_during(fn) -> list: + records = [] + diffusers_logger = logging.getLogger("diffusers.models.modeling_utils") + + class Collect(logging.Handler): + def emit(self, record): + records.append(record.getMessage()) + + h = Collect() + diffusers_logger.addHandler(h) + try: + fn() + finally: + diffusers_logger.removeHandler(h) + return records + + +def test_spurious_fp32_warning_is_dropped(): + vae = _tiny_vae() + msgs = _messages_during(lambda: vae.to(torch.float16)) + assert not any("should be kept in float32" in m for m in msgs), msgs + + +def test_real_fp32_warning_is_kept(): + vae = _tiny_vae() + vae._keep_in_fp32_modules = ["quant_conv"] + msgs = _messages_during(lambda: vae.to(torch.float16)) + assert any("should be kept in float32: ['quant_conv']" in m for m in msgs), msgs diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index c1cc9e3717..f6679d4b76 100644 --- a/tests/test_optimizer.py +++ b/tests/test_optimizer.py @@ -9,8 +9,14 @@ from lion_pytorch import lion_pytorch import schedulefree -import dadaptation -import dadaptation.experimental as dadapt_experimental +# dadaptation is not in requirements.txt (CI installs it separately); skip its +# cases when it is not available instead of failing the whole module. +try: + import dadaptation + import dadaptation.experimental as dadapt_experimental +except ImportError: + dadaptation = None + dadapt_experimental = None import prodigyopt import schedulefree as sf @@ -86,36 +92,6 @@ def test_all_supported_optimizers(): "instance": bnb.optim.PagedAdamW32bit, }, {"name": "torch.optim.sgd.SGD", "alias": "SGD", "instance": torch.optim.SGD}, - { - "name": "dadaptation.experimental.dadapt_adam_preprint.DAdaptAdamPreprint", - "alias": "DAdaptAdamPreprint", - "instance": dadapt_experimental.DAdaptAdamPreprint, - }, - { - "name": "dadaptation.dadapt_adagrad.DAdaptAdaGrad", - "alias": "DAdaptAdaGrad", - "instance": dadaptation.DAdaptAdaGrad, - }, - { - "name": "dadaptation.dadapt_adan.DAdaptAdan", - "alias": "DAdaptAdan", - "instance": dadaptation.DAdaptAdan, - }, - { - "name": "dadaptation.experimental.dadapt_adan_ip.DAdaptAdanIP", - "alias": "DAdaptAdanIP", - "instance": dadapt_experimental.DAdaptAdanIP, - }, - { - "name": "dadaptation.dadapt_lion.DAdaptLion", - "alias": "DAdaptLion", - "instance": dadaptation.DAdaptLion, - }, - { - "name": "dadaptation.dadapt_sgd.DAdaptSGD", - "alias": "DAdaptSGD", - "instance": dadaptation.DAdaptSGD, - }, { "name": "prodigyopt.prodigy.Prodigy", "alias": "Prodigy", @@ -137,6 +113,39 @@ def test_all_supported_optimizers(): "instance": sf.SGDScheduleFree, }, ] + if dadaptation is not None: + optimizers += [ + { + "name": "dadaptation.experimental.dadapt_adam_preprint.DAdaptAdamPreprint", + "alias": "DAdaptAdamPreprint", + "instance": dadapt_experimental.DAdaptAdamPreprint, + }, + { + "name": "dadaptation.dadapt_adagrad.DAdaptAdaGrad", + "alias": "DAdaptAdaGrad", + "instance": dadaptation.DAdaptAdaGrad, + }, + { + "name": "dadaptation.dadapt_adan.DAdaptAdan", + "alias": "DAdaptAdan", + "instance": dadaptation.DAdaptAdan, + }, + { + "name": "dadaptation.experimental.dadapt_adan_ip.DAdaptAdanIP", + "alias": "DAdaptAdanIP", + "instance": dadapt_experimental.DAdaptAdanIP, + }, + { + "name": "dadaptation.dadapt_lion.DAdaptLion", + "alias": "DAdaptLion", + "instance": dadaptation.DAdaptLion, + }, + { + "name": "dadaptation.dadapt_sgd.DAdaptSGD", + "alias": "DAdaptSGD", + "instance": dadaptation.DAdaptSGD, + }, + ] for opt in optimizers: with patch("sys.argv", ["", "--optimizer_type", opt.get("alias")]): diff --git a/tests/test_sampling_scheduler.py b/tests/test_sampling_scheduler.py new file mode 100644 index 0000000000..656547c5b6 --- /dev/null +++ b/tests/test_sampling_scheduler.py @@ -0,0 +1,80 @@ +"""``library.sampling.get_my_scheduler`` builds every ``--sample_sampler`` choice with the installed diffusers.""" + +import importlib.util + +import pytest + +from library import sampling +from library.sampling import check_sampler_requirements, get_my_scheduler + +HAS_SCIPY = importlib.util.find_spec("scipy") is not None + +# every choice of --sample_sampler (library/args.py) except the scipy-dependent lms / k_lms +SAMPLERS = [ + "ddim", + "pndm", + "euler", + "euler_a", + "heun", + "dpm_2", + "dpm_2_a", + "dpmsolver", + "dpmsolver++", + "dpmsingle", + "k_euler", + "k_euler_a", + "k_dpm_2", + "k_dpm_2_a", +] + + +@pytest.mark.parametrize("v_parameterization", [False, True]) +@pytest.mark.parametrize("sampler", SAMPLERS) +def test_every_sampler_builds_and_sets_timesteps(sampler, v_parameterization): + scheduler = get_my_scheduler(sample_sampler=sampler, v_parameterization=v_parameterization) + scheduler.set_timesteps(4) + assert len(scheduler.timesteps) >= 4 + if v_parameterization: + assert scheduler.config.prediction_type == "v_prediction" + + +def test_dpmsolver_uses_sigma_min_final_sigmas(): + # diffusers rejects final_sigmas_type="zero" (its default) for the non-++ algorithm + scheduler = get_my_scheduler(sample_sampler="dpmsolver", v_parameterization=False) + assert scheduler.config.algorithm_type == "dpmsolver" + assert scheduler.config.final_sigmas_type == "sigma_min" + # dpmsolver++ keeps the diffusers default so its output is unchanged + scheduler = get_my_scheduler(sample_sampler="dpmsolver++", v_parameterization=False) + assert scheduler.config.final_sigmas_type == "zero" + + +def test_dpmsingle_has_no_steps_offset(): + scheduler = get_my_scheduler(sample_sampler="dpmsingle", v_parameterization=False) + assert "steps_offset" not in scheduler.config + + +@pytest.mark.parametrize("sampler", [s for s in SAMPLERS if s != "dpmsingle"]) +def test_other_samplers_keep_steps_offset_one(sampler): + scheduler = get_my_scheduler(sample_sampler=sampler, v_parameterization=False) + assert scheduler.config.steps_offset == 1 + + +@pytest.mark.skipif(not HAS_SCIPY, reason="scipy is not installed") +def test_lms_builds_when_scipy_is_available(): + scheduler = get_my_scheduler(sample_sampler="lms", v_parameterization=False) + scheduler.set_timesteps(4) + + +@pytest.mark.parametrize("sampler", ["lms", "k_lms"]) +def test_lms_without_scipy_raises_a_clear_error(monkeypatch, sampler): + monkeypatch.setattr(sampling.importlib.util, "find_spec", lambda name: None if name == "scipy" else object()) + with pytest.raises(ImportError, match="pip install scipy"): + check_sampler_requirements(sampler) + with pytest.raises(ImportError, match="pip install scipy"): + get_my_scheduler(sample_sampler=sampler, v_parameterization=False) + + +def test_check_sampler_requirements_ignores_other_samplers(monkeypatch): + monkeypatch.setattr(sampling.importlib.util, "find_spec", lambda name: None) + for sampler in SAMPLERS: + check_sampler_requirements(sampler) diff --git a/tools/canny.py b/tools/canny.py index f2190975c1..84fcded78d 100644 --- a/tools/canny.py +++ b/tools/canny.py @@ -1,4 +1,8 @@ import argparse + +from library import cv2_compat + +cv2_compat.require_opencv("tools/canny.py") import cv2 import logging diff --git a/tools/detect_face_rotate.py b/tools/detect_face_rotate.py index 16fd7d0b79..f51e0fa2cd 100644 --- a/tools/detect_face_rotate.py +++ b/tools/detect_face_rotate.py @@ -9,6 +9,10 @@ import argparse import math + +from library import cv2_compat + +cv2_compat.require_opencv("tools/detect_face_rotate.py") import cv2 import glob import os diff --git a/tools/latent_upscaler.py b/tools/latent_upscaler.py index f05cf71942..0d8a1acbb5 100644 --- a/tools/latent_upscaler.py +++ b/tools/latent_upscaler.py @@ -4,6 +4,7 @@ import argparse import glob import os +from library import cv2_compat # noqa: F401 - must be imported before `import cv2` import cv2 from diffusers import AutoencoderKL diff --git a/tools/original_control_net.py b/tools/original_control_net.py index 5640d542d9..8274d94e5d 100644 --- a/tools/original_control_net.py +++ b/tools/original_control_net.py @@ -1,5 +1,6 @@ from typing import List, NamedTuple, Any import numpy as np +from library import cv2_compat # noqa: F401 - must be imported before `import cv2` import cv2 import torch from safetensors.torch import load_file @@ -110,6 +111,7 @@ def load_preprocess(prep_type: str): return None if prep_type.startswith("canny"): + cv2_compat.require_opencv("ControlNet canny preprocessor") args = prep_type.split("_") th1 = int(args[1]) if len(args) >= 2 else 63 th2 = int(args[2]) if len(args) >= 3 else 191 diff --git a/tools/resize_images_to_resolution.py b/tools/resize_images_to_resolution.py index f5fbae2bbd..d2d19f1570 100644 --- a/tools/resize_images_to_resolution.py +++ b/tools/resize_images_to_resolution.py @@ -1,6 +1,5 @@ import glob import os -import cv2 import argparse import shutil import math