Cover the DiT spec table, the QLoRA prequant heuristic, the Z-Image bf16-only
guard, the gated-repo name check, family resolution now that FLUX/Qwen/Z-Image
are trainable (and GGUF repos are rejected as inference-only), the families list
in /diffusion/info, and the gated-base 400 preflight that leaves the GPU
untouched.
Cover the trainer registry (get_trainer resolves SDXL, unknown family raises),
family resolution (explicit model_family validation, resolved_family on the config),
the metadata sidecar write + scan read with family gating, and the service loss-history
folding (append, bad-point skipping, decimation at cap, family/perf fields) plus the
status route nesting metric_history.
The trainer only supports the SDXL U-Net, but a FLUX / Qwen-Image / Z-Image
repo or a GGUF filename passed as base_model was accepted and then failed
minutes later inside StableDiffusionXLPipeline.from_pretrained with an
unrelated-looking error. Add a name-based guard in normalized() so known
DiT-family names and .gguf checkpoints are rejected up front, which the API
start route surfaces as an immediate 400 with a message that says exactly
which bases are trainable. Unrecognisable names still pass through so custom
local SDXL checkpoints keep working.
Addresses review findings on the SDXL LoRA trainer:
- Gate the base model with the same trust check as inference (unsloth/*, allowlisted
official bases, or a local path) before from_pretrained, so an untrusted remote repo
is never fetched or deserialised.
- Check the stop signal before the (slow) model load, not only between steps, so a
cancel during download is honoured; a stop may carry save=False to cancel without
leaving a partial adapter.
- Per-sample SDXL add_time_ids from the actual crop (original size + crop offset, with
the offset mirrored on horizontal flip) instead of a fixed uncropped-square tensor.
- Apply EXIF orientation before resize/crop so rotated photos train upright.
- Skip gradient clipping when max_grad_norm <= 0 (the Studio 'disable' value) instead
of scaling every gradient to zero.
- Coerce Studio config strings/blanks: learning_rate string to float, blank hf_token to
anonymous, gradient_checkpointing 'none'/'true'/'unsloth' to bool; reject a zero/negative
lora_alpha or learning_rate.
- Alias the generic Studio training payload keys (model_name/max_steps/batch_size/lora_r/
lr_scheduler_type/random_seed) onto the diffusion field names.
- Mirror the trained adapter into loras/diffusion so the Images LoRA picker discovers it.
- Report worker exceptions in both message and error keys so the failure is not lost.
Adds regression tests for the config coercion/validation and aliasing.
First diffusion training path in Studio: train a LoRA on the SDXL U-Net from an
image + caption dataset and export it as a diffusers .safetensors that the existing
diffusion LoRA loader (and any diffusers pipeline) can load.
core/training/diffusion_lora_trainer.py:
- DiffusionLoraConfig with validation/defaults (rank, alpha, targets, lr, steps, grad
accumulation, resolution, min-SNR gamma, gradient checkpointing, lr scheduler, seed,
mixed precision).
- discover_image_caption_pairs: captions from metadata.jsonl / captions.jsonl, per-image
.txt/.caption sidecars, or a dreambooth instance_prompt fallback (pure, unit-tested).
- run_diffusion_lora_training: the loop -- freeze base, PEFT-wrap the U-Net attention
projections, VAE-encode (fp32 VAE to avoid the SDXL fp16 overflow), sample noise +
timesteps, predict, MSE loss with optional min-SNR weighting (epsilon / v-prediction),
AdamW + get_scheduler + grad accumulation + grad clipping, then export via
save_lora_weights. Emits worker-protocol events (model_load_*, progress, complete) and
polls should_stop for a clean stop with a partial save.
- run_diffusion_training_process: mp.Queue subprocess adapter (event_queue / stop_queue),
so the training worker can spawn it; plus a CLI entry point.
Only SDXL (U-Net) is trained here; DiT families and the Studio UI form + route wiring are
follow-ups. The trainer is decoupled and worker-ready.
Tests: test_diffusion_lora_trainer.py covers caption discovery (metadata / sidecar /
instance prompt / skip-uncaptioned / errors), config normalisation + validation, the SDXL
add-time-ids, and the dict->config adapter. Verified live on GPU: a 60-step SDXL LoRA run
lowers the loss, exports a ~45 MB adapter, and loading it back shifts generation from
baseline (mean abs pixel diff ~55/255).
