Every LLM tutorial on this blog has a diffusion-model twin that nobody writes: teams that need images — product shots in a house style, marketing variants, synthetic training data, defect renderings for an inspection model — reach for a hosted API, discover it will not hold their brand's look, and then ask how hard it is to train their own. The answer in 2026 is: a few hundred images, one 24 GB GPU, and an afternoon. LoRA on a modern rectified-flow text-to-image model (FLUX.1-dev, SD 3.5, SDXL) is now as routine as LoRA on Llama.
This tutorial builds one end to end in PyTorch: dataset prep, a LoRA fine-tune of the transformer only, sane hyperparameters, evaluation that is not "squint at four pictures", and a serving path that loads adapters per request.
What a diffusion LoRA actually trains
A text-to-image model in the FLUX/SD3 family has three frozen parts you never touch — a text encoder (usually two: CLIP plus T5), a VAE that maps between pixels and a compressed latent space, and a scheduler — plus one part you do: the denoising transformer (DiT). Training injects low-rank adapters into the transformer's attention and MLP projections. The base weights stay frozen and quantized if you like; only the adapters get gradients.
The training objective is not next-token prediction. For rectified-flow models it is flow matching: sample an image latent, sample noise, interpolate between them at a random timestep, and ask the transformer to predict the velocity from noise toward data. Roughly twenty lines, and worth understanding because most training bugs live in the timestep sampling.
Licensing before compute. FLUX.1-dev is non-commercial; FLUX.1-schnell (Apache-2.0), SD 3.5 Medium (permissive community license) and SDXL (OpenRAIL++) are the usual commercial-friendly picks. Check what your client's legal team will accept before you build a pipeline around a checkpoint. The same goes for your training images: "we scraped them" is a discovery problem, not a modelling one.
1. Build the dataset
Two hundred to a thousand images beats fifty thousand scraped ones. Concretely, for a brand-style or product LoRA:
- 20–50 images for a single object or person; 300–800 for a broad house style.
- Consistent subject, varied context: angles, lighting, backgrounds, crops.
- Resolution at or above your training resolution (1024×1024 for FLUX/SD3.5, 1024 for SDXL). Do not upscale.
- Remove near-duplicates. Repeated frames teach the model to memorize.
Captions matter more than most people expect. Caption what varies and leave what you want the trigger token to absorb implicit. If every image is captioned "a photo of TOKENSTYLE product on a white background", the model binds the background to your token too.
Auto-caption with a VLM, then fix by hand:
# pip install transformers accelerate pillow
import torch, json, pathlib
from transformers import AutoProcessor, AutoModelForImageTextToText
MODEL = "Qwen/Qwen2.5-VL-7B-Instruct"
proc = AutoProcessor.from_pretrained(MODEL)
vlm = AutoModelForImageTextToText.from_pretrained(
MODEL, dtype=torch.bfloat16, device_map="cuda")
PROMPT = ("Describe this image in one sentence for an image-generation caption: "
"subject, pose, lighting, background, camera framing. No opinions.")
rows = []
for p in sorted(pathlib.Path("data/raw").glob("*.jpg")):
msgs = [{"role": "user", "content": [
{"type": "image", "image": str(p)}, {"type": "text", "text": PROMPT}]}]
inputs = proc.apply_chat_template(
msgs, add_generation_prompt=True, tokenize=True,
return_dict=True, return_tensors="pt").to(vlm.device)
out = vlm.generate(**inputs, max_new_tokens=64, do_sample=False)
cap = proc.decode(out[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
rows.append({"file_name": p.name, "text": f"ACMESTYLE photo, {cap.strip()}"})
pathlib.Path("data/train").mkdir(parents=True, exist_ok=True)
with open("data/train/metadata.jsonl", "w") as f:
for r in rows:
f.write(json.dumps(r) + "\n")
Pick a trigger token that is rare in the tokenizer's vocabulary (ACMESTYLE, not modern). Then read all the captions. Every hour spent here saves two training runs.
2. The training loop
You can drive this with diffusers' train_dreambooth_lora_flux.py script and be done. It is worth writing the core loop once, though, because that is what you will be debugging:
import torch
from diffusers import FluxPipeline
from peft import LoraConfig, get_peft_model
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
transformer = pipe.transformer
vae, sched = pipe.vae, pipe.scheduler
for m in (vae, pipe.text_encoder, pipe.text_encoder_2):
m.requires_grad_(False)
transformer = get_peft_model(transformer, LoraConfig(
r=16, lora_alpha=16, lora_dropout=0.0,
target_modules=["to_q", "to_k", "to_v", "to_out.0",
"ff.net.0.proj", "ff.net.2"],
))
transformer.print_trainable_parameters() # ~0.5% of parameters
transformer.enable_gradient_checkpointing()
transformer.to("cuda")
opt = torch.optim.AdamW(
[p for p in transformer.parameters() if p.requires_grad],
lr=1e-4, weight_decay=1e-4, fused=True)
for step, batch in enumerate(loader): # batch: latents + cached prompt embeds
latents = batch["latents"].to("cuda", torch.bfloat16)
noise = torch.randn_like(latents)
# logit-normal timestep sampling: concentrate on the middle of the schedule
u = torch.sigmoid(torch.randn(latents.shape[0], device="cuda"))
t = (u * sched.config.num_train_timesteps).long()
sigma = (t / sched.config.num_train_timesteps).view(-1, 1, 1, 1).to(latents.dtype)
noisy = (1.0 - sigma) * latents + sigma * noise
target = noise - latents # flow-matching velocity
pred = transformer(
hidden_states=noisy,
timestep=t.to(latents.dtype) / sched.config.num_train_timesteps,
encoder_hidden_states=batch["prompt_embeds"].to("cuda", torch.bfloat16),
pooled_projections=batch["pooled_embeds"].to("cuda", torch.bfloat16),
).sample
loss = torch.nn.functional.mse_loss(pred.float(), target.float())
loss.backward()
torch.nn.utils.clip_grad_norm_(
[p for p in transformer.parameters() if p.requires_grad], 1.0)
opt.step(); opt.zero_grad(set_to_none=True)
Three details do most of the work:
- Cache latents and prompt embeddings before training. The VAE and both text encoders are frozen, so encoding once to disk removes them from the loop entirely: several GB of VRAM freed and a large speedup. This is the single biggest win in the whole pipeline.
- Logit-normal timestep sampling, not uniform. Uniform sampling wastes steps at the extremes where the objective is trivial; the SD3 paper's logit-normal weighting is why your loss curve looks sane.
- Gradient checkpointing plus a frozen 4-bit base (
bitsandbytesNF4, ortorchao) fits a 12B FLUX LoRA on a 24 GB card. Adapters stay in BF16 — quantized adapters train badly.
Reasonable starting hyperparameters: rank 16 (style) or 32 (complex subject), LR 1e-4 with 10% warmup and cosine decay, batch size 1 with gradient accumulation to an effective 4, 1000–2000 steps for a subject, 3000–6000 for a style. Save an adapter checkpoint every 250 steps — image models overfit quietly, and the best checkpoint is rarely the last.
3. Evaluate like an engineer, not an art director
Loss tells you almost nothing here. Build a fixed evaluation grid instead: 10–15 held-out prompts × 3 fixed seeds, regenerated identically for every checkpoint. Then measure three things:
- Subject/style fidelity — CLIP or DINOv2 image-embedding cosine similarity between generations and held-out reference images.
- Prompt adherence — CLIPScore between generation and its prompt, so you catch the classic failure where the LoRA nails the style and ignores everything you asked for.
- Diversity collapse — mean pairwise embedding distance across seeds. If it plummets, you have memorized the training set.
# sketch: score a checkpoint directory against a fixed prompt/seed grid
import torch, itertools
from transformers import CLIPModel, CLIPProcessor
clip = CLIPModel.from_pretrained("openai/clip-vit-large-patch14").cuda().eval()
cp = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
def embed(images):
px = cp(images=images, return_tensors="pt").to("cuda")
with torch.no_grad():
e = clip.get_image_features(**px)
return torch.nn.functional.normalize(e, dim=-1)
gen = embed(generated_images); ref = embed(reference_images)
print("fidelity", (gen @ ref.T).mean().item())
print("diversity", 1 - (gen @ gen.T).triu(1).sum().item() /
max(1, len(list(itertools.combinations(range(len(gen)), 2)))))
Add one human pass at the end — a 20-image side-by-side against the base model, scored by whoever owns the brand. Automated metrics rank checkpoints; humans decide whether to ship.
Two failure modes and their fixes: fried/oversaturated output means LR too high or too many steps (halve the LR, take an earlier checkpoint, or lower the LoRA scale at inference with pipe.set_adapters(["style"], [0.7])); the trigger token does nothing means your captions described the style in words, so the token had nothing left to learn.
4. Serve it
Adapters are 20–200 MB, which makes multi-tenant serving genuinely cheap: one base model in VRAM, many adapters swapped per request.
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to("cuda")
pipe.load_lora_weights("out/acme-style", adapter_name="acme")
pipe.load_lora_weights("out/partner-style", adapter_name="partner")
pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune")
img = pipe("ACMESTYLE photo of a running shoe on wet asphalt, dawn light",
num_inference_steps=4, guidance_scale=0.0,
generator=torch.Generator("cuda").manual_seed(0)).images[0]
Production notes. pipe.fuse_lora() folds one adapter into the base for maximum speed but blocks per-request swapping — fuse for single-tenant, keep unfused for multi-tenant. torch.compile on the transformer gives a solid latency win but recompiles on resolution changes, so pin your output sizes to a small set and warm each one at startup. Batch requests by (resolution, step count); mixing shapes destroys throughput. And put a safety classifier plus an invisible watermark on the output path — for commercial deployments this is now a procurement requirement, not a nicety. For the deeper serving mechanics, see torch.compile in Practice and Serving PyTorch Without Python.
What to try next
Train a second LoRA on a different concept and compose them at inference with set_adapters(["style", "product"], [0.8, 0.6]) — composition is where diffusion LoRAs beat every hosted alternative. Add a ControlNet or IP-Adapter when layout, not just style, has to be controlled. And if generation is feeding a downstream model, close the loop: measure the downstream metric, because synthetic data that looks better does not always train better.
We build and deploy custom image-generation pipelines as part of our PyTorch computer vision and deep learning model development work — dataset curation, training, evaluation harness and a serving stack on your own infrastructure. Get in touch if you want a house-style model that actually holds the house style.