Train Your Own LoRA with ComfyUI: A Step-by-Step Guide

Machine learning models have become integral in various applications, from language understanding to image generation. With the rapid advancement of AI, there's a growing need for customization and fine-tuning models to cater to specific use cases. One such customization technique is Low-Rank Adaptation (LoRA), a popular method for fine-tuning large language models efficiently. In this blog post, we'll guide you through the process of training your own LoRA using ComfyUI, a user-friendly platform designed to simplify AI model training.

What is LoRA?

LoRA (Low-Rank Adaptation) is a method designed to efficiently fine-tune large neural networks by injecting trainable rank-decomposition matrices into every layer of the Transformer architecture. This allows for more efficient adaptation of models to new data, reducing the computational cost and time required compared to traditional fine-tuning methods.

Why Use ComfyUI?

ComfyUI is a lightweight, portable platform for building and training machine learning models. It offers a simplified interface, making it easier for developers and enthusiasts to focus on their models rather than the intricacies of configuring environments and managing dependencies. With its embedded Python environment and a set of pre-configured libraries, ComfyUI is perfect for anyone looking to train their own LoRA model without the hassle of setting up a complex development environment.

Prerequisites

Before we get started, make sure you have the following:

• ComfyUI installed on your machine. You can download it from the official website or GitHub repository here: comfyanonymous/ComfyUI: The most powerful and modular diffusion model GUI, api and backend with a graph/nodes interface. (github.com).
• A basic understanding of machine learning and neural network training.
• Python experience, as we will use Python to install additional packages.
• A dataset for training. This could be text, images, or any other data relevant to your use case. In our case I'll use a set of pictures of a model and use that to generate new images based on that. This is very useful if you want to create images of yourself by training your own LoRA with your own pictures. 

Step 1: Set Up ComfyUI Your Environment

First, ensure that your ComfyUI environment is correctly set up and running. If you haven't downloaded ComfyUI yet, grab it from the official website and extract it to a directory of your choice.

Navigate to the ComfyUI directory using the command line below to run the app:

C:\ComfyUI_windows_portable\run_nvidia_gpu.bat

This will launch the browser with ComfyUI. 

Step 2: Configure the ComfyUI Manager for the Portable Version

To properly use the ComfyUI portable version, you need to configure the ComfyUI Manager to handle environment variables and dependencies correctly.

1. Download scripts/install-manager-for-portable-version.bat into installed "ComfyUI_windows_portable" directory

2. double click install-manager-for-portable-version.bat batch file

3. Launch again ComfyUI.

Step 3: Install the additional packages we will need from the ComfyUI Manager

Navigate to the Manager -> Custom Node Manager, and install the following packages:

• ComfyUI WD 1.4 Tagger
• LoadLoraWithTags
• Lora-Training-in-Comfy
• Image-Captioning-in-ComfyUI

Step 4: Install Necessary Libraries

While ComfyUI comes pre-packaged with several libraries, we need to ensure that all required packages, like transformers, datasets, and accelerate, are installed. I had run into issues when running the Lora-Training-in-Comfy package and had to install many packages manually. Note that I'm using Python 3.10 here. Run the following commands:

First you need to see the version of Cuda you need for torch using the command below:

C:\Users\jordi>nvidia-smi

Sun Aug 25 14:35:40 2024

+---------------------------------------------------------------------------------------+

| NVIDIA-SMI 546.26                 Driver Version: 546.26       CUDA Version: 12.3     |

|-----------------------------------------+----------------------+----------------------+

| GPU  Name                     TCC/WDDM  | Bus-Id        Disp.A | Volatile Uncorr. ECC |

| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |

|                                         |                      |               MIG M. |

|=========================================+======================+======================|

|   0  NVIDIA GeForce RTX 2060      WDDM  | 00000000:01:00.0  On |                  N/A |

| N/A   63C    P0              28W /  80W |    942MiB /  6144MiB |     16%      Default |

|                                         |                      |                  N/A |

+-----------------------------------------+----------------------+----------------------+

The version of cuda available is 12.1 even though I'm running 12.3 and I'm going to use this one for my torch installation:

To check if your cuda is available, run the following:

C:\Python310>python

>>> import torch

>>> torch.cuda.is_available()

True

I ran the following packages to get it all working: 

• C:\Python310>python -m pip install --upgrade pip
• C:\Python310>python -m pip install accelerate
• C:\Python310>python -m pip install diffusers
• C:\Python310>python -m pip install transformers==4.25.1
• C:\Python310>python -m pip install einops
• C:\Python310>python -m pip install opencv-python
• C:\Python310>python -m pip install voluptuous
• C:\Python310>python -m pip install torch==2.4.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/test/cu121
• C:\Python310>python -m pip install xformers==0.0.27.post2
• C:\Python310>python -m pip install bitsandbytes=0.43.3

You can also use my requeriments.txt file to speed things up. Remember you'll need a cuda compatible GPU to try this locally.

Step 5: Prepare Your Data

All your images need to be in PNG format. If you need to transform them from jpeg, you can use this little function I created to translate them to PNG format. Also all my images are 800x530px to make things simpler.

Once your images are ready, copy them into a folder, e.g. C:\temp\1_jovanka. The 1_jovanka folder needs to start with a number and it has to be under another folder, in this case temp. It can be any name, you just need to follow the format described here (folder\1_name). 

This workflow can be downloaded here. Remember to define the prefix when you save this LoRA so the model can be triggered by the word you described. Adding a custom name might make things easier. in my case I just used the name of the model plus a number so it's easy to trigger: jovanka1.

Once this completes, you'll see all your images have a txt file with the tag information:

Now we are ready to train our LoRA.

Step 6: Configure Your LoRA Training Script

For this we'll use the Lora Training in ComfyUI function and use the folder where the dataset is

We'll use the dreamshaper_8 model to train our LoRA and then we'll define 400 epochs and that we will save a file every 100 epochs so we have different checkpoints during the training. the output name is the name of the LoRA and the output dir where the LoRA will be stored. This will take a few hours. It took me about 5h to train the LoRA.

Once finished you should see your new files under the loras folder:

Step 7: Evaluate Your Model

Once training is complete, you can evaluate your model's performance on test data to ensure it meets your expectations. Use the following script snippet to evaluate your model

After this, you'll have to try a bit to get something that resembles the original model:

Conclusion

Training your own LoRA model using ComfyUI is a powerful way to customize AI models for specific needs without the high computational cost of traditional fine-tuning. By following this guide, you've set up your environment, configured the ComfyUI Manager for the portable version, prepared your data, configured a training script, and trained a LoRA model effectively.

Feel free to experiment with different models, datasets, and training parameters to further enhance your model's performance. The world of AI is vast and ever-evolving—happy fine-tuning!