Introduction
AI-generated images have transformed how designers, artists, and content providers produce visual content today. Artists and designers are using AI models to produce high-quality visuals from text prompts in order to enhance creativity and streamline workflows. On the other hand, ecommerce businesses are using this technology to generate customized product visuals and improve engagement. The ability to quickly generate tailored images is helping companies streamline and scale their content-generation pipelines.
A key development is the use of AI models like Flux.1, which can produce detailed, context-based visuals from simple text prompts. This is particularly useful for ecommerce companies, allowing them to create personalized product displays that cater to individual customer preferences, offering a more dynamic and visually appealing shopping experience.
In this article, we will walk you through the process of fine-tuning Flux.1 to generate tailored images for any e-commerce platform.
AI Toolkit
The AI Toolkit is a comprehensive tool created to streamline AI model training, especially for utilizing LoRA and Flux. It has an easy-to-use configuration for activities such as generating images, managing datasets, and training models. This toolkit provides flexibility during training by allowing users to customize particular layers for training and control VRAM usage. Also, it can connect with Hugging Face for sharing models and has a Gradio UI for effortless interaction while training.
Steps to Fine-Tune Flux.1
Step 1: Set Up the Environment and Clone the AI Toolkit
For fine-tuning image generation models like Flux.1, you need a high-end GPU like NVIDIA A100 which you can rent on E2E Cloud.
To get started, sign up to E2E Cloud and launch a cloud GPU node. E2E Cloud offers the most price-performant cloud GPUs in the Indian market. You get blazing-fast performance at a price point that’s far cheaper than AWS, GCP, or Azure. Check the pricing page to learn more.
When launching the cloud GPU, make sure you add your public SSH key (id_rsa.pub). This will allow you to remote SSH into the node in the following way:
Once you have logged in, create a user using the adduser command, and add the user to the `sudo-ers` list using `visudo`.
You can now create a Python virtual environment.
Install Jupyter Lab and then use that to build this example:
Once you have your Jupyter environment set up, clone the ai-toolkit repository and install the following libraries:
Step 2: Provide a Hugging Face Token
Replace xxxxxxxxxxxxx with your own Hugging Face token and make sure it has access to the Flux.1-dev model.
Step 3: Prepare the Dataset
Prepare a dataset for fine-tuning the Flux model. The dataset should contain at least 15-20 images in JPG format; create a corresponding text file containing captions for each image with the same name. In this tutorial, we explore visual representations with an Indian touch.
Example:
Caption:
A girl with long black hair is sitting indoors on a chair, smiling with her mouth slightly open. She is holding a cup or mug in her hand, wearing a wristwatch and bracelet as part of her jewelry. Her outfit is a dress, and the scene has a realistic feel, capturing a casual moment of her enjoying a drink.
Step 4: Set Up the Configurations
After preparing the dataset, you need to set up the configurations for fine-tuning the model. Navigate to ai-tookit/config/examples. There, you will find the train_lora_flux_24gb.yaml configuration file according to your needs. Here are the configurations I used for this project.
Step 5: Start the Training Process
After the dataset and configurations are ready, start the training process by navigating the AI Toolkit directory and running the following command:
This process may take some time depending on your GPU and configurations.
Step 6: Test the Fine-Tuned Model
Now the fine-tuned model is ready to be tested. Test it first by loading the model through the Flux pipeline and adding the LoRA weights,
replace /home/vansh/ai-toolkit/output/Fine_tuned_Flux.1/Fine_tuned_Flux.1.safetensors with location and name of your saved model.
Add the prompt for the image you want to generate – its height, width, and name as you would like to save the image as.
Display the image using the `matplotlib` library.
Outputs:
Prompt 1:
A dark-skinned girl with black hair is wearing a red lehenga, including a red skirt. She is smiling and looking directly at the viewer. Her jewelry includes earrings, a necklace, a bracelet, and a ring. The scene has a realistic feel, with a slightly blurry background that emphasizes her makeup and lipstick, creating a vibrant yet soft appearance.
Prompt 2:
A girl with black hair is standing and looking at the viewer. She is wearing a white lehenga, and adorned with jewelry including a bracelet, necklace, and earrings.
Prompt 3:
A dark-skinned handsome boy with short black hair and a beard is standing against a brick wall. He is wearing a white collared shirt with sleeves rolled up, denim jeans, and sunglasses. His hands are clasped together, and he is adorned with jewelry, including a bracelet. The image is realistic, with his feet out of the frame and his lips slightly parted.
Conclusion
By following this guide, you can easily fine-tune the Flux.1 model to generate high-quality, customized images tailored to your needs. By using tools such as the AI Toolkit, you can simplify the tasks of preparing datasets, configuring models, and generating images for achieving effective and visually appealing outcomes. This can be a very useful tool for ecommerce companies to market products and create personalized shopping experiences.