Dreambooth in Stable Diffusion Web UI

In this tutorial, I will show you how to use DreamBooth to finetune pre-trained open-source text-to-image Stable Diffusion (SD) model. Our objective is to obtain a new checkpoint (CKPT) model which we create by training SD model with DreamBooth and images of Arnold Schwarzenegger.

Training is an intensive task and requires a powerful GPU. Luckily, we can use Collab notebook which is a Jupyter notebook used to run Python code. We will run notebook on Google Collab which is is a way to run the notebook on a virtual machine by Google. But most importantly, Google Collab provides Nvidia Tesla T4 GPU with 16GB vRAM so you do not need buy expensive hardware to finetune SD model.

TheLastBen from Github makes the entire process easy for us. I modified his Collab notebook (removed SD Web UI part) as I prefer to generate images local instance of Web UI. You just need to save my notebook to your Google drive and run the entire process on copy of my notebook in your Google Collab. You don't have to to care about the code, there is only like forms there that you can run. Of course, if you are interested in the code you can check it out if you want. But all you we need is a Google account to run Collab.

The good things is that your images will not be shared with any company or model. You just train a model and when it is created you will download it from Google drive to the local hard disk. Finally, you will run Stable Diffusion web UI on your PC and generate images with the text-to-image function integrated in Web UI.

Before you start working with Google Collab, you need to prepare your training images in advance. The images should have a resolution of 512x512 or 769x768 based on the Stable Diffusion model that we will be fine-tuning. In my case, I prepared 26 images, but DreamBooth requires 3 to 5 images (Figure 1).

Figure 1 - Training Images of Arnlod Schwarzenneger

1. Save notebook to your Google Drive

Open my notebook fast-DreambBoth.ipynb and save it to your Google drive (File -> Save a copy in Drive). The Directory "Collab Notebooks" is created in your Google drive. Got to the directory and right click on the file "Copy of fast-DreambBoth.ipynb". Rename the notebook if needed.

To open the notebook in your Google Collab, right click on the notebook and select Open in Application-> Google Colaboratory.

2. Mount your Google Drive

Mount your Google drive so the Google Collab can access your Google drive and store the files there. Just click the play button on the left to run the cell.

3. Install Dependencies

Click the Play button marked in the red circle to install the dependencies (Figure 2). The code inside cell is successfully.

Figure 2 - Cell Dependencies Executed Successfully

4. Model Download

The code in this cell downloads a pre-trained Stable Diffusion model, which we will fine-tune using DreamBooth. Based on the resolution of your training images, select the SD version 2.1 model (768x768) or keep the default 1.5 model (512x512).

5. DreamBooth

5.1 Create/Load Session

Enter session name Arnie and run the cell. A Fast-DreamBoth/Sessions/Arnie directory structure will be created on your Google Drive (Figure 3).

Figure 3 - New Directories Created on Google Drive

5.2 DreamBooth Instance Images

Now, you need upload your training images that are ready in your local drive to your Google drive. Before you click the Play button consider the following:

If you prefer to resize or crop images manually, uncheck the box "Smart_Crop_images" (Figure 4). Run the cell and click Browse button to navigate to your images located in your local PC. Images will be uploaded in to your Google drive in the directory: Fast-DreamBoth/Sessions/Arnie/instance-images.

Figure 4 - Configuration of Instance Images Cell

6. Training

The training takes about 20 minutes depending on the number of steps.

6.1 Start DreamBooth

Figure 5 - Training

When the training is complete, the new Arnie.cpkt model is saved to the Fast-DreamBoth/Sessions/Arnie folder on Google Drive. You can now disconnect from your Google Collab.

7. Testing model with SD Web UI and select the model

You can close the Collab now and download your new Checkpoint model Arnie.cpkt from Google drive to your computer. Start SD Web UI and select the model Arnie.cpkt (Figure 6).

Figure 6 - Selecting Arnie.cpkt Model in SD Web UI

Below are several images we created using the Stable Diffusion web UI using the "generate a photo of a young muscular bodybuilder" prompt. We used the Arnie.ckpt model, which we created by training a standard Stable Diffusion 1.5 model using DreamBooth and photos of Arnold Schwarzenegger. These are the best photos I could generate that somewhat resemble Arnold Schwarzenegger's photos. To achieve better results, I should have used higher-quality training photos with more details or adjusted certain parameters during the training process.