The most important evolutionary leap in AI image generation is the creation of custom models. By that I mean secondary packages that replicate specific styles or objects, but they can also be used to inject a person's face into unique scenarios. This is possible thanks to the Dreambooth platform, developed by Google Research and Boston University. The most popular application of Dreambooth these days is fine-tuning models compatible with Stable Diffusion, and today we'll explain how to do it.
Pushing the Limits
Of the three models available for AI image generation, Stable Diffusion has proven to be the most flexible. Why? Two reasons: on one hand, it's free and open source, and on the other, it plays well with alternative models. But where do these models come from? From the users themselves, who take advantage of their high-end hardware or free online options to train them.
The problem with custom training or fine-tuning of Stable Diffusion is its VRAM consumption. Training a model with little video memory is almost impossible, however, thanks to the magic of Google Colab, Google Drive, HuggingFace, and the DotCSV YouTube channel, today we can do it without too many hiccups, and without investing thousands of euros in hardware. The technical requirements will be explained at each step, so don't skip anything!
Train Your Own Image Generation Model with Dreambooth
- The first step is to go to the HuggingFace page and create an account. All it asks for is an email address and a username. This is important because we need one of the platform's access tokens. Once the account is created, go to Settings -> Access Tokens, and click New token. Give it a name, and specify the Write function. Don't go too far, because you'll need to copy the token from here later.
- Then you'll need to define what kind of model you want to train through a selection of images. The recommended average is 20... and they must be varied. The goal is to demonstrate a style, a concept, a format, an idea. For example, suppose you want a model inspired by Polaroid photos. If you only focus on portrait Polaroids, the model will never leave that mold, when the ideal is to reproduce the details, the errors, the 'warmth' and 'vintage' essence of those photos.
- The third step is simple: adjust the images to an exact size of 512 x 512 pixels. The Birme portal we recently talked about will help you do this in record time, adjusting the size and cropping any excess. Save the images in a temporary folder, we'll come back to them.
- Now it's time to enter the Google Colab environment configured by DotCSV. My first recommendation? Don't be afraid. At first it seems like something designed by aliens, but it's actually a more spartan 'step-by-step' interface than usual. Before anything, go to "Runtime", go into the "Change runtime type" section, and make sure "GPU" is the selected option. This will confirm that Google Colab should assign a suitable graphics card for the task.
- The Google Colab environment invites us to follow its steps. We press the Play button on step 1, and accept the warning to associate our Google Drive with the environment. Google Drive must have at least 4 GB free for this, and my suggestion is to use an alternative account.
- Then we press the Play button next to Dependencies. As its name indicates, this will install all necessary dependencies for the correct execution of Dreambooth. Fortunately, it doesn't take too long.
- Step 3 activates the field where we must enter the token we generated in HuggingFace. A click on the arrow will open the menu, we copy the token and paste it into "HuggingFace_Token". Ignore the rest of the options, and press Play. Wait until it finishes.
- The fourth step asks for data about the Dreambooth training, and this changes depending on your desires: Is the plan to train a person, or an object? Maybe an artist or a TV series? Choose the appropriate option from the list in the "Training_Subject" field, and then modify the expression in "Subject_Type" following the examples below. However, all your attention should go to "Instance_name". This is the wildcard, the keyword you'll use in prompts to 'activate' the content you've trained. It's recommended to use a rare word that won't be confused with anything else by Stable Diffusion.
- When you press Play on step 4, Google Colab will enable an upload button to upload the images that will serve as reference in training. Remember, they must be 512 x 512!
- The fifth step is optional, and especially useful when creating models of people. In our test it's not necessary, but there you'll find five different values. "person_ddim" is the default mode, don't hesitate to change it to experiment.
And Now, to Wait
The Play button on step 6 starts the training process, and from here... patience. How long does it take? For a standard package of 20 images with 1,600 steps, the average is about an hour, very reasonable considering we're accessing entire gigabytes of VRAM for free. The level of activity in Google Colab and the graphics card configured for the project also influence its final performance.
With the model trained, only two things remain: evaluate it online using step 7, or download the ckpt file from our Google Drive account for offline execution with Automatic1111. Don't forget to add your keyword to the prompts, or you'll never see correct results among the generated images.
In summary, Dreambooth is an extremely powerful tool, but it requires some hours of flight. On certain occasions, the training is not entirely satisfactory, and you may be forced to try again with a different set of images. Don't give up!
HuggingFace: Click here
Dreambooth in Google Colab: Click here