I tried Dreambooth LoRA fine-tuning of FLUX.1 on DGX Spark

Introduction

Hello, I'm Morishige from Classmethod's Manufacturing Business Technology Division.

In previous DGX Spark articles, we focused mainly on LLM inference and fine-tuning, but the 128GB unified memory isn't only beneficial for language models. This time, I tried Dreambooth LoRA fine-tuning with the image generation model FLUX.1.

FLUX.1 is a 12B parameter image generation model released by Black Forest Labs, characterized by its architecture where multiple large components work together. Since it's a model that consumes quite a bit of memory, I was curious about its compatibility with DGX Spark's large unified memory.

This time, using NVIDIA's official DGX Spark Playbook, I'll try training a Dreambooth LoRA with photos of the DGX Spark unit itself, all the way through generating "scenes with a DGX Spark."

What is FLUX.1

FLUX.1 is an image generation model released by Black Forest Labs (the original Stable Diffusion development team). It adopts a Diffusion Transformer (DiT) architecture rather than the conventional U-Net base, and has attracted attention as a model with a high balance between text instruction adherence and image quality.

Component Structure and Memory

FLUX.1 looks like a single "model," but in reality multiple components work together.

Loading all components simultaneously in FP16 totals over 30GB. In environments like the RTX 4090 (VRAM 24GB), it can't all fit at once, requiring techniques such as split-loading the model between CPU and GPU, or offloading unused components.

With DGX Spark's 128GB unified memory (UMA), the CPU and GPU share the same memory space, so such splitting is unnecessary. Even loading all components simultaneously uses only about 30GB, leaving nearly 100GB to spare.

Why 128GB UMA Is Effective

According to NVIDIA's official blog, DGX Spark can generate 1K images at 2.6 seconds/image using FLUX.1 (12B, FP4 precision). However, since this Playbook trains and infers in BF16 precision, the conditions differ from FP4. I'll touch on how this difference affects the measured values in the benchmark results later.

Setup

Requesting Access to HuggingFace Gated Models

FLUX.1-dev is a gated model, so you need to request access in advance.

1. Visit the FLUX.1-dev HuggingFace page
2. Agree to the terms of use and request access
3. Create a Fine-grained Access Token in HuggingFace settings
   • Specify black-forest-labs/FLUX.1-dev as the target repository
   • Check "Read access to contents of selected repos"

Cloning the Playbook and Downloading the Model

git clone https://github.com/NVIDIA/dgx-spark-playbooks.git
cd dgx-spark-playbooks/nvidia/flux-finetuning/assets

Downloading the model takes about 30–45 minutes.

# Set HuggingFace token
export HF_TOKEN=<your-token>

# Download model
sh download.sh

Preparing Training Data

This time, I'll train on the concept of "DGX Spark unit itself." Using photos I took of the DGX Spark as training data is a somewhat meta setup.

What I Kept in Mind When Shooting

The recommended number of images is 5–10 per concept. I took 9 photos this time. I shot from various angles — front, diagonal, and from above — and varied the lighting between a desk lamp and natural light. I also switched up the background between a desk surface, inside a shelf, and a solid-color background, being careful to ensure the subject was clearly visible.

Dataset Configuration

Place the captured images in flux_data/sparkgpu/ and define the concept in data.toml.

[general]
shuffle_caption = false
keep_tokens = 2

[[datasets]]
resolution = 1024
batch_size = 1

    [[datasets.subsets]]
    image_dir = "flux_data/sparkgpu"
    class_tokens = "sparkgpu gpu"
    num_repeats = 2
    is_reg = false
    flip_aug = true

The sparkgpu in class_tokens becomes the trigger word. By including it in the prompt during inference, such as "a photo of sparkgpu gpu on a desk," the learned DGX Spark characteristics are reflected.

Running Dreambooth LoRA Training

The Playbook includes a Docker image for training and a launch script.

sh launch_train.sh

Training Parameters

The Playbook's default settings are optimized for DGX Spark's unified memory.

The LoRA Rank of 256 stands out. Normally it's around 8–64, but with 128GB of unified memory there's no need to worry about memory when setting a high rank. The idea is that a higher rank increases the LoRA adapter's expressiveness, allowing it to capture the concept's characteristics in finer detail even with a small training dataset.

Measured Performance

I checked resource usage during training using nvidia-smi and free.

Training took about 3.4 hours. Initial compilation with torch_compile took about 30 minutes, so the actual training was about 3 hours. Pre-resizing the training images to 1536px (long side) likely also reduced preprocessing overhead.

What's interesting is how the memory was used. System memory reached 96GB during training. This is the result of all data needed for training — the model itself, LoRA parameters, optimizer state, etc. — being expanded in memory simultaneously. This is a configuration that only works because of the 128GB unified memory; it's a scale that cannot be achieved with the RTX 4090's 24GB VRAM.

Generating Images

When training completes, checkpoints in .safetensors format are saved to models/loras/. Let's first check the training results in ComfyUI.

Launching ComfyUI

sh launch_comfyui.sh

Accessing http://localhost:8188 in a browser opens the ComfyUI interface. For remote access, adding --listen 0.0.0.0 to the launch command in launch_comfyui.sh allows access via the DGX Spark's IP address.

Loading LoRA and Generating

Load the finetuned_flux.json workflow included with the Playbook and select the trained .safetensors in the "Load LoRA" node.

Generate with the trigger word sparkgpu gpu included in the prompt.

Generation Results

I tried several prompts. The first run takes time to load the model, but from the second run onward it stabilizes at around 100 seconds.

Something that looks like a DGX Spark came out!

Most of the initial ~5 minutes is model loading time. Loading FLUX.1's DiT (23GB), T5-XXL (9.2GB), and LoRA (4.8GB) into memory takes that long. From the second run onward, the model stays in memory, so generation takes about 100 seconds.

The official blog's "2.6 seconds/image" is for FP4 precision, but this Playbook uses BF16 precision, resulting in roughly a 40x difference. The tradeoff between quality and speed is clearly visible.

Comparison With and Without LoRA

I also generated using the same prompt without LoRA (the base_flux.json workflow).

Without LoRA, since there's no concept corresponding to the word sparkgpu, a generic graphics card was generated. With LoRA, DGX Spark's characteristic mesh front panel and box-like form factor are reflected, and even in unrealistic situations like outer space or a cyberpunk cityscape, something "recognizable" comes out. Particularly in the product photo (white background), the texture of the top panel and mesh pattern were reproduced with high accuracy. This confirms that Dreambooth LoRA can reliably capture a concept's characteristics even with a small training dataset (9 images).

ComfyUI Workflow

In ComfyUI, you incorporate trained LoRA into a workflow for use. LoRA intensity can be adjusted from 0.0 to 1.0, and multiple LoRAs can be applied in combination. The trained .safetensors file can be copied directly to other ComfyUI environments, making it very portable.

Comparison with LLM Fine-tuning

Having tried fine-tuning both an LLM and an image generation model on the same DGX Spark, the difference in how the 128GB unified memory is used becomes clear.

Image generation FT uses more memory because FLUX.1 needs to expand multiple large components simultaneously. Seeing memory reach 96GB during training shows that the 128GB unified memory is not barely sufficient but rather comfortably usable.

Summary

Using NVIDIA's FLUX.1 Dreambooth LoRA Playbook for DGX Spark, I fine-tuned an image generation model with photos of the DGX Spark unit itself.

By simply executing the Playbook scripts in order, everything from model download to training to inference in ComfyUI is covered. The procedure itself had no confusing steps. BF16 precision inference takes about 100 seconds/image, which is slower compared to the official FP4 value (2.6 seconds/image), but it's best understood as a quality-speed tradeoff.

The point where I most felt the benefit of the 128GB unified memory was seeing memory reach 96GB during training. With LLM fine-tuning, the UMA's strength was "being able to handle large models without quantization," but with image generation there's a different angle of benefit: "being able to load multiple large components all at once."

Verification Environment

Reference Links

• FLUX.1 Dreambooth LoRA Fine-tuning Playbook
• ComfyUI Playbook
• DGX Spark Playbooks (GitHub)
• FLUX.1-dev (HuggingFace)
• DGX Spark Performance (NVIDIA Blog)