The import-hoist lint flags newly re-exported names in the models/__init__.py hub as
unused (it does not treat __all__ membership as a use). Import the three diffusion
training schemas straight from models.training in routes/training.py, where they are
used in the route annotations and calls, and drop the __init__ re-export.
Surface the diffusion training API in the Images page. A "Train LoRA" button in the top
bar opens a self-contained dialog to fine-tune an SDXL LoRA on a folder of images: pick
the base model, dataset folder, output folder, an optional instance prompt, and the core
hyperparameters (steps, rank, resolution, batch, learning rate), then Start. The dialog
polls the training status while open and shows a progress bar, step count, live loss, and
the saved adapter path, with a Stop button for a clean stop.
The dialog is independent of the loaded generation model (training runs in its own
subprocess), and prefills the base model with the loaded checkpoint when it is SDXL, else
the SDXL base. api.ts gains startDiffusionTraining / stopDiffusionTraining /
getDiffusionTrainingStatus plus their types, matching the /api/train/diffusion routes.
Make the SDXL LoRA trainer reachable from the app with a small, self-contained job
service and JSON routes, deliberately separate from the LLM TrainingBackend (whose
lifecycle -- LLM config build, per-run SQLite rows, matplotlib plots, transfer-to-chat-
inference -- is text-training specific and would mis-handle a diffusion run).
core/training/diffusion_training_service.py: DiffusionTrainingService runs one job at a
time -- validate the config cheaply (before any spawn), spawn the trainer subprocess
(spawn context, parent-lifetime bound), pump its events (model_load_* / progress /
complete / error) into an in-memory status snapshot, and support a clean stop. The
subprocess context and target are injectable so the full start -> pump -> status ->
complete path is unit-tested without real multiprocessing or torch.
routes/training.py: POST /api/train/diffusion/start (400 on a bad config, 409 when a job
is already running), POST /api/train/diffusion/stop, GET /api/train/diffusion/status
(JSON poll). models/training.py: DiffusionTrainingStartRequest + response schemas
mirroring DiffusionLoraConfig, so model_dump() passes straight through.
Tests: test_diffusion_training.py -- service happy path, bad-config-before-spawn,
concurrent-job rejection, clean stop, crash-without-terminal-event, event transitions;
plus route wiring via the FastAPI TestClient (start / 422 / 400 / 409 / status / stop)
with a mocked service. The diffusion trainer's progress events already use the field
names this path expects.
The Studio training pump reads 'learning_rate' from progress events; the diffusion
trainer emitted 'lr'. Rename the field (and the CLI reader) so the trainer's events are
directly consumable by the existing training status/SSE machinery when it is wired into
the worker, without a translation shim.
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).
The backend has always accepted a bare Hugging Face repo id (owner/name, or
owner/name:weight-file.safetensors) as a LoRA, downloading and applying it. But the
picker only rendered when the curated catalog had entries, and the catalog is empty,
so there was no UI path to apply any LoRA. Show the LoRA section whenever the loaded
model supports LoRA, and replace the curated-only dropdown with a text input: type a
Hub repo id, or pick a discovered adapter from a datalist of suggestions when the
catalog is populated.
Also restore LoRAs when loading a recipe. restoreSettings now parses the recipe's
"id:weight" strings (splitting on the last colon, since the id itself may contain one
for a specific weight file) back into the selection, so replaying a saved image
reproduces its adapters. The generate payload trims hand-typed ids and drops empty /
zero-weight rows, and a model swap clears the selection (a LoRA is family-specific)
without discarding a free-text pick that is not in the curated list.
SDXL is the first U-Net family in the diffusion backend: its denoiser is
pipe.unet (UNet2DConditionModel), not a DiT pipe.transformer, and a single-file
.safetensors is the whole pipeline rather than a transformer-only file. The
backend previously assumed a DiT transformer everywhere, so add the two hooks a
U-Net family needs and register SDXL.
DiffusionFamily gains denoiser_attr ("transformer" for DiT, "unet" for SDXL) and
single_file_is_pipeline (SDXL loads a single file via pipeline_class.from_single_file
with the base repo as config, instead of transformer_class.from_single_file plus a
companion assembly). _align_vae_dtype now reads the denoiser generically so img2img
and inpaint keep the VAE and U-Net dtypes aligned.
The non-GGUF trust gate is extended with a short, exact-match, safetensors-only
allowlist of official base repos (the SDXL base/refiner and sdxl-turbo), because
SDXL ships only as a full pipeline and has no unsloth-hosted GGUF. Local paths stay
trusted as before; a random repo, even one that detects as SDXL, is still rejected.
The image-conditioned and ControlNet workflows are the standard SDXL pipelines,
built around the resident modules via from_pipe like every other family, so SDXL
gets txt2img, img2img, inpaint, outpaint, upscale, LoRA and ControlNet. There is no
native sd.cpp mapping yet, so the no-GPU route falls back to diffusers.
Frontend catalog gains SDXL Base 1.0 and SDXL Turbo entries with SDXL step/guidance
defaults (Turbo: few steps, no CFG; base: ~30 steps, real CFG).
Tests: new test_diffusion_sdxl.py (family shape, detection, trust allowlist, model
kind, U-Net VAE-dtype alignment, LoRA gate) plus loader-branch tests in
test_diffusion_backend.py (pipeline-kind from_pretrained, single-file whole-pipeline
from_single_file, allowlist accept/reject). Verified live on GPU: sdxl-turbo loads
both as a pipeline and as a single file and generates coherent txt2img + img2img.
