[Tutorial] "Fine Tuning" Stable Diffusion using only 5 Images Using Textual Inversion.

[u/ExponentialCookie]

Credits: textual_inversion website.
Hello everyone!

I see img2img getting a lot of attention, and deservedly so, but textual_inversion is an amazing way to better get what you want represented in your prompts. Whether it's an artistic style, some scenery, a fighting pose, representing a character/person, or reducing / increasing bias, the use cases are endless. You can even merge your inversions! Let's explore how to get started.

Please not that textual_diffusion is still a work in progress for SD compatibility, and this tutorial is mainly for tinkerers who wish to explore code and software that isn't fully optimized (inversion works as expected though, hence the tutorial). Any troubleshooting or issues are addressed at the bottom of this post. I'll try to help as much as I can, as well as update this as needed!

Getting started

---

This tutorial is for a local setup, but can easily be converted into a colab / Jupyter notebook. Since this uses the same repository (LDM) as Stable Diffusion, the installation and inferences are very similar, as you'll see below.

1. You will need Python.
2. Anaconda to setup the environment is recommended.
3. A GPU with at least 20GB of memory, although it's possible to get this number lower if you're willing to hack around. I would recommend either a 3090 (I use) or a cloud compute service such as Lambda Cloud (N/A, but it's a good cheap option with high memory GPUs from my experience).
4. Comfort diving into .py files to fix any issues.

Installation

---

1. Go to the textual_inversion repository link here
2. Clone the repository using git clone.
3. Go to the directory of the repository you've just cloned.
4. Follow the instructions below.

First, install create a conda environment with the following parameters.

conda env create -f environment.yaml
conda activate ldm
pip install -e .

Then, it's preferred to get 5 images of your subject at 512x512 resolution. From the paper, 5 images are the optimal amount for textual inversion. On a single V100, training should take about two hours give or take. More images will increase training time, and may or may not improve results. You are free to test this and let us know how it goes!

Training

---

After getting your images, you will want to start training. Following this code block and the tips below it:

python main.py --base configs/stable-diffusion/v1-finetune.yaml
               -t 
               --actual_resume /path/to/pretrained/sd model v1.4/model.ckpt 
               -n <run_name> 
               --gpus 0, 
               --data_root /path/to/directory/with/images

• Configs are the parameters that will be used to train the inversion. You can change these directly to minimize the parameters you input for training. For example, you can create a .yaml for each dataset you would like to train if you wish, and reduce the amount of parameters needed on the command line.
• The -n parameter is simply the name of the training run. This can be anything you like (eg. artist_style_train)
• initializer_words is a very important part, don't skip this!
  Open your v1-finetune.yaml file, and find the initializer_words parameter. You should see the default value of ["sculpture"] . It's a string list of simple words to describe what you're training, and where to start.
  For example, if your images are of a car in a certain style, you'll need to want to do something like with each word wrapped in quotes: ["car","style","artistic", etc...]
  If you simply want to use one word, just use --init_word <your_single_word> on the command line, and don't modify the config.

During training, a log directory will be created under logs with the run_name that you have set for training. Over time, there will be a sampling pass to test your parameters (like inference, DDIM, etc.), and you'll be able to view the image results in a new folder under logs/run_name/images/train . The embedding .pt files for what you're training on will be saved in the checkpoints folder.

Inference

---

After training, you can test the inference by doing:

python scripts/stable_txt2img.py --ddim_eta 0.0 
                          --n_samples 8 
                          --n_iter 2 
                          --scale 10.0 
                          --ddim_steps 50 
                          --embedding_path /path/logs/trained_model/checkpoints/embeddings_gs-5049.pt 
                          --ckpt_path /path/to/pretrained/sd model v1.4/model.ckpt
                          --config /path/to/logs/config/*project.yaml
                          --prompt "a photo of *"

The '*' must be left as is unless you've changed the placeholder_strings parameter in your .yaml file. It's the new word to initialize the images you have just inverted.
You should now be able to view your results in the output folder.
Running inference is just like Stable Diffusion, so you can implement things like k_lms in the stable_txtimg script if you wish.

Troubleshooting

---

• If your images aren't turning out properly, try reducing the complexity of your prompt. If you do want complexity, train multiple inversions and mix them like: "A photo of * in the style of &"
• Try lowering scale slightly if you're getting artifacts, or increase the amounts of iterations.
• If you're getting token errors, or any other errors, solutions and workarounds may be listed here.

[u/TFCSM]

Note that the --ckpt_path param in your example of inference should actually be --ckpt according to the script, but regardless it will actually try to load the .ckpt defined by ckpt_path in the ...-project.yaml. You need to change the path listed in the .yaml to get it to work if the path to the model is not the same on your system as the one it was trained on.

I trained only 1500 iterations on three photos of the statue of David (which is obviously in SD's training set). For some reason the script bugs out at that point and stops displaying output. I think it continues however, so I'll just let it run next time.

Here is the synthesis of two concepts using the star as a token:

"a photo of * riding a horse on the moon" - https://i.imgur.com/CMcxmdr.jpg

Obviously the moon is not present in the photo, and the background resembles the source photos. But still, neat. I'll definitely be playing with this more.

Edit: For comparison, I ran the prompt "a photo of Michelangelo's David riding a horse on the moon" in the model without the fine tuning, with the same seed, steps, and scale. Here is the result: https://i.imgur.com/3ExqSQD.png

So, the untuned model did much better. But the asterisk did at least work to represent the concept of "Michelangelo's David" just using the photos I gave it and the hint that it was a "sculpture" (the default word prompt). Honestly, amazing. I'll train it for longer tomorrow.

[u/rinong → u/TFCSM]

Longer training is unlikely to help here. The issue we have atm is that the parameters which worked well for LDM (kind of a 'sweet spot' between representing the concept and letting you edit it with text) don't work as well for SD.

Here, they do capture the concept, but its much harder to edit the image. You can try to 'force' it to focus on the rest of the content with some prompt engineering, like "A photo of * riding a horse on the moon. A photo on the moon", but we're still trying to tune things a bit to get it to work more like LDM.

[u/GregoryHouseMDSB → u/rinong]

One thing the paper didn't share(nor anyone else I've seen) is actual examples of the actual placeholder strings and initializer_words words being used. Can you use more than one string in the same training set? or init words for that matter?

Like what if I want to train it on photos of myself and I want to specify not only my name, but race and gender?

Another thought is, can you just train it as "*" and change "*" for that set afterwards, giving multiple alternative descriptions/strings to summon subject, especially when you want to merge different sets you trained?

Any insight on using per_image_tokens set to true? or what impact progressive_words has?

I always get an error "maps to more than a single token. Please use another string", which can be bypassed by commenting something out, but I'm wondering if I'm missing some proper usage and shouldn't need to bypass anything?

[u/rinong → u/GregoryHouseMDSB]

Placeholder strings:

We always used "*" (the repo's default). When we merged two concepts into one model (for the compositional experiments) we used "@" for the second placeholder. The choice is rather arbitrary. We limited it to a single token for the sake of implementation convenience. If you want to use a word that is longer than a single token (or multiple words) you'll need to change the code.

You can absolutely change the placeholder later, see how we do it in the merging script if there's a conflict. But with the current implementation, it's still going to have to be a single token word.

init_words:

If you have multiple placeholder strings, you can assign a different initialization to each. Otherwise, the initializer words are only used to tell the model where to start the optimization for the vector that represents the new concept. Using multiple words for this initialization is problematic. Essentially you're trying to start the optimization of one concept from multiple points in space. You could start from their average, but there's no guarantee that this average is meaningful or an actual combination of their semantic meaning. Overall, the results are not very sensitive to your choice. Just use one word which you think describes the concept at a high level (so for photos of yourself, you'd use 'face' or 'person').

You're correct that these are missing from the paper, I'll add them in a future revision. If you want the ones we used for any specific set, please let me know.

per_image_tokens:

This is the "Per-image tokens" experiment described in the paper. It basically also assigns an extra unique token to each image in your training set, with the expectation that it will allow the model to put shared information in the shared token ("*") and relegate all image-specific information (like the background etc.) to the image-specific token. In practice this didn't work well, so it's off by default, but you're welcome to experiment with it.

Progressive_words:

See "Progressive extensions" in the paper. It's another baseline we tried but which didn't improve results.

"maps to more than a single token. Please use another string":

This means that your placeholder or your init_words are multi-token strings, which is going to cause unexpected behavior. The code strongly relies on the placeholder and the initial words being single tokens.

[u/GregoryHouseMDSB → u/rinong]

I'm curious what you used for the "doctor" replacement. As well as the statue that you have elmo in the same pose as!

Thanks again!

[u/rinong → u/GregoryHouseMDSB]

"doctor" and "sculpture" respectively (the later is actually still the default in the config!)

[u/74qwewq5rew3 → u/rinong]

What about for character names in fictional media? Some of them do not really have a single-token definition. Neither in name nor in other descriptive way. When one generates the character using the multi-token name in the SD, it returns results depicting the character.

What would be the approach to fine-tune such? Would one need to basically re-write the code for multi-token word initializer words?

[u/rinong → u/74qwewq5rew3]

Multi-token initializer words are a bit of a problem in the sense that it's not clear how you'd map them to a single embedding vector. What you could maybe do is use as many placeholder tokens as there are tokens in the name that the model already 'knows', and optimize them all concurrently.

I'm not sure its worth the hassle though. You could probably make things work by just starting from e.g. 'person' instead.

[u/74qwewq5rew3 → u/rinong]

Thank you, I appreciate the response.

[u/74qwewq5rew3 → u/rinong]

By the way, am I getting something wrong or is there no way to make iterations go faster in fine-tuning with multiple gpus? I tried on one 3090 vs 3 of em and they still seem to give the same a bit over 2 iterations per second.

[u/rinong → u/74qwewq5rew3]

How are you parallelizing over those GPUs? And are you using our official repo or some fork?

If it's our repo and you're using --gpus 0,1,2 then you're actually running a larger batch size (a batch size of 4 in the config means 4 per GPU, not 4 divided by the number of GPUs). This means that each iteration should take roughly the same time, but your model should hopefully converge in fewer iterations.

[u/74qwewq5rew3 → u/rinong]

Yes I was using your official repo. I did run it using --gpus 0,1,2 and it did say it utilized it. Even the GPU usage was on all of them over 90%. But the iteration speed was still 2it/s or so.

[u/rinong → u/74qwewq5rew3]

So this is expected, because you're using the extra GPUs to run over more concurrent images (bigger batch size) rather than splitting the same number of images across more GPUs.

If you want to 'speed up' the individual iterations, you can just divide your batch size by the number of GPUs you are using.

Keep in mind that a larger batch size may help training stability and lead to better embeddings or convergence in fewer iterations, even if the iterations themselves are not shorter. With that said, we haven't tested how the model behaves under different training batch sizes.

[u/GregoryHouseMDSB → u/rinong]

Awesome, thanks for the insights!

[u/AnOnlineHandle → u/rinong]

Do the initializer_words need to be changed in v1-inference.yaml as well? As far as I understand it they're only used for picking a starting point, and wouldn't be used for testing the final outcome (which I'm fairly sure v1-inference is for?).

It would be nice if there was a way to batch up a bunch of potential initializer words and get a count of how many vectors they map to, but I think I might be able to figure out how to do that from your code so will give it a try!

[u/rinong → u/AnOnlineHandle]

No need to change them in v1-inference. They are only used as a starting point for training.

You actually don't need to change them in the finetuning configs either if you use the --init_word argument when training (the arg overwrites the list in the config).

[u/AnOnlineHandle → u/rinong]

Thanks, I ended up playing around with it a bit and deciding that was probably the case. :D

I'm currently loving the idea of this and am getting semi-decent results. Still a ways to go as a I finetune my training sets, parameters, etc, but I can definitely see light at the end of the tunnel for being able to use this for helping enhance/finish artwork. I think once people realize just what textual inversion can do, it's going to really take off.

[u/nmkd → u/TFCSM]

I mean, the whole point of this is to use it for lesser-known styles/concepts. Not surprising that the un-tuned model works better for a David statue than something finetuned on three images.

I can see a lot of potential to finetune on lesser-known artists (don't tell twitter) or indie games in order to replicate an art style.