Init

.gitignore

@@ -0,0 +1,181 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+.idea/
+
+/env.sh
+/models
+/custom/*
+!/custom/.gitkeep
+/.tmp
+/venv.bkp
+/venv.*
+/config/*
+!/config/examples
+!/config/_PUT_YOUR_CONFIGS_HERE).txt
+/output/*
+!/output/.gitkeep
+/extensions/*
+!/extensions/example
+/temp
+/wandb
+.vscode/settings.json
+.DS_Store
+._.DS_Store
+merge_file.py

.gitmodules

@@ -0,0 +1,12 @@
+[submodule "repositories/sd-scripts"]
+	path = repositories/sd-scripts
+	url = https://github.com/kohya-ss/sd-scripts.git
+[submodule "repositories/leco"]
+	path = repositories/leco
+	url = https://github.com/p1atdev/LECO
+[submodule "repositories/batch_annotator"]
+	path = repositories/batch_annotator
+	url = https://github.com/ostris/batch-annotator
+[submodule "repositories/ipadapter"]
+	path = repositories/ipadapter
+	url = https://github.com/tencent-ailab/IP-Adapter.git

FAQ.md

@@ -0,0 +1,10 @@
+# FAQ
+
+WIP. Will continue to add things as they are needed.
+
+## FLUX.1 Training
+
+#### How much VRAM is required to train a lora on FLUX.1?
+
+24GB minimum is required.
+

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Ostris, LLC
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,12 +1,40 @@
----
-title: Flux LoRA Trainning On Modal
-emoji: 🏢
-colorFrom: blue
-colorTo: red
-sdk: gradio
-sdk_version: 5.9.1
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# FLUX LoRA Training on Modal
+
+## IMPORTANT - READ THIS
+
+This Space allows you to train LoRA models for FLUX using Modal's GPU resources. You **must** use your own API tokens for this Space to function.
+
+## Setup Instructions:
+
+1.  **Create Accounts:** Create free accounts on [Hugging Face](https://huggingface.co), [Modal](https://modal.com), and [Weights & Biases](https://wandb.ai) (if you want to save training info).
+2.  **Get API Tokens:**
+    *   **Hugging Face:** Obtain a "write" access token from your Hugging Face [settings/tokens](https://huggingface.co/settings/tokens).
+    *   **Modal:** Get your Modal API token from your [Modal dashboard](https://modal.com/).
+    *   **WandB:** Generate a WandB API key from your Weights & Biases settings if you plan to use WandB.
+3.  **Duplicate This Space:** Duplicate (clone) this Hugging Face Space to your own account.
+4.  **Add API Tokens as Secrets:** In your duplicated space, navigate to "Settings" -> "Variables and Secrets" and add:
+    *   `HF_TOKEN`: Your Hugging Face write access token.
+    *   `WANDB_API_KEY`: Your Weights & Biases API key.
+5.  **Upload your dataset**
+    *  You must upload your dataset to the `/root/ai-toolkit` folder. The images can be in a subfolder but that folder's path needs to begin with `/root/ai-toolkit/`. For instance:  `/root/ai-toolkit/my-dataset`
+    *  Make sure the image file names match with a corresponding text caption in the same folder. `image1.jpg` and `image1.txt`
+    *  You can upload a zip file of your dataset, or just a collection of images and text files.
+6. **Customize and Train:**
+    *  Go to the `App` tab, and use the Gradio interface to train your LoRA.
+    *   Enter the required information, dataset, and configure the training parameters.
+        * Choose to upload a zip file or multiple images
+        * Make sure your image file names match with a corresponding `.txt` file.
+    *   Click the "Start Training" button and wait for the training to complete (check Modal for logs, WandB for training data).
+    * The UI will automatically upload the file(s) to the Modal compute environment
+7. **View Results:**
+    *   Trained LoRA models will be automatically pushed to your Hugging Face account if you enable the option and have the necessary write token set
+    *   Samples, Logs, optimizer and other training information will be stored on WandB if enabled.
+    *   Models, optimizer, and samples are always stored in `Storage > flux-lora-models` on Modal.
+
+## Notes
+
+*  Training data will always be located in a dataset folder under `/root/ai-toolkit/your-dataset` in your config, this is required to train correctly. If you upload a folder name `my-dataset` as a zip, the folder path to reference will be: `/root/ai-toolkit/my-dataset`.
+*  Make sure the images have a corresponding text file with the same name
+*  Training and downloading samples in this model will take a while be patient, especially in low vram mode
+
+If you encounter any problems, please open a new issue on the [ostris/ai-toolkit](https://github.com/ostris/ai-toolkit) github

README_origin.md

@@ -0,0 +1,468 @@
+# AI Toolkit by Ostris
+
+## IMPORTANT NOTE - READ THIS
+This is my research repo. I do a lot of experiments in it and it is possible that I will break things.
+If something breaks, checkout an earlier commit. This repo can train a lot of things, and it is
+hard to keep up with all of them.
+
+## Support my work
+
+<a href="https://glif.app" target="_blank">
+<img alt="glif.app" src="https://raw.githubusercontent.com/ostris/ai-toolkit/main/assets/glif.svg?v=1" width="256" height="auto">
+</a>
+
+
+My work on this project would not be possible without the amazing support of [Glif](https://glif.app/) and everyone on the 
+team. If you want to support me, support Glif. [Join the site](https://glif.app/), 
+[Join us on Discord](https://discord.com/invite/nuR9zZ2nsh), [follow us on Twitter](https://x.com/heyglif)
+and come make some cool stuff with us
+
+## Installation
+
+Requirements:
+- python >3.10
+- Nvidia GPU with enough ram to do what you need
+- python venv
+- git
+
+
+
+Linux:
+```bash
+git clone https://github.com/ostris/ai-toolkit.git
+cd ai-toolkit
+git submodule update --init --recursive
+python3 -m venv venv
+source venv/bin/activate
+# .\venv\Scripts\activate on windows
+# install torch first
+pip3 install torch
+pip3 install -r requirements.txt
+```
+
+Windows:
+```bash
+git clone https://github.com/ostris/ai-toolkit.git
+cd ai-toolkit
+git submodule update --init --recursive
+python -m venv venv
+.\venv\Scripts\activate
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
+pip install -r requirements.txt
+```
+
+## FLUX.1 Training
+
+### Tutorial
+
+To get started quickly, check out [@araminta_k](https://x.com/araminta_k) tutorial on [Finetuning Flux Dev on a 3090](https://www.youtube.com/watch?v=HzGW_Kyermg) with 24GB VRAM.
+
+
+### Requirements
+You currently need a GPU with **at least 24GB of VRAM** to train FLUX.1. If you are using it as your GPU to control 
+your monitors, you probably need to set the flag `low_vram: true` in the config file under `model:`. This will quantize
+the model on CPU and should allow it to train with monitors attached. Users have gotten it to work on Windows with WSL,
+but there are some reports of a bug when running on windows natively. 
+I have only tested on linux for now. This is still extremely experimental
+and a lot of quantizing and tricks had to happen to get it to fit on 24GB at all. 
+
+### FLUX.1-dev
+
+FLUX.1-dev has a non-commercial license. Which means anything you train will inherit the
+non-commercial license. It is also a gated model, so you need to accept the license on HF before using it.
+Otherwise, this will fail. Here are the required steps to setup a license.
+
+1. Sign into HF and accept the model access here [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev)
+2. Make a file named `.env` in the root on this folder
+3. [Get a READ key from huggingface](https://huggingface.co/settings/tokens/new?) and add it to the `.env` file like so `HF_TOKEN=your_key_here`
+
+### FLUX.1-schnell
+
+FLUX.1-schnell is Apache 2.0. Anything trained on it can be licensed however you want and it does not require a HF_TOKEN to train.
+However, it does require a special adapter to train with it, [ostris/FLUX.1-schnell-training-adapter](https://huggingface.co/ostris/FLUX.1-schnell-training-adapter).
+It is also highly experimental. For best overall quality, training on FLUX.1-dev is recommended.
+
+To use it, You just need to add the assistant to the `model` section of your config file like so:
+
+```yaml
+      model:
+        name_or_path: "black-forest-labs/FLUX.1-schnell"
+        assistant_lora_path: "ostris/FLUX.1-schnell-training-adapter"
+        is_flux: true
+        quantize: true
+```
+
+You also need to adjust your sample steps since schnell does not require as many
+
+```yaml
+      sample:
+        guidance_scale: 1  # schnell does not do guidance
+        sample_steps: 4  # 1 - 4 works well
+```
+
+### Training
+1. Copy the example config file located at `config/examples/train_lora_flux_24gb.yaml` (`config/examples/train_lora_flux_schnell_24gb.yaml` for schnell) to the `config` folder and rename it to `whatever_you_want.yml`
+2. Edit the file following the comments in the file
+3. Run the file like so `python run.py config/whatever_you_want.yml`
+
+A folder with the name and the training folder from the config file will be created when you start. It will have all 
+checkpoints and images in it. You can stop the training at any time using ctrl+c and when you resume, it will pick back up
+from the last checkpoint.
+
+IMPORTANT. If you press crtl+c while it is saving, it will likely corrupt that checkpoint. So wait until it is done saving
+
+### Need help?
+
+Please do not open a bug report unless it is a bug in the code. You are welcome to [Join my Discord](https://discord.gg/VXmU2f5WEU)
+and ask for help there. However, please refrain from PMing me directly with general question or support. Ask in the discord
+and I will answer when I can.
+
+## Gradio UI
+
+To get started training locally with a with a custom UI, once you followed the steps above and `ai-toolkit` is installed:
+
+```bash
+cd ai-toolkit #in case you are not yet in the ai-toolkit folder
+huggingface-cli login #provide a `write` token to publish your LoRA at the end
+python flux_train_ui.py
+```
+
+You will instantiate a UI that will let you upload your images, caption them, train and publish your LoRA
+![image](assets/lora_ease_ui.png)
+
+
+## Training in RunPod
+Example RunPod template: **runpod/pytorch:2.2.0-py3.10-cuda12.1.1-devel-ubuntu22.04**
+> You need a minimum of 24GB VRAM, pick a GPU by your preference.
+
+#### Example config ($0.5/hr):
+- 1x A40 (48 GB VRAM)
+- 19 vCPU 100 GB RAM
+
+#### Custom overrides (you need some storage to clone FLUX.1, store datasets, store trained models and samples):
+- ~120 GB Disk
+- ~120 GB Pod Volume
+- Start Jupyter Notebook
+
+### 1. Setup
+```
+git clone https://github.com/ostris/ai-toolkit.git
+cd ai-toolkit
+git submodule update --init --recursive
+python -m venv venv
+source venv/bin/activate
+pip install torch
+pip install -r requirements.txt
+pip install --upgrade accelerate transformers diffusers huggingface_hub #Optional, run it if you run into issues
+```
+### 2. Upload your dataset
+- Create a new folder in the root, name it `dataset` or whatever you like.
+- Drag and drop your .jpg, .jpeg, or .png images and .txt files inside the newly created dataset folder.
+
+### 3. Login into Hugging Face with an Access Token
+- Get a READ token from [here](https://huggingface.co/settings/tokens) and request access to Flux.1-dev model from [here](https://huggingface.co/black-forest-labs/FLUX.1-dev).
+- Run ```huggingface-cli login``` and paste your token.
+
+### 4. Training
+- Copy an example config file located at ```config/examples``` to the config folder and rename it to ```whatever_you_want.yml```.
+- Edit the config following the comments in the file.
+- Change ```folder_path: "/path/to/images/folder"``` to your dataset path like ```folder_path: "/workspace/ai-toolkit/your-dataset"```.
+- Run the file: ```python run.py config/whatever_you_want.yml```.
+
+### Screenshot from RunPod
+<img width="1728" alt="RunPod Training Screenshot" src="https://github.com/user-attachments/assets/53a1b8ef-92fa-4481-81a7-bde45a14a7b5">
+
+## Training in Modal
+
+### 1. Setup
+#### ai-toolkit:
+```
+git clone https://github.com/ostris/ai-toolkit.git
+cd ai-toolkit
+git submodule update --init --recursive
+python -m venv venv
+source venv/bin/activate
+pip install torch
+pip install -r requirements.txt
+pip install --upgrade accelerate transformers diffusers huggingface_hub #Optional, run it if you run into issues
+```
+#### Modal:
+- Run `pip install modal` to install the modal Python package.
+- Run `modal setup` to authenticate (if this doesn’t work, try `python -m modal setup`).
+
+#### Hugging Face:
+- Get a READ token from [here](https://huggingface.co/settings/tokens) and request access to Flux.1-dev model from [here](https://huggingface.co/black-forest-labs/FLUX.1-dev).
+- Run `huggingface-cli login` and paste your token.
+
+### 2. Upload your dataset
+- Drag and drop your dataset folder containing the .jpg, .jpeg, or .png images and .txt files in `ai-toolkit`.
+
+### 3. Configs
+- Copy an example config file located at ```config/examples/modal``` to the `config` folder and rename it to ```whatever_you_want.yml```.
+- Edit the config following the comments in the file, **<ins>be careful and follow the example `/root/ai-toolkit` paths</ins>**.
+
+### 4. Edit run_modal.py
+- Set your entire local `ai-toolkit` path at `code_mount = modal.Mount.from_local_dir` like:
+  
+   ```
+   code_mount = modal.Mount.from_local_dir("/Users/username/ai-toolkit", remote_path="/root/ai-toolkit")
+   ```
+- Choose a `GPU` and `Timeout` in `@app.function` _(default is A100 40GB and 2 hour timeout)_.
+
+### 5. Training
+- Run the config file in your terminal: `modal run run_modal.py --config-file-list-str=/root/ai-toolkit/config/whatever_you_want.yml`.
+- You can monitor your training in your local terminal, or on [modal.com](https://modal.com/).
+- Models, samples and optimizer will be stored in `Storage > flux-lora-models`.
+
+### 6. Saving the model
+- Check contents of the volume by running `modal volume ls flux-lora-models`. 
+- Download the content by running `modal volume get flux-lora-models your-model-name`.
+- Example: `modal volume get flux-lora-models my_first_flux_lora_v1`.
+
+### Screenshot from Modal
+
+<img width="1728" alt="Modal Traning Screenshot" src="https://github.com/user-attachments/assets/7497eb38-0090-49d6-8ad9-9c8ea7b5388b">
+
+---
+
+## Dataset Preparation
+
+Datasets generally need to be a folder containing images and associated text files. Currently, the only supported
+formats are jpg, jpeg, and png. Webp currently has issues. The text files should be named the same as the images
+but with a `.txt` extension. For example `image2.jpg` and `image2.txt`. The text file should contain only the caption.
+You can add the word `[trigger]` in the caption file and if you have `trigger_word` in your config, it will be automatically
+replaced. 
+
+Images are never upscaled but they are downscaled and placed in buckets for batching. **You do not need to crop/resize your images**.
+The loader will automatically resize them and can handle varying aspect ratios. 
+
+
+## Training Specific Layers
+
+To train specific layers with LoRA, you can use the `only_if_contains` network kwargs. For instance, if you want to train only the 2 layers
+used by The Last Ben, [mentioned in this post](https://x.com/__TheBen/status/1829554120270987740), you can adjust your
+network kwargs like so:
+
+```yaml
+      network:
+        type: "lora"
+        linear: 128
+        linear_alpha: 128
+        network_kwargs:
+          only_if_contains:
+            - "transformer.single_transformer_blocks.7.proj_out"
+            - "transformer.single_transformer_blocks.20.proj_out"
+```
+
+The naming conventions of the layers are in diffusers format, so checking the state dict of a model will reveal 
+the suffix of the name of the layers you want to train. You can also use this method to only train specific groups of weights.
+For instance to only train the `single_transformer` for FLUX.1, you can use the following:
+
+```yaml
+      network:
+        type: "lora"
+        linear: 128
+        linear_alpha: 128
+        network_kwargs:
+          only_if_contains:
+            - "transformer.single_transformer_blocks."
+```
+
+You can also exclude layers by their names by using `ignore_if_contains` network kwarg. So to exclude all the single transformer blocks,
+
+
+```yaml
+      network:
+        type: "lora"
+        linear: 128
+        linear_alpha: 128
+        network_kwargs:
+          ignore_if_contains:
+            - "transformer.single_transformer_blocks."
+```
+
+`ignore_if_contains` takes priority over `only_if_contains`. So if a weight is covered by both,
+if will be ignored.
+
+---
+
+## EVERYTHING BELOW THIS LINE IS OUTDATED 
+
+It may still work like that, but I have not tested it in a while.
+
+---
+
+### Batch Image Generation
+
+A image generator that can take frompts from a config file or form a txt file and generate them to a 
+folder. I mainly needed this for an SDXL test I am doing but added some polish to it so it can be used
+for generat batch image generation.
+It all runs off a config file, which you can find an example of in  `config/examples/generate.example.yaml`.
+Mere info is in the comments in the example
+
+---
+
+### LoRA (lierla), LoCON (LyCORIS) extractor
+
+It is based on the extractor in the [LyCORIS](https://github.com/KohakuBlueleaf/LyCORIS) tool, but adding some QOL features
+and LoRA (lierla) support. It can do multiple types of extractions in one run.
+It all runs off a config file, which you can find an example of in  `config/examples/extract.example.yml`.
+Just copy that file, into the `config` folder, and rename it to `whatever_you_want.yml`.
+Then you can edit the file to your liking. and call it like so:
+
+```bash
+python3 run.py config/whatever_you_want.yml
+```
+
+You can also put a full path to a config file, if you want to keep it somewhere else.
+
+```bash
+python3 run.py "/home/user/whatever_you_want.yml"
+```
+
+More notes on how it works are available in the example config file itself. LoRA and LoCON both support
+extractions of 'fixed', 'threshold', 'ratio', 'quantile'. I'll update what these do and mean later.
+Most people used fixed, which is traditional fixed dimension extraction.
+
+`process` is an array of different processes to run. You can add a few and mix and match. One LoRA, one LyCON, etc.
+
+---
+
+### LoRA Rescale
+
+Change `<lora:my_lora:4.6>` to `<lora:my_lora:1.0>` or whatever you want with the same effect. 
+A tool for rescaling a LoRA's weights. Should would with LoCON as well, but I have not tested it.
+It all runs off a config file, which you can find an example of in  `config/examples/mod_lora_scale.yml`.
+Just copy that file, into the `config` folder, and rename it to `whatever_you_want.yml`.
+Then you can edit the file to your liking. and call it like so:
+
+```bash
+python3 run.py config/whatever_you_want.yml
+```
+
+You can also put a full path to a config file, if you want to keep it somewhere else.
+
+```bash
+python3 run.py "/home/user/whatever_you_want.yml"
+```
+
+More notes on how it works are available in the example config file itself. This is useful when making 
+all LoRAs, as the ideal weight is rarely 1.0, but now you can fix that. For sliders, they can have weird scales form -2 to 2
+or even -15 to 15. This will allow you to dile it in so they all have your desired scale
+
+---
+
+### LoRA Slider Trainer
+
+<a target="_blank" href="https://colab.research.google.com/github/ostris/ai-toolkit/blob/main/notebooks/SliderTraining.ipynb">
+  <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
+</a>
+
+This is how I train most of the recent sliders I have on Civitai, you can check them out in my [Civitai profile](https://civitai.com/user/Ostris/models).
+It is based off the work by [p1atdev/LECO](https://github.com/p1atdev/LECO) and [rohitgandikota/erasing](https://github.com/rohitgandikota/erasing)
+But has been heavily modified to create sliders rather than erasing concepts. I have a lot more plans on this, but it is
+very functional as is. It is also very easy to use. Just copy the example config file in `config/examples/train_slider.example.yml`
+to the `config` folder and rename it to `whatever_you_want.yml`. Then you can edit the file to your liking. and call it like so:
+
+```bash
+python3 run.py config/whatever_you_want.yml
+```
+
+There is a lot more information in that example file. You can even run the example as is without any modifications to see
+how it works. It will create a slider that turns all animals into dogs(neg) or cats(pos). Just run it like so:
+
+```bash
+python3 run.py config/examples/train_slider.example.yml
+```
+
+And you will be able to see how it works without configuring anything. No datasets are required for this method.
+I will post an better tutorial soon. 
+
+---
+
+## Extensions!!
+
+You can now make and share custom extensions. That run within this framework and have all the inbuilt tools
+available to them. I will probably use this as the primary development method going
+forward so I dont keep adding and adding more and more features to this base repo. I will likely migrate a lot
+of the existing functionality as well to make everything modular. There is an example extension in the `extensions`
+folder that shows how to make a model merger extension. All of the code is heavily documented which is hopefully
+enough to get you started. To make an extension, just copy that example and replace all the things you need to.
+
+
+### Model Merger - Example Extension
+It is located in the `extensions` folder. It is a fully finctional model merger that can merge as many models together
+as you want. It is a good example of how to make an extension, but is also a pretty useful feature as well since most
+mergers can only do one model at a time and this one will take as many as you want to feed it. There is an 
+example config file in there, just copy that to your `config` folder and rename it to `whatever_you_want.yml`.
+and use it like any other config file.
+
+## WIP Tools
+
+
+### VAE (Variational Auto Encoder) Trainer
+
+This works, but is not ready for others to use and therefore does not have an example config. 
+I am still working on it. I will update this when it is ready.
+I am adding a lot of features for criteria that I have used in my image enlargement work. A Critic (discriminator),
+content loss, style loss, and a few more. If you don't know, the VAE
+for stable diffusion (yes even the MSE one, and SDXL), are horrible at smaller faces and it holds SD back. I will fix this.
+I'll post more about this later with better examples later, but here is a quick test of a run through with various VAEs.
+Just went in and out. It is much worse on smaller faces than shown here.
+
+<img src="https://raw.githubusercontent.com/ostris/ai-toolkit/main/assets/VAE_test1.jpg" width="768" height="auto"> 
+
+---
+
+## TODO
+- [X] Add proper regs on sliders
+- [X] Add SDXL support (base model only for now)
+- [ ] Add plain erasing
+- [ ] Make Textual inversion network trainer (network that spits out TI embeddings)
+
+---
+
+## Change Log
+
+#### 2023-08-05
+ - Huge memory rework and slider rework. Slider training is better thant ever with no more
+ram spikes. I also made it so all 4 parts of the slider algorythm run in one batch so they share gradient
+accumulation. This makes it much faster and more stable. 
+ - Updated the example config to be something more practical and more updated to current methods. It is now
+a detail slide and shows how to train one without a subject. 512x512 slider training for 1.5 should work on 
+6GB gpu now. Will test soon to verify. 
+
+
+#### 2021-10-20
+ - Windows support bug fixes
+ - Extensions! Added functionality to make and share custom extensions for training, merging, whatever.
+check out the example in the `extensions` folder. Read more about that above.
+ - Model Merging, provided via the example extension.
+
+#### 2023-08-03
+Another big refactor to make SD more modular.
+
+Made batch image generation script
+
+#### 2023-08-01
+Major changes and update. New LoRA rescale tool, look above for details. Added better metadata so
+Automatic1111 knows what the base model is. Added some experiments and a ton of updates. This thing is still unstable
+at the moment, so hopefully there are not breaking changes. 
+
+Unfortunately, I am too lazy to write a proper changelog with all the changes.
+
+I added SDXL training to sliders... but.. it does not work properly. 
+The slider training relies on a model's ability to understand that an unconditional (negative prompt)
+means you do not want that concept in the output. SDXL does not understand this for whatever reason, 
+which makes separating out
+concepts within the model hard. I am sure the community will find a way to fix this 
+over time, but for now, it is not 
+going to work properly. And if any of you are thinking "Could we maybe fix it by adding 1 or 2 more text
+encoders to the model as well as a few more entirely separate diffusion networks?" No. God no. It just needs a little
+training without every experimental new paper added to it. The KISS principal. 
+
+
+#### 2023-07-30
+Added "anchors" to the slider trainer. This allows you to set a prompt that will be used as a 
+regularizer. You can set the network multiplier to force spread consistency at high weights
+

app.py

@@ -0,0 +1,6 @@
+import os
+import sys
+from hf_ui import demo
+
+if __name__ == "__main__":
+    demo.launch() 

build_and_push_docker.yaml

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+
+echo "Docker builds from the repo, not this dir. Make sure changes are pushed to the repo."
+# wait 2 seconds
+sleep 2
+docker build --build-arg CACHEBUST=$(date +%s) -t aitoolkit:latest -f docker/Dockerfile .
+docker tag aitoolkit:latest ostris/aitoolkit:latest
+docker push ostris/aitoolkit:latest

config/examples/extract.example.yml

@@ -0,0 +1,75 @@
+---
+# this is in yaml format. You can use json if you prefer
+# I like both but yaml is easier to read and write
+# plus it has comments which is nice for documentation
+job: extract # tells the runner what to do
+config:
+  # the name will be used to create a folder in the output folder
+  # it will also replace any [name] token in the rest of this config
+  name: name_of_your_model
+  # can be hugging face model, a .ckpt, or a .safetensors
+  base_model: "/path/to/base/model.safetensors"
+  # can be hugging face model, a .ckpt, or a .safetensors
+  extract_model: "/path/to/model/to/extract/trained.safetensors"
+  # we will create folder here with name above so. This will create /path/to/output/folder/name_of_your_model
+  output_folder: "/path/to/output/folder"
+  is_v2: false
+  dtype: fp16 # saved dtype
+  device: cpu # cpu, cuda:0, etc
+
+  # processes can be chained like this to run multiple in a row
+  # they must all use same models above, but great for testing different
+  # sizes and typed of extractions. It is much faster as we already have the models loaded
+  process:
+  # process 1
+  - type: locon  # locon or lora (locon is lycoris)
+    filename: "[name]_64_32.safetensors" # will be put in output folder
+    dtype: fp16
+    mode: fixed
+    linear: 64
+    conv: 32
+
+  # process 2
+  - type: locon
+    output_path: "/absolute/path/for/this/output.safetensors" # can be absolute
+    mode: ratio
+    linear: 0.2
+    conv: 0.2
+
+  # process 3
+  - type: locon
+    filename: "[name]_ratio_02.safetensors"
+    mode: quantile
+    linear: 0.5
+    conv: 0.5
+
+  # process 4
+  - type: lora  # traditional lora extraction (lierla) with linear layers only
+    filename: "[name]_4.safetensors"
+    mode: fixed  # fixed, ratio, quantile supported for lora as well
+    linear: 4 # lora dim or rank
+    # no conv for lora
+
+  # process 5
+  - type: lora
+    filename: "[name]_q05.safetensors"
+    mode: quantile
+    linear: 0.5
+
+# you can put any information you want here, and it will be saved in the model
+# the below is an example. I recommend doing trigger words at a minimum
+# in the metadata. The software will include this plus some other information
+meta:
+  name: "[name]"  # [name] gets replaced with the name above
+  description: A short description of your model
+  trigger_words:
+  - put
+  - trigger
+  - words
+  - here
+  version: '0.1'
+  creator:
+    name: Your Name
+    email: [email protected]
+    website: https://yourwebsite.com
+  any: All meta data above is arbitrary, it can be whatever you want.

config/examples/generate.example.yaml

@@ -0,0 +1,60 @@
+---
+
+job: generate # tells the runner what to do
+config:
+  name: "generate" # this is not really used anywhere currently but required by runner
+  process:
+    # process 1
+    - type: to_folder  # process images to a folder
+      output_folder: "output/gen"
+      device: cuda:0 # cpu, cuda:0, etc
+      generate:
+        # these are your defaults you can override most of them with flags
+        sampler: "ddpm" # ignored for now, will add later though ddpm is used regardless for now
+        width: 1024
+        height: 1024
+        neg: "cartoon, fake, drawing, illustration, cgi, animated, anime"
+        seed: -1 # -1 is random
+        guidance_scale: 7
+        sample_steps: 20
+        ext: ".png" # .png, .jpg, .jpeg, .webp
+
+        # here ate the flags you can use for prompts. Always start with
+        # your prompt first then add these flags after. You can use as many
+        # like
+        # photo of a baseball --n painting, ugly --w 1024 --h 1024 --seed 42 --cfg 7 --steps 20
+        # we will try to support all sd-scripts flags where we can
+
+        # FROM SD-SCRIPTS
+        # --n Treat everything until the next option as a negative prompt.
+        # --w Specify the width of the generated image.
+        # --h Specify the height of the generated image.
+        # --d Specify the seed for the generated image.
+        # --l Specify the CFG scale for the generated image.
+        # --s Specify the number of steps during generation.
+
+        # OURS and some QOL additions
+        # --p2 Prompt for the second text encoder (SDXL only)
+        # --n2 Negative prompt for the second text encoder (SDXL only)
+        # --gr Specify the guidance rescale for the generated image (SDXL only)
+        # --seed Specify the seed for the generated image same as --d
+        # --cfg Specify the CFG scale for the generated image same as --l
+        # --steps Specify the number of steps during generation same as --s
+
+        prompt_file: false # if true a txt file will be created next to images with prompt strings used
+        # prompts can also be a path to a text file with one prompt per line
+        # prompts: "/path/to/prompts.txt"
+        prompts:
+          - "photo of batman"
+          - "photo of superman"
+          - "photo of spiderman"
+          - "photo of a superhero --n batman superman spiderman"
+
+      model:
+        # huggingface name, relative prom project path, or absolute path to .safetensors or .ckpt
+        #      name_or_path: "runwayml/stable-diffusion-v1-5"
+        name_or_path: "/mnt/Models/stable-diffusion/models/stable-diffusion/Ostris/Ostris_Real_v1.safetensors"
+        is_v2: false  # for v2 models
+        is_v_pred: false # for v-prediction models (most v2 models)
+        is_xl: false  # for SDXL models
+        dtype: bf16

config/examples/mod_lora_scale.yaml

@@ -0,0 +1,48 @@
+---
+job: mod
+config:
+  name: name_of_your_model_v1
+  process:
+    - type: rescale_lora
+      # path to your current lora model
+      input_path: "/path/to/lora/lora.safetensors"
+      # output path for your new lora model, can be the same as input_path to replace
+      output_path: "/path/to/lora/output_lora_v1.safetensors"
+      # replaces meta with the meta below (plus minimum meta fields)
+      # if false, we will leave the meta alone except for updating hashes (sd-script hashes)
+      replace_meta: true
+      # how to adjust, we can scale the up_down weights or the alpha
+      # up_down is the default and probably the best, they will both net the same outputs
+      # would only affect rare NaN cases and maybe merging with old merge tools
+      scale_target: 'up_down'
+      # precision to save, fp16 is the default and standard
+      save_dtype: fp16
+      # current_weight is the ideal weight you use as a multiplier when using the lora
+      # IE in automatic1111 <lora:my_lora:6.0> the 6.0 is the current_weight
+      # you can do negatives here too if you want to flip the lora
+      current_weight: 6.0
+      # target_weight is the ideal weight you use as a multiplier when using the lora
+      # instead of the one above. IE in automatic1111 instead of using <lora:my_lora:6.0>
+      # we want to use <lora:my_lora:1.0> so 1.0 is the target_weight
+      target_weight: 1.0
+
+      # base model for the lora
+      # this is just used to add meta so automatic111 knows which model it is for
+      # assume v1.5 if these are not set
+      is_xl: false
+      is_v2: false
+meta:
+  # this is only used if you set replace_meta to true above
+  name: "[name]"  # [name] gets replaced with the name above
+  description: A short description of your lora
+  trigger_words:
+  - put
+  - trigger
+  - words
+  - here
+  version: '0.1'
+  creator:
+    name: Your Name
+    email: [email protected]
+    website: https://yourwebsite.com
+  any: All meta data above is arbitrary, it can be whatever you want.

config/examples/modal/modal_train_lora_flux_24gb.yaml

@@ -0,0 +1,96 @@
+---
+job: extension
+config:
+  # this name will be the folder and filename name
+  name: "my_first_flux_lora_v1"
+  process:
+    - type: 'sd_trainer'
+      # root folder to save training sessions/samples/weights
+      training_folder: "/root/ai-toolkit/modal_output" # must match MOUNT_DIR from run_modal.py
+      # uncomment to see performance stats in the terminal every N steps
+#      performance_log_every: 1000
+      device: cuda:0
+      # if a trigger word is specified, it will be added to captions of training data if it does not already exist
+      # alternatively, in your captions you can add [trigger] and it will be replaced with the trigger word
+#      trigger_word: "p3r5on"
+      network:
+        type: "lora"
+        linear: 16
+        linear_alpha: 16
+      save:
+        dtype: float16 # precision to save
+        save_every: 250 # save every this many steps
+        max_step_saves_to_keep: 4 # how many intermittent saves to keep
+      datasets:
+        # datasets are a folder of images. captions need to be txt files with the same name as the image
+        # for instance image2.jpg and image2.txt. Only jpg, jpeg, and png are supported currently
+        # images will automatically be resized and bucketed into the resolution specified
+        # on windows, escape back slashes with another backslash so
+        # "C:\\path\\to\\images\\folder"
+        # your dataset must be placed in /ai-toolkit and /root is for modal to find the dir:
+        - folder_path: "/root/ai-toolkit/your-dataset"
+          caption_ext: "txt"
+          caption_dropout_rate: 0.05  # will drop out the caption 5% of time
+          shuffle_tokens: false  # shuffle caption order, split by commas
+          cache_latents_to_disk: true  # leave this true unless you know what you're doing
+          resolution: [ 512, 768, 1024 ]  # flux enjoys multiple resolutions
+      train:
+        batch_size: 1
+        steps: 2000  # total number of steps to train 500 - 4000 is a good range
+        gradient_accumulation_steps: 1
+        train_unet: true
+        train_text_encoder: false  # probably won't work with flux
+        gradient_checkpointing: true  # need the on unless you have a ton of vram
+        noise_scheduler: "flowmatch" # for training only
+        optimizer: "adamw8bit"
+        lr: 1e-4
+        # uncomment this to skip the pre training sample
+#        skip_first_sample: true
+        # uncomment to completely disable sampling
+#        disable_sampling: true
+        # uncomment to use new vell curved weighting. Experimental but may produce better results
+#        linear_timesteps: true
+
+        # ema will smooth out learning, but could slow it down. Recommended to leave on.
+        ema_config:
+          use_ema: true
+          ema_decay: 0.99
+
+        # will probably need this if gpu supports it for flux, other dtypes may not work correctly
+        dtype: bf16
+      model:
+        # huggingface model name or path
+        # if you get an error, or get stuck while downloading,
+        # check https://github.com/ostris/ai-toolkit/issues/84, download the model locally and
+        # place it like "/root/ai-toolkit/FLUX.1-dev"
+        name_or_path: "black-forest-labs/FLUX.1-dev"
+        is_flux: true
+        quantize: true  # run 8bit mixed precision
+#        low_vram: true  # uncomment this if the GPU is connected to your monitors. It will use less vram to quantize, but is slower.
+      sample:
+        sampler: "flowmatch" # must match train.noise_scheduler
+        sample_every: 250 # sample every this many steps
+        width: 1024
+        height: 1024
+        prompts:
+          # you can add [trigger] to the prompts here and it will be replaced with the trigger word
+#          - "[trigger] holding a sign that says 'I LOVE PROMPTS!'"\
+          - "woman with red hair, playing chess at the park, bomb going off in the background"
+          - "a woman holding a coffee cup, in a beanie, sitting at a cafe"
+          - "a horse is a DJ at a night club, fish eye lens, smoke machine, lazer lights, holding a martini"
+          - "a man showing off his cool new t shirt at the beach, a shark is jumping out of the water in the background"
+          - "a bear building a log cabin in the snow covered mountains"
+          - "woman playing the guitar, on stage, singing a song, laser lights, punk rocker"
+          - "hipster man with a beard, building a chair, in a wood shop"
+          - "photo of a man, white background, medium shot, modeling clothing, studio lighting, white backdrop"
+          - "a man holding a sign that says, 'this is a sign'"
+          - "a bulldog, in a post apocalyptic world, with a shotgun, in a leather jacket, in a desert, with a motorcycle"
+        neg: ""  # not used on flux
+        seed: 42
+        walk_seed: true
+        guidance_scale: 4
+        sample_steps: 20
+# you can add any additional meta info here. [name] is replaced with config name at top
+meta:
+  name: "[name]"
+  version: '1.0'

config/examples/modal/modal_train_lora_flux_schnell_24gb.yaml

@@ -0,0 +1,98 @@
+---
+job: extension
+config:
+  # this name will be the folder and filename name
+  name: "my_first_flux_lora_v1"
+  process:
+    - type: 'sd_trainer'
+      # root folder to save training sessions/samples/weights
+      training_folder: "/root/ai-toolkit/modal_output" # must match MOUNT_DIR from run_modal.py
+      # uncomment to see performance stats in the terminal every N steps
+#      performance_log_every: 1000
+      device: cuda:0
+      # if a trigger word is specified, it will be added to captions of training data if it does not already exist
+      # alternatively, in your captions you can add [trigger] and it will be replaced with the trigger word
+#      trigger_word: "p3r5on"
+      network:
+        type: "lora"
+        linear: 16
+        linear_alpha: 16
+      save:
+        dtype: float16 # precision to save
+        save_every: 250 # save every this many steps
+        max_step_saves_to_keep: 4 # how many intermittent saves to keep
+      datasets:
+        # datasets are a folder of images. captions need to be txt files with the same name as the image
+        # for instance image2.jpg and image2.txt. Only jpg, jpeg, and png are supported currently
+        # images will automatically be resized and bucketed into the resolution specified
+        # on windows, escape back slashes with another backslash so
+        # "C:\\path\\to\\images\\folder"
+        # your dataset must be placed in /ai-toolkit and /root is for modal to find the dir:
+        - folder_path: "/root/ai-toolkit/your-dataset"
+          caption_ext: "txt"
+          caption_dropout_rate: 0.05  # will drop out the caption 5% of time
+          shuffle_tokens: false  # shuffle caption order, split by commas
+          cache_latents_to_disk: true  # leave this true unless you know what you're doing
+          resolution: [ 512, 768, 1024 ]  # flux enjoys multiple resolutions
+      train:
+        batch_size: 1
+        steps: 2000  # total number of steps to train 500 - 4000 is a good range
+        gradient_accumulation_steps: 1
+        train_unet: true
+        train_text_encoder: false  # probably won't work with flux
+        gradient_checkpointing: true  # need the on unless you have a ton of vram
+        noise_scheduler: "flowmatch" # for training only
+        optimizer: "adamw8bit"
+        lr: 1e-4
+        # uncomment this to skip the pre training sample
+#        skip_first_sample: true
+        # uncomment to completely disable sampling
+#        disable_sampling: true
+        # uncomment to use new vell curved weighting. Experimental but may produce better results
+#        linear_timesteps: true
+
+        # ema will smooth out learning, but could slow it down. Recommended to leave on.
+        ema_config:
+          use_ema: true
+          ema_decay: 0.99
+
+        # will probably need this if gpu supports it for flux, other dtypes may not work correctly
+        dtype: bf16
+      model:
+        # huggingface model name or path
+        # if you get an error, or get stuck while downloading,
+        # check https://github.com/ostris/ai-toolkit/issues/84, download the models locally and
+        # place them like "/root/ai-toolkit/FLUX.1-schnell" and "/root/ai-toolkit/FLUX.1-schnell-training-adapter"
+        name_or_path: "black-forest-labs/FLUX.1-schnell"
+        assistant_lora_path: "ostris/FLUX.1-schnell-training-adapter" # Required for flux schnell training
+        is_flux: true
+        quantize: true  # run 8bit mixed precision
+        # low_vram is painfully slow to fuse in the adapter avoid it unless absolutely necessary
+#        low_vram: true  # uncomment this if the GPU is connected to your monitors. It will use less vram to quantize, but is slower.
+      sample:
+        sampler: "flowmatch" # must match train.noise_scheduler
+        sample_every: 250 # sample every this many steps
+        width: 1024
+        height: 1024
+        prompts:
+          # you can add [trigger] to the prompts here and it will be replaced with the trigger word
+#          - "[trigger] holding a sign that says 'I LOVE PROMPTS!'"\
+          - "woman with red hair, playing chess at the park, bomb going off in the background"
+          - "a woman holding a coffee cup, in a beanie, sitting at a cafe"
+          - "a horse is a DJ at a night club, fish eye lens, smoke machine, lazer lights, holding a martini"
+          - "a man showing off his cool new t shirt at the beach, a shark is jumping out of the water in the background"
+          - "a bear building a log cabin in the snow covered mountains"
+          - "woman playing the guitar, on stage, singing a song, laser lights, punk rocker"
+          - "hipster man with a beard, building a chair, in a wood shop"
+          - "photo of a man, white background, medium shot, modeling clothing, studio lighting, white backdrop"
+          - "a man holding a sign that says, 'this is a sign'"
+          - "a bulldog, in a post apocalyptic world, with a shotgun, in a leather jacket, in a desert, with a motorcycle"
+        neg: ""  # not used on flux
+        seed: 42
+        walk_seed: true
+        guidance_scale: 1  # schnell does not do guidance
+        sample_steps: 4  # 1 - 4 works well
+# you can add any additional meta info here. [name] is replaced with config name at top
+meta:
+  name: "[name]"
+  version: '1.0'

config/examples/train_lora_flux_24gb.yaml

@@ -0,0 +1,96 @@
+---
+job: extension
+config:
+  # this name will be the folder and filename name
+  name: "my_first_flux_lora_v1"
+  process:
+    - type: 'sd_trainer'
+      # root folder to save training sessions/samples/weights
+      training_folder: "output"
+      # uncomment to see performance stats in the terminal every N steps
+#      performance_log_every: 1000
+      device: cuda:0
+      # if a trigger word is specified, it will be added to captions of training data if it does not already exist
+      # alternatively, in your captions you can add [trigger] and it will be replaced with the trigger word
+#      trigger_word: "p3r5on"
+      network:
+        type: "lora"
+        linear: 16
+        linear_alpha: 16
+      save:
+        dtype: float16 # precision to save
+        save_every: 250 # save every this many steps
+        max_step_saves_to_keep: 4 # how many intermittent saves to keep
+        push_to_hub: false #change this to True to push your trained model to Hugging Face.
+        # You can either set up a HF_TOKEN env variable or you'll be prompted to log-in         
+#       hf_repo_id: your-username/your-model-slug
+#       hf_private: true #whether the repo is private or public
+      datasets:
+        # datasets are a folder of images. captions need to be txt files with the same name as the image
+        # for instance image2.jpg and image2.txt. Only jpg, jpeg, and png are supported currently
+        # images will automatically be resized and bucketed into the resolution specified
+        # on windows, escape back slashes with another backslash so
+        # "C:\\path\\to\\images\\folder"
+        - folder_path: "/path/to/images/folder"
+          caption_ext: "txt"
+          caption_dropout_rate: 0.05  # will drop out the caption 5% of time
+          shuffle_tokens: false  # shuffle caption order, split by commas
+          cache_latents_to_disk: true  # leave this true unless you know what you're doing
+          resolution: [ 512, 768, 1024 ]  # flux enjoys multiple resolutions
+      train:
+        batch_size: 1
+        steps: 2000  # total number of steps to train 500 - 4000 is a good range
+        gradient_accumulation_steps: 1
+        train_unet: true
+        train_text_encoder: false  # probably won't work with flux
+        gradient_checkpointing: true  # need the on unless you have a ton of vram
+        noise_scheduler: "flowmatch" # for training only
+        optimizer: "adamw8bit"
+        lr: 1e-4
+        # uncomment this to skip the pre training sample
+#        skip_first_sample: true
+        # uncomment to completely disable sampling
+#        disable_sampling: true
+        # uncomment to use new vell curved weighting. Experimental but may produce better results
+#        linear_timesteps: true
+
+        # ema will smooth out learning, but could slow it down. Recommended to leave on.
+        ema_config:
+          use_ema: true
+          ema_decay: 0.99
+
+        # will probably need this if gpu supports it for flux, other dtypes may not work correctly
+        dtype: bf16
+      model:
+        # huggingface model name or path
+        name_or_path: "black-forest-labs/FLUX.1-dev"
+        is_flux: true
+        quantize: true  # run 8bit mixed precision
+#        low_vram: true  # uncomment this if the GPU is connected to your monitors. It will use less vram to quantize, but is slower.
+      sample:
+        sampler: "flowmatch" # must match train.noise_scheduler
+        sample_every: 250 # sample every this many steps
+        width: 1024
+        height: 1024
+        prompts:
+          # you can add [trigger] to the prompts here and it will be replaced with the trigger word
+#          - "[trigger] holding a sign that says 'I LOVE PROMPTS!'"\
+          - "woman with red hair, playing chess at the park, bomb going off in the background"
+          - "a woman holding a coffee cup, in a beanie, sitting at a cafe"
+          - "a horse is a DJ at a night club, fish eye lens, smoke machine, lazer lights, holding a martini"
+          - "a man showing off his cool new t shirt at the beach, a shark is jumping out of the water in the background"
+          - "a bear building a log cabin in the snow covered mountains"
+          - "woman playing the guitar, on stage, singing a song, laser lights, punk rocker"
+          - "hipster man with a beard, building a chair, in a wood shop"
+          - "photo of a man, white background, medium shot, modeling clothing, studio lighting, white backdrop"
+          - "a man holding a sign that says, 'this is a sign'"
+          - "a bulldog, in a post apocalyptic world, with a shotgun, in a leather jacket, in a desert, with a motorcycle"
+        neg: ""  # not used on flux
+        seed: 42
+        walk_seed: true
+        guidance_scale: 4
+        sample_steps: 20
+# you can add any additional meta info here. [name] is replaced with config name at top
+meta:
+  name: "[name]"
+  version: '1.0'

config/examples/train_lora_flux_schnell_24gb.yaml

@@ -0,0 +1,98 @@
+---
+job: extension
+config:
+  # this name will be the folder and filename name
+  name: "my_first_flux_lora_v1"
+  process:
+    - type: 'sd_trainer'
+      # root folder to save training sessions/samples/weights
+      training_folder: "output"
+      # uncomment to see performance stats in the terminal every N steps
+#      performance_log_every: 1000
+      device: cuda:0
+      # if a trigger word is specified, it will be added to captions of training data if it does not already exist
+      # alternatively, in your captions you can add [trigger] and it will be replaced with the trigger word
+#      trigger_word: "p3r5on"
+      network:
+        type: "lora"
+        linear: 16
+        linear_alpha: 16
+      save:
+        dtype: float16 # precision to save
+        save_every: 250 # save every this many steps
+        max_step_saves_to_keep: 4 # how many intermittent saves to keep
+        push_to_hub: false #change this to True to push your trained model to Hugging Face.
+        # You can either set up a HF_TOKEN env variable or you'll be prompted to log-in         
+#       hf_repo_id: your-username/your-model-slug
+#       hf_private: true #whether the repo is private or public        
+      datasets:
+        # datasets are a folder of images. captions need to be txt files with the same name as the image
+        # for instance image2.jpg and image2.txt. Only jpg, jpeg, and png are supported currently
+        # images will automatically be resized and bucketed into the resolution specified
+        # on windows, escape back slashes with another backslash so
+        # "C:\\path\\to\\images\\folder"
+        - folder_path: "/path/to/images/folder"
+          caption_ext: "txt"
+          caption_dropout_rate: 0.05  # will drop out the caption 5% of time
+          shuffle_tokens: false  # shuffle caption order, split by commas
+          cache_latents_to_disk: true  # leave this true unless you know what you're doing
+          resolution: [ 512, 768, 1024 ]  # flux enjoys multiple resolutions
+      train:
+        batch_size: 1
+        steps: 2000  # total number of steps to train 500 - 4000 is a good range
+        gradient_accumulation_steps: 1
+        train_unet: true
+        train_text_encoder: false  # probably won't work with flux
+        gradient_checkpointing: true  # need the on unless you have a ton of vram
+        noise_scheduler: "flowmatch" # for training only
+        optimizer: "adamw8bit"
+        lr: 1e-4
+        # uncomment this to skip the pre training sample
+#        skip_first_sample: true
+        # uncomment to completely disable sampling
+#        disable_sampling: true
+        # uncomment to use new bell curved weighting. Experimental but may produce better results
+#        linear_timesteps: true
+
+        # ema will smooth out learning, but could slow it down. Recommended to leave on.
+        ema_config:
+          use_ema: true
+          ema_decay: 0.99
+
+        # will probably need this if gpu supports it for flux, other dtypes may not work correctly
+        dtype: bf16
+      model:
+        # huggingface model name or path
+        name_or_path: "black-forest-labs/FLUX.1-schnell"
+        assistant_lora_path: "ostris/FLUX.1-schnell-training-adapter" # Required for flux schnell training
+        is_flux: true
+        quantize: true  # run 8bit mixed precision
+        # low_vram is painfully slow to fuse in the adapter avoid it unless absolutely necessary
+#        low_vram: true  # uncomment this if the GPU is connected to your monitors. It will use less vram to quantize, but is slower.
+      sample:
+        sampler: "flowmatch" # must match train.noise_scheduler
+        sample_every: 250 # sample every this many steps
+        width: 1024
+        height: 1024
+        prompts:
+          # you can add [trigger] to the prompts here and it will be replaced with the trigger word
+#          - "[trigger] holding a sign that says 'I LOVE PROMPTS!'"\
+          - "woman with red hair, playing chess at the park, bomb going off in the background"
+          - "a woman holding a coffee cup, in a beanie, sitting at a cafe"
+          - "a horse is a DJ at a night club, fish eye lens, smoke machine, lazer lights, holding a martini"
+          - "a man showing off his cool new t shirt at the beach, a shark is jumping out of the water in the background"
+          - "a bear building a log cabin in the snow covered mountains"
+          - "woman playing the guitar, on stage, singing a song, laser lights, punk rocker"
+          - "hipster man with a beard, building a chair, in a wood shop"
+          - "photo of a man, white background, medium shot, modeling clothing, studio lighting, white backdrop"
+          - "a man holding a sign that says, 'this is a sign'"
+          - "a bulldog, in a post apocalyptic world, with a shotgun, in a leather jacket, in a desert, with a motorcycle"
+        neg: ""  # not used on flux
+        seed: 42
+        walk_seed: true
+        guidance_scale: 1  # schnell does not do guidance
+        sample_steps: 4  # 1 - 4 works well
+# you can add any additional meta info here. [name] is replaced with config name at top
+meta:
+  name: "[name]"
+  version: '1.0'

config/examples/train_lora_sd35_large_24gb.yaml

@@ -0,0 +1,97 @@
+---
+# NOTE!! THIS IS CURRENTLY EXPERIMENTAL AND UNDER DEVELOPMENT. SOME THINGS WILL CHANGE
+job: extension
+config:
+  # this name will be the folder and filename name
+  name: "my_first_sd3l_lora_v1"
+  process:
+    - type: 'sd_trainer'
+      # root folder to save training sessions/samples/weights
+      training_folder: "output"
+      # uncomment to see performance stats in the terminal every N steps
+#      performance_log_every: 1000
+      device: cuda:0
+      # if a trigger word is specified, it will be added to captions of training data if it does not already exist
+      # alternatively, in your captions you can add [trigger] and it will be replaced with the trigger word
+#      trigger_word: "p3r5on"
+      network:
+        type: "lora"
+        linear: 16
+        linear_alpha: 16
+      save:
+        dtype: float16 # precision to save
+        save_every: 250 # save every this many steps
+        max_step_saves_to_keep: 4 # how many intermittent saves to keep
+        push_to_hub: false #change this to True to push your trained model to Hugging Face.
+        # You can either set up a HF_TOKEN env variable or you'll be prompted to log-in         
+#       hf_repo_id: your-username/your-model-slug
+#       hf_private: true #whether the repo is private or public
+      datasets:
+        # datasets are a folder of images. captions need to be txt files with the same name as the image
+        # for instance image2.jpg and image2.txt. Only jpg, jpeg, and png are supported currently
+        # images will automatically be resized and bucketed into the resolution specified
+        # on windows, escape back slashes with another backslash so
+        # "C:\\path\\to\\images\\folder"
+        - folder_path: "/path/to/images/folder"
+          caption_ext: "txt"
+          caption_dropout_rate: 0.05  # will drop out the caption 5% of time
+          shuffle_tokens: false  # shuffle caption order, split by commas
+          cache_latents_to_disk: true  # leave this true unless you know what you're doing
+          resolution: [ 1024 ] 
+      train:
+        batch_size: 1
+        steps: 2000  # total number of steps to train 500 - 4000 is a good range
+        gradient_accumulation_steps: 1
+        train_unet: true
+        train_text_encoder: false  # May not fully work with SD3 yet
+        gradient_checkpointing: true  # need the on unless you have a ton of vram
+        noise_scheduler: "flowmatch"
+        timestep_type: "linear" # linear or sigmoid
+        optimizer: "adamw8bit"
+        lr: 1e-4
+        # uncomment this to skip the pre training sample
+#        skip_first_sample: true
+        # uncomment to completely disable sampling
+#        disable_sampling: true
+        # uncomment to use new vell curved weighting. Experimental but may produce better results
+#        linear_timesteps: true
+
+        # ema will smooth out learning, but could slow it down. Recommended to leave on.
+        ema_config:
+          use_ema: true
+          ema_decay: 0.99
+
+        # will probably need this if gpu supports it for sd3, other dtypes may not work correctly
+        dtype: bf16
+      model:
+        # huggingface model name or path
+        name_or_path: "stabilityai/stable-diffusion-3.5-large"
+        is_v3: true
+        quantize: true  # run 8bit mixed precision
+      sample:
+        sampler: "flowmatch" # must match train.noise_scheduler
+        sample_every: 250 # sample every this many steps
+        width: 1024
+        height: 1024
+        prompts:
+          # you can add [trigger] to the prompts here and it will be replaced with the trigger word
+#          - "[trigger] holding a sign that says 'I LOVE PROMPTS!'"\
+          - "woman with red hair, playing chess at the park, bomb going off in the background"
+          - "a woman holding a coffee cup, in a beanie, sitting at a cafe"
+          - "a horse is a DJ at a night club, fish eye lens, smoke machine, lazer lights, holding a martini"
+          - "a man showing off his cool new t shirt at the beach, a shark is jumping out of the water in the background"
+          - "a bear building a log cabin in the snow covered mountains"
+          - "woman playing the guitar, on stage, singing a song, laser lights, punk rocker"
+          - "hipster man with a beard, building a chair, in a wood shop"
+          - "photo of a man, white background, medium shot, modeling clothing, studio lighting, white backdrop"
+          - "a man holding a sign that says, 'this is a sign'"
+          - "a bulldog, in a post apocalyptic world, with a shotgun, in a leather jacket, in a desert, with a motorcycle"
+        neg: ""  
+        seed: 42
+        walk_seed: true
+        guidance_scale: 4
+        sample_steps: 25
+# you can add any additional meta info here. [name] is replaced with config name at top
+meta:
+  name: "[name]"
+  version: '1.0'

config/examples/train_slider.example.yml

@@ -0,0 +1,230 @@
+---
+# This is in yaml format. You can use json if you prefer
+# I like both but yaml is easier to write
+# Plus it has comments which is nice for documentation
+# This is the config I use on my sliders, It is solid and tested
+job: train
+config:
+  # the name will be used to create a folder in the output folder
+  # it will also replace any [name] token in the rest of this config
+  name: detail_slider_v1
+  # folder will be created with name above in folder below
+  # it can be relative to the project root or absolute
+  training_folder: "output/LoRA"
+  device: cuda:0 # cpu, cuda:0, etc
+  # for tensorboard logging, we will make a subfolder for this job
+  log_dir: "output/.tensorboard"
+  # you can stack processes for other jobs, It is not tested with sliders though
+  # just use one for now
+  process:
+    - type: slider # tells runner to run the slider process
+      # network is the LoRA network for a slider, I recommend to leave this be
+      network:
+        # network type lierla is traditional LoRA that works everywhere, only linear layers
+        type: "lierla"
+        # rank / dim of the network. Bigger is not always better. Especially for sliders. 8 is good
+        linear: 8
+        linear_alpha: 4 # Do about half of rank
+      # training config
+      train:
+        # this is also used in sampling. Stick with ddpm unless you know what you are doing
+        noise_scheduler: "ddpm" # or "ddpm", "lms", "euler_a"
+        # how many steps to train. More is not always better. I rarely go over 1000
+        steps: 500
+        # I have had good results with 4e-4 to 1e-4 at 500 steps
+        lr: 2e-4
+        # enables gradient checkpoint, saves vram, leave it on
+        gradient_checkpointing: true
+        # train the unet. I recommend leaving this true
+        train_unet: true
+        # train the text encoder. I don't recommend this unless you have a special use case
+        # for sliders we are adjusting representation of the concept (unet),
+        # not the description of it (text encoder)
+        train_text_encoder: false
+        # same as from sd-scripts, not fully tested but should speed up training
+        min_snr_gamma: 5.0
+        # just leave unless you know what you are doing
+        # also supports "dadaptation" but set lr to 1 if you use that,
+        # but it learns too fast and I don't recommend it
+        optimizer: "adamw"
+        # only constant for now
+        lr_scheduler: "constant"
+        # we randomly denoise random num of steps form 1 to this number
+        # while training. Just leave it
+        max_denoising_steps: 40
+        # works great at 1. I do 1 even with my 4090.
+        # higher may not work right with newer single batch stacking code anyway
+        batch_size: 1
+        # bf16 works best if your GPU supports it (modern)
+        dtype: bf16  # fp32, bf16, fp16
+        # if you have it, use it. It is faster and better
+        # torch 2.0 doesnt need xformers anymore, only use if you have lower version
+#        xformers: true
+        # I don't recommend using unless you are trying to make a darker lora. Then do 0.1 MAX
+        # although, the way we train sliders is comparative, so it probably won't work anyway
+        noise_offset: 0.0
+#        noise_offset: 0.0357  # SDXL was trained with offset of 0.0357. So use that when training on SDXL
+
+      # the model to train the LoRA network on
+      model:
+        # huggingface name, relative prom project path, or absolute path to .safetensors or .ckpt
+        name_or_path: "runwayml/stable-diffusion-v1-5"
+        is_v2: false  # for v2 models
+        is_v_pred: false # for v-prediction models (most v2 models)
+        # has some issues with the dual text encoder and the way we train sliders
+        # it works bit weights need to probably be higher to see it.
+        is_xl: false  # for SDXL models
+
+      # saving config
+      save:
+        dtype: float16 # precision to save. I recommend float16
+        save_every: 50 # save every this many steps
+        # this will remove step counts more than this number
+        # allows you to save more often in case of a crash without filling up your drive
+        max_step_saves_to_keep: 2
+
+      # sampling config
+      sample:
+        # must match train.noise_scheduler, this is not used here
+        # but may be in future and in other processes
+        sampler: "ddpm"
+        # sample every this many steps
+        sample_every: 20
+        # image size
+        width: 512
+        height: 512
+        # prompts to use for sampling. Do as many as you want, but it slows down training
+        # pick ones that will best represent the concept you are trying to adjust
+        # allows some flags after the prompt
+        #  --m [number]  # network multiplier. LoRA weight. -3 for the negative slide, 3 for the positive
+        #      slide are good tests. will inherit sample.network_multiplier if not set
+        #  --n [string]  # negative prompt, will inherit sample.neg if not set
+        # Only 75 tokens allowed currently
+        # I like to do a wide positive and negative spread so I can see a good range and stop
+        # early if the network is braking down
+        prompts:
+          - "a woman in a coffee shop, black hat, blonde hair, blue jacket --m -5"
+          - "a woman in a coffee shop, black hat, blonde hair, blue jacket --m -3"
+          - "a woman in a coffee shop, black hat, blonde hair, blue jacket --m 3"
+          - "a woman in a coffee shop, black hat, blonde hair, blue jacket --m 5"
+          - "a golden retriever sitting on a leather couch, --m -5"
+          - "a golden retriever sitting on a leather couch --m -3"
+          - "a golden retriever sitting on a leather couch --m 3"
+          - "a golden retriever sitting on a leather couch --m 5"
+          - "a man with a beard and red flannel shirt, wearing vr goggles, walking into traffic --m -5"
+          - "a man with a beard and red flannel shirt, wearing vr goggles, walking into traffic --m -3"
+          - "a man with a beard and red flannel shirt, wearing vr goggles, walking into traffic --m 3"
+          - "a man with a beard and red flannel shirt, wearing vr goggles, walking into traffic --m 5"
+        # negative prompt used on all prompts above as default if they don't have one
+        neg: "cartoon, fake, drawing, illustration, cgi, animated, anime, monochrome"
+        # seed for sampling. 42 is the answer for everything
+        seed: 42
+        # walks the seed so s1 is 42, s2 is 43, s3 is 44, etc
+        # will start over on next sample_every so s1 is always seed
+        # works well if you use same prompt but want different results
+        walk_seed: false
+        # cfg scale (4 to 10 is good)
+        guidance_scale: 7
+        # sampler steps (20 to 30 is good)
+        sample_steps: 20
+        # default network multiplier for all prompts
+        # since we are training a slider, I recommend overriding this with --m [number]
+        # in the prompts above to get both sides of the slider
+        network_multiplier: 1.0
+
+      # logging information
+      logging:
+        log_every: 10 # log every this many steps
+        use_wandb: false # not supported yet
+        verbose: false # probably done need unless you are debugging
+
+      # slider training config, best for last
+      slider:
+        # resolutions to train on. [ width, height ]. This is less important for sliders
+        # as we are not teaching the model anything it doesn't already know
+        # but must be a size it understands [ 512, 512 ] for sd_v1.5  and [ 768, 768 ] for sd_v2.1
+        # and [ 1024, 1024 ] for sd_xl
+        # you can do as many as you want here
+        resolutions:
+          - [ 512, 512 ]
+#          - [ 512, 768 ]
+#          - [ 768, 768 ]
+        # slider training uses 4 combined steps for a single round. This will do it in one gradient
+        # step. It is highly optimized and shouldn't take anymore vram than doing without it,
+        # since we break down batches for gradient accumulation now. so just leave it on.
+        batch_full_slide: true
+        # These are the concepts to train on. You can do as many as you want here,
+        # but they can conflict outweigh each other. Other than experimenting, I recommend
+        # just doing one for good results
+        targets:
+            # target_class is the base concept we are adjusting the representation of
+            # for example, if we are adjusting the representation of a person, we would use "person"
+            # if we are adjusting the representation of a cat, we would use "cat" It is not
+            # a keyword necessarily but what the model understands the concept to represent.
+            # "person" will affect men, women, children, etc but will not affect cats, dogs, etc
+            # it is the models base general understanding of the concept and everything it represents
+            # you can leave it blank to affect everything. In this example, we are adjusting
+            # detail, so we will leave it blank to affect everything
+          - target_class: ""
+            # positive is the prompt for the positive side of the slider.
+            # It is the concept that will be excited and amplified in the model when we slide the slider
+            # to the positive side and forgotten / inverted when we slide
+            # the slider to the negative side. It is generally best to include the target_class in
+            # the prompt. You want it to be the extreme of what you want to train on. For example,
+            # if you want to train on fat people, you would use "an extremely fat, morbidly obese person"
+            # as the prompt. Not just "fat person"
+            # max 75 tokens for now
+            positive: "high detail, 8k, intricate, detailed, high resolution, high res, high quality"
+            # negative is the prompt for the negative side of the slider and works the same as positive
+            # it does not necessarily work the same as a negative prompt when generating images
+            # these need to be polar opposites.
+            # max 76 tokens for now
+            negative: "blurry, boring, fuzzy, low detail, low resolution, low res, low quality"
+            # the loss for this target is multiplied by this number.
+            # if you are doing more than one target it may be good to set less important ones
+            # to a lower number like 0.1 so they don't outweigh the primary target
+            weight: 1.0
+            # shuffle the prompts split by the comma. We will run every combination randomly
+            # this will make the LoRA more robust. You probably want this on unless prompt order
+            # is important for some reason
+            shuffle: true
+
+
+        # anchors are prompts that we will try to hold on to while training the slider
+        # these are NOT necessary and can prevent the slider from converging if not done right
+        # leave them off if you are having issues, but they can help lock the network
+        # on certain concepts to help prevent catastrophic forgetting
+        # you want these to generate an image that is not your target_class, but close to it
+        # is fine as long as it does not directly overlap it.
+        # For example, if you are training on a person smiling,
+        # you could use "a person with a face mask" as an anchor. It is a person, the image is the same
+        # regardless if they are smiling or not, however, the closer the concept is to the target_class
+        # the less the multiplier needs to be. Keep multipliers less than 1.0 for anchors usually
+        # for close concepts, you want to be closer to 0.1 or 0.2
+        # these will slow down training. I am leaving them off for the demo
+
+#        anchors:
+#          - prompt: "a woman"
+#            neg_prompt: "animal"
+#            # the multiplier applied to the LoRA when this is run.
+#            # higher will give it more weight but also help keep the lora from collapsing
+#            multiplier: 1.0
+#          - prompt: "a man"
+#            neg_prompt: "animal"
+#            multiplier: 1.0
+#          - prompt: "a person"
+#            neg_prompt: "animal"
+#            multiplier: 1.0
+
+# You can put any information you want here, and it will be saved in the model.
+# The below is an example, but you can put your grocery list in it if you want.
+# It is saved in the model so be aware of that. The software will include this
+# plus some other information for you automatically
+meta:
+  # [name] gets replaced with the name above
+  name: "[name]"
+#  version: '1.0'
+#  creator:
+#    name: Your Name
+#    email: [email protected]
+#    website: https://your.website

docker/Dockerfile

@@ -0,0 +1,31 @@
+FROM runpod/base:0.6.2-cuda12.2.0
+
+LABEL authors="jaret"
+
+# Install dependencies
+RUN apt-get update
+
+WORKDIR /app
+ARG CACHEBUST=1
+RUN git clone https://github.com/ostris/ai-toolkit.git && \
+    cd ai-toolkit && \
+    git submodule update --init --recursive
+
+WORKDIR /app/ai-toolkit
+
+RUN ln -s /usr/bin/python3 /usr/bin/python
+RUN python -m pip install -r requirements.txt
+
+RUN apt-get install -y tmux nvtop htop
+
+RUN pip install jupyterlab
+
+# mask workspace
+RUN mkdir /workspace
+
+
+# symlink app to workspace
+RUN ln -s /app/ai-toolkit /workspace/ai-toolkit
+
+WORKDIR /
+CMD ["/start.sh"]

extensions/example/ExampleMergeModels.py

@@ -0,0 +1,129 @@
+import torch
+import gc
+from collections import OrderedDict
+from typing import TYPE_CHECKING
+from jobs.process import BaseExtensionProcess
+from toolkit.config_modules import ModelConfig
+from toolkit.stable_diffusion_model import StableDiffusion
+from toolkit.train_tools import get_torch_dtype
+from tqdm import tqdm
+
+# Type check imports. Prevents circular imports
+if TYPE_CHECKING:
+    from jobs import ExtensionJob
+
+
+# extend standard config classes to add weight
+class ModelInputConfig(ModelConfig):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.weight = kwargs.get('weight', 1.0)
+        # overwrite default dtype unless user specifies otherwise
+        # float 32 will give up better precision on the merging functions
+        self.dtype: str = kwargs.get('dtype', 'float32')
+
+
+def flush():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+# this is our main class process
+class ExampleMergeModels(BaseExtensionProcess):
+    def __init__(
+            self,
+            process_id: int,
+            job: 'ExtensionJob',
+            config: OrderedDict
+    ):
+        super().__init__(process_id, job, config)
+        # this is the setup process, do not do process intensive stuff here, just variable setup and
+        # checking requirements. This is called before the run() function
+        # no loading models or anything like that, it is just for setting up the process
+        # all of your process intensive stuff should be done in the run() function
+        # config will have everything from the process item in the config file
+
+        # convince methods exist on BaseProcess to get config values
+        # if required is set to true and the value is not found it will throw an error
+        # you can pass a default value to get_conf() as well if it was not in the config file
+        # as well as a type to cast the value to
+        self.save_path = self.get_conf('save_path', required=True)
+        self.save_dtype = self.get_conf('save_dtype', default='float16', as_type=get_torch_dtype)
+        self.device = self.get_conf('device', default='cpu', as_type=torch.device)
+
+        # build models to merge list
+        models_to_merge = self.get_conf('models_to_merge', required=True, as_type=list)
+        # build list of ModelInputConfig objects. I find it is a good idea to make a class for each config
+        # this way you can add methods to it and it is easier to read and code. There are a lot of
+        # inbuilt config classes located in toolkit.config_modules as well
+        self.models_to_merge = [ModelInputConfig(**model) for model in models_to_merge]
+        # setup is complete. Don't load anything else here, just setup variables and stuff
+
+    # this is the entire run process be sure to call super().run() first
+    def run(self):
+        # always call first
+        super().run()
+        print(f"Running process: {self.__class__.__name__}")
+
+        # let's adjust our weights first to normalize them so the total is 1.0
+        total_weight = sum([model.weight for model in self.models_to_merge])
+        weight_adjust = 1.0 / total_weight
+        for model in self.models_to_merge:
+            model.weight *= weight_adjust
+
+        output_model: StableDiffusion = None
+        # let's do the merge, it is a good idea to use tqdm to show progress
+        for model_config in tqdm(self.models_to_merge, desc="Merging models"):
+            # setup model class with our helper class
+            sd_model = StableDiffusion(
+                device=self.device,
+                model_config=model_config,
+                dtype="float32"
+            )
+            # load the model
+            sd_model.load_model()
+
+            # adjust the weight of the text encoder
+            if isinstance(sd_model.text_encoder, list):
+                # sdxl model
+                for text_encoder in sd_model.text_encoder:
+                    for key, value in text_encoder.state_dict().items():
+                        value *= model_config.weight
+            else:
+                # normal model
+                for key, value in sd_model.text_encoder.state_dict().items():
+                    value *= model_config.weight
+            # adjust the weights of the unet
+            for key, value in sd_model.unet.state_dict().items():
+                value *= model_config.weight
+
+            if output_model is None:
+                # use this one as the base
+                output_model = sd_model
+            else:
+                # merge the models
+                # text encoder
+                if isinstance(output_model.text_encoder, list):
+                    # sdxl model
+                    for i, text_encoder in enumerate(output_model.text_encoder):
+                        for key, value in text_encoder.state_dict().items():
+                            value += sd_model.text_encoder[i].state_dict()[key]
+                else:
+                    # normal model
+                    for key, value in output_model.text_encoder.state_dict().items():
+                        value += sd_model.text_encoder.state_dict()[key]
+                # unet
+                for key, value in output_model.unet.state_dict().items():
+                    value += sd_model.unet.state_dict()[key]
+
+                # remove the model to free memory
+                del sd_model
+                flush()
+
+        # merge loop is done, let's save the model
+        print(f"Saving merged model to {self.save_path}")
+        output_model.save(self.save_path, meta=self.meta, save_dtype=self.save_dtype)
+        print(f"Saved merged model to {self.save_path}")
+        # do cleanup here
+        del output_model
+        flush()

extensions/example/__init__.py

@@ -0,0 +1,25 @@
+# This is an example extension for custom training. It is great for experimenting with new ideas.
+from toolkit.extension import Extension
+
+
+# We make a subclass of Extension
+class ExampleMergeExtension(Extension):
+    # uid must be unique, it is how the extension is identified
+    uid = "example_merge_extension"
+
+    # name is the name of the extension for printing
+    name = "Example Merge Extension"
+
+    # This is where your process class is loaded
+    # keep your imports in here so they don't slow down the rest of the program
+    @classmethod
+    def get_process(cls):
+        # import your process class here so it is only loaded when needed and return it
+        from .ExampleMergeModels import ExampleMergeModels
+        return ExampleMergeModels
+
+
+AI_TOOLKIT_EXTENSIONS = [
+    # you can put a list of extensions here
+    ExampleMergeExtension
+]

extensions/example/config/config.example.yaml

@@ -0,0 +1,48 @@
+---
+# Always include at least one example config file to show how to use your extension.
+# use plenty of comments so users know how to use it and what everything does
+
+# all extensions will use this job name
+job: extension
+config:
+  name: 'my_awesome_merge'
+  process:
+    # Put your example processes here. This will be passed
+    # to your extension process in the config argument.
+    # the type MUST match your extension uid
+    - type: "example_merge_extension"
+      # save path for the merged model
+      save_path: "output/merge/[name].safetensors"
+      # save type
+      dtype: fp16
+      # device to run it on
+      device: cuda:0
+      # input models can only be SD1.x and SD2.x models for this example (currently)
+      models_to_merge:
+        # weights are relative, total weights will be normalized
+        # for example. If you have 2 models with weight 1.0, they will
+        # both be weighted 0.5. If you have 1 model with weight 1.0 and
+        # another with weight 2.0, the first will be weighted 1/3 and the
+        # second will be weighted 2/3
+        - name_or_path: "input/model1.safetensors"
+          weight: 1.0
+        - name_or_path: "input/model2.safetensors"
+          weight: 1.0
+        - name_or_path: "input/model3.safetensors"
+          weight: 0.3
+        - name_or_path: "input/model4.safetensors"
+          weight: 1.0
+
+
+# you can put any information you want here, and it will be saved in the model
+# the below is an example. I recommend doing trigger words at a minimum
+# in the metadata. The software will include this plus some other information
+meta:
+  name: "[name]"  # [name] gets replaced with the name above
+  description: A short description of your model
+  version: '0.1'
+  creator:
+    name: Your Name
+    email: [email protected]
+    website: https://yourwebsite.com
+  any: All meta data above is arbitrary, it can be whatever you want.

extensions_built_in/advanced_generator/Img2ImgGenerator.py

@@ -0,0 +1,256 @@
+import math
+import os
+import random
+from collections import OrderedDict
+from typing import List
+
+import numpy as np
+from PIL import Image
+from diffusers import T2IAdapter
+from diffusers.utils.torch_utils import randn_tensor
+from torch.utils.data import DataLoader
+from diffusers import StableDiffusionXLImg2ImgPipeline, PixArtSigmaPipeline
+from tqdm import tqdm
+
+from toolkit.config_modules import ModelConfig, GenerateImageConfig, preprocess_dataset_raw_config, DatasetConfig
+from toolkit.data_transfer_object.data_loader import FileItemDTO, DataLoaderBatchDTO
+from toolkit.sampler import get_sampler
+from toolkit.stable_diffusion_model import StableDiffusion
+import gc
+import torch
+from jobs.process import BaseExtensionProcess
+from toolkit.data_loader import get_dataloader_from_datasets
+from toolkit.train_tools import get_torch_dtype
+from controlnet_aux.midas import MidasDetector
+from diffusers.utils import load_image
+from torchvision.transforms import ToTensor
+
+
+def flush():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+
+
+
+class GenerateConfig:
+
+    def __init__(self, **kwargs):
+        self.prompts: List[str]
+        self.sampler = kwargs.get('sampler', 'ddpm')
+        self.neg = kwargs.get('neg', '')
+        self.seed = kwargs.get('seed', -1)
+        self.walk_seed = kwargs.get('walk_seed', False)
+        self.guidance_scale = kwargs.get('guidance_scale', 7)
+        self.sample_steps = kwargs.get('sample_steps', 20)
+        self.guidance_rescale = kwargs.get('guidance_rescale', 0.0)
+        self.ext = kwargs.get('ext', 'png')
+        self.denoise_strength = kwargs.get('denoise_strength', 0.5)
+        self.trigger_word = kwargs.get('trigger_word', None)
+
+
+class Img2ImgGenerator(BaseExtensionProcess):
+
+    def __init__(self, process_id: int, job, config: OrderedDict):
+        super().__init__(process_id, job, config)
+        self.output_folder = self.get_conf('output_folder', required=True)
+        self.copy_inputs_to = self.get_conf('copy_inputs_to', None)
+        self.device = self.get_conf('device', 'cuda')
+        self.model_config = ModelConfig(**self.get_conf('model', required=True))
+        self.generate_config = GenerateConfig(**self.get_conf('generate', required=True))
+        self.is_latents_cached = True
+        raw_datasets = self.get_conf('datasets', None)
+        if raw_datasets is not None and len(raw_datasets) > 0:
+            raw_datasets = preprocess_dataset_raw_config(raw_datasets)
+        self.datasets = None
+        self.datasets_reg = None
+        self.dtype = self.get_conf('dtype', 'float16')
+        self.torch_dtype = get_torch_dtype(self.dtype)
+        self.params = []
+        if raw_datasets is not None and len(raw_datasets) > 0:
+            for raw_dataset in raw_datasets:
+                dataset = DatasetConfig(**raw_dataset)
+                is_caching = dataset.cache_latents or dataset.cache_latents_to_disk
+                if not is_caching:
+                    self.is_latents_cached = False
+                if dataset.is_reg:
+                    if self.datasets_reg is None:
+                        self.datasets_reg = []
+                    self.datasets_reg.append(dataset)
+                else:
+                    if self.datasets is None:
+                        self.datasets = []
+                    self.datasets.append(dataset)
+
+        self.progress_bar = None
+        self.sd = StableDiffusion(
+            device=self.device,
+            model_config=self.model_config,
+            dtype=self.dtype,
+        )
+        print(f"Using device {self.device}")
+        self.data_loader: DataLoader = None
+        self.adapter: T2IAdapter = None
+
+    def to_pil(self, img):
+        # image comes in -1 to 1. convert to a PIL RGB image
+        img = (img + 1) / 2
+        img = img.clamp(0, 1)
+        img = img[0].permute(1, 2, 0).cpu().numpy()
+        img = (img * 255).astype(np.uint8)
+        image = Image.fromarray(img)
+        return image
+
+    def run(self):
+        with torch.no_grad():
+            super().run()
+            print("Loading model...")
+            self.sd.load_model()
+            device = torch.device(self.device)
+
+            if self.model_config.is_xl:
+                pipe = StableDiffusionXLImg2ImgPipeline(
+                    vae=self.sd.vae,
+                    unet=self.sd.unet,
+                    text_encoder=self.sd.text_encoder[0],
+                    text_encoder_2=self.sd.text_encoder[1],
+                    tokenizer=self.sd.tokenizer[0],
+                    tokenizer_2=self.sd.tokenizer[1],
+                    scheduler=get_sampler(self.generate_config.sampler),
+                ).to(device, dtype=self.torch_dtype)
+            elif self.model_config.is_pixart:
+                pipe = self.sd.pipeline.to(device, dtype=self.torch_dtype)
+            else:
+                raise NotImplementedError("Only XL models are supported")
+            pipe.set_progress_bar_config(disable=True)
+
+            # pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+            # midas_depth = torch.compile(midas_depth, mode="reduce-overhead", fullgraph=True)
+
+            self.data_loader = get_dataloader_from_datasets(self.datasets, 1, self.sd)
+
+            num_batches = len(self.data_loader)
+            pbar = tqdm(total=num_batches, desc="Generating images")
+            seed = self.generate_config.seed
+            # load images from datasets, use tqdm
+            for i, batch in enumerate(self.data_loader):
+                batch: DataLoaderBatchDTO = batch
+
+                gen_seed = seed if seed > 0 else random.randint(0, 2 ** 32 - 1)
+                generator = torch.manual_seed(gen_seed)
+
+                file_item: FileItemDTO = batch.file_items[0]
+                img_path = file_item.path
+                img_filename = os.path.basename(img_path)
+                img_filename_no_ext = os.path.splitext(img_filename)[0]
+                img_filename = img_filename_no_ext + '.' + self.generate_config.ext
+                output_path = os.path.join(self.output_folder, img_filename)
+                output_caption_path = os.path.join(self.output_folder, img_filename_no_ext + '.txt')
+
+                if self.copy_inputs_to is not None:
+                    output_inputs_path = os.path.join(self.copy_inputs_to, img_filename)
+                    output_inputs_caption_path = os.path.join(self.copy_inputs_to, img_filename_no_ext + '.txt')
+                else:
+                    output_inputs_path = None
+                    output_inputs_caption_path = None
+
+                caption = batch.get_caption_list()[0]
+                if self.generate_config.trigger_word is not None:
+                    caption = caption.replace('[trigger]', self.generate_config.trigger_word)
+
+                img: torch.Tensor = batch.tensor.clone()
+                image = self.to_pil(img)
+
+                # image.save(output_depth_path)
+                if self.model_config.is_pixart:
+                    pipe: PixArtSigmaPipeline = pipe
+
+                    # Encode the full image once
+                    encoded_image = pipe.vae.encode(
+                        pipe.image_processor.preprocess(image).to(device=pipe.device, dtype=pipe.dtype))
+                    if hasattr(encoded_image, "latent_dist"):
+                        latents = encoded_image.latent_dist.sample(generator)
+                    elif hasattr(encoded_image, "latents"):
+                        latents = encoded_image.latents
+                    else:
+                        raise AttributeError("Could not access latents of provided encoder_output")
+                    latents = pipe.vae.config.scaling_factor * latents
+
+                    # latents = self.sd.encode_images(img)
+
+                    # self.sd.noise_scheduler.set_timesteps(self.generate_config.sample_steps)
+                    # start_step = math.floor(self.generate_config.sample_steps * self.generate_config.denoise_strength)
+                    # timestep = self.sd.noise_scheduler.timesteps[start_step].unsqueeze(0)
+                    # timestep = timestep.to(device, dtype=torch.int32)
+                    # latent = latent.to(device, dtype=self.torch_dtype)
+                    # noise = torch.randn_like(latent, device=device, dtype=self.torch_dtype)
+                    # latent = self.sd.add_noise(latent, noise, timestep)
+                    # timesteps_to_use = self.sd.noise_scheduler.timesteps[start_step + 1:]
+                    batch_size = 1
+                    num_images_per_prompt = 1
+
+                    shape = (batch_size, pipe.transformer.config.in_channels, image.height // pipe.vae_scale_factor,
+                             image.width // pipe.vae_scale_factor)
+                    noise = randn_tensor(shape, generator=generator, device=pipe.device, dtype=pipe.dtype)
+
+                    # noise = torch.randn_like(latents, device=device, dtype=self.torch_dtype)
+                    num_inference_steps = self.generate_config.sample_steps
+                    strength = self.generate_config.denoise_strength
+                    # Get timesteps
+                    init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
+                    t_start = max(num_inference_steps - init_timestep, 0)
+                    pipe.scheduler.set_timesteps(num_inference_steps, device="cpu")
+                    timesteps = pipe.scheduler.timesteps[t_start:]
+                    timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
+                    latents = pipe.scheduler.add_noise(latents, noise, timestep)
+
+                    gen_images = pipe.__call__(
+                        prompt=caption,
+                        negative_prompt=self.generate_config.neg,
+                        latents=latents,
+                        timesteps=timesteps,
+                        width=image.width,
+                        height=image.height,
+                        num_inference_steps=num_inference_steps,
+                        num_images_per_prompt=num_images_per_prompt,
+                        guidance_scale=self.generate_config.guidance_scale,
+                        # strength=self.generate_config.denoise_strength,
+                        use_resolution_binning=False,
+                        output_type="np"
+                    ).images[0]
+                    gen_images = (gen_images * 255).clip(0, 255).astype(np.uint8)
+                    gen_images = Image.fromarray(gen_images)
+                else:
+                    pipe: StableDiffusionXLImg2ImgPipeline = pipe
+
+                    gen_images = pipe.__call__(
+                        prompt=caption,
+                        negative_prompt=self.generate_config.neg,
+                        image=image,
+                        num_inference_steps=self.generate_config.sample_steps,
+                        guidance_scale=self.generate_config.guidance_scale,
+                        strength=self.generate_config.denoise_strength,
+                    ).images[0]
+                os.makedirs(os.path.dirname(output_path), exist_ok=True)
+                gen_images.save(output_path)
+
+                # save caption
+                with open(output_caption_path, 'w') as f:
+                    f.write(caption)
+
+                if output_inputs_path is not None:
+                    os.makedirs(os.path.dirname(output_inputs_path), exist_ok=True)
+                    image.save(output_inputs_path)
+                    with open(output_inputs_caption_path, 'w') as f:
+                        f.write(caption)
+
+                pbar.update(1)
+                batch.cleanup()
+
+            pbar.close()
+            print("Done generating images")
+            # cleanup
+            del self.sd
+            gc.collect()
+            torch.cuda.empty_cache()

extensions_built_in/advanced_generator/PureLoraGenerator.py

@@ -0,0 +1,102 @@
+import os
+from collections import OrderedDict
+
+from toolkit.config_modules import ModelConfig, GenerateImageConfig, SampleConfig, LoRMConfig
+from toolkit.lorm import ExtractMode, convert_diffusers_unet_to_lorm
+from toolkit.sd_device_states_presets import get_train_sd_device_state_preset
+from toolkit.stable_diffusion_model import StableDiffusion
+import gc
+import torch
+from jobs.process import BaseExtensionProcess
+from toolkit.train_tools import get_torch_dtype
+
+
+def flush():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+class PureLoraGenerator(BaseExtensionProcess):
+
+    def __init__(self, process_id: int, job, config: OrderedDict):
+        super().__init__(process_id, job, config)
+        self.output_folder = self.get_conf('output_folder', required=True)
+        self.device = self.get_conf('device', 'cuda')
+        self.device_torch = torch.device(self.device)
+        self.model_config = ModelConfig(**self.get_conf('model', required=True))
+        self.generate_config = SampleConfig(**self.get_conf('sample', required=True))
+        self.dtype = self.get_conf('dtype', 'float16')
+        self.torch_dtype = get_torch_dtype(self.dtype)
+        lorm_config = self.get_conf('lorm', None)
+        self.lorm_config = LoRMConfig(**lorm_config) if lorm_config is not None else None
+
+        self.device_state_preset = get_train_sd_device_state_preset(
+            device=torch.device(self.device),
+        )
+
+        self.progress_bar = None
+        self.sd = StableDiffusion(
+            device=self.device,
+            model_config=self.model_config,
+            dtype=self.dtype,
+        )
+
+    def run(self):
+        super().run()
+        print("Loading model...")
+        with torch.no_grad():
+            self.sd.load_model()
+            self.sd.unet.eval()
+            self.sd.unet.to(self.device_torch)
+            if isinstance(self.sd.text_encoder, list):
+                for te in self.sd.text_encoder:
+                    te.eval()
+                    te.to(self.device_torch)
+            else:
+                self.sd.text_encoder.eval()
+                self.sd.to(self.device_torch)
+
+            print(f"Converting to LoRM UNet")
+            # replace the unet with LoRMUnet
+            convert_diffusers_unet_to_lorm(
+                self.sd.unet,
+                config=self.lorm_config,
+            )
+
+            sample_folder = os.path.join(self.output_folder)
+            gen_img_config_list = []
+
+            sample_config = self.generate_config
+            start_seed = sample_config.seed
+            current_seed = start_seed
+            for i in range(len(sample_config.prompts)):
+                if sample_config.walk_seed:
+                    current_seed = start_seed + i
+
+                filename = f"[time]_[count].{self.generate_config.ext}"
+                output_path = os.path.join(sample_folder, filename)
+                prompt = sample_config.prompts[i]
+                extra_args = {}
+                gen_img_config_list.append(GenerateImageConfig(
+                    prompt=prompt,  # it will autoparse the prompt
+                    width=sample_config.width,
+                    height=sample_config.height,
+                    negative_prompt=sample_config.neg,
+                    seed=current_seed,
+                    guidance_scale=sample_config.guidance_scale,
+                    guidance_rescale=sample_config.guidance_rescale,
+                    num_inference_steps=sample_config.sample_steps,
+                    network_multiplier=sample_config.network_multiplier,
+                    output_path=output_path,
+                    output_ext=sample_config.ext,
+                    adapter_conditioning_scale=sample_config.adapter_conditioning_scale,
+                    **extra_args
+                ))
+
+            # send to be generated
+            self.sd.generate_images(gen_img_config_list, sampler=sample_config.sampler)
+            print("Done generating images")
+            # cleanup
+            del self.sd
+            gc.collect()
+            torch.cuda.empty_cache()

extensions_built_in/advanced_generator/ReferenceGenerator.py

@@ -0,0 +1,212 @@
+import os
+import random
+from collections import OrderedDict
+from typing import List
+
+import numpy as np
+from PIL import Image
+from diffusers import T2IAdapter
+from torch.utils.data import DataLoader
+from diffusers import StableDiffusionXLAdapterPipeline, StableDiffusionAdapterPipeline
+from tqdm import tqdm
+
+from toolkit.config_modules import ModelConfig, GenerateImageConfig, preprocess_dataset_raw_config, DatasetConfig
+from toolkit.data_transfer_object.data_loader import FileItemDTO, DataLoaderBatchDTO
+from toolkit.sampler import get_sampler
+from toolkit.stable_diffusion_model import StableDiffusion
+import gc
+import torch
+from jobs.process import BaseExtensionProcess
+from toolkit.data_loader import get_dataloader_from_datasets
+from toolkit.train_tools import get_torch_dtype
+from controlnet_aux.midas import MidasDetector
+from diffusers.utils import load_image
+
+
+def flush():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+class GenerateConfig:
+
+    def __init__(self, **kwargs):
+        self.prompts: List[str]
+        self.sampler = kwargs.get('sampler', 'ddpm')
+        self.neg = kwargs.get('neg', '')
+        self.seed = kwargs.get('seed', -1)
+        self.walk_seed = kwargs.get('walk_seed', False)
+        self.t2i_adapter_path = kwargs.get('t2i_adapter_path', None)
+        self.guidance_scale = kwargs.get('guidance_scale', 7)
+        self.sample_steps = kwargs.get('sample_steps', 20)
+        self.prompt_2 = kwargs.get('prompt_2', None)
+        self.neg_2 = kwargs.get('neg_2', None)
+        self.prompts = kwargs.get('prompts', None)
+        self.guidance_rescale = kwargs.get('guidance_rescale', 0.0)
+        self.ext = kwargs.get('ext', 'png')
+        self.adapter_conditioning_scale = kwargs.get('adapter_conditioning_scale', 1.0)
+        if kwargs.get('shuffle', False):
+            # shuffle the prompts
+            random.shuffle(self.prompts)
+
+
+class ReferenceGenerator(BaseExtensionProcess):
+
+    def __init__(self, process_id: int, job, config: OrderedDict):
+        super().__init__(process_id, job, config)
+        self.output_folder = self.get_conf('output_folder', required=True)
+        self.device = self.get_conf('device', 'cuda')
+        self.model_config = ModelConfig(**self.get_conf('model', required=True))
+        self.generate_config = GenerateConfig(**self.get_conf('generate', required=True))
+        self.is_latents_cached = True
+        raw_datasets = self.get_conf('datasets', None)
+        if raw_datasets is not None and len(raw_datasets) > 0:
+            raw_datasets = preprocess_dataset_raw_config(raw_datasets)
+        self.datasets = None
+        self.datasets_reg = None
+        self.dtype = self.get_conf('dtype', 'float16')
+        self.torch_dtype = get_torch_dtype(self.dtype)
+        self.params = []
+        if raw_datasets is not None and len(raw_datasets) > 0:
+            for raw_dataset in raw_datasets:
+                dataset = DatasetConfig(**raw_dataset)
+                is_caching = dataset.cache_latents or dataset.cache_latents_to_disk
+                if not is_caching:
+                    self.is_latents_cached = False
+                if dataset.is_reg:
+                    if self.datasets_reg is None:
+                        self.datasets_reg = []
+                    self.datasets_reg.append(dataset)
+                else:
+                    if self.datasets is None:
+                        self.datasets = []
+                    self.datasets.append(dataset)
+
+        self.progress_bar = None
+        self.sd = StableDiffusion(
+            device=self.device,
+            model_config=self.model_config,
+            dtype=self.dtype,
+        )
+        print(f"Using device {self.device}")
+        self.data_loader: DataLoader = None
+        self.adapter: T2IAdapter = None
+
+    def run(self):
+        super().run()
+        print("Loading model...")
+        self.sd.load_model()
+        device = torch.device(self.device)
+
+        if self.generate_config.t2i_adapter_path is not None:
+            self.adapter = T2IAdapter.from_pretrained(
+                self.generate_config.t2i_adapter_path,
+                torch_dtype=self.torch_dtype,
+                varient="fp16"
+            ).to(device)
+
+        midas_depth = MidasDetector.from_pretrained(
+            "valhalla/t2iadapter-aux-models", filename="dpt_large_384.pt", model_type="dpt_large"
+        ).to(device)
+
+        if self.model_config.is_xl:
+            pipe = StableDiffusionXLAdapterPipeline(
+                vae=self.sd.vae,
+                unet=self.sd.unet,
+                text_encoder=self.sd.text_encoder[0],
+                text_encoder_2=self.sd.text_encoder[1],
+                tokenizer=self.sd.tokenizer[0],
+                tokenizer_2=self.sd.tokenizer[1],
+                scheduler=get_sampler(self.generate_config.sampler),
+                adapter=self.adapter,
+            ).to(device, dtype=self.torch_dtype)
+        else:
+            pipe = StableDiffusionAdapterPipeline(
+                vae=self.sd.vae,
+                unet=self.sd.unet,
+                text_encoder=self.sd.text_encoder,
+                tokenizer=self.sd.tokenizer,
+                scheduler=get_sampler(self.generate_config.sampler),
+                safety_checker=None,
+                feature_extractor=None,
+                requires_safety_checker=False,
+                adapter=self.adapter,
+            ).to(device, dtype=self.torch_dtype)
+        pipe.set_progress_bar_config(disable=True)
+
+        pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+        # midas_depth = torch.compile(midas_depth, mode="reduce-overhead", fullgraph=True)
+
+        self.data_loader = get_dataloader_from_datasets(self.datasets, 1, self.sd)
+
+        num_batches = len(self.data_loader)
+        pbar = tqdm(total=num_batches, desc="Generating images")
+        seed = self.generate_config.seed
+        # load images from datasets, use tqdm
+        for i, batch in enumerate(self.data_loader):
+            batch: DataLoaderBatchDTO = batch
+
+            file_item: FileItemDTO = batch.file_items[0]
+            img_path = file_item.path
+            img_filename = os.path.basename(img_path)
+            img_filename_no_ext = os.path.splitext(img_filename)[0]
+            output_path = os.path.join(self.output_folder, img_filename)
+            output_caption_path = os.path.join(self.output_folder, img_filename_no_ext + '.txt')
+            output_depth_path = os.path.join(self.output_folder, img_filename_no_ext + '.depth.png')
+
+            caption = batch.get_caption_list()[0]
+
+            img: torch.Tensor = batch.tensor.clone()
+            # image comes in -1 to 1. convert to a PIL RGB image
+            img = (img + 1) / 2
+            img = img.clamp(0, 1)
+            img = img[0].permute(1, 2, 0).cpu().numpy()
+            img = (img * 255).astype(np.uint8)
+            image = Image.fromarray(img)
+
+            width, height = image.size
+            min_res = min(width, height)
+
+            if self.generate_config.walk_seed:
+                seed = seed + 1
+
+            if self.generate_config.seed == -1:
+                # random
+                seed = random.randint(0, 1000000)
+
+            torch.manual_seed(seed)
+            torch.cuda.manual_seed(seed)
+
+            # generate depth map
+            image = midas_depth(
+                image,
+                detect_resolution=min_res,  # do 512 ?
+                image_resolution=min_res
+            )
+
+            # image.save(output_depth_path)
+
+            gen_images = pipe(
+                prompt=caption,
+                negative_prompt=self.generate_config.neg,
+                image=image,
+                num_inference_steps=self.generate_config.sample_steps,
+                adapter_conditioning_scale=self.generate_config.adapter_conditioning_scale,
+                guidance_scale=self.generate_config.guidance_scale,
+            ).images[0]
+            os.makedirs(os.path.dirname(output_path), exist_ok=True)
+            gen_images.save(output_path)
+
+            # save caption
+            with open(output_caption_path, 'w') as f:
+                f.write(caption)
+
+            pbar.update(1)
+            batch.cleanup()
+
+        pbar.close()
+        print("Done generating images")
+        # cleanup
+        del self.sd
+        gc.collect()
+        torch.cuda.empty_cache()

extensions_built_in/advanced_generator/__init__.py

@@ -0,0 +1,59 @@
+# This is an example extension for custom training. It is great for experimenting with new ideas.
+from toolkit.extension import Extension
+
+
+# This is for generic training (LoRA, Dreambooth, FineTuning)
+class AdvancedReferenceGeneratorExtension(Extension):
+    # uid must be unique, it is how the extension is identified
+    uid = "reference_generator"
+
+    # name is the name of the extension for printing
+    name = "Reference Generator"
+
+    # This is where your process class is loaded
+    # keep your imports in here so they don't slow down the rest of the program
+    @classmethod
+    def get_process(cls):
+        # import your process class here so it is only loaded when needed and return it
+        from .ReferenceGenerator import ReferenceGenerator
+        return ReferenceGenerator
+
+
+# This is for generic training (LoRA, Dreambooth, FineTuning)
+class PureLoraGenerator(Extension):
+    # uid must be unique, it is how the extension is identified
+    uid = "pure_lora_generator"
+
+    # name is the name of the extension for printing
+    name = "Pure LoRA Generator"
+
+    # This is where your process class is loaded
+    # keep your imports in here so they don't slow down the rest of the program
+    @classmethod
+    def get_process(cls):
+        # import your process class here so it is only loaded when needed and return it
+        from .PureLoraGenerator import PureLoraGenerator
+        return PureLoraGenerator
+
+
+# This is for generic training (LoRA, Dreambooth, FineTuning)
+class Img2ImgGeneratorExtension(Extension):
+    # uid must be unique, it is how the extension is identified
+    uid = "batch_img2img"
+
+    # name is the name of the extension for printing
+    name = "Img2ImgGeneratorExtension"
+
+    # This is where your process class is loaded
+    # keep your imports in here so they don't slow down the rest of the program
+    @classmethod
+    def get_process(cls):
+        # import your process class here so it is only loaded when needed and return it
+        from .Img2ImgGenerator import Img2ImgGenerator
+        return Img2ImgGenerator
+
+
+AI_TOOLKIT_EXTENSIONS = [
+    # you can put a list of extensions here
+    AdvancedReferenceGeneratorExtension, PureLoraGenerator, Img2ImgGeneratorExtension
+]

extensions_built_in/advanced_generator/config/train.example.yaml

@@ -0,0 +1,91 @@
+---
+job: extension
+config:
+  name: test_v1
+  process:
+    - type: 'textual_inversion_trainer'
+      training_folder: "out/TI"
+      device: cuda:0
+      # for tensorboard logging
+      log_dir: "out/.tensorboard"
+      embedding:
+        trigger: "your_trigger_here"
+        tokens: 12
+        init_words: "man with short brown hair"
+        save_format: "safetensors"  # 'safetensors' or 'pt'
+      save:
+        dtype: float16 # precision to save
+        save_every: 100 # save every this many steps
+        max_step_saves_to_keep: 5 # only affects step counts
+      datasets:
+        - folder_path: "/path/to/dataset"
+          caption_ext: "txt"
+          default_caption: "[trigger]"
+          buckets: true
+          resolution: 512
+      train:
+        noise_scheduler: "ddpm" # or "ddpm", "lms", "euler_a"
+        steps: 3000
+        weight_jitter: 0.0
+        lr: 5e-5
+        train_unet: false
+        gradient_checkpointing: true
+        train_text_encoder: false
+        optimizer: "adamw"
+#        optimizer: "prodigy"
+        optimizer_params:
+          weight_decay: 1e-2
+        lr_scheduler: "constant"
+        max_denoising_steps: 1000
+        batch_size: 4
+        dtype: bf16
+        xformers: true
+        min_snr_gamma: 5.0
+#        skip_first_sample: true
+        noise_offset: 0.0 # not needed for this
+      model:
+        # objective reality v2
+        name_or_path: "https://civitai.com/models/128453?modelVersionId=142465"
+        is_v2: false  # for v2 models
+        is_xl: false  # for SDXL models
+        is_v_pred: false # for v-prediction models (most v2 models)
+      sample:
+        sampler: "ddpm" # must match train.noise_scheduler
+        sample_every: 100 # sample every this many steps
+        width: 512
+        height: 512
+        prompts:
+          - "photo of [trigger] laughing"
+          - "photo of [trigger] smiling"
+          - "[trigger] close up"
+          - "dark scene [trigger] frozen"
+          - "[trigger] nighttime"
+          - "a painting of [trigger]"
+          - "a drawing of [trigger]"
+          - "a cartoon of [trigger]"
+          - "[trigger] pixar style"
+          - "[trigger] costume"
+        neg: ""
+        seed: 42
+        walk_seed: false
+        guidance_scale: 7
+        sample_steps: 20
+        network_multiplier: 1.0
+
+      logging:
+        log_every: 10 # log every this many steps
+        use_wandb: false # not supported yet
+        verbose: false
+
+# You can put any information you want here, and it will be saved in the model.
+# The below is an example, but you can put your grocery list in it if you want.
+# It is saved in the model so be aware of that. The software will include this
+# plus some other information for you automatically
+meta:
+  # [name] gets replaced with the name above
+  name: "[name]"
+#  version: '1.0'
+#  creator:
+#    name: Your Name
+#    email: [email protected]
+#    website: https://your.website

extensions_built_in/concept_replacer/ConceptReplacer.py

@@ -0,0 +1,151 @@
+import random
+from collections import OrderedDict
+from torch.utils.data import DataLoader
+from toolkit.prompt_utils import concat_prompt_embeds, split_prompt_embeds
+from toolkit.stable_diffusion_model import StableDiffusion, BlankNetwork
+from toolkit.train_tools import get_torch_dtype, apply_snr_weight
+import gc
+import torch
+from jobs.process import BaseSDTrainProcess
+
+
+def flush():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+class ConceptReplacementConfig:
+    def __init__(self, **kwargs):
+        self.concept: str = kwargs.get('concept', '')
+        self.replacement: str = kwargs.get('replacement', '')
+
+
+class ConceptReplacer(BaseSDTrainProcess):
+
+    def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
+        super().__init__(process_id, job, config, **kwargs)
+        replacement_list = self.config.get('replacements', [])
+        self.replacement_list = [ConceptReplacementConfig(**x) for x in replacement_list]
+
+    def before_model_load(self):
+        pass
+
+    def hook_before_train_loop(self):
+        self.sd.vae.eval()
+        self.sd.vae.to(self.device_torch)
+
+        # textual inversion
+        if self.embedding is not None:
+            # set text encoder to train. Not sure if this is necessary but diffusers example did it
+            self.sd.text_encoder.train()
+
+    def hook_train_loop(self, batch):
+        with torch.no_grad():
+            dtype = get_torch_dtype(self.train_config.dtype)
+            noisy_latents, noise, timesteps, conditioned_prompts, imgs = self.process_general_training_batch(batch)
+            network_weight_list = batch.get_network_weight_list()
+
+            # have a blank network so we can wrap it in a context and set multipliers without checking every time
+            if self.network is not None:
+                network = self.network
+            else:
+                network = BlankNetwork()
+
+            batch_replacement_list = []
+            # get a random replacement for each prompt
+            for prompt in conditioned_prompts:
+                replacement = random.choice(self.replacement_list)
+                batch_replacement_list.append(replacement)
+
+            # build out prompts
+            concept_prompts = []
+            replacement_prompts = []
+            for idx, replacement in enumerate(batch_replacement_list):
+                prompt = conditioned_prompts[idx]
+
+                # insert shuffled concept at beginning and end of prompt
+                shuffled_concept = [x.strip() for x in replacement.concept.split(',')]
+                random.shuffle(shuffled_concept)
+                shuffled_concept = ', '.join(shuffled_concept)
+                concept_prompts.append(f"{shuffled_concept}, {prompt}, {shuffled_concept}")
+
+                # insert replacement at beginning and end of prompt
+                shuffled_replacement = [x.strip() for x in replacement.replacement.split(',')]
+                random.shuffle(shuffled_replacement)
+                shuffled_replacement = ', '.join(shuffled_replacement)
+                replacement_prompts.append(f"{shuffled_replacement}, {prompt}, {shuffled_replacement}")
+
+            # predict the replacement without network
+            conditional_embeds = self.sd.encode_prompt(replacement_prompts).to(self.device_torch, dtype=dtype)
+
+            replacement_pred = self.sd.predict_noise(
+                latents=noisy_latents.to(self.device_torch, dtype=dtype),
+                conditional_embeddings=conditional_embeds.to(self.device_torch, dtype=dtype),
+                timestep=timesteps,
+                guidance_scale=1.0,
+            )
+
+            del conditional_embeds
+            replacement_pred = replacement_pred.detach()
+
+        self.optimizer.zero_grad()
+        flush()
+
+        # text encoding
+        grad_on_text_encoder = False
+        if self.train_config.train_text_encoder:
+            grad_on_text_encoder = True
+
+        if self.embedding:
+            grad_on_text_encoder = True
+
+        # set the weights
+        network.multiplier = network_weight_list
+
+        # activate network if it exits
+        with network:
+            with torch.set_grad_enabled(grad_on_text_encoder):
+                # embed the prompts
+                conditional_embeds = self.sd.encode_prompt(concept_prompts).to(self.device_torch, dtype=dtype)
+            if not grad_on_text_encoder:
+                # detach the embeddings
+                conditional_embeds = conditional_embeds.detach()
+                self.optimizer.zero_grad()
+                flush()
+
+            noise_pred = self.sd.predict_noise(
+                latents=noisy_latents.to(self.device_torch, dtype=dtype),
+                conditional_embeddings=conditional_embeds.to(self.device_torch, dtype=dtype),
+                timestep=timesteps,
+                guidance_scale=1.0,
+            )
+
+            loss = torch.nn.functional.mse_loss(noise_pred.float(), replacement_pred.float(), reduction="none")
+            loss = loss.mean([1, 2, 3])
+
+            if self.train_config.min_snr_gamma is not None and self.train_config.min_snr_gamma > 0.000001:
+                # add min_snr_gamma
+                loss = apply_snr_weight(loss, timesteps, self.sd.noise_scheduler, self.train_config.min_snr_gamma)
+
+            loss = loss.mean()
+
+            # back propagate loss to free ram
+            loss.backward()
+            flush()
+
+        # apply gradients
+        self.optimizer.step()
+        self.optimizer.zero_grad()
+        self.lr_scheduler.step()
+
+        if self.embedding is not None:
+            # Let's make sure we don't update any embedding weights besides the newly added token
+            self.embedding.restore_embeddings()
+
+        loss_dict = OrderedDict(
+            {'loss': loss.item()}
+        )
+        # reset network multiplier
+        network.multiplier = 1.0
+
+        return loss_dict

extensions_built_in/concept_replacer/__init__.py

@@ -0,0 +1,26 @@
+# This is an example extension for custom training. It is great for experimenting with new ideas.
+from toolkit.extension import Extension
+
+
+# This is for generic training (LoRA, Dreambooth, FineTuning)
+class ConceptReplacerExtension(Extension):
+    # uid must be unique, it is how the extension is identified
+    uid = "concept_replacer"
+
+    # name is the name of the extension for printing
+    name = "Concept Replacer"
+
+    # This is where your process class is loaded
+    # keep your imports in here so they don't slow down the rest of the program
+    @classmethod
+    def get_process(cls):
+        # import your process class here so it is only loaded when needed and return it
+        from .ConceptReplacer import ConceptReplacer
+        return ConceptReplacer
+
+
+
+AI_TOOLKIT_EXTENSIONS = [
+    # you can put a list of extensions here
+    ConceptReplacerExtension,
+]

extensions_built_in/concept_replacer/config/train.example.yaml

@@ -0,0 +1,91 @@
+---
+job: extension
+config:
+  name: test_v1
+  process:
+    - type: 'textual_inversion_trainer'
+      training_folder: "out/TI"
+      device: cuda:0
+      # for tensorboard logging
+      log_dir: "out/.tensorboard"
+      embedding:
+        trigger: "your_trigger_here"
+        tokens: 12
+        init_words: "man with short brown hair"
+        save_format: "safetensors"  # 'safetensors' or 'pt'
+      save:
+        dtype: float16 # precision to save
+        save_every: 100 # save every this many steps
+        max_step_saves_to_keep: 5 # only affects step counts
+      datasets:
+        - folder_path: "/path/to/dataset"
+          caption_ext: "txt"
+          default_caption: "[trigger]"
+          buckets: true
+          resolution: 512
+      train:
+        noise_scheduler: "ddpm" # or "ddpm", "lms", "euler_a"
+        steps: 3000
+        weight_jitter: 0.0
+        lr: 5e-5
+        train_unet: false
+        gradient_checkpointing: true
+        train_text_encoder: false
+        optimizer: "adamw"
+#        optimizer: "prodigy"
+        optimizer_params:
+          weight_decay: 1e-2
+        lr_scheduler: "constant"
+        max_denoising_steps: 1000
+        batch_size: 4
+        dtype: bf16
+        xformers: true
+        min_snr_gamma: 5.0
+#        skip_first_sample: true
+        noise_offset: 0.0 # not needed for this
+      model:
+        # objective reality v2
+        name_or_path: "https://civitai.com/models/128453?modelVersionId=142465"
+        is_v2: false  # for v2 models
+        is_xl: false  # for SDXL models
+        is_v_pred: false # for v-prediction models (most v2 models)
+      sample:
+        sampler: "ddpm" # must match train.noise_scheduler
+        sample_every: 100 # sample every this many steps
+        width: 512
+        height: 512
+        prompts:
+          - "photo of [trigger] laughing"
+          - "photo of [trigger] smiling"
+          - "[trigger] close up"
+          - "dark scene [trigger] frozen"
+          - "[trigger] nighttime"
+          - "a painting of [trigger]"
+          - "a drawing of [trigger]"
+          - "a cartoon of [trigger]"
+          - "[trigger] pixar style"
+          - "[trigger] costume"
+        neg: ""
+        seed: 42
+        walk_seed: false
+        guidance_scale: 7
+        sample_steps: 20
+        network_multiplier: 1.0
+
+      logging:
+        log_every: 10 # log every this many steps
+        use_wandb: false # not supported yet
+        verbose: false
+
+# You can put any information you want here, and it will be saved in the model.
+# The below is an example, but you can put your grocery list in it if you want.
+# It is saved in the model so be aware of that. The software will include this
+# plus some other information for you automatically
+meta:
+  # [name] gets replaced with the name above
+  name: "[name]"
+#  version: '1.0'
+#  creator:
+#    name: Your Name
+#    email: [email protected]
+#    website: https://your.website

extensions_built_in/dataset_tools/DatasetTools.py

@@ -0,0 +1,20 @@
+from collections import OrderedDict
+import gc
+import torch
+from jobs.process import BaseExtensionProcess
+
+
+def flush():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+class DatasetTools(BaseExtensionProcess):
+
+    def __init__(self, process_id: int, job, config: OrderedDict):
+        super().__init__(process_id, job, config)
+
+    def run(self):
+        super().run()
+
+        raise NotImplementedError("This extension is not yet implemented")

extensions_built_in/dataset_tools/SuperTagger.py

@@ -0,0 +1,196 @@
+import copy
+import json
+import os
+from collections import OrderedDict
+import gc
+import traceback
+import torch
+from PIL import Image, ImageOps
+from tqdm import tqdm
+
+from .tools.dataset_tools_config_modules import RAW_DIR, TRAIN_DIR, Step, ImgInfo
+from .tools.fuyu_utils import FuyuImageProcessor
+from .tools.image_tools import load_image, ImageProcessor, resize_to_max
+from .tools.llava_utils import LLaVAImageProcessor
+from .tools.caption import default_long_prompt, default_short_prompt, default_replacements
+from jobs.process import BaseExtensionProcess
+from .tools.sync_tools import get_img_paths
+
+img_ext = ['.jpg', '.jpeg', '.png', '.webp']
+
+
+def flush():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+VERSION = 2
+
+
+class SuperTagger(BaseExtensionProcess):
+
+    def __init__(self, process_id: int, job, config: OrderedDict):
+        super().__init__(process_id, job, config)
+        parent_dir = config.get('parent_dir', None)
+        self.dataset_paths: list[str] = config.get('dataset_paths', [])
+        self.device = config.get('device', 'cuda')
+        self.steps: list[Step] = config.get('steps', [])
+        self.caption_method = config.get('caption_method', 'llava:default')
+        self.caption_prompt = config.get('caption_prompt', default_long_prompt)
+        self.caption_short_prompt = config.get('caption_short_prompt', default_short_prompt)
+        self.force_reprocess_img = config.get('force_reprocess_img', False)
+        self.caption_replacements = config.get('caption_replacements', default_replacements)
+        self.caption_short_replacements = config.get('caption_short_replacements', default_replacements)
+        self.master_dataset_dict = OrderedDict()
+        self.dataset_master_config_file = config.get('dataset_master_config_file', None)
+        if parent_dir is not None and len(self.dataset_paths) == 0:
+            # find all folders in the patent_dataset_path
+            self.dataset_paths = [
+                os.path.join(parent_dir, folder)
+                for folder in os.listdir(parent_dir)
+                if os.path.isdir(os.path.join(parent_dir, folder))
+            ]
+        else:
+            # make sure they exist
+            for dataset_path in self.dataset_paths:
+                if not os.path.exists(dataset_path):
+                    raise ValueError(f"Dataset path does not exist: {dataset_path}")
+
+        print(f"Found {len(self.dataset_paths)} dataset paths")
+
+        self.image_processor: ImageProcessor = self.get_image_processor()
+
+    def get_image_processor(self):
+        if self.caption_method.startswith('llava'):
+            return LLaVAImageProcessor(device=self.device)
+        elif self.caption_method.startswith('fuyu'):
+            return FuyuImageProcessor(device=self.device)
+        else:
+            raise ValueError(f"Unknown caption method: {self.caption_method}")
+
+    def process_image(self, img_path: str):
+        root_img_dir = os.path.dirname(os.path.dirname(img_path))
+        filename = os.path.basename(img_path)
+        filename_no_ext = os.path.splitext(filename)[0]
+        train_dir = os.path.join(root_img_dir, TRAIN_DIR)
+        train_img_path = os.path.join(train_dir, filename)
+        json_path = os.path.join(train_dir, f"{filename_no_ext}.json")
+
+        # check if json exists, if it does load it as image info
+        if os.path.exists(json_path):
+            with open(json_path, 'r') as f:
+                img_info = ImgInfo(**json.load(f))
+        else:
+            img_info = ImgInfo()
+
+        # always send steps first in case other processes need them
+        img_info.add_steps(copy.deepcopy(self.steps))
+        img_info.set_version(VERSION)
+        img_info.set_caption_method(self.caption_method)
+
+        image: Image = None
+        caption_image: Image = None
+
+        did_update_image = False
+
+        # trigger reprocess of steps
+        if self.force_reprocess_img:
+            img_info.trigger_image_reprocess()
+
+        # set the image as updated if it does not exist on disk
+        if not os.path.exists(train_img_path):
+            did_update_image = True
+            image = load_image(img_path)
+        if img_info.force_image_process:
+            did_update_image = True
+            image = load_image(img_path)
+
+        # go through the needed steps
+        for step in copy.deepcopy(img_info.state.steps_to_complete):
+            if step == 'caption':
+                # load image
+                if image is None:
+                    image = load_image(img_path)
+                if caption_image is None:
+                    caption_image = resize_to_max(image, 1024, 1024)
+
+                if not self.image_processor.is_loaded:
+                    print('Loading Model. Takes a while, especially the first time')
+                    self.image_processor.load_model()
+
+                img_info.caption = self.image_processor.generate_caption(
+                    image=caption_image,
+                    prompt=self.caption_prompt,
+                    replacements=self.caption_replacements
+                )
+                img_info.mark_step_complete(step)
+            elif step == 'caption_short':
+                # load image
+                if image is None:
+                    image = load_image(img_path)
+
+                if caption_image is None:
+                    caption_image = resize_to_max(image, 1024, 1024)
+
+                if not self.image_processor.is_loaded:
+                    print('Loading Model. Takes a while, especially the first time')
+                    self.image_processor.load_model()
+                img_info.caption_short = self.image_processor.generate_caption(
+                    image=caption_image,
+                    prompt=self.caption_short_prompt,
+                    replacements=self.caption_short_replacements
+                )
+                img_info.mark_step_complete(step)
+            elif step == 'contrast_stretch':
+                # load image
+                if image is None:
+                    image = load_image(img_path)
+                image = ImageOps.autocontrast(image, cutoff=(0.1, 0), preserve_tone=True)
+                did_update_image = True
+                img_info.mark_step_complete(step)
+            else:
+                raise ValueError(f"Unknown step: {step}")
+
+        os.makedirs(os.path.dirname(train_img_path), exist_ok=True)
+        if did_update_image:
+            image.save(train_img_path)
+
+        if img_info.is_dirty:
+            with open(json_path, 'w') as f:
+                json.dump(img_info.to_dict(), f, indent=4)
+
+        if self.dataset_master_config_file:
+            # add to master dict
+            self.master_dataset_dict[train_img_path] = img_info.to_dict()
+
+    def run(self):
+        super().run()
+        imgs_to_process = []
+        # find all images
+        for dataset_path in self.dataset_paths:
+            raw_dir = os.path.join(dataset_path, RAW_DIR)
+            raw_image_paths = get_img_paths(raw_dir)
+            for raw_image_path in raw_image_paths:
+                imgs_to_process.append(raw_image_path)
+
+        if len(imgs_to_process) == 0:
+            print(f"No images to process")
+        else:
+            print(f"Found {len(imgs_to_process)} to process")
+
+            for img_path in tqdm(imgs_to_process, desc="Processing images"):
+                try:
+                    self.process_image(img_path)
+                except Exception:
+                    # print full stack trace
+                    print(traceback.format_exc())
+                    continue
+                # self.process_image(img_path)
+
+        if self.dataset_master_config_file is not None:
+            # save it as json
+            with open(self.dataset_master_config_file, 'w') as f:
+                json.dump(self.master_dataset_dict, f, indent=4)
+
+        del self.image_processor
+        flush()

extensions_built_in/dataset_tools/SyncFromCollection.py

@@ -0,0 +1,131 @@
+import os
+import shutil
+from collections import OrderedDict
+import gc
+from typing import List
+
+import torch
+from tqdm import tqdm
+
+from .tools.dataset_tools_config_modules import DatasetSyncCollectionConfig, RAW_DIR, NEW_DIR
+from .tools.sync_tools import get_unsplash_images, get_pexels_images, get_local_image_file_names, download_image, \
+    get_img_paths
+from jobs.process import BaseExtensionProcess
+
+
+def flush():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+class SyncFromCollection(BaseExtensionProcess):
+
+    def __init__(self, process_id: int, job, config: OrderedDict):
+        super().__init__(process_id, job, config)
+
+        self.min_width = config.get('min_width', 1024)
+        self.min_height = config.get('min_height', 1024)
+
+        # add our min_width and min_height to each dataset config if they don't exist
+        for dataset_config in config.get('dataset_sync', []):
+            if 'min_width' not in dataset_config:
+                dataset_config['min_width'] = self.min_width
+            if 'min_height' not in dataset_config:
+                dataset_config['min_height'] = self.min_height
+
+        self.dataset_configs: List[DatasetSyncCollectionConfig] = [
+            DatasetSyncCollectionConfig(**dataset_config)
+            for dataset_config in config.get('dataset_sync', [])
+        ]
+        print(f"Found {len(self.dataset_configs)} dataset configs")
+
+    def move_new_images(self, root_dir: str):
+        raw_dir = os.path.join(root_dir, RAW_DIR)
+        new_dir = os.path.join(root_dir, NEW_DIR)
+        new_images = get_img_paths(new_dir)
+
+        for img_path in new_images:
+            # move to raw
+            new_path = os.path.join(raw_dir, os.path.basename(img_path))
+            shutil.move(img_path, new_path)
+
+        # remove new dir
+        shutil.rmtree(new_dir)
+
+    def sync_dataset(self, config: DatasetSyncCollectionConfig):
+        if config.host == 'unsplash':
+            get_images = get_unsplash_images
+        elif config.host == 'pexels':
+            get_images = get_pexels_images
+        else:
+            raise ValueError(f"Unknown host: {config.host}")
+
+        results = {
+            'num_downloaded': 0,
+            'num_skipped': 0,
+            'bad': 0,
+            'total': 0,
+        }
+
+        photos = get_images(config)
+        raw_dir = os.path.join(config.directory, RAW_DIR)
+        new_dir = os.path.join(config.directory, NEW_DIR)
+        raw_images = get_local_image_file_names(raw_dir)
+        new_images = get_local_image_file_names(new_dir)
+
+        for photo in tqdm(photos, desc=f"{config.host}-{config.collection_id}"):
+            try:
+                if photo.filename not in raw_images and photo.filename not in new_images:
+                    download_image(photo, new_dir, min_width=self.min_width, min_height=self.min_height)
+                    results['num_downloaded'] += 1
+                else:
+                    results['num_skipped'] += 1
+            except Exception as e:
+                print(f" - BAD({photo.id}): {e}")
+                results['bad'] += 1
+                continue
+            results['total'] += 1
+
+        return results
+
+    def print_results(self, results):
+        print(
+            f" - new:{results['num_downloaded']}, old:{results['num_skipped']}, bad:{results['bad']} total:{results['total']}")
+
+    def run(self):
+        super().run()
+        print(f"Syncing {len(self.dataset_configs)} datasets")
+        all_results = None
+        failed_datasets = []
+        for dataset_config in tqdm(self.dataset_configs, desc="Syncing datasets", leave=True):
+            try:
+                results = self.sync_dataset(dataset_config)
+                if all_results is None:
+                    all_results = {**results}
+                else:
+                    for key, value in results.items():
+                        all_results[key] += value
+
+                self.print_results(results)
+            except Exception as e:
+                print(f" - FAILED: {e}")
+                if 'response' in e.__dict__:
+                    error = f"{e.response.status_code}: {e.response.text}"
+                    print(f"   - {error}")
+                    failed_datasets.append({'dataset': dataset_config, 'error': error})
+                else:
+                    failed_datasets.append({'dataset': dataset_config, 'error': str(e)})
+                continue
+
+        print("Moving new images to raw")
+        for dataset_config in self.dataset_configs:
+            self.move_new_images(dataset_config.directory)
+
+        print("Done syncing datasets")
+        self.print_results(all_results)
+
+        if len(failed_datasets) > 0:
+            print(f"Failed to sync {len(failed_datasets)} datasets")
+            for failed in failed_datasets:
+                print(f" - {failed['dataset'].host}-{failed['dataset'].collection_id}")
+                print(f"   - ERR: {failed['error']}")

extensions_built_in/dataset_tools/__init__.py

@@ -0,0 +1,43 @@
+from toolkit.extension import Extension
+
+
+class DatasetToolsExtension(Extension):
+    uid = "dataset_tools"
+
+    # name is the name of the extension for printing
+    name = "Dataset Tools"
+
+    # This is where your process class is loaded
+    # keep your imports in here so they don't slow down the rest of the program
+    @classmethod
+    def get_process(cls):
+        # import your process class here so it is only loaded when needed and return it
+        from .DatasetTools import DatasetTools
+        return DatasetTools
+
+
+class SyncFromCollectionExtension(Extension):
+    uid = "sync_from_collection"
+    name = "Sync from Collection"
+
+    @classmethod
+    def get_process(cls):
+        # import your process class here so it is only loaded when needed and return it
+        from .SyncFromCollection import SyncFromCollection
+        return SyncFromCollection
+    
+    
+class SuperTaggerExtension(Extension):
+    uid = "super_tagger"
+    name = "Super Tagger"
+
+    @classmethod
+    def get_process(cls):
+        # import your process class here so it is only loaded when needed and return it
+        from .SuperTagger import SuperTagger
+        return SuperTagger
+
+
+AI_TOOLKIT_EXTENSIONS = [
+    SyncFromCollectionExtension, DatasetToolsExtension, SuperTaggerExtension
+]

extensions_built_in/dataset_tools/tools/caption.py

@@ -0,0 +1,53 @@
+
+caption_manipulation_steps = ['caption', 'caption_short']
+
+default_long_prompt = 'caption this image. describe every single thing in the image in detail. Do not include any unnecessary words in your description for the sake of good grammar. I want many short statements that serve the single purpose of giving the most thorough description if items as possible in the smallest, comma separated way possible. be sure to describe people\'s moods, clothing, the environment, lighting, colors, and everything.'
+default_short_prompt = 'caption this image in less than ten words'
+
+default_replacements = [
+    ("the image features", ""),
+    ("the image shows", ""),
+    ("the image depicts", ""),
+    ("the image is", ""),
+    ("in this image", ""),
+    ("in the image", ""),
+]
+
+
+def clean_caption(cap, replacements=None):
+    if replacements is None:
+        replacements = default_replacements
+
+    # remove any newlines
+    cap = cap.replace("\n", ", ")
+    cap = cap.replace("\r", ", ")
+    cap = cap.replace(".", ",")
+    cap = cap.replace("\"", "")
+
+    # remove unicode characters
+    cap = cap.encode('ascii', 'ignore').decode('ascii')
+
+    # make lowercase
+    cap = cap.lower()
+    # remove any extra spaces
+    cap = " ".join(cap.split())
+
+    for replacement in replacements:
+        if replacement[0].startswith('*'):
+            # we are removing all text if it starts with this and the rest matches
+            search_text = replacement[0][1:]
+            if cap.startswith(search_text):
+                cap = ""
+        else:
+            cap = cap.replace(replacement[0].lower(), replacement[1].lower())
+
+    cap_list = cap.split(",")
+    # trim whitespace
+    cap_list = [c.strip() for c in cap_list]
+    # remove empty strings
+    cap_list = [c for c in cap_list if c != ""]
+    # remove duplicates
+    cap_list = list(dict.fromkeys(cap_list))
+    # join back together
+    cap = ", ".join(cap_list)
+    return cap

extensions_built_in/dataset_tools/tools/dataset_tools_config_modules.py

@@ -0,0 +1,187 @@
+import json
+from typing import Literal, Type, TYPE_CHECKING
+
+Host: Type = Literal['unsplash', 'pexels']
+
+RAW_DIR = "raw"
+NEW_DIR = "_tmp"
+TRAIN_DIR = "train"
+DEPTH_DIR = "depth"
+
+from .image_tools import Step, img_manipulation_steps
+from .caption import caption_manipulation_steps
+
+
+class DatasetSyncCollectionConfig:
+    def __init__(self, **kwargs):
+        self.host: Host = kwargs.get('host', None)
+        self.collection_id: str = kwargs.get('collection_id', None)
+        self.directory: str = kwargs.get('directory', None)
+        self.api_key: str = kwargs.get('api_key', None)
+        self.min_width: int = kwargs.get('min_width', 1024)
+        self.min_height: int = kwargs.get('min_height', 1024)
+
+        if self.host is None:
+            raise ValueError("host is required")
+        if self.collection_id is None:
+            raise ValueError("collection_id is required")
+        if self.directory is None:
+            raise ValueError("directory is required")
+        if self.api_key is None:
+            raise ValueError(f"api_key is required: {self.host}:{self.collection_id}")
+
+
+class ImageState:
+    def __init__(self, **kwargs):
+        self.steps_complete: list[Step] = kwargs.get('steps_complete', [])
+        self.steps_to_complete: list[Step] = kwargs.get('steps_to_complete', [])
+
+    def to_dict(self):
+        return {
+            'steps_complete': self.steps_complete
+        }
+
+
+class Rect:
+    def __init__(self, **kwargs):
+        self.x = kwargs.get('x', 0)
+        self.y = kwargs.get('y', 0)
+        self.width = kwargs.get('width', 0)
+        self.height = kwargs.get('height', 0)
+
+    def to_dict(self):
+        return {
+            'x': self.x,
+            'y': self.y,
+            'width': self.width,
+            'height': self.height
+        }
+
+
+class ImgInfo:
+    def __init__(self, **kwargs):
+        self.version: int = kwargs.get('version', None)
+        self.caption: str = kwargs.get('caption', None)
+        self.caption_short: str = kwargs.get('caption_short', None)
+        self.poi = [Rect(**poi) for poi in kwargs.get('poi', [])]
+        self.state = ImageState(**kwargs.get('state', {}))
+        self.caption_method = kwargs.get('caption_method', None)
+        self.other_captions = kwargs.get('other_captions', {})
+        self._upgrade_state()
+        self.force_image_process: bool = False
+        self._requested_steps: list[Step] = []
+
+        self.is_dirty: bool = False
+
+    def _upgrade_state(self):
+        # upgrades older states
+        if self.caption is not None and 'caption' not in self.state.steps_complete:
+            self.mark_step_complete('caption')
+            self.is_dirty = True
+        if self.caption_short is not None and 'caption_short' not in self.state.steps_complete:
+            self.mark_step_complete('caption_short')
+            self.is_dirty = True
+        if self.caption_method is None and self.caption is not None:
+            # added caption method in version 2. Was all llava before that
+            self.caption_method = 'llava:default'
+            self.is_dirty = True
+
+    def to_dict(self):
+        return {
+            'version': self.version,
+            'caption_method': self.caption_method,
+            'caption': self.caption,
+            'caption_short': self.caption_short,
+            'poi': [poi.to_dict() for poi in self.poi],
+            'state': self.state.to_dict(),
+            'other_captions': self.other_captions
+        }
+
+    def mark_step_complete(self, step: Step):
+        if step not in self.state.steps_complete:
+            self.state.steps_complete.append(step)
+        if step in self.state.steps_to_complete:
+            self.state.steps_to_complete.remove(step)
+        self.is_dirty = True
+
+    def add_step(self, step: Step):
+        if step not in self.state.steps_to_complete and step not in self.state.steps_complete:
+            self.state.steps_to_complete.append(step)
+
+    def trigger_image_reprocess(self):
+        if self._requested_steps is None:
+            raise Exception("Must call add_steps before trigger_image_reprocess")
+        steps = self._requested_steps
+        # remove all image manipulationf from steps_to_complete
+        for step in img_manipulation_steps:
+            if step in self.state.steps_to_complete:
+                self.state.steps_to_complete.remove(step)
+            if step in self.state.steps_complete:
+                self.state.steps_complete.remove(step)
+        self.force_image_process = True
+        self.is_dirty = True
+        # we want to keep the order passed in process file
+        for step in steps:
+            if step in img_manipulation_steps:
+                self.add_step(step)
+
+    def add_steps(self, steps: list[Step]):
+        self._requested_steps = [step for step in steps]
+        for stage in steps:
+            self.add_step(stage)
+
+        # update steps if we have any img processes not complete, we have to reprocess them all
+        # if any steps_to_complete are in img_manipulation_steps
+
+        is_manipulating_image = any([step in img_manipulation_steps for step in self.state.steps_to_complete])
+        order_has_changed = False
+
+        if not is_manipulating_image:
+            # check to see if order has changed. No need to if already redoing it. Will detect if ones are removed
+            target_img_manipulation_order = [step for step in steps if step in img_manipulation_steps]
+            current_img_manipulation_order = [step for step in self.state.steps_complete if
+                                              step in img_manipulation_steps]
+            if target_img_manipulation_order != current_img_manipulation_order:
+                order_has_changed = True
+
+        if is_manipulating_image or order_has_changed:
+            self.trigger_image_reprocess()
+
+    def set_caption_method(self, method: str):
+        if self._requested_steps is None:
+            raise Exception("Must call add_steps before set_caption_method")
+        if self.caption_method != method:
+            self.is_dirty = True
+            # move previous caption method to other_captions
+            if self.caption_method is not None and self.caption is not None or self.caption_short is not None:
+                self.other_captions[self.caption_method] = {
+                    'caption': self.caption,
+                    'caption_short': self.caption_short,
+                }
+            self.caption_method = method
+            self.caption = None
+            self.caption_short = None
+            # see if we have a caption from the new method
+            if method in self.other_captions:
+                self.caption = self.other_captions[method].get('caption', None)
+                self.caption_short = self.other_captions[method].get('caption_short', None)
+            else:
+                self.trigger_new_caption()
+
+    def trigger_new_caption(self):
+        self.caption = None
+        self.caption_short = None
+        self.is_dirty = True
+        # check to see if we have any steps in the complete list and move them to the to_complete list
+        for step in self.state.steps_complete:
+            if step in caption_manipulation_steps:
+                self.state.steps_complete.remove(step)
+                self.state.steps_to_complete.append(step)
+
+    def to_json(self):
+        return json.dumps(self.to_dict())
+
+    def set_version(self, version: int):
+        if self.version != version:
+            self.is_dirty = True
+        self.version = version

extensions_built_in/dataset_tools/tools/fuyu_utils.py

@@ -0,0 +1,66 @@
+from transformers import  CLIPImageProcessor, BitsAndBytesConfig, AutoTokenizer
+
+from .caption import default_long_prompt, default_short_prompt, default_replacements, clean_caption
+import torch
+from PIL import Image
+
+
+class FuyuImageProcessor:
+    def __init__(self, device='cuda'):
+        from transformers import FuyuProcessor, FuyuForCausalLM
+        self.device = device
+        self.model: FuyuForCausalLM = None
+        self.processor: FuyuProcessor = None
+        self.dtype = torch.bfloat16
+        self.tokenizer: AutoTokenizer
+        self.is_loaded = False
+
+    def load_model(self):
+        from transformers import FuyuProcessor, FuyuForCausalLM
+        model_path = "adept/fuyu-8b"
+        kwargs = {"device_map": self.device}
+        kwargs['load_in_4bit'] = True
+        kwargs['quantization_config'] = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=self.dtype,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4'
+        )
+        self.processor = FuyuProcessor.from_pretrained(model_path)
+        self.model = FuyuForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+        self.is_loaded = True
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = FuyuForCausalLM.from_pretrained(model_path, torch_dtype=self.dtype, **kwargs)
+        self.processor = FuyuProcessor(image_processor=FuyuImageProcessor(), tokenizer=self.tokenizer)
+
+    def generate_caption(
+            self, image: Image,
+            prompt: str = default_long_prompt,
+            replacements=default_replacements,
+            max_new_tokens=512
+    ):
+        # prepare inputs for the model
+        # text_prompt = f"{prompt}\n"
+
+        # image = image.convert('RGB')
+        model_inputs = self.processor(text=prompt, images=[image])
+        model_inputs = {k: v.to(dtype=self.dtype if torch.is_floating_point(v) else v.dtype, device=self.device) for k, v in
+                        model_inputs.items()}
+
+        generation_output = self.model.generate(**model_inputs, max_new_tokens=max_new_tokens)
+        prompt_len = model_inputs["input_ids"].shape[-1]
+        output = self.tokenizer.decode(generation_output[0][prompt_len:], skip_special_tokens=True)
+        output = clean_caption(output, replacements=replacements)
+        return output
+
+        # inputs = self.processor(text=text_prompt, images=image, return_tensors="pt")
+        # for k, v in inputs.items():
+        #     inputs[k] = v.to(self.device)
+
+        # # autoregressively generate text
+        # generation_output = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
+        # generation_text = self.processor.batch_decode(generation_output[:, -max_new_tokens:], skip_special_tokens=True)
+        # output = generation_text[0]
+        #
+        # return clean_caption(output, replacements=replacements)

extensions_built_in/dataset_tools/tools/image_tools.py

@@ -0,0 +1,49 @@
+from typing import Literal, Type, TYPE_CHECKING, Union
+
+import cv2
+import numpy as np
+from PIL import Image, ImageOps
+
+Step: Type = Literal['caption', 'caption_short', 'create_mask', 'contrast_stretch']
+
+img_manipulation_steps = ['contrast_stretch']
+
+img_ext = ['.jpg', '.jpeg', '.png', '.webp']
+
+if TYPE_CHECKING:
+    from .llava_utils import LLaVAImageProcessor
+    from .fuyu_utils import FuyuImageProcessor
+
+ImageProcessor = Union['LLaVAImageProcessor', 'FuyuImageProcessor']
+
+
+def pil_to_cv2(image):
+    """Convert a PIL image to a cv2 image."""
+    return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+
+
+def cv2_to_pil(image):
+    """Convert a cv2 image to a PIL image."""
+    return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+
+
+def load_image(img_path: str):
+    image = Image.open(img_path).convert('RGB')
+    try:
+        # transpose with exif data
+        image = ImageOps.exif_transpose(image)
+    except Exception as e:
+        pass
+    return image
+
+
+def resize_to_max(image, max_width=1024, max_height=1024):
+    width, height = image.size
+    if width <= max_width and height <= max_height:
+        return image
+
+    scale = min(max_width / width, max_height / height)
+    width = int(width * scale)
+    height = int(height * scale)
+
+    return image.resize((width, height), Image.LANCZOS)

extensions_built_in/dataset_tools/tools/llava_utils.py

@@ -0,0 +1,85 @@
+
+from .caption import default_long_prompt, default_short_prompt, default_replacements, clean_caption
+
+import torch
+from PIL import Image, ImageOps
+
+from transformers import AutoTokenizer, BitsAndBytesConfig, CLIPImageProcessor
+
+img_ext = ['.jpg', '.jpeg', '.png', '.webp']
+
+
+class LLaVAImageProcessor:
+    def __init__(self, device='cuda'):
+        try:
+            from llava.model import LlavaLlamaForCausalLM
+        except ImportError:
+            # print("You need to manually install llava -> pip install --no-deps  git+https://github.com/haotian-liu/LLaVA.git")
+            print(
+                "You need to manually install llava -> pip install --no-deps  git+https://github.com/haotian-liu/LLaVA.git")
+            raise
+        self.device = device
+        self.model: LlavaLlamaForCausalLM = None
+        self.tokenizer: AutoTokenizer = None
+        self.image_processor: CLIPImageProcessor = None
+        self.is_loaded = False
+
+    def load_model(self):
+        from llava.model import LlavaLlamaForCausalLM
+
+        model_path = "4bit/llava-v1.5-13b-3GB"
+        # kwargs = {"device_map": "auto"}
+        kwargs = {"device_map": self.device}
+        kwargs['load_in_4bit'] = True
+        kwargs['quantization_config'] = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4'
+        )
+        self.model = LlavaLlamaForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+        vision_tower = self.model.get_vision_tower()
+        if not vision_tower.is_loaded:
+            vision_tower.load_model()
+        vision_tower.to(device=self.device)
+        self.image_processor = vision_tower.image_processor
+        self.is_loaded = True
+
+    def generate_caption(
+            self, image:
+            Image, prompt: str = default_long_prompt,
+            replacements=default_replacements,
+            max_new_tokens=512
+    ):
+        from llava.conversation import conv_templates, SeparatorStyle
+        from llava.utils import disable_torch_init
+        from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+        from llava.mm_utils import tokenizer_image_token, KeywordsStoppingCriteria
+        # question = "how many dogs are in the picture?"
+        disable_torch_init()
+        conv_mode = "llava_v0"
+        conv = conv_templates[conv_mode].copy()
+        roles = conv.roles
+        image_tensor = self.image_processor.preprocess([image], return_tensors='pt')['pixel_values'].half().cuda()
+
+        inp = f"{roles[0]}: {prompt}"
+        inp = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + inp
+        conv.append_message(conv.roles[0], inp)
+        conv.append_message(conv.roles[1], None)
+        raw_prompt = conv.get_prompt()
+        input_ids = tokenizer_image_token(raw_prompt, self.tokenizer, IMAGE_TOKEN_INDEX,
+                                          return_tensors='pt').unsqueeze(0).cuda()
+        stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
+        keywords = [stop_str]
+        stopping_criteria = KeywordsStoppingCriteria(keywords, self.tokenizer, input_ids)
+        with torch.inference_mode():
+            output_ids = self.model.generate(
+                input_ids, images=image_tensor, do_sample=True, temperature=0.1,
+                max_new_tokens=max_new_tokens, use_cache=True, stopping_criteria=[stopping_criteria],
+                top_p=0.8
+            )
+        outputs = self.tokenizer.decode(output_ids[0, input_ids.shape[1]:]).strip()
+        conv.messages[-1][-1] = outputs
+        output = outputs.rsplit('</s>', 1)[0]
+        return clean_caption(output, replacements=replacements)

extensions_built_in/dataset_tools/tools/sync_tools.py

@@ -0,0 +1,279 @@
+import os
+import requests
+import tqdm
+from typing import List, Optional, TYPE_CHECKING
+
+
+def img_root_path(img_id: str):
+    return os.path.dirname(os.path.dirname(img_id))
+
+
+if TYPE_CHECKING:
+    from .dataset_tools_config_modules import DatasetSyncCollectionConfig
+
+img_exts = ['.jpg', '.jpeg', '.webp', '.png']
+
+class Photo:
+    def __init__(
+            self,
+            id,
+            host,
+            width,
+            height,
+            url,
+            filename
+    ):
+        self.id = str(id)
+        self.host = host
+        self.width = width
+        self.height = height
+        self.url = url
+        self.filename = filename
+
+
+def get_desired_size(img_width: int, img_height: int, min_width: int, min_height: int):
+    if img_width > img_height:
+        scale = min_height / img_height
+    else:
+        scale = min_width / img_width
+
+    new_width = int(img_width * scale)
+    new_height = int(img_height * scale)
+
+    return new_width, new_height
+
+
+def get_pexels_images(config: 'DatasetSyncCollectionConfig') -> List[Photo]:
+    all_images = []
+    next_page = f"https://api.pexels.com/v1/collections/{config.collection_id}?page=1&per_page=80&type=photos"
+
+    while True:
+        response = requests.get(next_page, headers={
+            "Authorization": f"{config.api_key}"
+        })
+        response.raise_for_status()
+        data = response.json()
+        all_images.extend(data['media'])
+        if 'next_page' in data and data['next_page']:
+            next_page = data['next_page']
+        else:
+            break
+
+    photos = []
+    for image in all_images:
+        new_width, new_height = get_desired_size(image['width'], image['height'], config.min_width, config.min_height)
+        url = f"{image['src']['original']}?auto=compress&cs=tinysrgb&h={new_height}&w={new_width}"
+        filename = os.path.basename(image['src']['original'])
+
+        photos.append(Photo(
+            id=image['id'],
+            host="pexels",
+            width=image['width'],
+            height=image['height'],
+            url=url,
+            filename=filename
+        ))
+
+    return photos
+
+
+def get_unsplash_images(config: 'DatasetSyncCollectionConfig') -> List[Photo]:
+    headers = {
+        # "Authorization": f"Client-ID {UNSPLASH_ACCESS_KEY}"
+        "Authorization": f"Client-ID {config.api_key}"
+    }
+    # headers['Authorization'] = f"Bearer {token}"
+
+    url = f"https://api.unsplash.com/collections/{config.collection_id}/photos?page=1&per_page=30"
+    response = requests.get(url, headers=headers)
+    response.raise_for_status()
+    res_headers = response.headers
+    # parse the link header to get the next page
+    # 'Link': '<https://api.unsplash.com/collections/mIPWwLdfct8/photos?page=82>; rel="last", <https://api.unsplash.com/collections/mIPWwLdfct8/photos?page=2>; rel="next"'
+    has_next_page = False
+    if 'Link' in res_headers:
+        has_next_page = True
+        link_header = res_headers['Link']
+        link_header = link_header.split(',')
+        link_header = [link.strip() for link in link_header]
+        link_header = [link.split(';') for link in link_header]
+        link_header = [[link[0].strip('<>'), link[1].strip().strip('"')] for link in link_header]
+        link_header = {link[1]: link[0] for link in link_header}
+
+        # get page number from last url
+        last_page = link_header['rel="last']
+        last_page = last_page.split('?')[1]
+        last_page = last_page.split('&')
+        last_page = [param.split('=') for param in last_page]
+        last_page = {param[0]: param[1] for param in last_page}
+        last_page = int(last_page['page'])
+
+    all_images = response.json()
+
+    if has_next_page:
+        # assume we start on page 1, so we don't need to get it again
+        for page in tqdm.tqdm(range(2, last_page + 1)):
+            url = f"https://api.unsplash.com/collections/{config.collection_id}/photos?page={page}&per_page=30"
+            response = requests.get(url, headers=headers)
+            response.raise_for_status()
+            all_images.extend(response.json())
+
+    photos = []
+    for image in all_images:
+        new_width, new_height = get_desired_size(image['width'], image['height'], config.min_width, config.min_height)
+        url = f"{image['urls']['raw']}&w={new_width}"
+        filename = f"{image['id']}.jpg"
+
+        photos.append(Photo(
+            id=image['id'],
+            host="unsplash",
+            width=image['width'],
+            height=image['height'],
+            url=url,
+            filename=filename
+        ))
+
+    return photos
+
+
+def get_img_paths(dir_path: str):
+    os.makedirs(dir_path, exist_ok=True)
+    local_files = os.listdir(dir_path)
+    # remove non image files
+    local_files = [file for file in local_files if os.path.splitext(file)[1].lower() in img_exts]
+    # make full path
+    local_files = [os.path.join(dir_path, file) for file in local_files]
+    return local_files
+
+
+def get_local_image_ids(dir_path: str):
+    os.makedirs(dir_path, exist_ok=True)
+    local_files = get_img_paths(dir_path)
+    # assuming local files are named after Unsplash IDs, e.g., 'abc123.jpg'
+    return set([os.path.basename(file).split('.')[0] for file in local_files])
+
+
+def get_local_image_file_names(dir_path: str):
+    os.makedirs(dir_path, exist_ok=True)
+    local_files = get_img_paths(dir_path)
+    # assuming local files are named after Unsplash IDs, e.g., 'abc123.jpg'
+    return set([os.path.basename(file) for file in local_files])
+
+
+def download_image(photo: Photo, dir_path: str, min_width: int = 1024, min_height: int = 1024):
+    img_width = photo.width
+    img_height = photo.height
+
+    if img_width < min_width or img_height < min_height:
+        raise ValueError(f"Skipping {photo.id} because it is too small: {img_width}x{img_height}")
+
+    img_response = requests.get(photo.url)
+    img_response.raise_for_status()
+    os.makedirs(dir_path, exist_ok=True)
+
+    filename = os.path.join(dir_path, photo.filename)
+    with open(filename, 'wb') as file:
+        file.write(img_response.content)
+
+
+def update_caption(img_path: str):
+    # if the caption is a txt file, convert it to a json file
+    filename_no_ext = os.path.splitext(os.path.basename(img_path))[0]
+    # see if it exists
+    if os.path.exists(os.path.join(os.path.dirname(img_path), f"{filename_no_ext}.json")):
+        # todo add poi and what not
+        return  # we have a json file
+    caption = ""
+    # see if txt file exists
+    if os.path.exists(os.path.join(os.path.dirname(img_path), f"{filename_no_ext}.txt")):
+        # read it
+        with open(os.path.join(os.path.dirname(img_path), f"{filename_no_ext}.txt"), 'r') as file:
+            caption = file.read()
+    # write json file
+    with open(os.path.join(os.path.dirname(img_path), f"{filename_no_ext}.json"), 'w') as file:
+        file.write(f'{{"caption": "{caption}"}}')
+
+    # delete txt file
+    os.remove(os.path.join(os.path.dirname(img_path), f"{filename_no_ext}.txt"))
+
+
+# def equalize_img(img_path: str):
+#     input_path = img_path
+#     output_path = os.path.join(img_root_path(img_path), COLOR_CORRECTED_DIR, os.path.basename(img_path))
+#     os.makedirs(os.path.dirname(output_path), exist_ok=True)
+#     process_img(
+#         img_path=input_path,
+#         output_path=output_path,
+#         equalize=True,
+#         max_size=2056,
+#         white_balance=False,
+#         gamma_correction=False,
+#         strength=0.6,
+#     )
+
+
+# def annotate_depth(img_path: str):
+#     # make fake args
+#     args = argparse.Namespace()
+#     args.annotator = "midas"
+#     args.res = 1024
+#
+#     img = cv2.imread(img_path)
+#     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+#
+#     output = annotate(img, args)
+#
+#     output = output.astype('uint8')
+#     output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
+#
+#     os.makedirs(os.path.dirname(img_path), exist_ok=True)
+#     output_path = os.path.join(img_root_path(img_path), DEPTH_DIR, os.path.basename(img_path))
+#
+#     cv2.imwrite(output_path, output)
+
+
+# def invert_depth(img_path: str):
+#     img = cv2.imread(img_path)
+#     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+#     # invert the colors
+#     img = cv2.bitwise_not(img)
+#
+#     os.makedirs(os.path.dirname(img_path), exist_ok=True)
+#     output_path = os.path.join(img_root_path(img_path), INVERTED_DEPTH_DIR, os.path.basename(img_path))
+#     cv2.imwrite(output_path, img)
+
+
+    #
+    # # update our list of raw images
+    # raw_images = get_img_paths(raw_dir)
+    #
+    # # update raw captions
+    # for image_id in tqdm.tqdm(raw_images, desc="Updating raw captions"):
+    #     update_caption(image_id)
+    #
+    # # equalize images
+    # for img_path in tqdm.tqdm(raw_images, desc="Equalizing images"):
+    #     if img_path not in eq_images:
+    #         equalize_img(img_path)
+    #
+    # # update our list of eq images
+    # eq_images = get_img_paths(eq_dir)
+    # # update eq captions
+    # for image_id in tqdm.tqdm(eq_images, desc="Updating eq captions"):
+    #     update_caption(image_id)
+    #
+    # # annotate depth
+    # depth_dir = os.path.join(root_dir, DEPTH_DIR)
+    # depth_images = get_img_paths(depth_dir)
+    # for img_path in tqdm.tqdm(eq_images, desc="Annotating depth"):
+    #     if img_path not in depth_images:
+    #         annotate_depth(img_path)
+    #
+    # depth_images = get_img_paths(depth_dir)
+    #
+    # # invert depth
+    # inv_depth_dir = os.path.join(root_dir, INVERTED_DEPTH_DIR)
+    # inv_depth_images = get_img_paths(inv_depth_dir)
+    # for img_path in tqdm.tqdm(depth_images, desc="Inverting depth"):
+    #     if img_path not in inv_depth_images:
+    #         invert_depth(img_path)

extensions_built_in/image_reference_slider_trainer/ImageReferenceSliderTrainerProcess.py

@@ -0,0 +1,235 @@
+import copy
+import random
+from collections import OrderedDict
+import os
+from contextlib import nullcontext
+from typing import Optional, Union, List
+from torch.utils.data import ConcatDataset, DataLoader
+
+from toolkit.config_modules import ReferenceDatasetConfig
+from toolkit.data_loader import PairedImageDataset
+from toolkit.prompt_utils import concat_prompt_embeds, split_prompt_embeds
+from toolkit.stable_diffusion_model import StableDiffusion, PromptEmbeds
+from toolkit.train_tools import get_torch_dtype, apply_snr_weight
+import gc
+from toolkit import train_tools
+import torch
+from jobs.process import BaseSDTrainProcess
+import random
+from toolkit.basic import value_map
+
+
+def flush():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+class ReferenceSliderConfig:
+    def __init__(self, **kwargs):
+        self.additional_losses: List[str] = kwargs.get('additional_losses', [])
+        self.weight_jitter: float = kwargs.get('weight_jitter', 0.0)
+        self.datasets: List[ReferenceDatasetConfig] = [ReferenceDatasetConfig(**d) for d in kwargs.get('datasets', [])]
+
+
+class ImageReferenceSliderTrainerProcess(BaseSDTrainProcess):
+    sd: StableDiffusion
+    data_loader: DataLoader = None
+
+    def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
+        super().__init__(process_id, job, config, **kwargs)
+        self.prompt_txt_list = None
+        self.step_num = 0
+        self.start_step = 0
+        self.device = self.get_conf('device', self.job.device)
+        self.device_torch = torch.device(self.device)
+        self.slider_config = ReferenceSliderConfig(**self.get_conf('slider', {}))
+
+    def load_datasets(self):
+        if self.data_loader is None:
+            print(f"Loading datasets")
+            datasets = []
+            for dataset in self.slider_config.datasets:
+                print(f" - Dataset: {dataset.pair_folder}")
+                config = {
+                    'path': dataset.pair_folder,
+                    'size': dataset.size,
+                    'default_prompt': dataset.target_class,
+                    'network_weight': dataset.network_weight,
+                    'pos_weight': dataset.pos_weight,
+                    'neg_weight': dataset.neg_weight,
+                    'pos_folder': dataset.pos_folder,
+                    'neg_folder': dataset.neg_folder,
+                }
+                image_dataset = PairedImageDataset(config)
+                datasets.append(image_dataset)
+
+            concatenated_dataset = ConcatDataset(datasets)
+            self.data_loader = DataLoader(
+                concatenated_dataset,
+                batch_size=self.train_config.batch_size,
+                shuffle=True,
+                num_workers=2
+            )
+
+    def before_model_load(self):
+        pass
+
+    def hook_before_train_loop(self):
+        self.sd.vae.eval()
+        self.sd.vae.to(self.device_torch)
+        self.load_datasets()
+
+        pass
+
+    def hook_train_loop(self, batch):
+        with torch.no_grad():
+            imgs, prompts, network_weights = batch
+            network_pos_weight, network_neg_weight = network_weights
+
+            if isinstance(network_pos_weight, torch.Tensor):
+                network_pos_weight = network_pos_weight.item()
+            if isinstance(network_neg_weight, torch.Tensor):
+                network_neg_weight = network_neg_weight.item()
+
+            # get an array of random floats between -weight_jitter and weight_jitter
+            loss_jitter_multiplier = 1.0
+            weight_jitter = self.slider_config.weight_jitter
+            if weight_jitter > 0.0:
+                jitter_list = random.uniform(-weight_jitter, weight_jitter)
+                orig_network_pos_weight = network_pos_weight
+                network_pos_weight += jitter_list
+                network_neg_weight += (jitter_list * -1.0)
+                # penalize the loss for its distance from network_pos_weight
+                # a jitter_list of abs(3.0) on a weight of 5.0 is a 60% jitter
+                # so the loss_jitter_multiplier needs to be 0.4
+                loss_jitter_multiplier = value_map(abs(jitter_list), 0.0, weight_jitter, 1.0, 0.0)
+
+
+            # if items in network_weight list are tensors, convert them to floats
+
+            dtype = get_torch_dtype(self.train_config.dtype)
+            imgs: torch.Tensor = imgs.to(self.device_torch, dtype=dtype)
+            # split batched images in half so left is negative and right is positive
+            negative_images, positive_images = torch.chunk(imgs, 2, dim=3)
+
+            positive_latents = self.sd.encode_images(positive_images)
+            negative_latents = self.sd.encode_images(negative_images)
+
+            height = positive_images.shape[2]
+            width = positive_images.shape[3]
+            batch_size = positive_images.shape[0]
+
+            if self.train_config.gradient_checkpointing:
+                # may get disabled elsewhere
+                self.sd.unet.enable_gradient_checkpointing()
+
+            noise_scheduler = self.sd.noise_scheduler
+            optimizer = self.optimizer
+            lr_scheduler = self.lr_scheduler
+
+            self.sd.noise_scheduler.set_timesteps(
+                self.train_config.max_denoising_steps, device=self.device_torch
+            )
+
+            timesteps = torch.randint(0, self.train_config.max_denoising_steps, (1,), device=self.device_torch)
+            timesteps = timesteps.long()
+
+            # get noise
+            noise_positive = self.sd.get_latent_noise(
+                pixel_height=height,
+                pixel_width=width,
+                batch_size=batch_size,
+                noise_offset=self.train_config.noise_offset,
+            ).to(self.device_torch, dtype=dtype)
+
+            noise_negative = noise_positive.clone()
+
+            # Add noise to the latents according to the noise magnitude at each timestep
+            # (this is the forward diffusion process)
+            noisy_positive_latents = noise_scheduler.add_noise(positive_latents, noise_positive, timesteps)
+            noisy_negative_latents = noise_scheduler.add_noise(negative_latents, noise_negative, timesteps)
+
+            noisy_latents = torch.cat([noisy_positive_latents, noisy_negative_latents], dim=0)
+            noise = torch.cat([noise_positive, noise_negative], dim=0)
+            timesteps = torch.cat([timesteps, timesteps], dim=0)
+            network_multiplier = [network_pos_weight * 1.0, network_neg_weight * -1.0]
+
+        self.optimizer.zero_grad()
+        noisy_latents.requires_grad = False
+
+        # if training text encoder enable grads, else do context of no grad
+        with torch.set_grad_enabled(self.train_config.train_text_encoder):
+            # fix issue with them being tuples sometimes
+            prompt_list = []
+            for prompt in prompts:
+                if isinstance(prompt, tuple):
+                    prompt = prompt[0]
+                prompt_list.append(prompt)
+            conditional_embeds = self.sd.encode_prompt(prompt_list).to(self.device_torch, dtype=dtype)
+            conditional_embeds = concat_prompt_embeds([conditional_embeds, conditional_embeds])
+
+        # if self.model_config.is_xl:
+        #     # todo also allow for setting this for low ram in general, but sdxl spikes a ton on back prop
+        #     network_multiplier_list = network_multiplier
+        #     noisy_latent_list = torch.chunk(noisy_latents, 2, dim=0)
+        #     noise_list = torch.chunk(noise, 2, dim=0)
+        #     timesteps_list = torch.chunk(timesteps, 2, dim=0)
+        #     conditional_embeds_list = split_prompt_embeds(conditional_embeds)
+        # else:
+        network_multiplier_list = [network_multiplier]
+        noisy_latent_list = [noisy_latents]
+        noise_list = [noise]
+        timesteps_list = [timesteps]
+        conditional_embeds_list = [conditional_embeds]
+
+        losses = []
+        # allow to chunk it out to save vram
+        for network_multiplier, noisy_latents, noise, timesteps, conditional_embeds in zip(
+                network_multiplier_list, noisy_latent_list, noise_list, timesteps_list, conditional_embeds_list
+        ):
+            with self.network:
+                assert self.network.is_active
+
+                self.network.multiplier = network_multiplier
+
+                noise_pred = self.sd.predict_noise(
+                    latents=noisy_latents.to(self.device_torch, dtype=dtype),
+                    conditional_embeddings=conditional_embeds.to(self.device_torch, dtype=dtype),
+                    timestep=timesteps,
+                )
+                noise = noise.to(self.device_torch, dtype=dtype)
+
+                if self.sd.prediction_type == 'v_prediction':
+                    # v-parameterization training
+                    target = noise_scheduler.get_velocity(noisy_latents, noise, timesteps)
+                else:
+                    target = noise
+
+                loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="none")
+                loss = loss.mean([1, 2, 3])
+
+                if self.train_config.min_snr_gamma is not None and self.train_config.min_snr_gamma > 0.000001:
+                    # add min_snr_gamma
+                    loss = apply_snr_weight(loss, timesteps, noise_scheduler, self.train_config.min_snr_gamma)
+
+                loss = loss.mean() * loss_jitter_multiplier
+
+                loss_float = loss.item()
+                losses.append(loss_float)
+
+                # back propagate loss to free ram
+                loss.backward()
+
+        # apply gradients
+        optimizer.step()
+        lr_scheduler.step()
+
+        # reset network
+        self.network.multiplier = 1.0
+
+        loss_dict = OrderedDict(
+            {'loss': sum(losses) / len(losses) if len(losses) > 0 else 0.0}
+        )
+
+        return loss_dict
+        # end hook_train_loop

extensions_built_in/image_reference_slider_trainer/__init__.py

@@ -0,0 +1,25 @@
+# This is an example extension for custom training. It is great for experimenting with new ideas.
+from toolkit.extension import Extension
+
+
+# We make a subclass of Extension
+class ImageReferenceSliderTrainer(Extension):
+    # uid must be unique, it is how the extension is identified
+    uid = "image_reference_slider_trainer"
+
+    # name is the name of the extension for printing
+    name = "Image Reference Slider Trainer"
+
+    # This is where your process class is loaded
+    # keep your imports in here so they don't slow down the rest of the program
+    @classmethod
+    def get_process(cls):
+        # import your process class here so it is only loaded when needed and return it
+        from .ImageReferenceSliderTrainerProcess import ImageReferenceSliderTrainerProcess
+        return ImageReferenceSliderTrainerProcess
+
+
+AI_TOOLKIT_EXTENSIONS = [
+    # you can put a list of extensions here
+    ImageReferenceSliderTrainer
+]

extensions_built_in/image_reference_slider_trainer/config/train.example.yaml

@@ -0,0 +1,107 @@
+---
+job: extension
+config:
+  name: example_name
+  process:
+    - type: 'image_reference_slider_trainer'
+      training_folder: "/mnt/Train/out/LoRA"
+      device: cuda:0
+      # for tensorboard logging
+      log_dir: "/home/jaret/Dev/.tensorboard"
+      network:
+        type: "lora"
+        linear: 8
+        linear_alpha: 8
+      train:
+        noise_scheduler: "ddpm" # or "ddpm", "lms", "euler_a"
+        steps: 5000
+        lr: 1e-4
+        train_unet: true
+        gradient_checkpointing: true
+        train_text_encoder: true
+        optimizer: "adamw"
+        optimizer_params:
+          weight_decay: 1e-2
+        lr_scheduler: "constant"
+        max_denoising_steps: 1000
+        batch_size: 1
+        dtype: bf16
+        xformers: true
+        skip_first_sample: true
+        noise_offset: 0.0
+      model:
+        name_or_path: "/path/to/model.safetensors"
+        is_v2: false  # for v2 models
+        is_xl: false  # for SDXL models
+        is_v_pred: false # for v-prediction models (most v2 models)
+      save:
+        dtype: float16 # precision to save
+        save_every: 1000 # save every this many steps
+        max_step_saves_to_keep: 2 # only affects step counts
+      sample:
+        sampler: "ddpm" # must match train.noise_scheduler
+        sample_every: 100 # sample every this many steps
+        width: 512
+        height: 512
+        prompts:
+          - "photo of a woman with red hair taking a selfie --m -3"
+          - "photo of a woman with red hair taking a selfie --m -1"
+          - "photo of a woman with red hair taking a selfie --m 1"
+          - "photo of a woman with red hair taking a selfie --m 3"
+          - "close up photo of a man smiling at the camera, in a tank top --m -3"
+          - "close up photo of a man smiling at the camera, in a tank top--m -1"
+          - "close up photo of a man smiling at the camera, in a tank top --m 1"
+          - "close up photo of a man smiling at the camera, in a tank top --m 3"
+          - "photo of a blonde woman smiling, barista --m -3"
+          - "photo of a blonde woman smiling, barista --m -1"
+          - "photo of a blonde woman smiling, barista --m 1"
+          - "photo of a blonde woman smiling, barista --m 3"
+          - "photo of a Christina Hendricks --m -1"
+          - "photo of a Christina Hendricks --m -1"
+          - "photo of a Christina Hendricks --m 1"
+          - "photo of a Christina Hendricks --m 3"
+          - "photo of a Christina Ricci --m -3"
+          - "photo of a Christina Ricci --m -1"
+          - "photo of a Christina Ricci --m 1"
+          - "photo of a Christina Ricci --m 3"
+        neg: "cartoon, fake, drawing, illustration, cgi, animated, anime"
+        seed: 42
+        walk_seed: false
+        guidance_scale: 7
+        sample_steps: 20
+        network_multiplier: 1.0
+
+      logging:
+        log_every: 10 # log every this many steps
+        use_wandb: false # not supported yet
+        verbose: false
+
+      slider:
+        datasets:
+          - pair_folder: "/path/to/folder/side/by/side/images"
+            network_weight: 2.0
+            target_class: "" # only used as default if caption txt are not present
+            size: 512
+          - pair_folder: "/path/to/folder/side/by/side/images"
+            network_weight: 4.0
+            target_class: "" # only used as default if caption txt are not present
+            size: 512
+
+
+# you can put any information you want here, and it will be saved in the model
+# the below is an example. I recommend doing trigger words at a minimum
+# in the metadata. The software will include this plus some other information
+meta:
+  name: "[name]"  # [name] gets replaced with the name above
+  description: A short description of your model
+  trigger_words:
+  - put
+  - trigger
+  - words
+  - here
+  version: '0.1'
+  creator:
+    name: Your Name
+    email: [email protected]
+    website: https://yourwebsite.com
+  any: All meta data above is arbitrary, it can be whatever you want.

extensions_built_in/sd_trainer/SDTrainer.py

@@ -0,0 +1,1679 @@
+import os
+import random
+from collections import OrderedDict
+from typing import Union, Literal, List, Optional
+
+import numpy as np
+from diffusers import T2IAdapter, AutoencoderTiny, ControlNetModel
+
+import torch.functional as F
+from safetensors.torch import load_file
+from torch.utils.data import DataLoader, ConcatDataset
+
+from toolkit import train_tools
+from toolkit.basic import value_map, adain, get_mean_std
+from toolkit.clip_vision_adapter import ClipVisionAdapter
+from toolkit.config_modules import GuidanceConfig
+from toolkit.data_loader import get_dataloader_datasets
+from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO, FileItemDTO
+from toolkit.guidance import get_targeted_guidance_loss, get_guidance_loss, GuidanceType
+from toolkit.image_utils import show_tensors, show_latents
+from toolkit.ip_adapter import IPAdapter
+from toolkit.custom_adapter import CustomAdapter
+from toolkit.prompt_utils import PromptEmbeds, concat_prompt_embeds
+from toolkit.reference_adapter import ReferenceAdapter
+from toolkit.stable_diffusion_model import StableDiffusion, BlankNetwork
+from toolkit.train_tools import get_torch_dtype, apply_snr_weight, add_all_snr_to_noise_scheduler, \
+    apply_learnable_snr_gos, LearnableSNRGamma
+import gc
+import torch
+from jobs.process import BaseSDTrainProcess
+from torchvision import transforms
+from diffusers import EMAModel
+import math
+from toolkit.train_tools import precondition_model_outputs_flow_match
+
+
+def flush():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+adapter_transforms = transforms.Compose([
+    transforms.ToTensor(),
+])
+
+
+class SDTrainer(BaseSDTrainProcess):
+
+    def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
+        super().__init__(process_id, job, config, **kwargs)
+        self.assistant_adapter: Union['T2IAdapter', 'ControlNetModel', None]
+        self.do_prior_prediction = False
+        self.do_long_prompts = False
+        self.do_guided_loss = False
+        self.taesd: Optional[AutoencoderTiny] = None
+
+        self._clip_image_embeds_unconditional: Union[List[str], None] = None
+        self.negative_prompt_pool: Union[List[str], None] = None
+        self.batch_negative_prompt: Union[List[str], None] = None
+
+        self.scaler = torch.cuda.amp.GradScaler()
+
+        self.is_bfloat = self.train_config.dtype == "bfloat16" or self.train_config.dtype == "bf16"
+
+        self.do_grad_scale = True
+        if self.is_fine_tuning and self.is_bfloat:
+            self.do_grad_scale = False
+        if self.adapter_config is not None:
+            if self.adapter_config.train:
+                self.do_grad_scale = False
+
+        if self.train_config.dtype in ["fp16", "float16"]:
+            # patch the scaler to allow fp16 training
+            org_unscale_grads = self.scaler._unscale_grads_
+            def _unscale_grads_replacer(optimizer, inv_scale, found_inf, allow_fp16):
+                return org_unscale_grads(optimizer, inv_scale, found_inf, True)
+            self.scaler._unscale_grads_ = _unscale_grads_replacer
+
+        self.cached_blank_embeds: Optional[PromptEmbeds] = None
+        self.cached_trigger_embeds: Optional[PromptEmbeds] = None
+
+
+    def before_model_load(self):
+        pass
+
+    def before_dataset_load(self):
+        self.assistant_adapter = None
+        # get adapter assistant if one is set
+        if self.train_config.adapter_assist_name_or_path is not None:
+            adapter_path = self.train_config.adapter_assist_name_or_path
+
+            if self.train_config.adapter_assist_type == "t2i":
+                # dont name this adapter since we are not training it
+                self.assistant_adapter = T2IAdapter.from_pretrained(
+                    adapter_path, torch_dtype=get_torch_dtype(self.train_config.dtype)
+                ).to(self.device_torch)
+            elif self.train_config.adapter_assist_type == "control_net":
+                self.assistant_adapter = ControlNetModel.from_pretrained(
+                    adapter_path, torch_dtype=get_torch_dtype(self.train_config.dtype)
+                ).to(self.device_torch, dtype=get_torch_dtype(self.train_config.dtype))
+            else:
+                raise ValueError(f"Unknown adapter assist type {self.train_config.adapter_assist_type}")
+
+            self.assistant_adapter.eval()
+            self.assistant_adapter.requires_grad_(False)
+            flush()
+        if self.train_config.train_turbo and self.train_config.show_turbo_outputs:
+            if self.model_config.is_xl:
+                self.taesd = AutoencoderTiny.from_pretrained("madebyollin/taesdxl",
+                                                             torch_dtype=get_torch_dtype(self.train_config.dtype))
+            else:
+                self.taesd = AutoencoderTiny.from_pretrained("madebyollin/taesd",
+                                                             torch_dtype=get_torch_dtype(self.train_config.dtype))
+            self.taesd.to(dtype=get_torch_dtype(self.train_config.dtype), device=self.device_torch)
+            self.taesd.eval()
+            self.taesd.requires_grad_(False)
+
+    def hook_before_train_loop(self):
+        super().hook_before_train_loop()
+        
+        if self.train_config.do_prior_divergence:
+            self.do_prior_prediction = True
+        # move vae to device if we did not cache latents
+        if not self.is_latents_cached:
+            self.sd.vae.eval()
+            self.sd.vae.to(self.device_torch)
+        else:
+            # offload it. Already cached
+            self.sd.vae.to('cpu')
+            flush()
+        add_all_snr_to_noise_scheduler(self.sd.noise_scheduler, self.device_torch)
+        if self.adapter is not None:
+            self.adapter.to(self.device_torch)
+
+            # check if we have regs and using adapter and caching clip embeddings
+            has_reg = self.datasets_reg is not None and len(self.datasets_reg) > 0
+            is_caching_clip_embeddings = self.datasets is not None and any([self.datasets[i].cache_clip_vision_to_disk for i in range(len(self.datasets))])
+
+            if has_reg and is_caching_clip_embeddings:
+                # we need a list of unconditional clip image embeds from other datasets to handle regs
+                unconditional_clip_image_embeds = []
+                datasets = get_dataloader_datasets(self.data_loader)
+                for i in range(len(datasets)):
+                    unconditional_clip_image_embeds += datasets[i].clip_vision_unconditional_cache
+
+                if len(unconditional_clip_image_embeds) == 0:
+                    raise ValueError("No unconditional clip image embeds found. This should not happen")
+
+                self._clip_image_embeds_unconditional = unconditional_clip_image_embeds
+
+        if self.train_config.negative_prompt is not None:
+            if os.path.exists(self.train_config.negative_prompt):
+                with open(self.train_config.negative_prompt, 'r') as f:
+                    self.negative_prompt_pool = f.readlines()
+                    # remove empty
+                    self.negative_prompt_pool = [x.strip() for x in self.negative_prompt_pool if x.strip() != ""]
+            else:
+                # single prompt
+                self.negative_prompt_pool = [self.train_config.negative_prompt]
+
+        # handle unload text encoder
+        if self.train_config.unload_text_encoder:
+            with torch.no_grad():
+                if self.train_config.train_text_encoder:
+                    raise ValueError("Cannot unload text encoder if training text encoder")
+                # cache embeddings
+
+                print("\n***** UNLOADING TEXT ENCODER *****")
+                print("This will train only with a blank prompt or trigger word, if set")
+                print("If this is not what you want, remove the unload_text_encoder flag")
+                print("***********************************")
+                print("")
+                self.sd.text_encoder_to(self.device_torch)
+                self.cached_blank_embeds = self.sd.encode_prompt("")
+                if self.trigger_word is not None:
+                    self.cached_trigger_embeds = self.sd.encode_prompt(self.trigger_word)
+
+                # move back to cpu
+                self.sd.text_encoder_to('cpu')
+                flush()
+
+
+    def process_output_for_turbo(self, pred, noisy_latents, timesteps, noise, batch):
+        # to process turbo learning, we make one big step from our current timestep to the end
+        # we then denoise the prediction on that remaining step and target our loss to our target latents
+        # this currently only works on euler_a (that I know of). Would work on others, but needs to be coded to do so.
+        # needs to be done on each item in batch as they may all have different timesteps
+        batch_size = pred.shape[0]
+        pred_chunks = torch.chunk(pred, batch_size, dim=0)
+        noisy_latents_chunks = torch.chunk(noisy_latents, batch_size, dim=0)
+        timesteps_chunks = torch.chunk(timesteps, batch_size, dim=0)
+        latent_chunks = torch.chunk(batch.latents, batch_size, dim=0)
+        noise_chunks = torch.chunk(noise, batch_size, dim=0)
+
+        with torch.no_grad():
+            # set the timesteps to 1000 so we can capture them to calculate the sigmas
+            self.sd.noise_scheduler.set_timesteps(
+                self.sd.noise_scheduler.config.num_train_timesteps,
+                device=self.device_torch
+            )
+            train_timesteps = self.sd.noise_scheduler.timesteps.clone().detach()
+
+            train_sigmas = self.sd.noise_scheduler.sigmas.clone().detach()
+
+            # set the scheduler to one timestep, we build the step and sigmas for each item in batch for the partial step
+            self.sd.noise_scheduler.set_timesteps(
+                1,
+                device=self.device_torch
+            )
+
+        denoised_pred_chunks = []
+        target_pred_chunks = []
+
+        for i in range(batch_size):
+            pred_item = pred_chunks[i]
+            noisy_latents_item = noisy_latents_chunks[i]
+            timesteps_item = timesteps_chunks[i]
+            latents_item = latent_chunks[i]
+            noise_item = noise_chunks[i]
+            with torch.no_grad():
+                timestep_idx = [(train_timesteps == t).nonzero().item() for t in timesteps_item][0]
+                single_step_timestep_schedule = [timesteps_item.squeeze().item()]
+                # extract the sigma idx for our midpoint timestep
+                sigmas = train_sigmas[timestep_idx:timestep_idx + 1].to(self.device_torch)
+
+                end_sigma_idx = random.randint(timestep_idx, len(train_sigmas) - 1)
+                end_sigma = train_sigmas[end_sigma_idx:end_sigma_idx + 1].to(self.device_torch)
+
+                # add noise to our target
+
+                # build the big sigma step. The to step will now be to 0 giving it a full remaining denoising half step
+                # self.sd.noise_scheduler.sigmas = torch.cat([sigmas, torch.zeros_like(sigmas)]).detach()
+                self.sd.noise_scheduler.sigmas = torch.cat([sigmas, end_sigma]).detach()
+                # set our single timstep
+                self.sd.noise_scheduler.timesteps = torch.from_numpy(
+                    np.array(single_step_timestep_schedule, dtype=np.float32)
+                ).to(device=self.device_torch)
+
+                # set the step index to None so it will be recalculated on first step
+                self.sd.noise_scheduler._step_index = None
+
+            denoised_latent = self.sd.noise_scheduler.step(
+                pred_item, timesteps_item, noisy_latents_item.detach(), return_dict=False
+            )[0]
+
+            residual_noise = (noise_item * end_sigma.flatten()).detach().to(self.device_torch, dtype=get_torch_dtype(
+                self.train_config.dtype))
+            # remove the residual noise from the denoised latents. Output should be a clean prediction (theoretically)
+            denoised_latent = denoised_latent - residual_noise
+
+            denoised_pred_chunks.append(denoised_latent)
+
+        denoised_latents = torch.cat(denoised_pred_chunks, dim=0)
+        # set the scheduler back to the original timesteps
+        self.sd.noise_scheduler.set_timesteps(
+            self.sd.noise_scheduler.config.num_train_timesteps,
+            device=self.device_torch
+        )
+
+        output = denoised_latents / self.sd.vae.config['scaling_factor']
+        output = self.sd.vae.decode(output).sample
+
+        if self.train_config.show_turbo_outputs:
+            # since we are completely denoising, we can show them here
+            with torch.no_grad():
+                show_tensors(output)
+
+        # we return our big partial step denoised latents as our pred and our untouched latents as our target.
+        # you can do mse against the two here  or run the denoised through the vae for pixel space loss against the
+        # input tensor images.
+
+        return output, batch.tensor.to(self.device_torch, dtype=get_torch_dtype(self.train_config.dtype))
+
+    # you can expand these in a child class to make customization easier
+    def calculate_loss(
+            self,
+            noise_pred: torch.Tensor,
+            noise: torch.Tensor,
+            noisy_latents: torch.Tensor,
+            timesteps: torch.Tensor,
+            batch: 'DataLoaderBatchDTO',
+            mask_multiplier: Union[torch.Tensor, float] = 1.0,
+            prior_pred: Union[torch.Tensor, None] = None,
+            **kwargs
+    ):
+        loss_target = self.train_config.loss_target
+        is_reg = any(batch.get_is_reg_list())
+
+        prior_mask_multiplier = None
+        target_mask_multiplier = None
+        dtype = get_torch_dtype(self.train_config.dtype)
+
+        has_mask = batch.mask_tensor is not None
+
+        with torch.no_grad():
+            loss_multiplier = torch.tensor(batch.loss_multiplier_list).to(self.device_torch, dtype=torch.float32)
+
+        if self.train_config.match_noise_norm:
+            # match the norm of the noise
+            noise_norm = torch.linalg.vector_norm(noise, ord=2, dim=(1, 2, 3), keepdim=True)
+            noise_pred_norm = torch.linalg.vector_norm(noise_pred, ord=2, dim=(1, 2, 3), keepdim=True)
+            noise_pred = noise_pred * (noise_norm / noise_pred_norm)
+
+        if self.train_config.pred_scaler != 1.0:
+            noise_pred = noise_pred * self.train_config.pred_scaler
+
+        target = None
+
+        if self.train_config.target_noise_multiplier != 1.0:
+            noise = noise * self.train_config.target_noise_multiplier
+
+        if self.train_config.correct_pred_norm or (self.train_config.inverted_mask_prior and prior_pred is not None and has_mask):
+            if self.train_config.correct_pred_norm and not is_reg:
+                with torch.no_grad():
+                    # this only works if doing a prior pred
+                    if prior_pred is not None:
+                        prior_mean = prior_pred.mean([2,3], keepdim=True)
+                        prior_std = prior_pred.std([2,3], keepdim=True)
+                        noise_mean = noise_pred.mean([2,3], keepdim=True)
+                        noise_std = noise_pred.std([2,3], keepdim=True)
+
+                        mean_adjust = prior_mean - noise_mean
+                        std_adjust = prior_std - noise_std
+
+                        mean_adjust = mean_adjust * self.train_config.correct_pred_norm_multiplier
+                        std_adjust = std_adjust * self.train_config.correct_pred_norm_multiplier
+
+                        target_mean = noise_mean + mean_adjust
+                        target_std = noise_std + std_adjust
+
+                        eps = 1e-5
+                        # match the noise to the prior
+                        noise = (noise - noise_mean) / (noise_std + eps)
+                        noise = noise * (target_std + eps) + target_mean
+                        noise = noise.detach()
+
+            if self.train_config.inverted_mask_prior and prior_pred is not None and has_mask:
+                assert not self.train_config.train_turbo
+                with torch.no_grad():
+                    # we need to make the noise prediction be a masked blending of noise and prior_pred
+                    stretched_mask_multiplier = value_map(
+                        mask_multiplier,
+                        batch.file_items[0].dataset_config.mask_min_value,
+                        1.0,
+                        0.0,
+                        1.0
+                    )
+
+                    prior_mask_multiplier = 1.0 - stretched_mask_multiplier
+
+
+                # target_mask_multiplier = mask_multiplier
+                # mask_multiplier = 1.0
+                target = noise
+                # target = (noise * mask_multiplier) + (prior_pred * prior_mask_multiplier)
+                # set masked multiplier to 1.0 so we dont double apply it
+                # mask_multiplier = 1.0
+        elif prior_pred is not None and not self.train_config.do_prior_divergence:
+            assert not self.train_config.train_turbo
+            # matching adapter prediction
+            target = prior_pred
+        elif self.sd.prediction_type == 'v_prediction':
+            # v-parameterization training
+            target = self.sd.noise_scheduler.get_velocity(batch.tensor, noise, timesteps)
+
+        elif self.sd.is_flow_matching:
+            target = (noise - batch.latents).detach()
+        else:
+            target = noise
+
+        if target is None:
+            target = noise
+
+        pred = noise_pred
+
+        if self.train_config.train_turbo:
+            pred, target = self.process_output_for_turbo(pred, noisy_latents, timesteps, noise, batch)
+
+        ignore_snr = False
+
+        if loss_target == 'source' or loss_target == 'unaugmented':
+            assert not self.train_config.train_turbo
+            # ignore_snr = True
+            if batch.sigmas is None:
+                raise ValueError("Batch sigmas is None. This should not happen")
+
+            # src https://github.com/huggingface/diffusers/blob/324d18fba23f6c9d7475b0ff7c777685f7128d40/examples/t2i_adapter/train_t2i_adapter_sdxl.py#L1190
+            denoised_latents = noise_pred * (-batch.sigmas) + noisy_latents
+            weighing = batch.sigmas ** -2.0
+            if loss_target == 'source':
+                # denoise the latent and compare to the latent in the batch
+                target = batch.latents
+            elif loss_target == 'unaugmented':
+                # we have to encode images into latents for now
+                # we also denoise as the unaugmented tensor is not a noisy diffirental
+                with torch.no_grad():
+                    unaugmented_latents = self.sd.encode_images(batch.unaugmented_tensor).to(self.device_torch, dtype=dtype)
+                    unaugmented_latents = unaugmented_latents * self.train_config.latent_multiplier
+                    target = unaugmented_latents.detach()
+
+                # Get the target for loss depending on the prediction type
+                if self.sd.noise_scheduler.config.prediction_type == "epsilon":
+                    target = target  # we are computing loss against denoise latents
+                elif self.sd.noise_scheduler.config.prediction_type == "v_prediction":
+                    target = self.sd.noise_scheduler.get_velocity(target, noise, timesteps)
+                else:
+                    raise ValueError(f"Unknown prediction type {self.sd.noise_scheduler.config.prediction_type}")
+
+            # mse loss without reduction
+            loss_per_element = (weighing.float() * (denoised_latents.float() - target.float()) ** 2)
+            loss = loss_per_element
+        else:
+
+            if self.train_config.loss_type == "mae":
+                loss = torch.nn.functional.l1_loss(pred.float(), target.float(), reduction="none")
+            else:
+                loss = torch.nn.functional.mse_loss(pred.float(), target.float(), reduction="none")
+
+            # handle linear timesteps and only adjust the weight of the timesteps
+            if self.sd.is_flow_matching and (self.train_config.linear_timesteps or self.train_config.linear_timesteps2):
+                # calculate the weights for the timesteps
+                timestep_weight = self.sd.noise_scheduler.get_weights_for_timesteps(
+                    timesteps,
+                    v2=self.train_config.linear_timesteps2
+                ).to(loss.device, dtype=loss.dtype)
+                timestep_weight = timestep_weight.view(-1, 1, 1, 1).detach()
+                loss = loss * timestep_weight
+
+        if self.train_config.do_prior_divergence and prior_pred is not None:
+            loss = loss + (torch.nn.functional.mse_loss(pred.float(), prior_pred.float(), reduction="none") * -1.0)
+
+        if self.train_config.train_turbo:
+            mask_multiplier = mask_multiplier[:, 3:, :, :]
+            # resize to the size of the loss
+            mask_multiplier = torch.nn.functional.interpolate(mask_multiplier, size=(pred.shape[2], pred.shape[3]), mode='nearest')
+
+        # multiply by our mask
+        loss = loss * mask_multiplier
+
+        prior_loss = None
+        if self.train_config.inverted_mask_prior and prior_pred is not None and prior_mask_multiplier is not None:
+            assert not self.train_config.train_turbo
+            if self.train_config.loss_type == "mae":
+                prior_loss = torch.nn.functional.l1_loss(pred.float(), prior_pred.float(), reduction="none")
+            else:
+                prior_loss = torch.nn.functional.mse_loss(pred.float(), prior_pred.float(), reduction="none")
+
+            prior_loss = prior_loss * prior_mask_multiplier * self.train_config.inverted_mask_prior_multiplier
+            if torch.isnan(prior_loss).any():
+                print("Prior loss is nan")
+                prior_loss = None
+            else:
+                prior_loss = prior_loss.mean([1, 2, 3])
+                # loss = loss + prior_loss
+                # loss = loss + prior_loss
+            # loss = loss + prior_loss
+        loss = loss.mean([1, 2, 3])
+        # apply loss multiplier before prior loss
+        loss = loss * loss_multiplier
+        if prior_loss is not None:
+            loss = loss + prior_loss
+
+        if not self.train_config.train_turbo:
+            if self.train_config.learnable_snr_gos:
+                # add snr_gamma
+                loss = apply_learnable_snr_gos(loss, timesteps, self.snr_gos)
+            elif self.train_config.snr_gamma is not None and self.train_config.snr_gamma > 0.000001 and not ignore_snr:
+                # add snr_gamma
+                loss = apply_snr_weight(loss, timesteps, self.sd.noise_scheduler, self.train_config.snr_gamma,
+                                        fixed=True)
+            elif self.train_config.min_snr_gamma is not None and self.train_config.min_snr_gamma > 0.000001 and not ignore_snr:
+                # add min_snr_gamma
+                loss = apply_snr_weight(loss, timesteps, self.sd.noise_scheduler, self.train_config.min_snr_gamma)
+
+        loss = loss.mean()
+
+        # check for additional losses
+        if self.adapter is not None and hasattr(self.adapter, "additional_loss") and self.adapter.additional_loss is not None:
+
+            loss = loss + self.adapter.additional_loss.mean()
+            self.adapter.additional_loss = None
+
+        if self.train_config.target_norm_std:
+            # seperate out the batch and channels
+            pred_std = noise_pred.std([2, 3], keepdim=True)
+            norm_std_loss = torch.abs(self.train_config.target_norm_std_value - pred_std).mean()
+            loss = loss + norm_std_loss
+
+
+        return loss
+
+    def preprocess_batch(self, batch: 'DataLoaderBatchDTO'):
+        return batch
+
+    def get_guided_loss(
+            self,
+            noisy_latents: torch.Tensor,
+            conditional_embeds: PromptEmbeds,
+            match_adapter_assist: bool,
+            network_weight_list: list,
+            timesteps: torch.Tensor,
+            pred_kwargs: dict,
+            batch: 'DataLoaderBatchDTO',
+            noise: torch.Tensor,
+            unconditional_embeds: Optional[PromptEmbeds] = None,
+            **kwargs
+    ):
+        loss = get_guidance_loss(
+            noisy_latents=noisy_latents,
+            conditional_embeds=conditional_embeds,
+            match_adapter_assist=match_adapter_assist,
+            network_weight_list=network_weight_list,
+            timesteps=timesteps,
+            pred_kwargs=pred_kwargs,
+            batch=batch,
+            noise=noise,
+            sd=self.sd,
+            unconditional_embeds=unconditional_embeds,
+            scaler=self.scaler,
+            **kwargs
+        )
+
+        return loss
+
+    def get_guided_loss_targeted_polarity(
+            self,
+            noisy_latents: torch.Tensor,
+            conditional_embeds: PromptEmbeds,
+            match_adapter_assist: bool,
+            network_weight_list: list,
+            timesteps: torch.Tensor,
+            pred_kwargs: dict,
+            batch: 'DataLoaderBatchDTO',
+            noise: torch.Tensor,
+            **kwargs
+    ):
+        with torch.no_grad():
+            # Perform targeted guidance (working title)
+            dtype = get_torch_dtype(self.train_config.dtype)
+
+            conditional_latents = batch.latents.to(self.device_torch, dtype=dtype).detach()
+            unconditional_latents = batch.unconditional_latents.to(self.device_torch, dtype=dtype).detach()
+
+            mean_latents = (conditional_latents + unconditional_latents) / 2.0
+
+            unconditional_diff = (unconditional_latents - mean_latents)
+            conditional_diff = (conditional_latents - mean_latents)
+
+            # we need to determine the amount of signal and noise that would be present at the current timestep
+            # conditional_signal = self.sd.add_noise(conditional_diff, torch.zeros_like(noise), timesteps)
+            # unconditional_signal = self.sd.add_noise(torch.zeros_like(noise), unconditional_diff, timesteps)
+            # unconditional_signal = self.sd.add_noise(unconditional_diff, torch.zeros_like(noise), timesteps)
+            # conditional_blend = self.sd.add_noise(conditional_latents, unconditional_latents, timesteps)
+            # unconditional_blend = self.sd.add_noise(unconditional_latents, conditional_latents, timesteps)
+
+            # target_noise = noise + unconditional_signal
+
+            conditional_noisy_latents = self.sd.add_noise(
+                mean_latents,
+                noise,
+                timesteps
+            ).detach()
+
+            unconditional_noisy_latents = self.sd.add_noise(
+                mean_latents,
+                noise,
+                timesteps
+            ).detach()
+
+            # Disable the LoRA network so we can predict parent network knowledge without it
+            self.network.is_active = False
+            self.sd.unet.eval()
+
+            # Predict noise to get a baseline of what the parent network wants to do with the latents + noise.
+            # This acts as our control to preserve the unaltered parts of the image.
+            baseline_prediction = self.sd.predict_noise(
+                latents=unconditional_noisy_latents.to(self.device_torch, dtype=dtype).detach(),
+                conditional_embeddings=conditional_embeds.to(self.device_torch, dtype=dtype).detach(),
+                timestep=timesteps,
+                guidance_scale=1.0,
+                **pred_kwargs  # adapter residuals in here
+            ).detach()
+
+            # double up everything to run it through all at once
+            cat_embeds = concat_prompt_embeds([conditional_embeds, conditional_embeds])
+            cat_latents = torch.cat([conditional_noisy_latents, conditional_noisy_latents], dim=0)
+            cat_timesteps = torch.cat([timesteps, timesteps], dim=0)
+
+            # since we are dividing the polarity from the middle out, we need to double our network
+            # weights on training since the convergent point will be at half network strength
+
+            negative_network_weights = [weight * -2.0 for weight in network_weight_list]
+            positive_network_weights = [weight * 2.0 for weight in network_weight_list]
+            cat_network_weight_list = positive_network_weights + negative_network_weights
+
+            # turn the LoRA network back on.
+            self.sd.unet.train()
+            self.network.is_active = True
+
+            self.network.multiplier = cat_network_weight_list
+
+        # do our prediction with LoRA active on the scaled guidance latents
+        prediction = self.sd.predict_noise(
+            latents=cat_latents.to(self.device_torch, dtype=dtype).detach(),
+            conditional_embeddings=cat_embeds.to(self.device_torch, dtype=dtype).detach(),
+            timestep=cat_timesteps,
+            guidance_scale=1.0,
+            **pred_kwargs  # adapter residuals in here
+        )
+
+        pred_pos, pred_neg = torch.chunk(prediction, 2, dim=0)
+
+        pred_pos = pred_pos - baseline_prediction
+        pred_neg = pred_neg - baseline_prediction
+
+        pred_loss = torch.nn.functional.mse_loss(
+            pred_pos.float(),
+            unconditional_diff.float(),
+            reduction="none"
+        )
+        pred_loss = pred_loss.mean([1, 2, 3])
+
+        pred_neg_loss = torch.nn.functional.mse_loss(
+            pred_neg.float(),
+            conditional_diff.float(),
+            reduction="none"
+        )
+        pred_neg_loss = pred_neg_loss.mean([1, 2, 3])
+
+        loss = (pred_loss + pred_neg_loss) / 2.0
+
+        # loss = self.apply_snr(loss, timesteps)
+        loss = loss.mean()
+        loss.backward()
+
+        # detach it so parent class can run backward on no grads without throwing error
+        loss = loss.detach()
+        loss.requires_grad_(True)
+
+        return loss
+
+    def get_guided_loss_masked_polarity(
+            self,
+            noisy_latents: torch.Tensor,
+            conditional_embeds: PromptEmbeds,
+            match_adapter_assist: bool,
+            network_weight_list: list,
+            timesteps: torch.Tensor,
+            pred_kwargs: dict,
+            batch: 'DataLoaderBatchDTO',
+            noise: torch.Tensor,
+            **kwargs
+    ):
+        with torch.no_grad():
+            # Perform targeted guidance (working title)
+            dtype = get_torch_dtype(self.train_config.dtype)
+
+            conditional_latents = batch.latents.to(self.device_torch, dtype=dtype).detach()
+            unconditional_latents = batch.unconditional_latents.to(self.device_torch, dtype=dtype).detach()
+            inverse_latents = unconditional_latents - (conditional_latents - unconditional_latents)
+
+            mean_latents = (conditional_latents + unconditional_latents) / 2.0
+
+            # unconditional_diff = (unconditional_latents - mean_latents)
+            # conditional_diff = (conditional_latents - mean_latents)
+
+            # we need to determine the amount of signal and noise that would be present at the current timestep
+            # conditional_signal = self.sd.add_noise(conditional_diff, torch.zeros_like(noise), timesteps)
+            # unconditional_signal = self.sd.add_noise(torch.zeros_like(noise), unconditional_diff, timesteps)
+            # unconditional_signal = self.sd.add_noise(unconditional_diff, torch.zeros_like(noise), timesteps)
+            # conditional_blend = self.sd.add_noise(conditional_latents, unconditional_latents, timesteps)
+            # unconditional_blend = self.sd.add_noise(unconditional_latents, conditional_latents, timesteps)
+
+            # make a differential mask
+            differential_mask = torch.abs(conditional_latents - unconditional_latents)
+            max_differential = \
+                differential_mask.max(dim=1, keepdim=True)[0].max(dim=2, keepdim=True)[0].max(dim=3, keepdim=True)[0]
+            differential_scaler = 1.0 / max_differential
+            differential_mask = differential_mask * differential_scaler
+            spread_point = 0.1
+            # adjust mask to amplify the differential at 0.1
+            differential_mask = ((differential_mask - spread_point) * 10.0) + spread_point
+            # clip it
+            differential_mask = torch.clamp(differential_mask, 0.0, 1.0)
+
+            # target_noise = noise + unconditional_signal
+
+            conditional_noisy_latents = self.sd.add_noise(
+                conditional_latents,
+                noise,
+                timesteps
+            ).detach()
+
+            unconditional_noisy_latents = self.sd.add_noise(
+                unconditional_latents,
+                noise,
+                timesteps
+            ).detach()
+
+            inverse_noisy_latents = self.sd.add_noise(
+                inverse_latents,
+                noise,
+                timesteps
+            ).detach()
+
+            # Disable the LoRA network so we can predict parent network knowledge without it
+            self.network.is_active = False
+            self.sd.unet.eval()
+
+            # Predict noise to get a baseline of what the parent network wants to do with the latents + noise.
+            # This acts as our control to preserve the unaltered parts of the image.
+            # baseline_prediction = self.sd.predict_noise(
+            #     latents=unconditional_noisy_latents.to(self.device_torch, dtype=dtype).detach(),
+            #     conditional_embeddings=conditional_embeds.to(self.device_torch, dtype=dtype).detach(),
+            #     timestep=timesteps,
+            #     guidance_scale=1.0,
+            #     **pred_kwargs  # adapter residuals in here
+            # ).detach()
+
+            # double up everything to run it through all at once
+            cat_embeds = concat_prompt_embeds([conditional_embeds, conditional_embeds])
+            cat_latents = torch.cat([conditional_noisy_latents, unconditional_noisy_latents], dim=0)
+            cat_timesteps = torch.cat([timesteps, timesteps], dim=0)
+
+            # since we are dividing the polarity from the middle out, we need to double our network
+            # weights on training since the convergent point will be at half network strength
+
+            negative_network_weights = [weight * -1.0 for weight in network_weight_list]
+            positive_network_weights = [weight * 1.0 for weight in network_weight_list]
+            cat_network_weight_list = positive_network_weights + negative_network_weights
+
+            # turn the LoRA network back on.
+            self.sd.unet.train()
+            self.network.is_active = True
+
+            self.network.multiplier = cat_network_weight_list
+
+        # do our prediction with LoRA active on the scaled guidance latents
+        prediction = self.sd.predict_noise(
+            latents=cat_latents.to(self.device_torch, dtype=dtype).detach(),
+            conditional_embeddings=cat_embeds.to(self.device_torch, dtype=dtype).detach(),
+            timestep=cat_timesteps,
+            guidance_scale=1.0,
+            **pred_kwargs  # adapter residuals in here
+        )
+
+        pred_pos, pred_neg = torch.chunk(prediction, 2, dim=0)
+
+        # create a loss to balance the mean to 0 between the two predictions
+        differential_mean_pred_loss = torch.abs(pred_pos - pred_neg).mean([1, 2, 3]) ** 2.0
+
+        # pred_pos = pred_pos - baseline_prediction
+        # pred_neg = pred_neg - baseline_prediction
+
+        pred_loss = torch.nn.functional.mse_loss(
+            pred_pos.float(),
+            noise.float(),
+            reduction="none"
+        )
+        # apply mask
+        pred_loss = pred_loss * (1.0 + differential_mask)
+        pred_loss = pred_loss.mean([1, 2, 3])
+
+        pred_neg_loss = torch.nn.functional.mse_loss(
+            pred_neg.float(),
+            noise.float(),
+            reduction="none"
+        )
+        # apply inverse mask
+        pred_neg_loss = pred_neg_loss * (1.0 - differential_mask)
+        pred_neg_loss = pred_neg_loss.mean([1, 2, 3])
+
+        # make a loss to balance to losses of the pos and neg so they are equal
+        # differential_mean_loss_loss = torch.abs(pred_loss - pred_neg_loss)
+        #
+        # differential_mean_loss = differential_mean_pred_loss + differential_mean_loss_loss
+        #
+        # # add a multiplier to balancing losses to make them the top priority
+        # differential_mean_loss = differential_mean_loss
+
+        # remove the grads from the negative as it is only a balancing loss
+        # pred_neg_loss = pred_neg_loss.detach()
+
+        # loss = pred_loss + pred_neg_loss + differential_mean_loss
+        loss = pred_loss + pred_neg_loss
+
+        # loss = self.apply_snr(loss, timesteps)
+        loss = loss.mean()
+        loss.backward()
+
+        # detach it so parent class can run backward on no grads without throwing error
+        loss = loss.detach()
+        loss.requires_grad_(True)
+
+        return loss
+
+    def get_prior_prediction(
+            self,
+            noisy_latents: torch.Tensor,
+            conditional_embeds: PromptEmbeds,
+            match_adapter_assist: bool,
+            network_weight_list: list,
+            timesteps: torch.Tensor,
+            pred_kwargs: dict,
+            batch: 'DataLoaderBatchDTO',
+            noise: torch.Tensor,
+            unconditional_embeds: Optional[PromptEmbeds] = None,
+            conditioned_prompts=None,
+            **kwargs
+    ):
+        # todo for embeddings, we need to run without trigger words
+        was_unet_training = self.sd.unet.training
+        was_network_active = False
+        if self.network is not None:
+            was_network_active = self.network.is_active
+            self.network.is_active = False
+        can_disable_adapter = False
+        was_adapter_active = False
+        if self.adapter is not None and (isinstance(self.adapter, IPAdapter) or
+                                         isinstance(self.adapter, ReferenceAdapter) or
+                                         (isinstance(self.adapter, CustomAdapter))
+        ):
+            can_disable_adapter = True
+            was_adapter_active = self.adapter.is_active
+            self.adapter.is_active = False
+
+        if self.train_config.unload_text_encoder:
+            raise ValueError("Prior predictions currently do not support unloading text encoder")
+        # do a prediction here so we can match its output with network multiplier set to 0.0
+        with torch.no_grad():
+            dtype = get_torch_dtype(self.train_config.dtype)
+
+            embeds_to_use = conditional_embeds.clone().detach()
+            # handle clip vision adapter by removing triggers from prompt and replacing with the class name
+            if (self.adapter is not None and isinstance(self.adapter, ClipVisionAdapter)) or self.embedding is not None:
+                prompt_list = batch.get_caption_list()
+                class_name = ''
+
+                triggers = ['[trigger]', '[name]']
+                remove_tokens = []
+
+                if self.embed_config is not None:
+                    triggers.append(self.embed_config.trigger)
+                    for i in range(1, self.embed_config.tokens):
+                        remove_tokens.append(f"{self.embed_config.trigger}_{i}")
+                    if self.embed_config.trigger_class_name is not None:
+                        class_name = self.embed_config.trigger_class_name
+
+                if self.adapter is not None:
+                    triggers.append(self.adapter_config.trigger)
+                    for i in range(1, self.adapter_config.num_tokens):
+                        remove_tokens.append(f"{self.adapter_config.trigger}_{i}")
+                    if self.adapter_config.trigger_class_name is not None:
+                        class_name = self.adapter_config.trigger_class_name
+
+                for idx, prompt in enumerate(prompt_list):
+                    for remove_token in remove_tokens:
+                        prompt = prompt.replace(remove_token, '')
+                    for trigger in triggers:
+                        prompt = prompt.replace(trigger, class_name)
+                    prompt_list[idx] = prompt
+
+                embeds_to_use = self.sd.encode_prompt(
+                    prompt_list,
+                    long_prompts=self.do_long_prompts).to(
+                    self.device_torch,
+                    dtype=dtype).detach()
+
+            # dont use network on this
+            # self.network.multiplier = 0.0
+            self.sd.unet.eval()
+
+            if self.adapter is not None and isinstance(self.adapter, IPAdapter) and not self.sd.is_flux:
+                # we need to remove the image embeds from the prompt except for flux
+                embeds_to_use: PromptEmbeds = embeds_to_use.clone().detach()
+                end_pos = embeds_to_use.text_embeds.shape[1] - self.adapter_config.num_tokens
+                embeds_to_use.text_embeds = embeds_to_use.text_embeds[:, :end_pos, :]
+                if unconditional_embeds is not None:
+                    unconditional_embeds = unconditional_embeds.clone().detach()
+                    unconditional_embeds.text_embeds = unconditional_embeds.text_embeds[:, :end_pos]
+
+            if unconditional_embeds is not None:
+                unconditional_embeds = unconditional_embeds.to(self.device_torch, dtype=dtype).detach()
+
+            prior_pred = self.sd.predict_noise(
+                latents=noisy_latents.to(self.device_torch, dtype=dtype).detach(),
+                conditional_embeddings=embeds_to_use.to(self.device_torch, dtype=dtype).detach(),
+                unconditional_embeddings=unconditional_embeds,
+                timestep=timesteps,
+                guidance_scale=self.train_config.cfg_scale,
+                rescale_cfg=self.train_config.cfg_rescale,
+                **pred_kwargs  # adapter residuals in here
+            )
+            if was_unet_training:
+                self.sd.unet.train()
+            prior_pred = prior_pred.detach()
+            # remove the residuals as we wont use them on prediction when matching control
+            if match_adapter_assist and 'down_intrablock_additional_residuals' in pred_kwargs:
+                del pred_kwargs['down_intrablock_additional_residuals']
+            if match_adapter_assist and 'down_block_additional_residuals' in pred_kwargs:
+                del pred_kwargs['down_block_additional_residuals']
+            if match_adapter_assist and 'mid_block_additional_residual' in pred_kwargs:
+                del pred_kwargs['mid_block_additional_residual']
+
+            if can_disable_adapter:
+                self.adapter.is_active = was_adapter_active
+            # restore network
+            # self.network.multiplier = network_weight_list
+            if self.network is not None:
+                self.network.is_active = was_network_active
+        return prior_pred
+
+    def before_unet_predict(self):
+        pass
+
+    def after_unet_predict(self):
+        pass
+
+    def end_of_training_loop(self):
+        pass
+
+    def predict_noise(
+            self,
+            noisy_latents: torch.Tensor,
+            timesteps: Union[int, torch.Tensor] = 1,
+            conditional_embeds: Union[PromptEmbeds, None] = None,
+            unconditional_embeds: Union[PromptEmbeds, None] = None,
+            **kwargs,
+    ):
+        dtype = get_torch_dtype(self.train_config.dtype)
+        return self.sd.predict_noise(
+            latents=noisy_latents.to(self.device_torch, dtype=dtype),
+            conditional_embeddings=conditional_embeds.to(self.device_torch, dtype=dtype),
+            unconditional_embeddings=unconditional_embeds,
+            timestep=timesteps,
+            guidance_scale=self.train_config.cfg_scale,
+            guidance_embedding_scale=self.train_config.cfg_scale,
+            detach_unconditional=False,
+            rescale_cfg=self.train_config.cfg_rescale,
+            bypass_guidance_embedding=self.train_config.bypass_guidance_embedding,
+            **kwargs
+        )
+
+    def train_single_accumulation(self, batch: DataLoaderBatchDTO):
+        self.timer.start('preprocess_batch')
+        batch = self.preprocess_batch(batch)
+        dtype = get_torch_dtype(self.train_config.dtype)
+        # sanity check
+        if self.sd.vae.dtype != self.sd.vae_torch_dtype:
+            self.sd.vae = self.sd.vae.to(self.sd.vae_torch_dtype)
+        if isinstance(self.sd.text_encoder, list):
+            for encoder in self.sd.text_encoder:
+                if encoder.dtype != self.sd.te_torch_dtype:
+                    encoder.to(self.sd.te_torch_dtype)
+        else:
+            if self.sd.text_encoder.dtype != self.sd.te_torch_dtype:
+                self.sd.text_encoder.to(self.sd.te_torch_dtype)
+
+        noisy_latents, noise, timesteps, conditioned_prompts, imgs = self.process_general_training_batch(batch)
+        if self.train_config.do_cfg or self.train_config.do_random_cfg:
+            # pick random negative prompts
+            if self.negative_prompt_pool is not None:
+                negative_prompts = []
+                for i in range(noisy_latents.shape[0]):
+                    num_neg = random.randint(1, self.train_config.max_negative_prompts)
+                    this_neg_prompts = [random.choice(self.negative_prompt_pool) for _ in range(num_neg)]
+                    this_neg_prompt = ', '.join(this_neg_prompts)
+                    negative_prompts.append(this_neg_prompt)
+                self.batch_negative_prompt = negative_prompts
+            else:
+                self.batch_negative_prompt = ['' for _ in range(batch.latents.shape[0])]
+
+        if self.adapter and isinstance(self.adapter, CustomAdapter):
+            # condition the prompt
+            # todo handle more than one adapter image
+            self.adapter.num_control_images = 1
+            conditioned_prompts = self.adapter.condition_prompt(conditioned_prompts)
+
+        network_weight_list = batch.get_network_weight_list()
+        if self.train_config.single_item_batching:
+            network_weight_list = network_weight_list + network_weight_list
+
+        has_adapter_img = batch.control_tensor is not None
+        has_clip_image = batch.clip_image_tensor is not None
+        has_clip_image_embeds = batch.clip_image_embeds is not None
+        # force it to be true if doing regs as we handle those differently
+        if any([batch.file_items[idx].is_reg for idx in range(len(batch.file_items))]):
+            has_clip_image = True
+            if self._clip_image_embeds_unconditional is not None:
+                has_clip_image_embeds = True  # we are caching embeds, handle that differently
+                has_clip_image = False
+
+        if self.adapter is not None and isinstance(self.adapter, IPAdapter) and not has_clip_image and has_adapter_img:
+            raise ValueError(
+                "IPAdapter control image is now 'clip_image_path' instead of 'control_path'. Please update your dataset config ")
+
+        match_adapter_assist = False
+
+        # check if we are matching the adapter assistant
+        if self.assistant_adapter:
+            if self.train_config.match_adapter_chance == 1.0:
+                match_adapter_assist = True
+            elif self.train_config.match_adapter_chance > 0.0:
+                match_adapter_assist = torch.rand(
+                    (1,), device=self.device_torch, dtype=dtype
+                ) < self.train_config.match_adapter_chance
+
+        self.timer.stop('preprocess_batch')
+
+        is_reg = False
+        with torch.no_grad():
+            loss_multiplier = torch.ones((noisy_latents.shape[0], 1, 1, 1), device=self.device_torch, dtype=dtype)
+            for idx, file_item in enumerate(batch.file_items):
+                if file_item.is_reg:
+                    loss_multiplier[idx] = loss_multiplier[idx] * self.train_config.reg_weight
+                    is_reg = True
+
+            adapter_images = None
+            sigmas = None
+            if has_adapter_img and (self.adapter or self.assistant_adapter):
+                with self.timer('get_adapter_images'):
+                    # todo move this to data loader
+                    if batch.control_tensor is not None:
+                        adapter_images = batch.control_tensor.to(self.device_torch, dtype=dtype).detach()
+                        # match in channels
+                        if self.assistant_adapter is not None:
+                            in_channels = self.assistant_adapter.config.in_channels
+                            if adapter_images.shape[1] != in_channels:
+                                # we need to match the channels
+                                adapter_images = adapter_images[:, :in_channels, :, :]
+                    else:
+                        raise NotImplementedError("Adapter images now must be loaded with dataloader")
+
+            clip_images = None
+            if has_clip_image:
+                with self.timer('get_clip_images'):
+                    # todo move this to data loader
+                    if batch.clip_image_tensor is not None:
+                        clip_images = batch.clip_image_tensor.to(self.device_torch, dtype=dtype).detach()
+
+            mask_multiplier = torch.ones((noisy_latents.shape[0], 1, 1, 1), device=self.device_torch, dtype=dtype)
+            if batch.mask_tensor is not None:
+                with self.timer('get_mask_multiplier'):
+                    # upsampling no supported for bfloat16
+                    mask_multiplier = batch.mask_tensor.to(self.device_torch, dtype=torch.float16).detach()
+                    # scale down to the size of the latents, mask multiplier shape(bs, 1, width, height), noisy_latents shape(bs, channels, width, height)
+                    mask_multiplier = torch.nn.functional.interpolate(
+                        mask_multiplier, size=(noisy_latents.shape[2], noisy_latents.shape[3])
+                    )
+                    # expand to match latents
+                    mask_multiplier = mask_multiplier.expand(-1, noisy_latents.shape[1], -1, -1)
+                    mask_multiplier = mask_multiplier.to(self.device_torch, dtype=dtype).detach()
+
+        def get_adapter_multiplier():
+            if self.adapter and isinstance(self.adapter, T2IAdapter):
+                # training a t2i adapter, not using as assistant.
+                return 1.0
+            elif match_adapter_assist:
+                # training a texture. We want it high
+                adapter_strength_min = 0.9
+                adapter_strength_max = 1.0
+            else:
+                # training with assistance, we want it low
+                # adapter_strength_min = 0.4
+                # adapter_strength_max = 0.7
+                adapter_strength_min = 0.5
+                adapter_strength_max = 1.1
+
+            adapter_conditioning_scale = torch.rand(
+                (1,), device=self.device_torch, dtype=dtype
+            )
+
+            adapter_conditioning_scale = value_map(
+                adapter_conditioning_scale,
+                0.0,
+                1.0,
+                adapter_strength_min,
+                adapter_strength_max
+            )
+            return adapter_conditioning_scale
+
+        # flush()
+        with self.timer('grad_setup'):
+
+            # text encoding
+            grad_on_text_encoder = False
+            if self.train_config.train_text_encoder:
+                grad_on_text_encoder = True
+
+            if self.embedding is not None:
+                grad_on_text_encoder = True
+
+            if self.adapter and isinstance(self.adapter, ClipVisionAdapter):
+                grad_on_text_encoder = True
+
+            if self.adapter_config and self.adapter_config.type == 'te_augmenter':
+                grad_on_text_encoder = True
+
+            # have a blank network so we can wrap it in a context and set multipliers without checking every time
+            if self.network is not None:
+                network = self.network
+            else:
+                network = BlankNetwork()
+
+            # set the weights
+            network.multiplier = network_weight_list
+
+        # activate network if it exits
+
+        prompts_1 = conditioned_prompts
+        prompts_2 = None
+        if self.train_config.short_and_long_captions_encoder_split and self.sd.is_xl:
+            prompts_1 = batch.get_caption_short_list()
+            prompts_2 = conditioned_prompts
+
+            # make the batch splits
+        if self.train_config.single_item_batching:
+            if self.model_config.refiner_name_or_path is not None:
+                raise ValueError("Single item batching is not supported when training the refiner")
+            batch_size = noisy_latents.shape[0]
+            # chunk/split everything
+            noisy_latents_list = torch.chunk(noisy_latents, batch_size, dim=0)
+            noise_list = torch.chunk(noise, batch_size, dim=0)
+            timesteps_list = torch.chunk(timesteps, batch_size, dim=0)
+            conditioned_prompts_list = [[prompt] for prompt in prompts_1]
+            if imgs is not None:
+                imgs_list = torch.chunk(imgs, batch_size, dim=0)
+            else:
+                imgs_list = [None for _ in range(batch_size)]
+            if adapter_images is not None:
+                adapter_images_list = torch.chunk(adapter_images, batch_size, dim=0)
+            else:
+                adapter_images_list = [None for _ in range(batch_size)]
+            if clip_images is not None:
+                clip_images_list = torch.chunk(clip_images, batch_size, dim=0)
+            else:
+                clip_images_list = [None for _ in range(batch_size)]
+            mask_multiplier_list = torch.chunk(mask_multiplier, batch_size, dim=0)
+            if prompts_2 is None:
+                prompt_2_list = [None for _ in range(batch_size)]
+            else:
+                prompt_2_list = [[prompt] for prompt in prompts_2]
+
+        else:
+            noisy_latents_list = [noisy_latents]
+            noise_list = [noise]
+            timesteps_list = [timesteps]
+            conditioned_prompts_list = [prompts_1]
+            imgs_list = [imgs]
+            adapter_images_list = [adapter_images]
+            clip_images_list = [clip_images]
+            mask_multiplier_list = [mask_multiplier]
+            if prompts_2 is None:
+                prompt_2_list = [None]
+            else:
+                prompt_2_list = [prompts_2]
+
+        for noisy_latents, noise, timesteps, conditioned_prompts, imgs, adapter_images, clip_images, mask_multiplier, prompt_2 in zip(
+                noisy_latents_list,
+                noise_list,
+                timesteps_list,
+                conditioned_prompts_list,
+                imgs_list,
+                adapter_images_list,
+                clip_images_list,
+                mask_multiplier_list,
+                prompt_2_list
+        ):
+
+            # if self.train_config.negative_prompt is not None:
+            #     # add negative prompt
+            #     conditioned_prompts = conditioned_prompts + [self.train_config.negative_prompt for x in
+            #                                                  range(len(conditioned_prompts))]
+            #     if prompt_2 is not None:
+            #         prompt_2 = prompt_2 + [self.train_config.negative_prompt for x in range(len(prompt_2))]
+
+            with (network):
+                # encode clip adapter here so embeds are active for tokenizer
+                if self.adapter and isinstance(self.adapter, ClipVisionAdapter):
+                    with self.timer('encode_clip_vision_embeds'):
+                        if has_clip_image:
+                            conditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors(
+                                clip_images.detach().to(self.device_torch, dtype=dtype),
+                                is_training=True,
+                                has_been_preprocessed=True
+                            )
+                        else:
+                            # just do a blank one
+                            conditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors(
+                                torch.zeros(
+                                    (noisy_latents.shape[0], 3, 512, 512),
+                                    device=self.device_torch, dtype=dtype
+                                ),
+                                is_training=True,
+                                has_been_preprocessed=True,
+                                drop=True
+                            )
+                        # it will be injected into the tokenizer when called
+                        self.adapter(conditional_clip_embeds)
+
+                # do the custom adapter after the prior prediction
+                if self.adapter and isinstance(self.adapter, CustomAdapter) and (has_clip_image or is_reg):
+                    quad_count = random.randint(1, 4)
+                    self.adapter.train()
+                    self.adapter.trigger_pre_te(
+                        tensors_0_1=clip_images if not is_reg else None,  # on regs we send none to get random noise
+                        is_training=True,
+                        has_been_preprocessed=True,
+                        quad_count=quad_count,
+                        batch_size=noisy_latents.shape[0]
+                    )
+
+                with self.timer('encode_prompt'):
+                    unconditional_embeds = None
+                    if self.train_config.unload_text_encoder:
+                        with torch.set_grad_enabled(False):
+                            embeds_to_use = self.cached_blank_embeds.clone().detach().to(
+                                self.device_torch, dtype=dtype
+                            )
+                            if self.cached_trigger_embeds is not None and not is_reg:
+                                embeds_to_use = self.cached_trigger_embeds.clone().detach().to(
+                                    self.device_torch, dtype=dtype
+                                )
+                            conditional_embeds = concat_prompt_embeds(
+                                [embeds_to_use] * noisy_latents.shape[0]
+                            )
+                            if self.train_config.do_cfg:
+                                unconditional_embeds = self.cached_blank_embeds.clone().detach().to(
+                                    self.device_torch, dtype=dtype
+                                )
+                                unconditional_embeds = concat_prompt_embeds(
+                                    [unconditional_embeds] * noisy_latents.shape[0]
+                                )
+
+                            if isinstance(self.adapter, CustomAdapter):
+                                self.adapter.is_unconditional_run = False
+
+                    elif grad_on_text_encoder:
+                        with torch.set_grad_enabled(True):
+                            if isinstance(self.adapter, CustomAdapter):
+                                self.adapter.is_unconditional_run = False
+                            conditional_embeds = self.sd.encode_prompt(
+                                conditioned_prompts, prompt_2,
+                                dropout_prob=self.train_config.prompt_dropout_prob,
+                                long_prompts=self.do_long_prompts).to(
+                                self.device_torch,
+                                dtype=dtype)
+
+                            if self.train_config.do_cfg:
+                                if isinstance(self.adapter, CustomAdapter):
+                                    self.adapter.is_unconditional_run = True
+                                # todo only do one and repeat it
+                                unconditional_embeds = self.sd.encode_prompt(
+                                    self.batch_negative_prompt,
+                                    self.batch_negative_prompt,
+                                    dropout_prob=self.train_config.prompt_dropout_prob,
+                                    long_prompts=self.do_long_prompts).to(
+                                    self.device_torch,
+                                    dtype=dtype)
+                                if isinstance(self.adapter, CustomAdapter):
+                                    self.adapter.is_unconditional_run = False
+                    else:
+                        with torch.set_grad_enabled(False):
+                            # make sure it is in eval mode
+                            if isinstance(self.sd.text_encoder, list):
+                                for te in self.sd.text_encoder:
+                                    te.eval()
+                            else:
+                                self.sd.text_encoder.eval()
+                            if isinstance(self.adapter, CustomAdapter):
+                                self.adapter.is_unconditional_run = False
+                            conditional_embeds = self.sd.encode_prompt(
+                                conditioned_prompts, prompt_2,
+                                dropout_prob=self.train_config.prompt_dropout_prob,
+                                long_prompts=self.do_long_prompts).to(
+                                self.device_torch,
+                                dtype=dtype)
+                            if self.train_config.do_cfg:
+                                if isinstance(self.adapter, CustomAdapter):
+                                    self.adapter.is_unconditional_run = True
+                                unconditional_embeds = self.sd.encode_prompt(
+                                    self.batch_negative_prompt,
+                                    dropout_prob=self.train_config.prompt_dropout_prob,
+                                    long_prompts=self.do_long_prompts).to(
+                                    self.device_torch,
+                                    dtype=dtype)
+                                if isinstance(self.adapter, CustomAdapter):
+                                    self.adapter.is_unconditional_run = False
+
+                        # detach the embeddings
+                        conditional_embeds = conditional_embeds.detach()
+                        if self.train_config.do_cfg:
+                            unconditional_embeds = unconditional_embeds.detach()
+                    
+                    if self.decorator:
+                        conditional_embeds.text_embeds = self.decorator(
+                            conditional_embeds.text_embeds
+                        )
+                        if self.train_config.do_cfg:
+                            unconditional_embeds.text_embeds = self.decorator(
+                                unconditional_embeds.text_embeds, 
+                                is_unconditional=True
+                            )
+
+                # flush()
+                pred_kwargs = {}
+
+                if has_adapter_img:
+                    if (self.adapter and isinstance(self.adapter, T2IAdapter)) or (
+                            self.assistant_adapter and isinstance(self.assistant_adapter, T2IAdapter)):
+                        with torch.set_grad_enabled(self.adapter is not None):
+                            adapter = self.assistant_adapter if self.assistant_adapter is not None else self.adapter
+                            adapter_multiplier = get_adapter_multiplier()
+                            with self.timer('encode_adapter'):
+                                down_block_additional_residuals = adapter(adapter_images)
+                                if self.assistant_adapter:
+                                    # not training. detach
+                                    down_block_additional_residuals = [
+                                        sample.to(dtype=dtype).detach() * adapter_multiplier for sample in
+                                        down_block_additional_residuals
+                                    ]
+                                else:
+                                    down_block_additional_residuals = [
+                                        sample.to(dtype=dtype) * adapter_multiplier for sample in
+                                        down_block_additional_residuals
+                                    ]
+
+                                pred_kwargs['down_intrablock_additional_residuals'] = down_block_additional_residuals
+
+                if self.adapter and isinstance(self.adapter, IPAdapter):
+                    with self.timer('encode_adapter_embeds'):
+                        # number of images to do if doing a quad image
+                        quad_count = random.randint(1, 4)
+                        image_size = self.adapter.input_size
+                        if has_clip_image_embeds:
+                            # todo handle reg images better than this
+                            if is_reg:
+                                # get unconditional image embeds from cache
+                                embeds = [
+                                    load_file(random.choice(batch.clip_image_embeds_unconditional)) for i in
+                                    range(noisy_latents.shape[0])
+                                ]
+                                conditional_clip_embeds = self.adapter.parse_clip_image_embeds_from_cache(
+                                    embeds,
+                                    quad_count=quad_count
+                                )
+
+                                if self.train_config.do_cfg:
+                                    embeds = [
+                                        load_file(random.choice(batch.clip_image_embeds_unconditional)) for i in
+                                        range(noisy_latents.shape[0])
+                                    ]
+                                    unconditional_clip_embeds = self.adapter.parse_clip_image_embeds_from_cache(
+                                        embeds,
+                                        quad_count=quad_count
+                                    )
+
+                            else:
+                                conditional_clip_embeds = self.adapter.parse_clip_image_embeds_from_cache(
+                                    batch.clip_image_embeds,
+                                    quad_count=quad_count
+                                )
+                                if self.train_config.do_cfg:
+                                    unconditional_clip_embeds = self.adapter.parse_clip_image_embeds_from_cache(
+                                        batch.clip_image_embeds_unconditional,
+                                        quad_count=quad_count
+                                    )
+                        elif is_reg:
+                            # we will zero it out in the img embedder
+                            clip_images = torch.zeros(
+                                (noisy_latents.shape[0], 3, image_size, image_size),
+                                device=self.device_torch, dtype=dtype
+                            ).detach()
+                            # drop will zero it out
+                            conditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors(
+                                clip_images,
+                                drop=True,
+                                is_training=True,
+                                has_been_preprocessed=False,
+                                quad_count=quad_count
+                            )
+                            if self.train_config.do_cfg:
+                                unconditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors(
+                                    torch.zeros(
+                                        (noisy_latents.shape[0], 3, image_size, image_size),
+                                        device=self.device_torch, dtype=dtype
+                                    ).detach(),
+                                    is_training=True,
+                                    drop=True,
+                                    has_been_preprocessed=False,
+                                    quad_count=quad_count
+                                )
+                        elif has_clip_image:
+                            conditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors(
+                                clip_images.detach().to(self.device_torch, dtype=dtype),
+                                is_training=True,
+                                has_been_preprocessed=True,
+                                quad_count=quad_count,
+                                # do cfg on clip embeds to normalize the embeddings for when doing cfg
+                                # cfg_embed_strength=3.0 if not self.train_config.do_cfg else None
+                                # cfg_embed_strength=3.0 if not self.train_config.do_cfg else None
+                            )
+                            if self.train_config.do_cfg:
+                                unconditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors(
+                                    clip_images.detach().to(self.device_torch, dtype=dtype),
+                                    is_training=True,
+                                    drop=True,
+                                    has_been_preprocessed=True,
+                                    quad_count=quad_count
+                                )
+                        else:
+                            print("No Clip Image")
+                            print([file_item.path for file_item in batch.file_items])
+                            raise ValueError("Could not find clip image")
+
+                    if not self.adapter_config.train_image_encoder:
+                        # we are not training the image encoder, so we need to detach the embeds
+                        conditional_clip_embeds = conditional_clip_embeds.detach()
+                        if self.train_config.do_cfg:
+                            unconditional_clip_embeds = unconditional_clip_embeds.detach()
+
+                    with self.timer('encode_adapter'):
+                        self.adapter.train()
+                        conditional_embeds = self.adapter(
+                            conditional_embeds.detach(),
+                            conditional_clip_embeds,
+                            is_unconditional=False
+                        )
+                        if self.train_config.do_cfg:
+                            unconditional_embeds = self.adapter(
+                                unconditional_embeds.detach(),
+                                unconditional_clip_embeds,
+                                is_unconditional=True
+                            )
+                        else:
+                            # wipe out unconsitional
+                            self.adapter.last_unconditional = None
+
+                if self.adapter and isinstance(self.adapter, ReferenceAdapter):
+                    # pass in our scheduler
+                    self.adapter.noise_scheduler = self.lr_scheduler
+                    if has_clip_image or has_adapter_img:
+                        img_to_use = clip_images if has_clip_image else adapter_images
+                        # currently 0-1 needs to be -1 to 1
+                        reference_images = ((img_to_use - 0.5) * 2).detach().to(self.device_torch, dtype=dtype)
+                        self.adapter.set_reference_images(reference_images)
+                        self.adapter.noise_scheduler = self.sd.noise_scheduler
+                    elif is_reg:
+                        self.adapter.set_blank_reference_images(noisy_latents.shape[0])
+                    else:
+                        self.adapter.set_reference_images(None)
+
+                prior_pred = None
+
+                do_reg_prior = False
+                # if is_reg and (self.network is not None or self.adapter is not None):
+                #     # we are doing a reg image and we have a network or adapter
+                #     do_reg_prior = True
+
+                do_inverted_masked_prior = False
+                if self.train_config.inverted_mask_prior and batch.mask_tensor is not None:
+                    do_inverted_masked_prior = True
+
+                do_correct_pred_norm_prior = self.train_config.correct_pred_norm
+
+                do_guidance_prior = False
+
+                if batch.unconditional_latents is not None:
+                    # for this not that, we need a prior pred to normalize
+                    guidance_type: GuidanceType = batch.file_items[0].dataset_config.guidance_type
+                    if guidance_type == 'tnt':
+                        do_guidance_prior = True
+
+                if ((
+                        has_adapter_img and self.assistant_adapter and match_adapter_assist) or self.do_prior_prediction or do_guidance_prior or do_reg_prior or do_inverted_masked_prior or self.train_config.correct_pred_norm):
+                    with self.timer('prior predict'):
+                        prior_pred = self.get_prior_prediction(
+                            noisy_latents=noisy_latents,
+                            conditional_embeds=conditional_embeds,
+                            match_adapter_assist=match_adapter_assist,
+                            network_weight_list=network_weight_list,
+                            timesteps=timesteps,
+                            pred_kwargs=pred_kwargs,
+                            noise=noise,
+                            batch=batch,
+                            unconditional_embeds=unconditional_embeds,
+                            conditioned_prompts=conditioned_prompts
+                        )
+                        if prior_pred is not None:
+                            prior_pred = prior_pred.detach()
+
+                # do the custom adapter after the prior prediction
+                if self.adapter and isinstance(self.adapter, CustomAdapter) and has_clip_image:
+                    quad_count = random.randint(1, 4)
+                    self.adapter.train()
+                    conditional_embeds = self.adapter.condition_encoded_embeds(
+                        tensors_0_1=clip_images,
+                        prompt_embeds=conditional_embeds,
+                        is_training=True,
+                        has_been_preprocessed=True,
+                        quad_count=quad_count
+                    )
+                    if self.train_config.do_cfg and unconditional_embeds is not None:
+                        unconditional_embeds = self.adapter.condition_encoded_embeds(
+                            tensors_0_1=clip_images,
+                            prompt_embeds=unconditional_embeds,
+                            is_training=True,
+                            has_been_preprocessed=True,
+                            is_unconditional=True,
+                            quad_count=quad_count
+                        )
+
+                if self.adapter and isinstance(self.adapter, CustomAdapter) and batch.extra_values is not None:
+                    self.adapter.add_extra_values(batch.extra_values.detach())
+
+                    if self.train_config.do_cfg:
+                        self.adapter.add_extra_values(torch.zeros_like(batch.extra_values.detach()),
+                                                      is_unconditional=True)
+
+                if has_adapter_img:
+                    if (self.adapter and isinstance(self.adapter, ControlNetModel)) or (
+                            self.assistant_adapter and isinstance(self.assistant_adapter, ControlNetModel)):
+                        if self.train_config.do_cfg:
+                            raise ValueError("ControlNetModel is not supported with CFG")
+                        with torch.set_grad_enabled(self.adapter is not None):
+                            adapter: ControlNetModel = self.assistant_adapter if self.assistant_adapter is not None else self.adapter
+                            adapter_multiplier = get_adapter_multiplier()
+                            with self.timer('encode_adapter'):
+                                # add_text_embeds is pooled_prompt_embeds for sdxl
+                                added_cond_kwargs = {}
+                                if self.sd.is_xl:
+                                    added_cond_kwargs["text_embeds"] = conditional_embeds.pooled_embeds
+                                    added_cond_kwargs['time_ids'] = self.sd.get_time_ids_from_latents(noisy_latents)
+                                down_block_res_samples, mid_block_res_sample = adapter(
+                                    noisy_latents,
+                                    timesteps,
+                                    encoder_hidden_states=conditional_embeds.text_embeds,
+                                    controlnet_cond=adapter_images,
+                                    conditioning_scale=1.0,
+                                    guess_mode=False,
+                                    added_cond_kwargs=added_cond_kwargs,
+                                    return_dict=False,
+                                )
+                                pred_kwargs['down_block_additional_residuals'] = down_block_res_samples
+                                pred_kwargs['mid_block_additional_residual'] = mid_block_res_sample
+
+                self.before_unet_predict()
+                # do a prior pred if we have an unconditional image, we will swap out the giadance later
+                if batch.unconditional_latents is not None or self.do_guided_loss:
+                    # do guided loss
+                    loss = self.get_guided_loss(
+                        noisy_latents=noisy_latents,
+                        conditional_embeds=conditional_embeds,
+                        match_adapter_assist=match_adapter_assist,
+                        network_weight_list=network_weight_list,
+                        timesteps=timesteps,
+                        pred_kwargs=pred_kwargs,
+                        batch=batch,
+                        noise=noise,
+                        unconditional_embeds=unconditional_embeds,
+                        mask_multiplier=mask_multiplier,
+                        prior_pred=prior_pred,
+                    )
+
+                else:
+                    with self.timer('predict_unet'):
+                        if unconditional_embeds is not None:
+                            unconditional_embeds = unconditional_embeds.to(self.device_torch, dtype=dtype).detach()
+                        noise_pred = self.predict_noise(
+                            noisy_latents=noisy_latents.to(self.device_torch, dtype=dtype),
+                            timesteps=timesteps,
+                            conditional_embeds=conditional_embeds.to(self.device_torch, dtype=dtype),
+                            unconditional_embeds=unconditional_embeds,
+                            **pred_kwargs
+                        )
+                    self.after_unet_predict()
+
+                    with self.timer('calculate_loss'):
+                        noise = noise.to(self.device_torch, dtype=dtype).detach()
+                        loss = self.calculate_loss(
+                            noise_pred=noise_pred,
+                            noise=noise,
+                            noisy_latents=noisy_latents,
+                            timesteps=timesteps,
+                            batch=batch,
+                            mask_multiplier=mask_multiplier,
+                            prior_pred=prior_pred,
+                        )
+                # check if nan
+                if torch.isnan(loss):
+                    print("loss is nan")
+                    loss = torch.zeros_like(loss).requires_grad_(True)
+
+                with self.timer('backward'):
+                    # todo we have multiplier seperated. works for now as res are not in same batch, but need to change
+                    loss = loss * loss_multiplier.mean()
+                    # IMPORTANT if gradient checkpointing do not leave with network when doing backward
+                    # it will destroy the gradients. This is because the network is a context manager
+                    # and will change the multipliers back to 0.0 when exiting. They will be
+                    # 0.0 for the backward pass and the gradients will be 0.0
+                    # I spent weeks on fighting this. DON'T DO IT
+                    # with fsdp_overlap_step_with_backward():
+                    # if self.is_bfloat:
+                    # loss.backward()
+                    # else:
+                    if not self.do_grad_scale:
+                        loss.backward()
+                    else:
+                        self.scaler.scale(loss).backward()
+
+        return loss.detach()
+        # flush()
+
+    def hook_train_loop(self, batch: Union[DataLoaderBatchDTO, List[DataLoaderBatchDTO]]):
+        if isinstance(batch, list):
+            batch_list = batch
+        else:
+            batch_list = [batch]
+        total_loss = None
+        self.optimizer.zero_grad()
+        for batch in batch_list:
+            loss = self.train_single_accumulation(batch)
+            if total_loss is None:
+                total_loss = loss
+            else:
+                total_loss += loss
+            if len(batch_list) > 1 and self.model_config.low_vram:
+                torch.cuda.empty_cache()
+
+
+        if not self.is_grad_accumulation_step:
+            # fix this for multi params
+            if self.train_config.optimizer != 'adafactor':
+                if self.do_grad_scale:
+                    self.scaler.unscale_(self.optimizer)
+                if isinstance(self.params[0], dict):
+                    for i in range(len(self.params)):
+                        torch.nn.utils.clip_grad_norm_(self.params[i]['params'], self.train_config.max_grad_norm)
+                else:
+                    torch.nn.utils.clip_grad_norm_(self.params, self.train_config.max_grad_norm)
+            # only step if we are not accumulating
+            with self.timer('optimizer_step'):
+                # self.optimizer.step()
+                if not self.do_grad_scale:
+                    self.optimizer.step()
+                else:
+                    self.scaler.step(self.optimizer)
+                    self.scaler.update()
+
+                self.optimizer.zero_grad(set_to_none=True)
+                if self.adapter and isinstance(self.adapter, CustomAdapter):
+                    self.adapter.post_weight_update()
+            if self.ema is not None:
+                with self.timer('ema_update'):
+                    self.ema.update()
+        else:
+            # gradient accumulation. Just a place for breakpoint
+            pass
+
+        # TODO Should we only step scheduler on grad step? If so, need to recalculate last step
+        with self.timer('scheduler_step'):
+            self.lr_scheduler.step()
+
+        if self.embedding is not None:
+            with self.timer('restore_embeddings'):
+                # Let's make sure we don't update any embedding weights besides the newly added token
+                self.embedding.restore_embeddings()
+        if self.adapter is not None and isinstance(self.adapter, ClipVisionAdapter):
+            with self.timer('restore_adapter'):
+                # Let's make sure we don't update any embedding weights besides the newly added token
+                self.adapter.restore_embeddings()
+
+        loss_dict = OrderedDict(
+            {'loss': loss.item()}
+        )
+
+        self.end_of_training_loop()
+
+        return loss_dict

extensions_built_in/sd_trainer/__init__.py

@@ -0,0 +1,30 @@
+# This is an example extension for custom training. It is great for experimenting with new ideas.
+from toolkit.extension import Extension
+
+
+# This is for generic training (LoRA, Dreambooth, FineTuning)
+class SDTrainerExtension(Extension):
+    # uid must be unique, it is how the extension is identified
+    uid = "sd_trainer"
+
+    # name is the name of the extension for printing
+    name = "SD Trainer"
+
+    # This is where your process class is loaded
+    # keep your imports in here so they don't slow down the rest of the program
+    @classmethod
+    def get_process(cls):
+        # import your process class here so it is only loaded when needed and return it
+        from .SDTrainer import SDTrainer
+        return SDTrainer
+
+
+# for backwards compatability
+class TextualInversionTrainer(SDTrainerExtension):
+    uid = "textual_inversion_trainer"
+
+
+AI_TOOLKIT_EXTENSIONS = [
+    # you can put a list of extensions here
+    SDTrainerExtension, TextualInversionTrainer
+]

extensions_built_in/sd_trainer/config/train.example.yaml

@@ -0,0 +1,91 @@
+---
+job: extension
+config:
+  name: test_v1
+  process:
+    - type: 'textual_inversion_trainer'
+      training_folder: "out/TI"
+      device: cuda:0
+      # for tensorboard logging
+      log_dir: "out/.tensorboard"
+      embedding:
+        trigger: "your_trigger_here"
+        tokens: 12
+        init_words: "man with short brown hair"
+        save_format: "safetensors"  # 'safetensors' or 'pt'
+      save:
+        dtype: float16 # precision to save
+        save_every: 100 # save every this many steps
+        max_step_saves_to_keep: 5 # only affects step counts
+      datasets:
+        - folder_path: "/path/to/dataset"
+          caption_ext: "txt"
+          default_caption: "[trigger]"
+          buckets: true
+          resolution: 512
+      train:
+        noise_scheduler: "ddpm" # or "ddpm", "lms", "euler_a"
+        steps: 3000
+        weight_jitter: 0.0
+        lr: 5e-5
+        train_unet: false
+        gradient_checkpointing: true
+        train_text_encoder: false
+        optimizer: "adamw"
+#        optimizer: "prodigy"
+        optimizer_params:
+          weight_decay: 1e-2
+        lr_scheduler: "constant"
+        max_denoising_steps: 1000
+        batch_size: 4
+        dtype: bf16
+        xformers: true
+        min_snr_gamma: 5.0
+#        skip_first_sample: true
+        noise_offset: 0.0 # not needed for this
+      model:
+        # objective reality v2
+        name_or_path: "https://civitai.com/models/128453?modelVersionId=142465"
+        is_v2: false  # for v2 models
+        is_xl: false  # for SDXL models
+        is_v_pred: false # for v-prediction models (most v2 models)
+      sample:
+        sampler: "ddpm" # must match train.noise_scheduler
+        sample_every: 100 # sample every this many steps
+        width: 512
+        height: 512
+        prompts:
+          - "photo of [trigger] laughing"
+          - "photo of [trigger] smiling"
+          - "[trigger] close up"
+          - "dark scene [trigger] frozen"
+          - "[trigger] nighttime"
+          - "a painting of [trigger]"
+          - "a drawing of [trigger]"
+          - "a cartoon of [trigger]"
+          - "[trigger] pixar style"
+          - "[trigger] costume"
+        neg: ""
+        seed: 42
+        walk_seed: false
+        guidance_scale: 7
+        sample_steps: 20
+        network_multiplier: 1.0
+
+      logging:
+        log_every: 10 # log every this many steps
+        use_wandb: false # not supported yet
+        verbose: false
+
+# You can put any information you want here, and it will be saved in the model.
+# The below is an example, but you can put your grocery list in it if you want.
+# It is saved in the model so be aware of that. The software will include this
+# plus some other information for you automatically
+meta:
+  # [name] gets replaced with the name above
+  name: "[name]"
+#  version: '1.0'
+#  creator:
+#    name: Your Name
+#    email: [email protected]
+#    website: https://your.website

extensions_built_in/ultimate_slider_trainer/UltimateSliderTrainerProcess.py

@@ -0,0 +1,533 @@
+import copy
+import random
+from collections import OrderedDict
+import os
+from contextlib import nullcontext
+from typing import Optional, Union, List
+from torch.utils.data import ConcatDataset, DataLoader
+
+from toolkit.config_modules import ReferenceDatasetConfig
+from toolkit.data_loader import PairedImageDataset
+from toolkit.prompt_utils import concat_prompt_embeds, split_prompt_embeds, build_latent_image_batch_for_prompt_pair
+from toolkit.stable_diffusion_model import StableDiffusion, PromptEmbeds
+from toolkit.train_tools import get_torch_dtype, apply_snr_weight
+import gc
+from toolkit import train_tools
+import torch
+from jobs.process import BaseSDTrainProcess
+import random
+
+import random
+from collections import OrderedDict
+from tqdm import tqdm
+
+from toolkit.config_modules import SliderConfig
+from toolkit.train_tools import get_torch_dtype, apply_snr_weight
+import gc
+from toolkit import train_tools
+from toolkit.prompt_utils import \
+    EncodedPromptPair, ACTION_TYPES_SLIDER, \
+    EncodedAnchor, concat_prompt_pairs, \
+    concat_anchors, PromptEmbedsCache, encode_prompts_to_cache, build_prompt_pair_batch_from_cache, split_anchors, \
+    split_prompt_pairs
+
+import torch
+
+
+def flush():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+class UltimateSliderConfig(SliderConfig):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.additional_losses: List[str] = kwargs.get('additional_losses', [])
+        self.weight_jitter: float = kwargs.get('weight_jitter', 0.0)
+        self.img_loss_weight: float = kwargs.get('img_loss_weight', 1.0)
+        self.cfg_loss_weight: float = kwargs.get('cfg_loss_weight', 1.0)
+        self.datasets: List[ReferenceDatasetConfig] = [ReferenceDatasetConfig(**d) for d in kwargs.get('datasets', [])]
+
+
+class UltimateSliderTrainerProcess(BaseSDTrainProcess):
+    sd: StableDiffusion
+    data_loader: DataLoader = None
+
+    def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
+        super().__init__(process_id, job, config, **kwargs)
+        self.prompt_txt_list = None
+        self.step_num = 0
+        self.start_step = 0
+        self.device = self.get_conf('device', self.job.device)
+        self.device_torch = torch.device(self.device)
+        self.slider_config = UltimateSliderConfig(**self.get_conf('slider', {}))
+
+        self.prompt_cache = PromptEmbedsCache()
+        self.prompt_pairs: list[EncodedPromptPair] = []
+        self.anchor_pairs: list[EncodedAnchor] = []
+        # keep track of prompt chunk size
+        self.prompt_chunk_size = 1
+
+        # store a list of all the prompts from the dataset so we can cache it
+        self.dataset_prompts = []
+        self.train_with_dataset = self.slider_config.datasets is not None and len(self.slider_config.datasets) > 0
+
+    def load_datasets(self):
+        if self.data_loader is None and \
+                self.slider_config.datasets is not None and len(self.slider_config.datasets) > 0:
+            print(f"Loading datasets")
+            datasets = []
+            for dataset in self.slider_config.datasets:
+                print(f" - Dataset: {dataset.pair_folder}")
+                config = {
+                    'path': dataset.pair_folder,
+                    'size': dataset.size,
+                    'default_prompt': dataset.target_class,
+                    'network_weight': dataset.network_weight,
+                    'pos_weight': dataset.pos_weight,
+                    'neg_weight': dataset.neg_weight,
+                    'pos_folder': dataset.pos_folder,
+                    'neg_folder': dataset.neg_folder,
+                }
+                image_dataset = PairedImageDataset(config)
+                datasets.append(image_dataset)
+
+                # capture all the prompts from it so we can cache the embeds
+                self.dataset_prompts += image_dataset.get_all_prompts()
+
+            concatenated_dataset = ConcatDataset(datasets)
+            self.data_loader = DataLoader(
+                concatenated_dataset,
+                batch_size=self.train_config.batch_size,
+                shuffle=True,
+                num_workers=2
+            )
+
+    def before_model_load(self):
+        pass
+
+    def hook_before_train_loop(self):
+        # load any datasets if they were passed
+        self.load_datasets()
+
+        # read line by line from file
+        if self.slider_config.prompt_file:
+            self.print(f"Loading prompt file from {self.slider_config.prompt_file}")
+            with open(self.slider_config.prompt_file, 'r', encoding='utf-8') as f:
+                self.prompt_txt_list = f.readlines()
+                # clean empty lines
+                self.prompt_txt_list = [line.strip() for line in self.prompt_txt_list if len(line.strip()) > 0]
+
+            self.print(f"Found {len(self.prompt_txt_list)} prompts.")
+
+            if not self.slider_config.prompt_tensors:
+                print(f"Prompt tensors not found. Building prompt tensors for {self.train_config.steps} steps.")
+                # shuffle
+                random.shuffle(self.prompt_txt_list)
+                # trim to max steps
+                self.prompt_txt_list = self.prompt_txt_list[:self.train_config.steps]
+                # trim list to our max steps
+
+        cache = PromptEmbedsCache()
+
+        # get encoded latents for our prompts
+        with torch.no_grad():
+            # list of neutrals. Can come from file or be empty
+            neutral_list = self.prompt_txt_list if self.prompt_txt_list is not None else [""]
+
+            # build the prompts to cache
+            prompts_to_cache = []
+            for neutral in neutral_list:
+                for target in self.slider_config.targets:
+                    prompt_list = [
+                        f"{target.target_class}",  # target_class
+                        f"{target.target_class} {neutral}",  # target_class with neutral
+                        f"{target.positive}",  # positive_target
+                        f"{target.positive} {neutral}",  # positive_target with neutral
+                        f"{target.negative}",  # negative_target
+                        f"{target.negative} {neutral}",  # negative_target with neutral
+                        f"{neutral}",  # neutral
+                        f"{target.positive} {target.negative}",  # both targets
+                        f"{target.negative} {target.positive}",  # both targets reverse
+                    ]
+                    prompts_to_cache += prompt_list
+
+            # remove duplicates
+            prompts_to_cache = list(dict.fromkeys(prompts_to_cache))
+
+            # trim to max steps if max steps is lower than prompt count
+            prompts_to_cache = prompts_to_cache[:self.train_config.steps]
+
+            if len(self.dataset_prompts) > 0:
+                # add the prompts from the dataset
+                prompts_to_cache += self.dataset_prompts
+
+            # encode them
+            cache = encode_prompts_to_cache(
+                prompt_list=prompts_to_cache,
+                sd=self.sd,
+                cache=cache,
+                prompt_tensor_file=self.slider_config.prompt_tensors
+            )
+
+            prompt_pairs = []
+            prompt_batches = []
+            for neutral in tqdm(neutral_list, desc="Building Prompt Pairs", leave=False):
+                for target in self.slider_config.targets:
+                    prompt_pair_batch = build_prompt_pair_batch_from_cache(
+                        cache=cache,
+                        target=target,
+                        neutral=neutral,
+
+                    )
+                    if self.slider_config.batch_full_slide:
+                        # concat the prompt pairs
+                        # this allows us to run the entire 4 part process in one shot (for slider)
+                        self.prompt_chunk_size = 4
+                        concat_prompt_pair_batch = concat_prompt_pairs(prompt_pair_batch).to('cpu')
+                        prompt_pairs += [concat_prompt_pair_batch]
+                    else:
+                        self.prompt_chunk_size = 1
+                        # do them one at a time (probably not necessary after new optimizations)
+                        prompt_pairs += [x.to('cpu') for x in prompt_pair_batch]
+
+        # move to cpu to save vram
+        # We don't need text encoder anymore, but keep it on cpu for sampling
+        # if text encoder is list
+        if isinstance(self.sd.text_encoder, list):
+            for encoder in self.sd.text_encoder:
+                encoder.to("cpu")
+        else:
+            self.sd.text_encoder.to("cpu")
+        self.prompt_cache = cache
+        self.prompt_pairs = prompt_pairs
+        # end hook_before_train_loop
+
+        # move vae to device so we can encode on the fly
+        # todo cache latents
+        self.sd.vae.to(self.device_torch)
+        self.sd.vae.eval()
+        self.sd.vae.requires_grad_(False)
+
+        if self.train_config.gradient_checkpointing:
+            # may get disabled elsewhere
+            self.sd.unet.enable_gradient_checkpointing()
+
+        flush()
+        # end hook_before_train_loop
+
+    def hook_train_loop(self, batch):
+        dtype = get_torch_dtype(self.train_config.dtype)
+
+        with torch.no_grad():
+            ### LOOP SETUP ###
+            noise_scheduler = self.sd.noise_scheduler
+            optimizer = self.optimizer
+            lr_scheduler = self.lr_scheduler
+
+            ### TARGET_PROMPTS ###
+            # get a random pair
+            prompt_pair: EncodedPromptPair = self.prompt_pairs[
+                torch.randint(0, len(self.prompt_pairs), (1,)).item()
+            ]
+            # move to device and dtype
+            prompt_pair.to(self.device_torch, dtype=dtype)
+
+            ### PREP REFERENCE IMAGES ###
+
+            imgs, prompts, network_weights = batch
+            network_pos_weight, network_neg_weight = network_weights
+
+            if isinstance(network_pos_weight, torch.Tensor):
+                network_pos_weight = network_pos_weight.item()
+            if isinstance(network_neg_weight, torch.Tensor):
+                network_neg_weight = network_neg_weight.item()
+
+            # get an array of random floats between -weight_jitter and weight_jitter
+            weight_jitter = self.slider_config.weight_jitter
+            if weight_jitter > 0.0:
+                jitter_list = random.uniform(-weight_jitter, weight_jitter)
+                network_pos_weight += jitter_list
+                network_neg_weight += (jitter_list * -1.0)
+
+            # if items in network_weight list are tensors, convert them to floats
+            imgs: torch.Tensor = imgs.to(self.device_torch, dtype=dtype)
+            # split batched images in half so left is negative and right is positive
+            negative_images, positive_images = torch.chunk(imgs, 2, dim=3)
+
+            height = positive_images.shape[2]
+            width = positive_images.shape[3]
+            batch_size = positive_images.shape[0]
+
+            positive_latents = self.sd.encode_images(positive_images)
+            negative_latents = self.sd.encode_images(negative_images)
+
+            self.sd.noise_scheduler.set_timesteps(
+                self.train_config.max_denoising_steps, device=self.device_torch
+            )
+
+            timesteps = torch.randint(0, self.train_config.max_denoising_steps, (1,), device=self.device_torch)
+            current_timestep_index = timesteps.item()
+            current_timestep = noise_scheduler.timesteps[current_timestep_index]
+            timesteps = timesteps.long()
+
+            # get noise
+            noise_positive = self.sd.get_latent_noise(
+                pixel_height=height,
+                pixel_width=width,
+                batch_size=batch_size,
+                noise_offset=self.train_config.noise_offset,
+            ).to(self.device_torch, dtype=dtype)
+
+            noise_negative = noise_positive.clone()
+
+            # Add noise to the latents according to the noise magnitude at each timestep
+            # (this is the forward diffusion process)
+            noisy_positive_latents = noise_scheduler.add_noise(positive_latents, noise_positive, timesteps)
+            noisy_negative_latents = noise_scheduler.add_noise(negative_latents, noise_negative, timesteps)
+
+            ### CFG SLIDER TRAINING PREP ###
+
+            # get CFG txt latents
+            noisy_cfg_latents = build_latent_image_batch_for_prompt_pair(
+                pos_latent=noisy_positive_latents,
+                neg_latent=noisy_negative_latents,
+                prompt_pair=prompt_pair,
+                prompt_chunk_size=self.prompt_chunk_size,
+            )
+            noisy_cfg_latents.requires_grad = False
+
+            assert not self.network.is_active
+
+            # 4.20 GB RAM for 512x512
+            positive_latents = self.sd.predict_noise(
+                latents=noisy_cfg_latents,
+                text_embeddings=train_tools.concat_prompt_embeddings(
+                    prompt_pair.positive_target,  # negative prompt
+                    prompt_pair.negative_target,  # positive prompt
+                    self.train_config.batch_size,
+                ),
+                timestep=current_timestep,
+                guidance_scale=1.0
+            )
+            positive_latents.requires_grad = False
+
+            neutral_latents = self.sd.predict_noise(
+                latents=noisy_cfg_latents,
+                text_embeddings=train_tools.concat_prompt_embeddings(
+                    prompt_pair.positive_target,  # negative prompt
+                    prompt_pair.empty_prompt,  # positive prompt (normally neutral
+                    self.train_config.batch_size,
+                ),
+                timestep=current_timestep,
+                guidance_scale=1.0
+            )
+            neutral_latents.requires_grad = False
+
+            unconditional_latents = self.sd.predict_noise(
+                latents=noisy_cfg_latents,
+                text_embeddings=train_tools.concat_prompt_embeddings(
+                    prompt_pair.positive_target,  # negative prompt
+                    prompt_pair.positive_target,  # positive prompt
+                    self.train_config.batch_size,
+                ),
+                timestep=current_timestep,
+                guidance_scale=1.0
+            )
+            unconditional_latents.requires_grad = False
+
+            positive_latents_chunks = torch.chunk(positive_latents, self.prompt_chunk_size, dim=0)
+            neutral_latents_chunks = torch.chunk(neutral_latents, self.prompt_chunk_size, dim=0)
+            unconditional_latents_chunks = torch.chunk(unconditional_latents, self.prompt_chunk_size, dim=0)
+            prompt_pair_chunks = split_prompt_pairs(prompt_pair, self.prompt_chunk_size)
+            noisy_cfg_latents_chunks = torch.chunk(noisy_cfg_latents, self.prompt_chunk_size, dim=0)
+            assert len(prompt_pair_chunks) == len(noisy_cfg_latents_chunks)
+
+            noisy_latents = torch.cat([noisy_positive_latents, noisy_negative_latents], dim=0)
+            noise = torch.cat([noise_positive, noise_negative], dim=0)
+            timesteps = torch.cat([timesteps, timesteps], dim=0)
+            network_multiplier = [network_pos_weight * 1.0, network_neg_weight * -1.0]
+
+        flush()
+
+        loss_float = None
+        loss_mirror_float = None
+
+        self.optimizer.zero_grad()
+        noisy_latents.requires_grad = False
+
+        # TODO allow both processed to train text encoder, for now, we just to unet and cache all text encodes
+        # if training text encoder enable grads, else do context of no grad
+        # with torch.set_grad_enabled(self.train_config.train_text_encoder):
+        #     # text encoding
+        #     embedding_list = []
+        #     # embed the prompts
+        #     for prompt in prompts:
+        #         embedding = self.sd.encode_prompt(prompt).to(self.device_torch, dtype=dtype)
+        #         embedding_list.append(embedding)
+        #     conditional_embeds = concat_prompt_embeds(embedding_list)
+        #     conditional_embeds = concat_prompt_embeds([conditional_embeds, conditional_embeds])
+
+        if self.train_with_dataset:
+            embedding_list = []
+            with torch.set_grad_enabled(self.train_config.train_text_encoder):
+                for prompt in prompts:
+                    # get embedding form cache
+                    embedding = self.prompt_cache[prompt]
+                    embedding = embedding.to(self.device_torch, dtype=dtype)
+                    embedding_list.append(embedding)
+                conditional_embeds = concat_prompt_embeds(embedding_list)
+                # double up so we can do both sides of the slider
+                conditional_embeds = concat_prompt_embeds([conditional_embeds, conditional_embeds])
+        else:
+            # throw error. Not supported yet
+            raise Exception("Datasets and targets required for ultimate slider")
+
+        if self.model_config.is_xl:
+            # todo also allow for setting this for low ram in general, but sdxl spikes a ton on back prop
+            network_multiplier_list = network_multiplier
+            noisy_latent_list = torch.chunk(noisy_latents, 2, dim=0)
+            noise_list = torch.chunk(noise, 2, dim=0)
+            timesteps_list = torch.chunk(timesteps, 2, dim=0)
+            conditional_embeds_list = split_prompt_embeds(conditional_embeds)
+        else:
+            network_multiplier_list = [network_multiplier]
+            noisy_latent_list = [noisy_latents]
+            noise_list = [noise]
+            timesteps_list = [timesteps]
+            conditional_embeds_list = [conditional_embeds]
+
+        ## DO REFERENCE IMAGE TRAINING ##
+
+        reference_image_losses = []
+        # allow to chunk it out to save vram
+        for network_multiplier, noisy_latents, noise, timesteps, conditional_embeds in zip(
+                network_multiplier_list, noisy_latent_list, noise_list, timesteps_list, conditional_embeds_list
+        ):
+            with self.network:
+                assert self.network.is_active
+
+                self.network.multiplier = network_multiplier
+
+                noise_pred = self.sd.predict_noise(
+                    latents=noisy_latents.to(self.device_torch, dtype=dtype),
+                    conditional_embeddings=conditional_embeds.to(self.device_torch, dtype=dtype),
+                    timestep=timesteps,
+                )
+                noise = noise.to(self.device_torch, dtype=dtype)
+
+                if self.sd.prediction_type == 'v_prediction':
+                    # v-parameterization training
+                    target = noise_scheduler.get_velocity(noisy_latents, noise, timesteps)
+                else:
+                    target = noise
+
+                loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="none")
+                loss = loss.mean([1, 2, 3])
+
+                # todo add snr gamma here
+                if self.train_config.min_snr_gamma is not None and self.train_config.min_snr_gamma > 0.000001:
+                    # add min_snr_gamma
+                    loss = apply_snr_weight(loss, timesteps, noise_scheduler, self.train_config.min_snr_gamma)
+
+                loss = loss.mean()
+                loss = loss * self.slider_config.img_loss_weight
+                loss_slide_float = loss.item()
+
+                loss_float = loss.item()
+                reference_image_losses.append(loss_float)
+
+                # back propagate loss to free ram
+                loss.backward()
+                flush()
+
+        ## DO CFG SLIDER TRAINING ##
+
+        cfg_loss_list = []
+
+        with self.network:
+            assert self.network.is_active
+            for prompt_pair_chunk, \
+                    noisy_cfg_latent_chunk, \
+                    positive_latents_chunk, \
+                    neutral_latents_chunk, \
+                    unconditional_latents_chunk \
+                    in zip(
+                prompt_pair_chunks,
+                noisy_cfg_latents_chunks,
+                positive_latents_chunks,
+                neutral_latents_chunks,
+                unconditional_latents_chunks,
+            ):
+                self.network.multiplier = prompt_pair_chunk.multiplier_list
+
+                target_latents = self.sd.predict_noise(
+                    latents=noisy_cfg_latent_chunk,
+                    text_embeddings=train_tools.concat_prompt_embeddings(
+                        prompt_pair_chunk.positive_target,  # negative prompt
+                        prompt_pair_chunk.target_class,  # positive prompt
+                        self.train_config.batch_size,
+                    ),
+                    timestep=current_timestep,
+                    guidance_scale=1.0
+                )
+
+                guidance_scale = 1.0
+
+                offset = guidance_scale * (positive_latents_chunk - unconditional_latents_chunk)
+
+                # make offset multiplier based on actions
+                offset_multiplier_list = []
+                for action in prompt_pair_chunk.action_list:
+                    if action == ACTION_TYPES_SLIDER.ERASE_NEGATIVE:
+                        offset_multiplier_list += [-1.0]
+                    elif action == ACTION_TYPES_SLIDER.ENHANCE_NEGATIVE:
+                        offset_multiplier_list += [1.0]
+
+                offset_multiplier = torch.tensor(offset_multiplier_list).to(offset.device, dtype=offset.dtype)
+                # make offset multiplier match rank of offset
+                offset_multiplier = offset_multiplier.view(offset.shape[0], 1, 1, 1)
+                offset *= offset_multiplier
+
+                offset_neutral = neutral_latents_chunk
+                # offsets are already adjusted on a per-batch basis
+                offset_neutral += offset
+
+                # 16.15 GB RAM for 512x512 -> 4.20GB RAM for 512x512 with new grad_checkpointing
+                loss = torch.nn.functional.mse_loss(target_latents.float(), offset_neutral.float(), reduction="none")
+                loss = loss.mean([1, 2, 3])
+
+                if self.train_config.min_snr_gamma is not None and self.train_config.min_snr_gamma > 0.000001:
+                    # match batch size
+                    timesteps_index_list = [current_timestep_index for _ in range(target_latents.shape[0])]
+                    # add min_snr_gamma
+                    loss = apply_snr_weight(loss, timesteps_index_list, noise_scheduler,
+                                            self.train_config.min_snr_gamma)
+
+                loss = loss.mean() * prompt_pair_chunk.weight * self.slider_config.cfg_loss_weight
+
+                loss.backward()
+                cfg_loss_list.append(loss.item())
+                del target_latents
+                del offset_neutral
+                del loss
+                flush()
+
+        # apply gradients
+        optimizer.step()
+        lr_scheduler.step()
+
+        # reset network
+        self.network.multiplier = 1.0
+
+        reference_image_loss = sum(reference_image_losses) / len(reference_image_losses) if len(
+            reference_image_losses) > 0 else 0.0
+        cfg_loss = sum(cfg_loss_list) / len(cfg_loss_list) if len(cfg_loss_list) > 0 else 0.0
+
+        loss_dict = OrderedDict({
+            'loss/img': reference_image_loss,
+            'loss/cfg': cfg_loss,
+        })
+
+        return loss_dict
+        # end hook_train_loop

extensions_built_in/ultimate_slider_trainer/__init__.py

@@ -0,0 +1,25 @@
+# This is an example extension for custom training. It is great for experimenting with new ideas.
+from toolkit.extension import Extension
+
+
+# We make a subclass of Extension
+class UltimateSliderTrainer(Extension):
+    # uid must be unique, it is how the extension is identified
+    uid = "ultimate_slider_trainer"
+
+    # name is the name of the extension for printing
+    name = "Ultimate Slider Trainer"
+
+    # This is where your process class is loaded
+    # keep your imports in here so they don't slow down the rest of the program
+    @classmethod
+    def get_process(cls):
+        # import your process class here so it is only loaded when needed and return it
+        from .UltimateSliderTrainerProcess import UltimateSliderTrainerProcess
+        return UltimateSliderTrainerProcess
+
+
+AI_TOOLKIT_EXTENSIONS = [
+    # you can put a list of extensions here
+    UltimateSliderTrainer
+]

extensions_built_in/ultimate_slider_trainer/config/train.example.yaml

@@ -0,0 +1,107 @@
+---
+job: extension
+config:
+  name: example_name
+  process:
+    - type: 'image_reference_slider_trainer'
+      training_folder: "/mnt/Train/out/LoRA"
+      device: cuda:0
+      # for tensorboard logging
+      log_dir: "/home/jaret/Dev/.tensorboard"
+      network:
+        type: "lora"
+        linear: 8
+        linear_alpha: 8
+      train:
+        noise_scheduler: "ddpm" # or "ddpm", "lms", "euler_a"
+        steps: 5000
+        lr: 1e-4
+        train_unet: true
+        gradient_checkpointing: true
+        train_text_encoder: true
+        optimizer: "adamw"
+        optimizer_params:
+          weight_decay: 1e-2
+        lr_scheduler: "constant"
+        max_denoising_steps: 1000
+        batch_size: 1
+        dtype: bf16
+        xformers: true
+        skip_first_sample: true
+        noise_offset: 0.0
+      model:
+        name_or_path: "/path/to/model.safetensors"
+        is_v2: false  # for v2 models
+        is_xl: false  # for SDXL models
+        is_v_pred: false # for v-prediction models (most v2 models)
+      save:
+        dtype: float16 # precision to save
+        save_every: 1000 # save every this many steps
+        max_step_saves_to_keep: 2 # only affects step counts
+      sample:
+        sampler: "ddpm" # must match train.noise_scheduler
+        sample_every: 100 # sample every this many steps
+        width: 512
+        height: 512
+        prompts:
+          - "photo of a woman with red hair taking a selfie --m -3"
+          - "photo of a woman with red hair taking a selfie --m -1"
+          - "photo of a woman with red hair taking a selfie --m 1"
+          - "photo of a woman with red hair taking a selfie --m 3"
+          - "close up photo of a man smiling at the camera, in a tank top --m -3"
+          - "close up photo of a man smiling at the camera, in a tank top--m -1"
+          - "close up photo of a man smiling at the camera, in a tank top --m 1"
+          - "close up photo of a man smiling at the camera, in a tank top --m 3"
+          - "photo of a blonde woman smiling, barista --m -3"
+          - "photo of a blonde woman smiling, barista --m -1"
+          - "photo of a blonde woman smiling, barista --m 1"
+          - "photo of a blonde woman smiling, barista --m 3"
+          - "photo of a Christina Hendricks --m -1"
+          - "photo of a Christina Hendricks --m -1"
+          - "photo of a Christina Hendricks --m 1"
+          - "photo of a Christina Hendricks --m 3"
+          - "photo of a Christina Ricci --m -3"
+          - "photo of a Christina Ricci --m -1"
+          - "photo of a Christina Ricci --m 1"
+          - "photo of a Christina Ricci --m 3"
+        neg: "cartoon, fake, drawing, illustration, cgi, animated, anime"
+        seed: 42
+        walk_seed: false
+        guidance_scale: 7
+        sample_steps: 20
+        network_multiplier: 1.0
+
+      logging:
+        log_every: 10 # log every this many steps
+        use_wandb: false # not supported yet
+        verbose: false
+
+      slider:
+        datasets:
+          - pair_folder: "/path/to/folder/side/by/side/images"
+            network_weight: 2.0
+            target_class: "" # only used as default if caption txt are not present
+            size: 512
+          - pair_folder: "/path/to/folder/side/by/side/images"
+            network_weight: 4.0
+            target_class: "" # only used as default if caption txt are not present
+            size: 512
+
+
+# you can put any information you want here, and it will be saved in the model
+# the below is an example. I recommend doing trigger words at a minimum
+# in the metadata. The software will include this plus some other information
+meta:
+  name: "[name]"  # [name] gets replaced with the name above
+  description: A short description of your model
+  trigger_words:
+  - put
+  - trigger
+  - words
+  - here
+  version: '0.1'
+  creator:
+    name: Your Name
+    email: [email protected]
+    website: https://yourwebsite.com
+  any: All meta data above is arbitrary, it can be whatever you want.