Diffusion training base precision modes: bf16 speed mode (2.3-2.6x), int8, fp8, auto (#6839)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add SDXL diffusion family (U-Net pipeline support)

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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Images: LoRA free-text Hugging Face entry + recipe round-trip

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.

* Add diffusion LoRA training (SDXL text-to-image)

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).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* diffusion trainer: emit learning_rate in progress events (Studio pump compatibility)

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.

* Wire diffusion LoRA training into the Studio API

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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Images: add a Train LoRA (SDXL) dialog

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.

* Import diffusion training schemas from models.training directly

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.

* Remove stray async task scratch outputs committed by mistake

* ControlNet: reject filesystem-like ids and do not cache a model past an unload race

Two review findings on the ControlNet path:
- resolve_controlnet's bare-repo fallback accepted any id with a slash, so a
  path-shaped id (/tmp/x, ../x) reached from_pretrained as a local directory.
  Restrict the fallback to a strict owner/name HF repo id shape.
- _controlnet_pipe now re-checks the cancel event after the blocking
  from_pretrained: an unload that raced the download had already cleared the
  caches, so caching the late module would pin it past the unload.

* Pipeline prefetch: fetch only the default torch weights

A full-pipeline prefetch kept every repo file outside assets/, so an official
repo that ships multiple formats (SDXL Base: fp16 variants, ONNX, OpenVINO,
Flax, a top-level single-file twin) downloaded tens of GB from_pretrained never
loads. Skip non-torch exports and dtype-variant twins in
_pipeline_file_downloaded, and drop a component .bin when the same directory
carries a picked safetensors weight (diffusers' own preference).

* Diffusion LoRA training: fall back to fp16 when CUDA lacks bf16

The default mixed_precision=bf16 hard-fails on pre-Ampere GPUs (T4 / V100 /
RTX 20xx) which have no bf16 compute; check torch.cuda.is_bf16_supported()
and drop to fp16 there.

* Diffusion training service: join the old pump outside the lock

start() joined a finished job's pump thread while holding the service lock,
but the pump's final state writes need that same lock, so the join always
burned its full timeout and a stale pump could then overwrite the new job's
state. Join outside the lock (with a re-check after), and fence _apply_event
and the exit handler by process identity so a superseded pump can never touch
the current job's state. Adds regression tests for both.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Diffusion LoRA training: harden config handling, cancellation, SDXL conditioning, and safety

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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* ControlNet: address review findings on the diffusers path

- resolve_controlnet enforces catalog family compatibility so a direct API call
  cannot load a ControlNet built for another family through the wrong pipeline.
- Unknown ControlNet ids now surface as a 400 (call site maps FileNotFoundError
  to ValueError) instead of a generic 500.
- strength 0 disables ControlNet entirely, so a no-op selection never pays the
  download / VRAM cost; the control image is decoded and validated BEFORE the
  ControlNet is resolved or built, so a malformed image fails fast for the same reason.
- ControlNet loads use the base compute dtype (state.dtype is a display string,
  not a torch.dtype, so it silently fell back to float32) and honor the base
  offload policy via group offloading instead of forcing the module resident.
- Empty/malformed HF token coerced to anonymous access.
- Flux Union ControlNet control_mode mapped from the selected control type.
- resolve_controlnet drops the unused hf_token/cancel_event params.
- ControlNetSpec validates guidance_start <= guidance_end (clean 422).
- Images UI ControlNet Select shows its placeholder when nothing is selected.

Adds regression tests for family enforcement and the union control-mode map.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Diffusion training API: LLM interlock, pre-spawn VRAM free, path containment, no dropped knobs

Four review findings on the diffusion training start path:
- It spawned the SDXL trainer without checking the LLM TrainingBackend, so a
  start while an LLM run was active put two trainers on the same GPU. Add a
  symmetric interlock: diffusion start returns 409 when LLM training is active,
  and LLM start refuses while a diffusion job is active.
- It went straight to service.start() without freeing GPU residents. Add a
  pre-spawn free of the export subprocess, the resident Images pipeline (with an
  arbiter release), and chat models, mirroring the LLM start path.
- data_dir / output_dir were passed through unresolved, so Studio-relative names
  failed and absolute paths bypassed containment. Resolve them with
  resolve_dataset_path / resolve_output_dir before spawn (400 on an uncontained
  path).
- The request model dropped max_grad_norm and lora_target_modules, so runs that
  set them trained with defaults. Add both fields.

The gemini pump-join deadlock was already fixed earlier (join outside the lock +
proc-identity fence). Note: honoring a stop DURING model load is a trainer-loop
change owned by the diffusion training engine PR (should_stop polled before the
first optimizer step). Adds route + model regression tests.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Diffusion LoRA: harden resolution, native tag precedence, and diffusers teardown

Address review findings on the LoRA path:
- resolve_one: normalise a blank/whitespace hf_token to None (anonymous access)
  and reject a client-supplied weight file with traversal / absolute path.
- resolve_specs: convert FileNotFoundError from an unknown/stale id to ValueError
  so the route returns 400 instead of a generic 500.
- _scan_local: disambiguate local adapters that share a stem (foo.safetensors vs
  foo.gguf) so each is uniquely addressable.
- inject_prompt_tags: the backend-validated weight now wins over a user-typed
  <lora:ALIAS:...> for a selected adapter; unselected user tags are left alone.
- diffusers _apply_loras: reject a .gguf adapter with a clear error before touching
  the pipe (diffusers loads safetensors only).
- _unload_locked: drop the explicit unload_lora_weights() on teardown; the pipe is
  dropped wholesale (freeing adapters), so the previous call could race an in-flight
  denoise on the same pipe.
- Images page: use a stable LoRA key and clear the selection (not just the options)
  when the catalog refresh fails.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Diffusion: guard trust check against OSError and validate conditioning inputs

- _is_trusted_diffusion_repo: wrap Path.exists() so a repo id with invalid
  characters (or a bare owner/name id) can't raise OSError; treat any failure as
  not-a-local-path and fall through to the unsloth/ allowlist. validate_load_request
  still raises the clear FileNotFoundError for a genuinely missing local pick.
- generate(): reject mask_image / upscale / reference_images supplied without an
  input image, and reject reference_images on a family that does not support
  reference conditioning, instead of silently degrading to txt2img / img2img.

* SDXL: reject GGUF up front, skip unused base weights, drop refiner, and harden helpers

Addresses review findings on the SDXL family:
- Reject a GGUF load for single_file_is_pipeline families (SDXL) in validate_load_request,
  before the route evicts the current model; SDXL has no transformer-only GGUF variant.
- Skip base-repo weight files when a whole-pipeline single file is loaded: from_single_file
  (config=base) needs only the base config/tokenizer/scheduler, so a local .safetensors no
  longer triggers a multi-GB base download.
- Remove the SDXL refiner from the non-GGUF trust allowlist: it is an img2img-only pipeline
  but this backend loads every sdxl repo as the base txt2img pipeline.
- Normalize a blank/whitespace hf_token to None once in load_pipeline so every load branch
  degrades to anonymous instead of erroring on a malformed token.
- Read the denoiser dtype from a parameter (compile-wrapped modules may lack .dtype) and
  access state.family.denoiser_attr directly.

Adds/updates regression tests for the trust allowlist, GGUF rejection, and base-config filter.

* Images: preserve restored LoRAs through model load and never send hidden LoRAs

- The LoRA effect cleared the selection on every load->capable transition, which
  wiped adapters restored from a gallery recipe before the model finished loading.
  Track the previously-loaded family in a ref and clear only on a real family swap;
  keep the selection on the initial load and on unload.
- Gate the generate payload's loras on loraCapable so a restored selection that is
  hidden (loaded model does not support LoRA) is never sent to the backend.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Images Train LoRA dialog: token, validation, precision, base-repo prefill, gating, refresh

Nine review findings on the SDXL training dialog:
- Forward the saved Hub token so a gated/private SDXL base can be trained (the
  image load flow already sends it).
- Re-seed the base-model field from the current default each time the dialog
  opens; the keep-alive dialog otherwise kept its mount-time default after a
  model loaded.
- Prefill from base_repo (the diffusers pipeline) rather than repo_id, which for
  a GGUF/single-file SDXL load is the checkpoint path from_pretrained can't open.
- Add client-side validation of steps/rank/resolution/batch/learning-rate before
  the request.
- Expose a precision selector (bf16/fp16/fp32) so non-bf16 GPUs can train from
  the UI, not only the API.
- Gate the dialog on the active Images route (active && trainOpen) so switching
  tabs closes it and stops its polling.
- Rescan the LoRA picker when a run completes, so a freshly-trained adapter
  appears without a model reload.
- Cap the dialog height and scroll the body so the Start/Stop footer stays
  reachable on short viewports.
- Correct the copy to not over-promise picker auto-discovery.

Freeing the resident Images pipeline before training is handled backend-side in
the diffusion training start route.

* Merge diffusion-sdxl into diffusion-lora-ux; keep options-only LoRA catch

The catalog-refresh .catch from the lower branch clears the selected adapters
too, which is right for its catalog-only picker but wrong here: this picker
holds free-text HF repo ids that are valid without being in the catalog, so a
transient refresh failure must not wipe them. Family swaps still clear the
selection and hidden LoRAs are never sent.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Train LoRA dialog: stop suggesting absolute paths the backend rejects

The dataset and output placeholders showed /path/to/... examples, but the
training routes resolve those fields inside the Studio home and reject
absolute paths outside the approved roots, so following the placeholder
produced a 400. Use folder-name placeholders and say in the labels and the
dialog description where each folder resolves.

* Align the VAE to the denoiser's first FLOATING dtype, not its first parameter

A GGUF-quantized transformer's leading parameters are packed uint8 storage,
so reading next(parameters()).dtype handed nn.Module.to() an integer dtype
and every image-conditioned generation on a GGUF model (Qwen-Image-Edit)
failed with a 500. Probe the parameters for the first floating dtype, treat
an all-integer module as a no-op, and also catch TypeError so an unexpected
dtype can never break generation. Regression test included.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Count LR scheduler warmup/decay in optimizer steps, not micro-steps

lr_sched.step() runs once per outer optimizer step (after the gradient
accumulation inner loop), for train_steps total. The scheduler was
configured with num_warmup_steps and num_training_steps multiplied by
gradient_accumulation_steps, so with accumulation > 1 a warmup or
non-constant schedule stretched past the run and never reached the
intended decay. Count both in optimizer steps.

* Address Codex review findings on the image-workflows PR

Keep diffusion.py importable without torch: the compile/arch patch modules
import torch at module level, so import them lazily at their load/unload
call sites instead of at module load. This restores the torchless contract
so get_diffusion_backend() works on a CPU/native sd.cpp install.

Match family reject keywords and aliases as whole path/name segments, not
raw substrings, so an unrelated word like edited, edition, or kontextual no
longer misroutes or hides a valid base image model, while supported edit
families (Qwen-Image-Edit, FLUX Kontext) still resolve. Mirror the same
segment matching in the picker task filter.

Route FLUX.2-dev native guidance through --guidance like the other FLUX
families rather than --cfg-scale. Reject native upscale requests that have
no input image. Read image header dimensions and reject over-limit inputs
before decoding pixels, so a crafted small-payload image cannot spike
memory. Reject an upscale that would shrink the source below its input
size. Validate the model_kind against the filename extension before the
GPU handoff. Estimate a local diffusers pipeline's size from its on-disk
weights so auto memory planning does not skip offload and OOM. Report
workflows: [txt2img] from the native backend status so the Create tab
stays enabled for a loaded native model. Clamp the outpaint canvas to the
backend's 4096px decode limit.

Adds regression tests for segment matching and kind/extension validation.

* Guard inference loads and worker lifetime against diffusion training

Teach the chat and image load guards about an active diffusion (SDXL) LoRA
job: a chat load is refused (its footprint cannot be fit-checked against the
trainer) and an image load is refused outright, mirroring the existing LLM
training guards, so a load can no longer allocate GPU memory alongside the
trainer and undo the pre-start cleanup.

Bind the diffusion trainer subprocess to the parent's lifetime and scrub the
native path lease secret from it by running the child through
run_without_native_path_secret, matching the inference/export/LLM workers, so
a Studio crash or kill no longer leaves the trainer holding the GPU.

Reset in_model_load on the complete and error terminal events: a stop or
failure during model loading otherwise leaves the status reporting a stale
loading indicator after the job has ended.

* Harden diffusion LoRA handling on the diffusers and native paths

Reject LoRA on a torch.compile'd diffusers transformer (Speed=default/max):
diffusers requires the adapter loaded before compilation, so applying one to
the already-compiled module fails with adapter-key mismatches. The status
gate now hides the picker and generate raises a clear message instead.

Convert a cancelled Hub LoRA download (RuntimeError Cancelled) to the
diffusion cancellation sentinel in resolve_specs, so an unload/superseding
load during resolution maps to a 409 instead of a generic server error.

Drop weight-0 LoRA rows before the native support gate so a request carrying
only disabled adapters stays a no-op on families where native LoRA is
unsupported, matching the diffusers path.

Reject duplicate LoRA ids in the request model: both apply paths suffix
colliding names, so a repeated id would stack the same adapter past its
per-adapter weight bound.

Strip all user-typed <lora:...> prompt tags on the native path (only the
selected adapters are materialized in the managed lora-model-dir, so an
unselected tag can never resolve), and restore saved LoRA selections from a
gallery recipe so restore reproduces a LoRA image.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Harden ControlNet resolve, gallery metadata, and the control-type picker

Check cancellation immediately after a ControlNet from_pretrained and before
any device placement, so an unload/eviction that raced the download does not
allocate several GB onto the GPU after the load was already cleared.

Require a loadable weight or shard index (not just config.json) before a local
ControlNet folder is advertised, so an interrupted copy is hidden instead of
failing deep in from_pretrained as a generic 500.

Do not record a strength-0 ControlNet in the gallery recipe: it is treated as
disabled and skipped, so the image is unconditioned and the metadata must not
claim a ControlNet was applied.

Build the control-type picker from the selected ControlNet's advertised
control_types instead of a hardcoded passthrough/canny pair, so a union model
with a precomputed depth or pose map sends the correct control_mode.

* Address further Codex findings on the image-workflows PR

- Persist the actual output image size in the gallery recipe instead of the
  request sliders: Transform/Inpaint/Edit derive the size from the uploaded
  image, Extend grows the canvas, and Upscale resizes it, so the sliders
  recorded (and later restored) the wrong dimensions for those workflows.
- Reject a remote '*-GGUF' repo loaded as a full pipeline (no single-file
  name) in validate_load_request, so the unloadable pick fails before chat is
  evicted rather than deep in from_pretrained.
- Only publish an image-conditioned from_pipe wrapper to the shared aux cache
  when the load is still current: from_pipe runs under the generate lock but
  not the state lock, so an unload racing its construction could otherwise
  cache a wrapper over torn-down modules that a later load would reuse.
- Verify the Windows CUDA runtime archive checksum before extracting it, like
  the main sd-cli archive, so a corrupt or tampered runtime is rejected rather
  than extracted next to the binary.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Refuse non-SDXL base models at diffusion training start

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.

* Add diffusion dataset upload and training info endpoints

Training an image LoRA required knowing the Studio home layout and copying
files onto the server by hand, which is the most confusing step of the whole
flow. Two small endpoints fix that:

- GET /api/train/diffusion/info reports the datasets and outputs roots plus
  every dataset folder that contains images (with image/caption counts), so
  the UI can offer a picker instead of a blind free-text path.
- POST /api/train/diffusion/dataset uploads images and optional caption
  .txt / metadata.jsonl files into a named folder under the datasets root,
  creating it on first use and accumulating on repeat uploads so large sets
  can arrive in batches. Names are validated to a single path component and
  files stream to disk under the same per-upload size cap as LLM dataset
  uploads. The returned name is a valid data_dir for /diffusion/start.

* Rework the Train LoRA dialog into a guided SDXL flow

The dialog assumed users knew the Studio home layout and that only SDXL is
trainable, and hid both facts behind free-text fields. Restructure it around
the three real decisions:

- Base model is a dropdown of the trainable SDXL picks (Base 1.0, Turbo, the
  loaded SDXL pipeline when there is one) with a custom repo/path escape
  hatch, instead of a bare text field defaulting to a repo id.
- Training images come from an in-browser upload (new dataset endpoints) or
  a picker over existing dataset folders with image/caption counts. No shell
  access or knowledge of the datasets root is needed any more, and the
  captioning rules are explained inline.
- The output field is now Adapter name and the instance prompt is labelled
  as the trigger prompt, with a no-captions warning wired to the selected
  dataset's actual caption count.

Hyperparameters collapse behind a training settings toggle since the
defaults suit a first run. A completed run says where the adapter went and
offers Done / Train another, and the top-bar button gets an icon and a
plainer description. The dialog title states the SDXL-only scope and that
other families load LoRAs but cannot train them yet.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Validate diffusion training config before freeing the GPU

The start route freed resident GPU workloads (export, Images pipeline, chat)
before the service validated the config, so a start that was then refused,
now including a non-SDXL base model, tore down the user's loaded model for
nothing. Run the same cheap normalise pass first; the LLM path already
follows this rule via its before_spawn hook.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Refactor diffusion LoRA training into a family-aware platform

Split the SDXL trainer into a shared, architecture-agnostic layer so more model
families can be trained without duplicating the plumbing:

- New core/training/diffusion_train_common.py holds the config + validation, dataset
  discovery, event emission, stop protocol, adapter publishing, and a lazy trainer
  registry (get_trainer). diffusion_lora_trainer.py keeps the SDXL-specific loop and
  re-exports the moved names so existing imports are unchanged.
- The SDXL-only base-model blocklist becomes a positive check: the family is resolved
  from the base model (or an explicit model_family) via the diffusion family registry,
  and a known-but-not-yet-trainable family is refused with a clear message. Unknown
  custom names still default to the SDXL trainer.
- DiffusionFamily gains a trainable flag and train_base_repos; SDXL is marked trainable.
  DiT families flip on when their trainers land.
- Trained adapters now write a <name>.json metadata sidecar (family, base model, rank,
  trigger prompt, ...) that the LoRA scanner reads to family-gate the adapter in the
  picker instead of showing it as unknown for every model.
- The training base-model trust allowlist adds the official FLUX.1-dev, Z-Image-Turbo,
  and Qwen-Image repos (safetensors-only, no remote code).

* Retain diffusion training loss history and expose it in status

The training service kept only the latest loss, so a live loss chart could show a
single point. Fold each progress event into bounded (step, loss, lr) history arrays
(capped at 4000 points, decimated when full) plus the latest throughput and peak VRAM,
and record the family / base model / catalog path on completion. The status endpoint
returns these as a nested metric_history object the UI can chart directly, and the
start request accepts an optional model_family override.

* Tests for the diffusion training platform

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.

* Add diffusion dataset labeling and example-import endpoints

The Train tab needs to let users caption small datasets in the browser and
pull in a ready-made set to see training work end to end, neither of which
the upload-only endpoint supported.

Add, under /api/train/diffusion/dataset:
- GET {name}/images lists every image with its resolved caption (metadata
  beats a per-image sidecar, matching the trainer's discovery order) so
  uncaptioned images are visible and flaggable.
- GET {name}/image/{filename} serves an image, with ?thumb=<px> returning a
  cached downscaled JPEG kept in a hidden .thumbs subdir (regenerated when
  the source is newer) so the labeling grid stays light.
- PUT {name}/caption/{filename} writes, or when blank clears, the .txt
  sidecar; DELETE {name}/image/{filename} removes the image plus its
  sidecars and thumbnails.
- GET dataset-examples lists a curated, license-labelled registry, and
  POST dataset/import-example materializes one into a dataset folder as
  numbered images + .txt captions. Two loaders cover the shapes seen in the
  wild: streaming rows from datasets.load_dataset (dog-example, Tuxemon) and
  a snapshot + jsonl walk for imagefolder repos whose captions live in a
  non-standard *.jsonl (the public-domain tarot set). Imports are idempotent
  and cap the image count.

Filenames and dataset names are validated against path traversal and pinned
inside the datasets root.

* Test diffusion dataset labeling and example-import endpoints

Cover caption precedence, thumbnail generation and .thumbs exclusion,
caption write/clear, image delete cleanup, path-traversal rejection on
names and filenames, and example import with a mocked datasets.load_dataset
(files plus sidecars written, idempotent second call, cap respected, load
failure mapped to 502).

* Add flow-matching DiT LoRA trainers (FLUX.1-dev, Qwen-Image, Z-Image)

Extends diffusion LoRA training beyond SDXL to the three popular DiT families
via a single shared flow-matching loop parameterised by small per-family specs
(loading, prompt/latent encoding, transformer forward, save). Verified against
diffusers 0.38.0:

- FLUX.1-dev: 2x2 latent packing + image ids, guidance-embed forward, on-the-fly
  nf4 QLoRA of the 12B transformer (the dev repo is gated, so training needs the
  user's HF token).
- Qwen-Image: 5D VAE latents normalised by the per-channel latents_mean/std,
  img_shapes forward, prequant nf4 base by default (on-the-fly nf4 for the bf16
  base).
- Z-Image: list I/O with the reversed timestep convention and a negated
  prediction, bf16 only.

The registry (get_trainer) and DiffusionFamily.trainable / train_base_repos now
route these families to the DiT trainer; the SDXL blocklist guard is replaced by
a positive family resolution that also rejects GGUF repos (inference-only) and
still-unsupported families. Per-family defaults + labels + VRAM notes are exposed
via family_train_infos for the Train UI.

Memory: caption embeddings are precomputed once and the text encoders freed
before the loop; gradient checkpointing (non-reentrant, required for bnb 4-bit)
and 8-bit AdamW are on by default.

* Speed up + shrink SDXL LoRA training (precompute text embeds, 8-bit AdamW)

SDXL re-encoded every caption with both CLIP text encoders on every step (pure
waste, since captions are constant) and kept the encoders resident. Precompute
each unique caption's embeddings once, then free the text encoders before the
loop: numerically identical (embeddings are deterministic and this consumes no
torch RNG, so the noise/timestep stream is unchanged) but faster and ~1.5 GB
lighter. Default the optimizer to 8-bit AdamW (bitsandbytes) with an fp32
fallback, halving optimizer state with no meaningful LoRA quality cost. Env
toggles (UNSLOTH_DIFFUSION_NO_PRECOMPUTE / _FP32_OPTIM) let the accuracy guard
A/B the paths.

* Expose trainable families in /diffusion/info and preflight gated bases

The training info endpoint now returns the trainable model families (name,
label, default + allowed base repos, recommended defaults, and a VRAM/access
note) so the Train UI can offer a base picker with realistic guidance. The start
route preflights a gated base repo (HEAD model_index.json with the user's token)
BEFORE freeing resident GPU workloads, so a missing FLUX.1-dev license/token
fails fast with an actionable 400 instead of evicting the loaded model and then
hitting a confusing mid-load 401.

* Tests for DiT trainers, family resolution, info families, gated preflight

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.

* Add diffusion training API client: metrics, families, dataset labeling, examples

Extends the Images training client for the Train tab: the status type now carries
metric_history (step/loss/lr) plus catalog_path/family/base_model/samples_per_second/
peak_memory_gb; the start request gains model_family; and info gains an optional
families list (per-family bases + defaults). Adds typed calls for the dataset
labeling and one-click example endpoints: list images with captions, thumbnail URL,
write/clear a caption, delete an image, list example datasets, and import an example.

* Add diffusion Train panel: config, dataset labeling, live charts, deploy

New full-page training workspace for the Images tab. Left column configures the run:
model family (FLUX.1-dev, Qwen-Image, Z-Image, SDXL in popularity order, with per-family
VRAM/license notes and defaults, backfilled from the backend families list when present),
base repo, dataset (existing folder, browser upload, or one-click example import), an
in-browser caption labeling grid (per-image thumbnail + caption saved on blur, delete,
uncaptioned highlight), adapter name, trigger prompt, and collapsed training settings.
Right column shows the live run: progress + loss/avg/speed/peak-VRAM readouts, the reused
training loss/LR charts fed from metric_history, and a completion card that deploys the
adapter into Create or starts another run.

* Wire Create/Train tab switch into the Images page and deploy flow

Replaces the Train LoRA dialog with a top-bar Create | Train segmented control next to
the model selector. Create renders the existing generation workspace unchanged; Train
renders the full-page training panel (unmounted in Create so its polling stops while the
backend run and its retained metric history survive a tab switch). Adds a deploy handler:
loading the trained adapter's base as a pipeline, queueing the adapter so the LoRA
discovery effect applies it once the base is loaded and LoRA-capable for the matching
family (with a mismatch warning), seeding the prompt with the trigger, and switching back
to Create. Removes the now-unused dialog.

* Wrap the DiT training forward in bf16 autocast

The fp32 LoRA parameters and the bnb 4-bit base matmuls need a single
compute dtype during the forward, exactly like the diffusers dreambooth
scripts run under accelerator.autocast. Without it the 4-bit backward on
FLUX.1-dev fails with an illegal-address CUBLAS error partway into the
first step. Z-Image and Qwen-Image smokes are unaffected and the SDXL
path (its own trainer) is untouched.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix Train tab example cards and Create/Train tab layout

The example-dataset cards used a two-column grid in the ~340px config
column, which wrapped titles one word per line and let the long license
text overrun into the neighbouring card. Switch to one card per row with a
horizontal layout: title with a compact truncated license badge (full text
in the tooltip), a two-line clamped description, and the Import button on
the right.

The Create/Train switch had an icon inside the Train trigger that overhung
the pill corner. Drop the icon, make both triggers a fixed equal width so
the active pill sits flush in the top bar.

* Show only loss and learning-rate charts for diffusion training

The Train tab reused the LLM charts section, which also rendered an empty
Grad Norm card and an Eval Loss card showing an Evaluation not configured
placeholder with a red smear. Neither applies to diffusion LoRA training.
Add a diffusion-only two-card view that reuses the loss and learning-rate
cards directly with fixed presentation defaults, and note under the loss
chart that per-step loss is noisy by design so users read the smoothed
line for the trend rather than the raw jitter.

* Add a dataset preview strip to the Train tab

When a dataset with images is selected, show a strip of up to 8 sampled
thumbnails with a +N more tile, so users can see what is in the folder
before training. Clicking the strip opens the existing caption review
grid. Samples are drawn evenly across the folder and refresh on dataset
change or after an upload/import.

* Stop example cards from overflowing the Train config column

The example-dataset cards still overran the ~340px config column: the
license used the Badge component whose baked-in w-fit and whitespace-nowrap
ignored the max-width and truncate, and the grid children had the default
min-width auto so wide content pushed past the column edge and clipped the
Import buttons. Replace the badge with a plain truncating pill span, and
give the config column min-w-0 with overflow-x-hidden so nothing escapes
its width.

* Add Smithsonian Butterflies and Nouns example datasets

Two permissive ~100-image sets for the Train tab: huggan/smithsonian_butterflies_subset
(CC0, the classic diffusers-docs training set, imported as a subject set with a trigger
prompt since its metadata columns are species names not captions) and m1guelpf/nouns
(CC0, captioned pixel-art avatars via the text column). Both cap at 100 images.

* Paginate the Train tab caption grid with prev/next controls

Large example datasets (100+ images) rendered every tile at once, so the
caption review grid grew unbounded. Show 24 images per page with < >
chevrons and an x-y of N indicator; a new dataset or refresh resets to
the first page.

* Offer example datasets in the Train dropdown with previews

Add an Examples group to the training-images dropdown that imports a
curated dataset in one pick, alongside the existing cards. Cards now show
up to three preview thumbnails pulled from the public HF datasets-server
so the set is visible before download. Hide the trigger prompt when every
image already has a caption (a captioned style set needs no trigger), and
turn the training-settings toggle into a ghost button with a rotating
chevron.

* Clamp the training base repo to the selected family

The base-model select's state could briefly hold the previous family's
repo after a family switch (the reseed effect runs a beat later, and a
value with no matching option makes the browser display the first option
anyway). The request then carried the stale repo: picking Qwen or Z-Image
still sent black-forest-labs/FLUX.1-dev and surfaced FLUX's gated-repo
error under the wrong family. Derive an effectiveBase clamped to the
current family's repos and use it for the select value, the start
request, and the deploy fallback.

Also move the Trigger prompt above Adapter name: the trigger describes
the dataset, the name only labels the output.

* Speed up diffusion LoRA training and cut DiT peak VRAM by a third

Perf core for the diffusion trainers, defaults preserving the training math:

- Phased model loading: the pipeline now loads without its transformer
  (conditioning only), captions are encoded and the text encoders freed,
  the VAE latent cache is built and the VAE freed, and only then does the
  transformer load. The multi-GB denoiser never shares VRAM with the
  encoders, cutting measured peak VRAM on B200: FLUX 17.1 -> 10.4 GB,
  Qwen-Image 19.1 -> 12.8 GB, Z-Image 7.3 -> 4.7 GB.
- Latent cache (cache_latents, default on): per-image crop/flip variants
  (cache_variants, default 4 vs the single frozen variant of the diffusers
  --cache_latents) store the VAE posterior's affine parameters, so every
  step still draws a fresh VAE sample; a cached center-crop Z-Image run
  matches the uncached one at the bf16 nondeterminism floor.
- True batching: train_batch_size now actually batches the transformer
  forward (it was silently 1). nf4 dequant dominates the step cost, so
  batch 4 lands near batch-1 step time: 4.0x samples/s on Qwen-Image,
  3.1x on FLUX, 2.1x on Z-Image, with multi-seed loss envelopes
  overlapping batch-1.
- LR scheduler support in the DiT loop (lr_scheduler / lr_warmup_steps
  were accepted but ignored); progress events now report the real
  per-step LR.
- TF32 + high fp32 matmul precision under enable_tf32 (default on),
  snapshot/restored around the run. cudnn.benchmark is scoped to a
  caller opt-in only: autotuning the fp32 VAE convs doubled peak VRAM
  on the DiT families for zero steady-state gain.
- Vectorized sigma gathering (drops a per-step Python search loop),
  cached FLUX img_ids/guidance, fused torch AdamW fallback, steady-state
  samples_per_second (excludes the first-step warmup).
- Regional torch.compile plumbing (compile_transformer off/on/auto with
  eager fallback): auto stays off over a bitsandbytes base where compile
  is a net loss (27 s warmup, slightly slower steady on Z-Image); it
  arms automatically for the dense/quantized speed modes that follow.
- Stop parity with the LLM trainer: /api/train/diffusion/stop accepts an
  optional {save} body and the service forwards save=False as a
  no-save cancel; a new preparing event surfaces cache-build progress.
- SDXL trainer gets the same latent cache, perf flags, and fused
  fallback; its batching, LR schedule, and min-SNR stay as they were.

Verified: 83 backend tests green; per-family 30-40 step runs with
adapter round-trip generation through the normal LoRA path (FLUX,
Qwen-Image, Z-Image all pass).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add base_precision speed modes to DiT training: bf16 2.3-2.6x, int8, fp8

New base_precision config for the DiT trainers: nf4 (unchanged default) |
bf16 | int8 | fp8 | auto, advertised per family + per machine through
/api/train/diffusion/info (precision_modes, recommended_precision,
supports_compile) so the UI can gate the selector.

- bf16: dense transformer + regional torch.compile (auto-armed). The
  measured speed mode: 2.3x nf4 on FLUX (1.81 -> 4.12 steps/s), 2.6x on
  Z-Image (2.5 -> 6.38 steps/s) on B200, at dense-weight VRAM
  (FLUX 24.7 GB / Z-Image 13.6 GB peak vs 10.4 / 4.7 for nf4).
- int8: torchao weight-only int8 on the frozen base, quantized AFTER
  add_adapter (quantizing first trips peft 0.18's TorchaoLoraLinear,
  which is incompatible with the torchao 0.16 config API). Runs eager:
  inductor rejects the int8 subclass training graph (aliased subclass
  outputs), so compile is force-disabled for it.
- fp8: torchao convert_to_float8_training on the frozen linears
  (filter skips lora_ modules, proj_out, non-divisible-by-16 dims,
  pad_inner_dim), applied after add_adapter, compile auto-armed.
  Works and round-trips, but measured SLOWER than compiled bf16 at
  LoRA-training shapes (FLUX 3.15 vs 4.12 steps/s; Z-Image similar),
  so it is an explicit opt-in and auto never picks it.
- auto: free VRAM (measured before load) + dense-size table -> bf16
  when it fits with headroom, int8 in the middle band, else nf4.
  Prequant bnb repos always resolve to nf4; dense modes on them are
  rejected at validation with a pointer to the family's dense base.

Two crashes found and fixed along the way:
- The cuDNN SDPA backend's training graph fails on the FLUX attention
  shapes (torch 2.10 + cu130, B200): mha_graph.execute errors, then the
  context degrades into illegal memory accesses. The perf-flag guard now
  pins flash/mem-efficient SDPA for the run (mathematically equivalent,
  snapshot/restored). nf4 escaped it by routing attention differently.
- Regional compile now uses dynamic=True (the inference layer's proven
  default): dynamic=False specialisation fused a gemm_and_bias epilogue
  that failed with CUBLAS_STATUS_EXECUTION_FAILED on the FLUX training
  graph; dynamic=True is also faster (Z-Image 3.84 -> 6.38 steps/s).

Verified: 98 backend tests green (new test_diffusion_base_precision.py:
validation, auto policy table, fp8 filter, compile gating, /info fields);
per-mode 40-step runs on FLUX + Z-Image with loss means inside the nf4
envelope and adapter round-trip generation through the normal LoRA path
for bf16-, fp8-, and int8-trained adapters.

* Clear TF32 flags when enable_tf32 is off so the opt-out is strict fp32

* Address review: auto int8 requires the dense-load transient to fit, dense modes are CUDA-only, auto respects bf16 compute, exact cudnn SDPA restore

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address review: fp32 latent cache stats + strict-JSON-safe progress floats

- Latent caches (DiT + SDXL) now hold the posterior mean/std in fp32 and draw the
  per-step sample in fp32, casting only the result to the training dtype. This
  matches the in-loop path (encode fp32 -> sample fp32 -> cast) exactly instead of
  sampling in bf16; the cache is tiny so the doubled RAM is negligible.
- The training service nulls non-finite floats (NaN/Inf loss, avg_loss,
  learning_rate) at its single ingestion point so status snapshots and persisted
  run records stay strict-JSON serializable; the metric history skips non-finite
  loss points. Test covers NaN/Inf progress followed by a finite point.

* Address review: gate auto int8 on torchao, scope dense validation to DiT

- base_precision="auto" only picks int8 when torchao is importable (the int8
  quantize has no runtime fallback, unlike fp8); otherwise the middle band falls
  back to nf4. Threaded as a parameter so the policy stays pure.
- The dense-mode validation (prequant base / bf16 compute) now applies only to
  DiT families: sdxl ignores base_precision entirely, so a leftover value can no
  longer fail an SDXL run. The mode-name validity check still runs everywhere.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Gate int8 and fp8 on a functional torchao import, not find_spec

The Windows ROCm torchao import stub satisfies find_spec and even lets
from torchao.quantization import quantize_ succeed, but its quantize_ is a
no-op: auto would pick int8, leave the transformer dense, and disable
compile as if it were quantized. has_functional_torchao imports the exact
symbols the int8 path uses and rejects the stub via its sentinel; both the
auto picker and the /info advertised modes now use it

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Skip the sigma-gather test when diffusers is not installed

CI runs the backend suite without diffusers; the test checks our index math against
the scheduler's own gather, so it skips rather than fails there.

* Coerce cache_latents and enable_tf32 string flags in the config dict path

The generic Studio config dict path can deliver these flags as strings, and a
non-empty string like "false" is truthy, so an opt-out silently no-ops (the
latent cache still builds, TF32 stays on). Coerce them the same way
gradient_checkpointing already is.

* Remove committed runtime scratch artifacts and ignore their dirs

logs/ (a 1.3 MB ComfyUI object_info dump plus stale PID files), temp/ (PR body
and commit message scratch), and async_task_outputs/ (agent task transcripts)
are environment specific runtime artifacts that were committed by accident and
carry stale local state into every checkout. Remove them and gitignore the
directories so they cannot be re-added.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Restore pre-Ampere bf16 fail-fast in the DiT trainer

The perf rewrite dropped the bf16 capability guard, so a pre-Ampere CUDA
device (T4/V100/RTX 20xx) would die deep in model load with an opaque dtype
error instead of a clear message. Restores parity with the SDXL trainer.

* Size-gate the automatic diffusion latent cache

The latent cache holds two fp32 posterior tensors per crop/flip variant per
image, pinned on CUDA hosts, so datasets with thousands of images can exhaust
host or pinned memory with no fallback. Estimate the cache size from the first
real encoded latent and fall back to per-step VAE encoding when it exceeds a
4 GiB budget. UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE bypasses the gate; the
existing UNSLOTH_DIFFUSION_NO_LATENT_CACHE opt-out is unchanged.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Gate DiT training precision: deny fp8 for Qwen, gate explicit int8 on torchao, gate advertised dense modes + route on bf16

- normalized() + family_train_infos() mirror the inference fp8 deny for
  Qwen-Image (activation outliers exceed fp8's range and corrupt the trained
  result); int8 stays allowed and the UI no longer advertises fp8 for it.
- _resolve_base_precision() gates an explicit int8 on a FUNCTIONAL torchao, the
  same gate auto and /info already apply, so a missing/stub torchao fails fast
  instead of silently loading dense with compile disabled.
- train_precision_modes() gates the dense modes (bf16/int8/fp8/auto) on
  torch.cuda.is_bf16_supported(), so a non-bf16 CUDA GPU (T4/V100/RTX 20xx) is
  offered only nf4 instead of a start that evicts resident models and then fails.
- start_diffusion_training preflights bf16 support for the DiT families BEFORE
  _free_gpu_for_diffusion_training(), so any DiT start (nf4 included, since the
  trainer requires bf16 unconditionally on CUDA) fails fast without eviction.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Gate DiT training on functional torchao for explicit int8; hide always-400 DiT modes on non-bf16 GPUs

The start route preflight only rejected non-bf16 GPUs; an explicit int8 request on
a host with a missing or stub torchao passed the preflight, evicted resident GPU
workloads, then died in the trainer child (its int8 base quantizer has no fallback).
Fold both gates into training_precision_preflight_error so int8-without-torchao fails
fast before eviction. Also empty the advertised DiT precision_modes (and surface the
reason in vram_note, drop compile) whenever the bf16 preflight would reject the family,
so /info never offers an nf4 DiT option the route always 400s.

* Reject dense DiT precisions on a CUDA-absent host before eviction; stabilize family-info tests

The start-route preflight caught the bf16-GPU and int8-torchao requirements but not the dense
precisions' CUDA requirement: on a GPU-less host bf16_unsupported_reason exempts CPU-only, so a
bf16/fp8 (or int8-with-torchao) DiT request passed the preflight, evicted resident workloads, then
raised only in the trainer child. Add the dense-mode CUDA gate mirroring _resolve_base_precision so
the doomed run is rejected up front. Also pin bf16_unsupported_reason in the two positive-path
family-info tests so they are deterministic across GPU types (a non-bf16 CUDA box would otherwise
empty every DiT family's advertised modes).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>

.gitignore

@@ -237,3 +237,8 @@ package-lock.json
 llama.cpp/
 # Stray "~" dir some tools create from a literal ~ TMPDIR; never part of the repo.
 /~/
+
+# Agent workspace runtime/scratch artifacts, never part of the repo.
+async_task_outputs/
+temp/
+logs/

plans/controlnet-workflow.md

@@ -0,0 +1,113 @@
+# Plan: ControlNet for the Studio Images workflow (stacked on #6769)
+
+## Context & research
+
+ControlNet is the **#2 most-used "beyond text-to-image" diffusion workflow** after LoRA (now shipped
+in #6771). It conditions generation on a spatial control map so the output follows a structure. Web
+research (this session) on ComfyUI / Forge / A1111 usage:
+
+- The dominant control types are **Depth, Canny (edges), and OpenPose (human pose)**.
+- The biggest 2025 shift is toward **Union / all-in-one ControlNets** that bundle many control modes in
+  one model: **InstantX / Shakker-Labs `FLUX.1-dev-ControlNet-Union-Pro`** for FLUX, and **`xinsir/
+  controlnet-union-sdxl-1.0`** for SDXL. SDXL has no official ControlNet (community: xinsir, TheMistoAI,
+  BRIA). SD1.5 has the original lllyasviel set.
+- Sources: comfyui-wiki.com ControlNet collections (flux-1 / sdxl), stable-diffusion-art.com ControlNet
+  ComfyUI, education.civitai.com ControlNet guide, stablediffusiontutorials.com Qwen-Image ControlNets.
+
+**Both Studio backends can do ControlNet** (verified against the live tree):
+- diffusers has `FluxControlNetPipeline` / `FluxControlNetModel`, `StableDiffusionXLControlNetPipeline` /
+  `ControlNetModel` / `ControlNetUnionModel`, `QwenImageControlNetPipeline`, `FluxControlNetInpaintPipeline`
+  (diffusers 0.38 in the studio venv). Also `FluxControlPipeline` (the Flux.1-Canny/Depth "Control"
+  in-model variants).
+- native sd.cpp (`stable-diffusion.cpp`, the #6769 sd-cli base) has `--control-net <path>`,
+  `--control-image <path>`, `--control-strength <f>`, `--control-net-cpu`, and a built-in
+  `preprocess_canny` (examples/cli/main.cpp:704, examples/common/common.cpp:422+).
+
+Studio has **no ControlNet wiring today**: `diffusion_families.py` has no controlnet fields and
+`sd_cpp_args.py` has no `--control-*`. This adds it, mirroring the LoRA architecture (#6771) and the
+existing img2img/inpaint/reference workflow patterns.
+
+## Scope (this PR = diffusers ControlNet, Union-first)
+
+Ship the highest-value slice first; keep it shippable and consistent with the shipped LoRA design.
+
+- **In scope:** diffusers ControlNet for the families with the strongest ecosystems and pipeline support:
+  **FLUX** (FluxControlNetPipeline + Union Pro), **SDXL** (StableDiffusionXLControlNetPipeline + xinsir
+  Union), **Qwen-Image** (QwenImageControlNetPipeline). Single ControlNet per generation. A `control_type`
+  hint (canny / depth / pose / tile / passthrough). **Canny preprocessing built in** (cheap, cv2/PIL) plus
+  **passthrough** for user-supplied control maps (depth/pose maps made elsewhere, matching ComfyUI where
+  preprocessing is separate). Strength + guidance start/end. Discovery endpoint + family-gated picker.
+- **Out of scope (follow-ups):** native sd.cpp ControlNet (`--control-net`, needs GGUF ControlNet assets +
+  family support probe); heavy preprocessors (Depth-Anything, OpenPose detector) as optional server-side
+  auto-preprocess; multi-ControlNet stacking; ControlNet + inpaint combo.
+
+## Key facts (verified)
+
+- diffusers ControlNet pipelines are built with `Pipeline.from_pipe(base_pipe, controlnet=cn_model)` (or
+  `from_pretrained(base, controlnet=...)`), so the resident base modules are reused with **no reload** --
+  same `from_pipe` machinery the img2img/inpaint/edit workflows already use (`diffusion.py` ~:1104-1130,
+  `_workflow_pipe`). The ControlNet model (`ControlNetModel` / `FluxControlNetModel` / `ControlNetUnionModel`)
+  is a small extra module loaded once and cached on `_LoadState`.
+- ControlNet models are **family-specific** (SD1.5 CN != FLUX CN != SDXL CN). So discovery must be
+  **family-gated**, exactly like the LoRA picker's `supports_lora`/family filter.
+- Generate-time contract mirrors reference/inpaint: a control image (b64) + params, threaded through
+  `routes/inference.py` into both backends (diffusers serves it; native rejects clearly until the
+  follow-up wires `--control-*`).
+- Reuse: `diffusion_lora.py` discovery/resolve/family-gate patterns; the reference-image upload component +
+  the LoRA picker UI shape; `_workflow_pipe` from_pipe; `hf_hub_download_with_xet_fallback`.
+
+## Approach
+
+### Families (`core/inference/diffusion_families.py`)
+- Add per-family ControlNet declaration: `controlnet_pipeline_class` (e.g. "FluxControlNetPipeline",
+  "StableDiffusionXLControlNetPipeline", "QwenImageControlNetPipeline"), `controlnet_model_class`
+  ("FluxControlNetModel" / "ControlNetModel" / a union class), and a small curated list of recommended
+  ControlNet repos tagged by control type. Expose a `controlnet: bool` capability (like `reference`).
+
+### Discovery (`core/inference/diffusion_controlnet.py`, new -- mirrors diffusion_lora.py)
+- `list_controlnets(family)` = curated family-tagged repos + a local scan; `resolve_controlnet(id, family,
+  hf_token)` downloads via the xet-fallback helper; `preprocess_control(image, control_type)` (canny via
+  cv2/PIL; passthrough otherwise); `supports_controlnet(engine, family, model_kind, quant)` gate
+  (diffusers bf16 / bnb-4bit yes; GGUF-via-diffusers + torchao fp8/int8 dense = no, same rule as LoRA;
+  native = follow-up).
+
+### Backend -- diffusers (`core/inference/diffusion.py`)
+- A ControlNet manager parallel to `_apply_loras`: load the requested `ControlNetModel` once (cache on
+  `_LoadState`, reset on unload/model change), build the CN pipeline via `from_pipe(base, controlnet=...)`
+  in `_workflow_pipe`, and pass `control_image` + `controlnet_conditioning_scale` +
+  `control_guidance_start/end` at generate time. Never fuse; CN model stays bf16.
+
+### Backend -- native (`core/inference/sd_cpp_backend.py`, sd_cpp_args.py) -- FOLLOW-UP
+- Add `--control-net` / `--control-image` / `--control-strength` to the arg builder and a GGUF-ControlNet
+  resolve; gate to families sd.cpp supports. Deferred out of this PR.
+
+### Routes + request models (`models/inference.py`, `routes/inference.py`)
+- Add optional `controlnet: ControlNetSpec` to `DiffusionGenerateRequest` (`{id, image, control_type,
+  strength (0..2, default 1), guidance_start (0..1), guidance_end (0..1)}`); thread into `backend.generate`;
+  surface `supports_controlnet` in status; persist the chosen CN + type in gallery recipe metadata.
+- New `GET /api/models/diffusion-controlnets?family=` (mirror the LoRA discovery route).
+
+### Frontend (`features/images/images-page.tsx`, `api.ts`)
+- A "ControlNet" control in the left rail (reuse the reference-image uploader + a control-type Select +
+  ControlNet-model Select gated by `supports_controlnet`/family + a strength SliderField). Show a small
+  preview of the preprocessed control map. Thread `controlnet` into `generateDiffusionImage`; omit when no
+  control image.
+
+## Verification
+- **Unit:** request validation (optional/empty unchanged; bad strength rejected; unsupported family/quant
+  rejected); discovery (family filter, resolve, canny preprocess shape); diffusers manager (loads CN once,
+  from_pipe built, scale threaded, reset on model change) with a fake pipe; routes (no-CN path unchanged).
+- **Live smoke (critical):** on GPU 4, drive the real diffusers backend with a real family + Union CN and a
+  canny control image; same prompt/seed at strength 0 vs 0.8; assert (a) output DIFFERS from no-control and
+  (b) the strong-control output structurally follows the control map (edge-overlap / SSIM vs the control).
+- **Playwright (`unsloth_studio_workflow`):** upload a control image, pick type + model + strength,
+  generate; capture screenshots/GIF against the live secure studio.
+- Full `pytest studio/backend/tests/` green; frontend `vite build` clean; ruff clean.
+
+## Delivery
+- New branch `diffusion-controlnet` off `diffusion-image-workflows` (#6769 head) in an isolated worktree,
+  sibling to #6771 (LoRA) and #6772 (fp8 fix). PR on `unslothai/unsloth`, base = diffusion-image-workflows,
+  part of the single logical stack rooted on #6763 (continuation of #6658). Commit as Daniel Han; no AI/bot
+  mentions, no emojis, no em dashes.
+- Follow-ups: native sd.cpp ControlNet; server-side Depth/OpenPose auto-preprocessors; multi-ControlNet;
+  ControlNet+inpaint.

plans/diffusion-popularity-findings.md

@@ -0,0 +1,113 @@
+# Diffusion workflow popularity findings (HF download data)
+
+Read-only HF metadata pull (`scripts/investigate_popularity.py`), to ground the Studio
+Images scope against what people actually download. Downloads are HF's 30-day count and
+all-time count; pulled 2026-06-30.
+
+## Qwen-Image-Edit vs Qwen-Image-Layered (the explicit "determine popularity" question)
+
+| Model | dl / 30d | dl all-time | likes | pipeline |
+|---|---:|---:|---:|---|
+| Qwen/Qwen-Image-Edit-2509 | 511,996 | 2,942,966 | 1,185 | image-to-image |
+| Qwen/Qwen-Image-Edit-2511 | 162,185 | 1,088,015 | 1,087 | image-to-image |
+| Qwen/Qwen-Image-Edit (base) | 70,728 | 1,161,044 | 2,440 | image-to-image |
+| **Qwen-Image-Edit (all variants)** | **~745,000** | **~5,192,000** | - | - |
+| Qwen/Qwen-Image-Layered | 51,303 | 234,785 | 1,112 | image-text-to-image |
+| unsloth/Qwen-Image-Edit-2511-GGUF | 218,313 | - | - | image-to-image |
+
+**Conclusion:** Qwen-Image-Edit is ~10-14x more downloaded than Layered (combined 745K/30d
+vs 51K, 5.2M vs 235K all-time). Shipping Edit (2511 + the unsloth GGUF, which alone pulls
+218K/30d) and rejecting/deferring Layered is the correct, data-backed call. Layered also
+needs a dedicated pipeline (`additional_t_cond=True`) the standard QwenImagePipeline can't
+drive, so it would be both niche AND extra engineering. Reject stands.
+
+## ControlNet is niche on the modern (diffusers/FLUX/Qwen) stack
+
+| Model | dl / 30d | dl all-time | likes |
+|---|---:|---:|---:|
+| InstantX/FLUX.1-dev-Controlnet-Canny | 2,891 | 136,727 | 194 |
+| lllyasviel/ControlNet (SD1.5-era) | 0 | 14 | 3,820 |
+| stabilityai/stable-diffusion-x4-upscaler | 10,040 | 2,976,405 | 725 |
+
+**Conclusion:** ControlNet's large user base lives in the older SD1.5 / A1111 ecosystem, not
+the diffusers/FLUX/Qwen stack Studio targets (the modern FLUX ControlNet is ~3K/30d). It is
+NOT part of the "most popular ~80%" for current-gen models, so deferring it is justified by
+the data, not just by effort. The dedicated x4 upscaler is also low 30-day (10K) though high
+all-time; our generic hires-fix upscale (img2img re-detail) covers the use case for any
+loaded family without an extra model.
+
+## The shipped six cover the popular workflows
+
+Top text-to-image (HF list, 30-day): SD1.5 (1.78M), SDXL (1.32M), FLUX.1-dev (1.09M),
+dreamshaper-7 (1.03M), **Tongyi-MAI/Z-Image-Turbo (886K)**, sd-turbo (684K), SD3.5-medium
+(606K), sdxl-turbo (598K), Qwen-Image-Lightning (483K). All are plain txt2img -> our Create
+tab; the GGUF/bnb families + Z-Image cover the modern ones.
+
+Top image-to-image (HF list, 30-day): Qwen-Image-Edit-2509 (512K) -> Edit tab; SDXL-refiner
+(162K) -> Upscale/Transform; Kontext (150K) -> Edit tab.
+
+So Create / Transform / Inpaint / Extend / Upscale / Edit map onto the head of both
+distributions.
+
+## SHIPPED: FLUX.2-klein image (reference) conditioning
+
+> Status: IMPLEMENTED + verified live (2026-06-30). `flux.2-klein` now has `reference=True`;
+> the backend exposes a "reference" workflow that passes the image to the loaded
+> Flux2KleinPipeline directly (no from_pipe, no strength, output at the requested size); the
+> frontend has a "Reference" tab. Verified with `scripts/verify_reference_http.py` on
+> `unsloth/FLUX.2-klein-4B-GGUF` (Q4_K_M): a reference-conditioned 1024x1024 result is
+> non-blank, correctly sized, and DIFFERS from the identical-seed plain txt2img.
+>
+> FLUX.2-klein ALSO gained inpaint (`Flux2KleinInpaintPipeline` via from_pipe; verified with
+> `scripts/verify_klein_inpaint.py`). It does NOT get outpaint/extend: FLUX.2 scales any >1MP
+> input down to ~1MP, so a padded outpaint canvas shrinks back. "outpaint" is now a distinct
+> capability advertised only for size-preserving inpaint families (`inpaint_preserves_size`).
+> Multi-reference is shipped too (the pipeline accepts a list; the Reference tab has add/remove
+> slots, backend caps at 3 extra; verified with `scripts/verify_multiref_http.py`: two
+> references differ from one at the same seed). The analysis that motivated the work follows.
+
+The data surfaced this gap (now closed):
+
+| Model | dl / 30d | pipeline |
+|---|---:|---|
+| black-forest-labs/FLUX.2-klein-4B | 470,482 | image-to-image (#2 overall) |
+| black-forest-labs/FLUX.2-dev | 271,037 | image-to-image |
+| **unsloth/FLUX.2-klein-4B-GGUF** | **243,307** | image-to-image |
+| black-forest-labs/FLUX.2-klein-9B | 178,964 | image-to-image |
+
+`flux.2-klein` is ALREADY a registered family in `diffusion_families.py` (txt2img only,
+base `FLUX.2-klein-4B`, open repo). But `Flux2KleinPipeline.__call__` natively accepts an
+`image` argument (verified in diffusers 0.38.0; params: image, prompt, height, width,
+num_inference_steps, guidance_scale -- NOTE: no `strength`). FLUX.2 is a unified
+text-to-image + reference/edit model: the SAME loaded pipe does both, depending on whether
+`image` is passed. Today Studio exposes only txt2img for it, so the popular image-editing
+mode of the #2 image-to-image model is unreachable.
+
+### Why it's a separate PR, not a tail-of-session add
+FLUX.2 reference conditioning is a DIFFERENT semantic from the shipped workflows:
+- No `strength` (it is reference-conditioning, not a denoise blend like img2img).
+- Output size comes from width/height (txt2img-style), not from the input image size, so the
+  image-conditioned width/height rule we added for img2img/inpaint/upscale does NOT apply.
+- FLUX.2 supports MULTIPLE reference images; single-image is the common case but the UX
+  should not preclude multi-ref.
+This needs: read the Flux2KleinPipeline source for exact `image` semantics (list vs single,
+how it is resized/tiled, recommended guidance), decide the UX (a "Reference" workflow that is
+available alongside Create for `reference=True` families, distinct from the strength-based
+Transform tab), then verify on the open FLUX.2-klein-4B base (and the unsloth GGUF) with a
+reference image before/after.
+
+### Sketch (for the follow-up PR)
+- `diffusion_families.py`: add `reference: bool = False`; set `reference=True` on flux.2-klein.
+- `_family_workflows`: when `fam.reference`, expose `"reference"` (in addition to txt2img).
+- `generate()`: a `reference` branch that passes `image` to `state.pipe` directly (no
+  from_pipe, no strength), with width/height = the requested size (NOT the input size).
+- Frontend: a "Reference" tab (image dropzone + prompt), gated to `reference` families;
+  Create stays pure txt2img for the same model.
+- Verify: load unsloth/FLUX.2-klein-4B-GGUF, pass a reference image, confirm the output is
+  conditioned on it and differs from a no-image run at the same seed.
+
+## Net
+The seven shipped workflows (create, transform, inpaint, extend, upscale, reference, edit)
+cover the popular ~80% across both the txt2img and image-to-image distributions, including the
+#1 image-to-image model (Qwen-Image-Edit) and the #2 (FLUX.2-klein, now via the reference tab).
+ControlNet / SD1.5-era ControlNet remain deferred with data backing (niche on the modern stack).

plans/diffusion-workflows-pr-plan.md

@@ -0,0 +1,156 @@
+# Stacked-PR plan: Studio diffusion workflows (Images redesign)
+
+Branch tip: `diffusion-eager-and-compile-cache` (latest commit "Phase 16 review fixes").
+Remote: `oobabooga/unsloth`. New PRs stack on top of the existing diffusion stack
+(ultimately on top of unslothai/unsloth#6658), treated as one logical change.
+
+Nothing here is committed yet (commit/push only on explicit request).
+
+## CRITICAL: the working tree holds TWO uncommitted streams, and three core files INTERMINGLE them
+
+A full `git status` / marker audit (branch `diffusion-eager-and-compile-cache`, tip "Phase 16
+review fixes") shows the uncommitted tree is NOT a clean single feature. There are two streams:
+
+A) **The eager/compile-cache phase** (the branch's own in-progress work, NOT this session's —
+   zero of this feature's markers). Purely-its files, safe to NOT touch in the workflow PRs:
+   - new modules: `diffusion_arch_patches.py`, `diffusion_compile_cache.py`,
+     `diffusion_eager_patches.py`, `diffusion_gguf_compile.py`, `diffusion_patch_backend.py`
+   - new tests: `test_diffusion_arch_patches.py`, `test_diffusion_compile_cache.py`,
+     `test_diffusion_eager_patches.py`, `test_diffusion_gguf_compile.py`
+   - modified: `diffusion_speed.py`, `test_diffusion_speed.py`, `conftest.py`,
+     `scripts/diffusion_bench.py`, and ~25 untracked `scripts/*bench*/*probe*/*orchestrator*`.
+
+B) **The Images workflows feature** (this session): the workflow engine + frontend + installer.
+
+**The two streams INTERMINGLE inside three shared files and are NOT separable by file:**
+  - `studio/backend/core/inference/diffusion.py` — this feature's workflow hunks are interleaved
+    with the eager/compile wiring (imports at L67-75; `install_arch_patches`/`compile_cache.begin`/
+    `.restore`/`.save` and the `eager_patched`/`compile_cache_ctx` state throughout
+    `load_pipeline`/`generate`/`unload`). A single `diffusion.py` cannot go into one PR without the
+    other stream's hunks.
+  - `studio/backend/models/inference.py` — this feature's `init_image`/`mask_image`/
+    `reference_images`/`upscale`/`model_kind` fields sit next to the pre-existing `speed_mode`/
+    `transformer_prequant_path` fields in the same request models.
+  - `studio/backend/tests/test_diffusion_backend.py` — this feature's workflow tests sit next to
+    the pre-existing `test_failed_load_rolls_back_eager_patches` (imports `diffusion_eager_patches`).
+
+**Implication / options (USER DECIDES — it is their branch + their eager/compile work):**
+  - CLEANLY separable now (purely this feature, can be committed/PR'd on their own anytime):
+    frontend `images-page.tsx` + `api.ts` + `pickers.tsx`, and the sd.cpp installer
+    `install_sd_cpp_prebuilt.py` + `test_sd_cpp_install.py`. (These are PR 2 and PR 3 below.)
+  - The backend engine (PR 1) CANNOT be cleanly split from the eager/compile phase via files.
+    Realistic paths: (a) finalize + commit the eager/compile phase first, then this feature's
+    backend lands as a clean diff on top; or (b) commit both streams together as the branch's
+    next chunk (consistent with treating the stack as one logical change); or (c) a manual
+    `git add -p` hunk split of the three shared files (tedious, risks a non-compiling
+    intermediate). NOT auto-doable safely without the owner's intent for the eager/compile work.
+
+## Proposed stack (3 PRs, bottom to top)
+
+### PR 1 - Backend: diffusion workflow engine (safetensors + image-conditioned + editing)
+Files:
+- `studio/backend/core/inference/diffusion.py` (the feature hunks: three load "kinds"
+  gguf/single_file/pipeline; `_workflow_pipe` via `from_pipe(torch_dtype=None)`;
+  `_align_vae_dtype`; `generate()` routing for reference/img2img/inpaint/upscale/edit;
+  image-conditioned width/height from the input image (but reference + txt2img use the slider
+  size); `upscale` (hires fix) branch on the img2img pipe; `reference` (FLUX.2) branch that
+  passes the image(s) to the loaded pipe directly (no from_pipe, no strength) incl. multi-
+  reference (`reference_images` combined into a list, capped at 3 extra); branch ORDER
+  inpaint/upscale before reference so a mask/upscale request on a reference family still routes
+  right; `_family_workflows` (adds "upscale" wherever img2img is supported, "reference" for
+  reference families, "outpaint" only for size-preserving inpaint families); `kind` on state +
+  `model_kind` in status; `load_progress` double-count fix). NOTE: this file ALSO carries
+  pre-existing speed hunks if any landed here - review per-hunk and exclude non-feature hunks.
+- `studio/backend/core/inference/diffusion_families.py` (img2img/inpaint pipeline slots;
+  `edit` flag + `reference` flag + `inpaint_preserves_size` flag; `qwen-image-edit` +
+  `flux.1-kontext` families; flux.2-klein gains reference + inpaint (no outpaint: FLUX.2
+  normalizes to ~1MP); `detect_family` longest-match + leftover-reject; `layered` reject).
+- `studio/backend/core/inference/diffusion_engine_router.py` (model_kind -> diffusers for
+  non-gguf kinds).
+- `studio/backend/core/inference/diffusion_memory.py` (`estimate_safetensors_dense_mib`).
+- `studio/backend/core/inference/sd_cpp_backend.py` (model_kind passthrough; reject
+  img2img/inpaint on the native engine).
+- `studio/backend/models/inference.py` (load request: optional gguf_filename, model_kind,
+  init/mask/strength, advanced knobs; status: workflows, model_kind).
+- `studio/backend/routes/inference.py` (model_kind forwarding; ValueError -> 400;
+  exc_info logging).
+- Tests: `test_diffusion_backend.py`, `test_diffusion_routes.py`.
+
+Title: `Studio diffusion: safetensors + image-conditioned + instruction-editing workflows`
+Summary: Adds non-GGUF safetensors loading (full bnb-4bit pipelines + single-file fp8,
+gated to unsloth/*), the image-conditioned workflows (img2img, inpaint, outpaint via the
+inpaint path) built with `Pipeline.from_pipe` for zero-extra-VRAM component reuse, and
+instruction editing as its own family kind (Qwen-Image-Edit-2511 + FLUX.1-Kontext-dev).
+Fixes two bugs: `from_pipe` defaulting to a float32 recast that crashed torchao-quantized
+transformers, and image-conditioned calls forcing the slider size onto the input image.
+
+### PR 2 - Frontend: redesigned Images page (workflow tabs + Advanced Options)
+Files:
+- `studio/frontend/src/features/images/images-page.tsx` (workflow tabs Create/Transform/
+  Inpaint/Extend/Upscale/Reference/Edit; capability gating + auto-switch; `MaskCanvas`;
+  `buildOutpaint`; Upscale tab with Scale + Detail-strength sliders; Reference tab (FLUX.2,
+  reference image + add/remove extra references, no strength); Advanced Options accordion gated
+  to GGUF for transformer-quant; spinner-overlap fix).
+- `studio/frontend/src/features/images/api.ts` (request/status types incl. model_kind, upscale,
+  reference_images).
+- `studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx` (curated
+  safetensors + edit GGUF rows; `SUPPORTED_EDIT_KEYWORDS` un-hide; layered hide).
+
+Title: `Studio Images: workflow tabs (create/transform/inpaint/extend/upscale/reference/edit) + Advanced Options`
+Summary: Redesigns the Images page around capability-gated workflow tabs with a brush mask
+editor, client-side outpaint, a hires-fix upscale tab, a FLUX.2 reference tab, an instruction-edit
+tab, and an Advanced Options panel (speed/quant/attention/memory/step-cache/offload), plus the
+number-input spinner fix.
+
+### PR 3 - sd.cpp prebuilt installer hardening
+Files: `studio/install_sd_cpp_prebuilt.py`, `studio/backend/tests/test_sd_cpp_install.py`.
+Title: `Studio sd.cpp: pin release + verify sha256 + mirror-ready source`
+Summary: Pins the stable-diffusion.cpp release (was tracking `latest`), verifies each
+download's sha256 against GitHub's published asset digest before extract/execute, adds a
+download timeout, and makes the source repo configurable (`UNSLOTH_SD_CPP_REPO`) so a
+future unslothai mirror needs no code change. Cleanly separable from the rest.
+
+## Pre-PR review (done)
+An independent 3-angle review (backend correctness, frontend/UX, security/robustness) ran over
+the full session diff. No High findings; the load-gating to `unsloth/*` and the multi-reference
+count caps were verified intact end-to-end. Fixes applied before the PRs:
+- Frontend [Med]: multi-reference slots no longer renumber mid-edit (dropped the eager
+  `filter(Boolean)` in the per-slot onChange; empties dropped only at send).
+- Backend [Med]: upscale now caps the ABSOLUTE output (longest side <= 2048), not just the
+  factor, so a large upload * 4x can't OOM.
+- Backend/security [Med]: `_decode_b64_image` rejects images > 4096px/side (uniform guard for
+  init/mask/reference vs decompression-bomb / OOM inputs); base64 image fields capped at 32 MiB.
+- Security [Low]: the native sd.cpp engine guard also rejects `reference_images`.
+All covered by new tests (82 backend pass) and a post-fix e2e (all five workflows still pass).
+
+## Post-deploy user feedback fixes (done)
+From live use of the deployed studio:
+- Backend [Med]: image-conditioned workflows passed the raw upload size to pipelines that
+  require multiples of 16 (Z-Image/Qwen/FLUX), so an odd upload (e.g. 186px) failed with
+  "Height must be divisible by 16". Added `_snap_to_multiple` and auto-resize init (and the
+  matched mask) to the nearest /16 for img2img/inpaint/extend/edit. Verified live: a 186x250
+  Transform and Inpaint now return 200 at 192x256. Tests added.
+- Frontend [Med]: the Advanced options (FP8/INT8 quant, speed, attention, memory) were a
+  collapsed, muted accordion at the bottom of the left rail that users missed (HF screenshots
+  discussion #25). Moved them into a RIGHT-DOCKED panel mirroring Chat's settings panel.
+  Per follow-up (discussion #26): CLOSED by default, toggled by a SINGLE fixed top-bar button
+  using Chat's `LayoutAlignRightIcon` that stays in the exact same position in both states
+  (verified x/y identical open vs closed) and highlights when open. Controls extracted to a
+  render-local `advancedControls`; unused Accordion import removed.
+
+## Constraints for execution (when authorized)
+- Write as the user; no AI/bot mentions, no emojis, no em dashes; well-formatted bodies.
+- `gh auth status` first. Push to `oobabooga/unsloth`, stack on the current branch.
+- Keep PR 3 independent; PR 2 depends on PR 1 (frontend needs the backend contract).
+- Re-run `pytest studio/backend/tests/test_diffusion_*.py test_sd_cpp_install.py` +
+  frontend `tsc`/`build` before each PR.
+
+## Out of scope / follow-ups
+Scope decisions are backed by HF download data in `plans/diffusion-popularity-findings.md`.
+The seven shipped workflows are create, transform, inpaint, extend, upscale, reference, edit.
+- Publish the unslothai/stable-diffusion.cpp mirror + macOS/Windows staging (#152).
+- ControlNet / style-transfer: the goal's "most popular" set is covered by the seven shipped
+  workflows. ControlNet on the modern diffusers/FLUX/Qwen stack is niche by downloads
+  (~3K/30d), so deferring it is data-backed, not just an effort call.
+- FLUX.2-klein inpaint and multi-reference are DONE (shipped). Outpaint is intentionally not
+  offered for FLUX.2 (it scales >1MP inputs to ~1MP). No further FLUX.2 follow-ups outstanding.

plans/diffusion-workflows-studio.md

@@ -0,0 +1,125 @@
+# Plan: Unsloth Studio diffusion workflows + Images UI redesign
+
+Stacked as NEW PRs on top of the existing 14-PR diffusion stack above unslothai/unsloth#6658
+(treated as one logical base). Goal: cover ~80% of common real-world diffusion workflows,
+across macOS/Linux/Windows/CPU, optimizing performance, accuracy, and memory.
+
+## Current state (verified by recon)
+
+- Backend is **text-to-image only** end-to-end. No image/mask/control plumbing in
+  `DiffusionGenerateRequest`, `/images/generate`, or `DiffusionBackend.generate()`.
+- **diffusers 0.38.0 already imports every pipeline we need**: `*Img2ImgPipeline`,
+  `*InpaintPipeline`, `FluxFillPipeline`, `*ControlNetPipeline`/`ControlNetModel`,
+  `FluxKontextPipeline`, `QwenImageEditPipeline`/`QwenImageEditPlusPipeline`,
+  `StableDiffusion(Latent)UpscalePipeline`. No diffusers upgrade required.
+- The native **sd.cpp engine already has dormant fields** (`init_img`, `strength`, `mask`,
+  `ref_images`) and a complete `upscale()` path — never wired to the request/route.
+- **All advanced LOAD options are already wired** end-to-end (speed_mode, transformer_quant
+  fp8/int8/nvfp4/mxfp8, attention_backend, memory_mode, cpu_offload, transformer_cache,
+  vae_tiling status). The Advanced panel is mostly a FRONTEND surfacing job.
+- Frontend has `Tabs` (`components/ui/tabs.tsx`) and `Accordion` ready. No dropzone and no
+  mask/brush canvas — both greenfield. `SliderField` is the customizer number input.
+- `_EDIT_KEYWORDS = ("edit","kontext","inpaint","layered")` rejects edit/inpaint/Kontext/
+  Layered repos at family detection.
+- sd.cpp binary is downloaded prebuilt from **upstream leejet/stable-diffusion.cpp** (no
+  unsloth mirror, no checksum/manifest/version pin, not wired into setup.sh). llama.cpp uses
+  an `unslothai/llama.cpp` mirror with manifest+sha256+version pin+source fallback.
+- Chat "Images" pill = provider-side (OpenAI/Gemini) hosted tool, separate from local diffusion.
+
+## Workflow popularity ranking (what to build for 80% coverage)
+
+1. txt2img (keep, polished) — done
+2. img2img / variations — **P0**
+3. inpainting (mask edit) — **P0**
+4. upscaling / hires fix — **P0**
+5. ControlNet (Canny/Depth/Pose/Lineart/Tile) — **P0/P1**
+6. outpainting (canvas extend) — **P1**
+7. instruction image editing (Qwen-Image-Edit, FLUX Kontext) — **P1**
+8. style transfer / reference — **P1** (via img2img/edit/control)
+9. batch generation/edit/upscale — **P1**
+10. LoRA/style packs — **P2**
+
+## Editing-model decisions (researched)
+
+- **Qwen-Image-Edit / Edit-2511**: popular, best-in-class clean targeted edits + multilingual
+  text. **Support** (instruction edit, mask-optional).
+- **FLUX.1 Kontext**: popular, character-consistent in-context editing. **Support** (note:
+  Kontext-dev is non-commercial/gated — surface license, don't block local custom models).
+- **Qwen-Image-Layered**: newer, niche (Photoshop RGBA layer decomposition). Needs a dedicated
+  pipeline (`additional_t_cond`) — **defer** (keep rejected for now; optional later behind a
+  layered-specific view). This already crashed the standard path (the earlier bug).
+
+## UI design — workflow tabs (inside ImagesPage, route/nav unchanged)
+
+`Tabs` across the top of the controls area. Combine related workflows:
+- **Create** — txt2img (current behavior preserved)
+- **Transform** — img2img + style transfer (upload + strength/denoise + presets)
+- **Edit** — inpaint (mask brush/upload/invert/feather, masked-vs-whole) + instruction edit
+  (Qwen-Image-Edit / FLUX Kontext, mask-optional)
+- **Extend** — outpainting (directional handles, aspect presets, overlap/feather)
+- **Control** — ControlNet (one control slot first: Canny/Depth/Pose/Lineart/Tile + preview)
+- **Enhance** — upscaling (ESRGAN/RealESRGAN + latent/tiled) 
+- **Advanced Options** — Accordion surfacing existing load knobs (speed/compile/attention/
+  quant fp8/int8/nvfp4/memory/offload/vae tiling/cache) with Auto defaults + resolved values.
+
+Capability gating: a workflow/control is shown enabled only when the selected engine+family+
+device+quant supports it; otherwise disabled with a plain-language "why".
+
+## Backend architecture
+
+- Extend `DiffusionGenerateRequest`: optional `workflow` (txt2img|img2img|inpaint|outpaint|
+  control|edit|upscale), `init_image` (b64), `mask_image` (b64), `control_image` (b64),
+  `strength`, `controlnet_conditioning_scale`, `control_start/end`, `upscale_factor`,
+  `ref_images`. Add an image-decode (b64→PIL) helper (none exists).
+- `DiffusionFamily`: add optional pipeline-class slots (`img2img_pipeline_class`,
+  `inpaint_pipeline_class`, `edit_pipeline_class`, `controlnet_pipeline_class` + control repos).
+  Build the right pipeline around the already-loaded `transformer=` (reuse `_assemble_pipe`
+  shape); swap/cache pipeline class per workflow without reloading the transformer where
+  possible.
+- `generate()` kwarg builder must branch: img2img/edit pipelines take `image=`/`strength=` and
+  reject `width/height`; inpaint adds `mask_image=`; control adds `control_image=`. Gate each
+  kwarg via `inspect.signature`.
+- Capability resolver: maps engine+family+device+quant → supported workflows + reasons; echoed
+  in run metadata so the UI shows what actually ran.
+- Memory planner must account for input/latent size, control models, VAE decode, upscale.
+
+## PR breakdown (stacked, small, capability-gated)
+
+- **PR-1 UI fixes + workflow shell**: fix number-input spinner overlap (DONE in tree), tab
+  scaffold (Create/Transform/Edit/Extend/Control/Enhance/Advanced), Advanced Options accordion
+  surfacing existing load knobs, capability banner, loading/empty/error states.
+- **PR-2 Backend workflow contract + capability registry**: extend request/response, decode
+  helper, per-family pipeline slots, resolver. No new behavior yet beyond txt2img.
+- **PR-3 img2img (Transform)**: backend + Transform tab + dropzone (adapt from
+  shared-composer `addFiles`/`PendingImageThumb`). Smoke test low vs high denoise.
+- **PR-4 inpaint + instruction edit (Edit)**: mask canvas (greenfield), inpaint pipeline,
+  Qwen-Image-Edit/FLUX Kontext edit; relax `_EDIT_KEYWORDS` → route to edit family.
+- **PR-5 outpaint (Extend)**: expanded-canvas inpaint, directional handles, feather/overlap.
+- **PR-6 ControlNet (Control)**: one control slot + preprocessor preview + strength/start/end.
+- **PR-7 upscaling (Enhance)**: wire dormant sd.cpp `upscale()` + diffusers upscale + `/images/upscale`.
+- **PR-8 Advanced panel polish + FP8/INT8 verification matrix**.
+- **PR-9 sd.cpp prebuilt packaging**: mirror to `unslothai/stable-diffusion.cpp`, manifest+
+  sha256+version pin+`--published-repo`+source fallback, wire into setup.sh (ref
+  install_llama_prebuilt.py).
+- **PR-10 cross-platform staging validation** (danielhanchen staging repos, small GGUFs).
+- **PR-11 Playwright tests + screenshots/GIFs per tab** (studio_test_kit / unsloth_studio_workflow).
+- **PR-12 batch + multi-control + reproducibility polish** (later).
+
+## Done so far
+
+- Fixed the customizer number-input spinner overlap (`SliderField` in images-page.tsx): native
+  spinners covered the value on the narrow field; now fully suppressed (webkit inner+outer +
+  Firefox `appearance:textfield`) and field widened to `w-14`. Frontend rebuilt clean.
+
+## Verification
+
+- Playwright (studio_test_kit) per tab: screenshots + GIFs, capability gating, upload/mask,
+  progress/cancel/error, gallery.
+- B200 functional: load + generate one image per workflow per representative family.
+- FP8 + INT8 verified (build matrix: SDXL/FLUX/Qwen-Image/Qwen-Image-Edit/GGUF; measure
+  black-image/NaN rate, peak VRAM, time-to-first-image, prompt adherence, source preservation).
+- Cross-platform staging (Linux CUDA/CPU, Windows CUDA/CPU, macOS MPS) with small GGUFs.
+
+## Delivery
+
+New branch(es) off the current tip; new stacked PRs. Commit/push only when asked.

plans/wobbly-jumping-narwhal.md

@@ -0,0 +1,141 @@
+# Plan: Publish unslothai/stable-diffusion.cpp mirror + our own CPU/Apple prebuilts
+
+## Context
+
+The Unsloth Studio native diffusion engine downloads a prebuilt `sd-cli` / `sd-server`
+(stable-diffusion.cpp) via `studio/install_sd_cpp_prebuilt.py`. Today it pulls from
+**leejet/stable-diffusion.cpp** upstream releases. We want to own this like we own
+**unslothai/llama.cpp**: a fork that builds and publishes OUR OWN prebuilt binaries on a
+schedule, so we control reproducibility, integrity, and the macOS load floor.
+
+**Why native is CPU/Apple-only.** On a GPU host, diffusers + our optimizations (regional
+`torch.compile` ~2.2x, cuDNN/flash attention, FP8/INT8/NVFP4 quant, first-block-cache) is
+faster than sd.cpp's CUDA path, which has none of those levers — so GPU hosts route to
+diffusers. Native sd.cpp only wins where diffusers is weak: **CPU and Apple**. Therefore we
+build native binaries ONLY for the platforms where native is actually the faster engine, and
+skip CUDA/ROCm/Vulkan entirely (GPU = diffusers/torch). This also makes the CI far cheaper.
+
+The Studio side is already prepared: `install_sd_cpp_prebuilt.py` reads `UNSLOTH_SD_CPP_REPO`
+(repo override) + `UNSLOTH_SD_CPP_TAG` (pinned tag) and verifies the GitHub asset `digest`
+(`_verify_sha256`). So the bulk of the work is the mirror repo + release CI; the Studio change
+is a small default flip + resolver tweak.
+
+## Coverage (user-confirmed): CPU / Apple ONLY
+
+| Platform | Arch | Build | Runner | Notes |
+|---|---|---|---|---|
+| macOS | arm64 | Metal (`-DSD_METAL=ON`) | macos-26, `OSX_DEPLOYMENT_TARGET=14.0` | Apple fast path (diffusers/MPS weak) |
+| macOS | x86_64 | CPU | macos-15-intel, `OSX_DEPLOYMENT_TARGET=13.3` | Intel Macs |
+| Linux | x86_64 | CPU | ubuntu-22.04 (glibc 2.35) | **also covers WSL** (WSL = Linux x64) |
+| Linux | aarch64 | CPU | ubuntu-24.04-arm | ARM servers |
+| Windows | x86_64 | CPU | windows-2022 (MSVC+Ninja) | |
+
+**Explicitly out of scope:** CUDA, ROCm, Vulkan native builds; GPU runners; cudart bundling;
+per-gfx matrices. GPU stays on diffusers/torch.
+
+## Reference pattern (verified this session)
+
+`unslothai/llama.cpp` builds via `.github/workflows/unsloth-prebuilt.yml` (orchestrator) + six
+per-accel children + `scripts/unsloth/` helpers (`assemble_metadata.py`, `package_bundle.py`,
+`assert_macho_minos.sh`). Mechanisms to mirror: `resolve` (supply-chain aging — only build a
+release public >=6h; stamp build-info + Unsloth fingerprint; upload ONE source artifact all
+children extract) -> per-platform children (build from the source artifact, load-gate, package,
+upload) -> `assemble` (fingerprint gate + manifest/sha256 index + coverage gate + **atomic
+draft->publish**, no partial releases). Template files fetched to `workspace_81/temp/llamacpp_workflows/`.
+
+## Key facts (verified)
+
+- leejet builds both `sd-cli` and `sd-server` (`examples/cli`, `examples/server`) — the mirror
+  ships both (sd-server is used by PR #6768's persistent server).
+- leejet naming: `sd-<tag>-bin-<Darwin-macOS-…-arm64 | Linux-Ubuntu-…-x86_64 | win-cpu-x64>.zip`.
+  leejet already ships macOS arm64, Linux x64 CPU, Windows CPU — we ADD macOS x86_64 and Linux
+  aarch64 (the gaps in our target set), and rebuild the rest under our own fingerprint/integrity.
+- Studio resolver (`resolve_release_asset`, `install_sd_cpp_prebuilt.py:88`): filters to `.zip`;
+  macOS = darwin/macos + arch token; Linux = `linux` + arch + (no accel marker for auto/cpu);
+  Windows = `bin-win` + `avx2` else any. For a CPU-only mirror the resolver needs essentially NO
+  change — macOS x86_64 and Linux aarch64 already match by arch token; just confirm the Windows
+  CPU asset resolves (contains `bin-win`, falls back to the plain build).
+
+## Design
+
+### A. Mirror repo (fork of leejet/stable-diffusion.cpp)
+
+Fork so upstream C++ stays intact; add only `.github/workflows/` + `scripts/unsloth/`. Adapt the
+llama.cpp orchestrator, heavily simplified (no CUDA/ROCm/Vulkan, no PR-mix):
+
+- **`resolve`**: pick the upstream leejet tag with the >=6h aging window; reuse leejet's
+  `master-<count>-<sha>` as the mirror tag (keeps `UNSLOTH_SD_CPP_TAG` comparable to upstream);
+  stamp a source tarball with build-info + the "Compiled by the Unsloth team" fingerprint; upload
+  the source artifact. Skip-if-already-published like llama.cpp.
+- **Build children** (reusable `workflow_call`), each `cmake -DSD_BUILD_EXAMPLES=ON` (cli+server):
+  - `macos` (arm64 Metal + x64 CPU): pinned `CMAKE_OSX_DEPLOYMENT_TARGET`, `@loader_path` rpath,
+    load-gate via `assert_macho_minos.sh` (adapted for `sd-cli`/`sd-server`).
+  - `cpu-linux` (x64 + arm64) and `cpu-windows` (x64, MSVC+Ninja).
+- **Asset naming = leejet-compatible**, all `.zip`:
+  `sd-<tag>-bin-Darwin-macOS-arm64.zip`, `sd-<tag>-bin-Darwin-macOS-x86_64.zip`,
+  `sd-<tag>-bin-Linux-Ubuntu-24.04-x86_64.zip`, `sd-<tag>-bin-Linux-Ubuntu-24.04-aarch64.zip`,
+  `sd-<tag>-bin-win-cpu-x64.zip`.
+- **`assemble`**: fingerprint gate (verify the mark in every archive), generate
+  `sd-prebuilt-manifest.json` + `sd-prebuilt-sha256.json`, coverage gate (all 5 assets present),
+  atomic draft->publish. GitHub sets each asset `digest`, which the Studio already verifies.
+- **Signing/notarization:** none. The Studio downloads via `urllib` (not a browser), so no macOS
+  quarantine xattr is set and Gatekeeper does not block CLI-run binaries (matches llama.cpp).
+
+### B. Studio-side switch (PR on the diffusion stack, after the mirror's first green release)
+
+Small, in `studio/install_sd_cpp_prebuilt.py` + its test:
+1. `DEFAULT_REPO = "unslothai/stable-diffusion.cpp"`; `DEFAULT_TAG` = the mirror's first tag.
+2. Confirm `resolve_release_asset` picks correctly for all 5 CPU/Apple hosts (add a Windows CPU
+   token only if the plain-`bin-win` fallback proves insufficient; likely no change needed).
+   Keep the leejet fallback (env override still points back upstream).
+3. Extend `test_sd_cpp_install.py` `_ASSETS` to the mirror's 5-asset set; assert host->pick for
+   macOS arm64/x64, Linux x64/arm64, Windows x64; assert GPU hosts are unaffected (still diffusers).
+
+## Critical files
+
+- New (mirror repo): `.github/workflows/unsloth-sd-prebuilt.yml` (+ `-macos.yml`, `-cpu-linux.yml`,
+  `-cpu-windows.yml`), `scripts/unsloth/{assemble_metadata.py,package_bundle.py,assert_macho_minos.sh}`.
+- Studio: `studio/install_sd_cpp_prebuilt.py`, `studio/backend/tests/test_sd_cpp_install.py`.
+- Local templates to adapt: `workspace_81/temp/llamacpp_workflows/{unsloth-prebuilt.yml,unsloth-prebuilt-macos.yml,unsloth-prebuilt-cpu.yml}`.
+
+## Sequencing (chicken-and-egg)
+
+1. Build the mirror repo + CI; `publish=false` dry run to validate the 5-way matrix (~10-20 min,
+   no GPU runners so cheap).
+2. First green **published** release with all 5 assets + manifest/sha256.
+3. THEN the Studio PR flips `DEFAULT_REPO`/`DEFAULT_TAG` + resolver test (on the diffusion stack).
+
+## Verification
+
+- **Resolver unit tests** (hermetic): feed the mirror's 5 asset names to `resolve_release_asset`
+  for macOS arm64/x64, Linux x64/arm64, Windows x64 -> correct pick; and CUDA/GPU host -> still
+  routes to diffusers (native not selected).
+- **CI dry run**: `publish=false` artifact-only run; inspect the 5 archives each contain `sd-cli`
+  (+ `sd-server`) and carry the fingerprint.
+- **Live install smoke** (this Linux box): `UNSLOTH_SD_CPP_REPO=unslothai/stable-diffusion.cpp
+  python studio/install_sd_cpp_prebuilt.py --print-asset` then real `install()`, confirm
+  `sd-cli --version` + `sd-server` launch, and drive one native CPU generation via the Studio.
+- **Integrity**: each published archive matches its manifest sha256 and the GitHub asset digest.
+
+## Staging (user-confirmed): fork + push CI now
+
+Execution order:
+1. **Preflight permissions**: `gh auth status`; confirm the token can create/fork under the
+   `unslothai` org and enable Actions. If it CANNOT, stop and report (fall back to scaffold-only,
+   or a private fork under danielhanchen), rather than pushing somewhere unintended.
+2. **Fork** leejet/stable-diffusion.cpp -> `unslothai/stable-diffusion.cpp` (clone into the
+   workspace to add files). Keep upstream C++ intact.
+3. **Add CI + scripts** on a branch: `.github/workflows/unsloth-sd-prebuilt.yml` +
+   `-macos.yml`/`-cpu-linux.yml`/`-cpu-windows.yml`, `scripts/unsloth/*`. Commit as Daniel Han
+   (no AI/bot mentions, no emojis, no em dashes). `unset GH_TOKEN`/use `gh` creds for pushes that
+   touch `.github/workflows/*` (needs `workflow` scope).
+4. **Dry run**: trigger the orchestrator with `publish=false` (artifact-only), confirm all 5
+   archives build + carry `sd-cli`/`sd-server` + the fingerprint. Iterate until green.
+5. **First publish**: `publish=true` (or let the schedule run) -> a real release with the 5
+   assets + manifest/sha256.
+6. **Studio PR** (section B) on the diffusion stack once the release tag exists.
+
+## Follow-ups (not this task)
+
+- Nightly schedule + auto-bump of the Studio `DEFAULT_TAG` (PR bot), like llama.cpp.
+- Add GPU native builds later ONLY if a real need appears (today: GPU = diffusers/torch).

studio/backend/core/training/diffusion_lora_trainer.py

@@ -20,6 +20,13 @@ Design:
 - ``run_diffusion_training_process`` is the thin mp.Queue adapter; it dispatches to the
   trainer registered for the resolved family (SDXL here, DiT families in a follow-up).
   ``main`` is a CLI.
+
+Memory/perf: captions are encoded once up front and the CLIP text encoders freed; VAE
+latents are likewise precomputed into a small CPU cache (``cache_latents``) and the VAE
+freed. The cache stores the posterior's affine pair (mean/std, scale folded in), so every
+step still draws a fresh VAE sample -- distribution-identical to encoding in the loop,
+without keeping the VAE resident or paying a per-step encode. TF32 matmuls + cudnn
+autotuning are enabled for the run under ``cfg.enable_tf32``.
 """
 
 from __future__ import annotations
@@ -40,12 +47,18 @@ from core.training.diffusion_train_common import (  # noqa: F401
     EventCb,
     StopCb,
     DiffusionLoraConfig,
+    LATENT_CACHE_OVER_BUDGET,
+    _apply_perf_flags,
     _assert_trusted_base_model,
     _coerce_gradient_checkpointing,
     _config_from_dict,
     _CONFIG_ALIASES,
     _emit,
+    _latent_cache_forced,
+    _latent_cache_over_budget,
+    _plan_cache_variants,
     _publish_to_lora_catalog,
+    _restore_perf_flags,
     discover_image_caption_pairs,
     get_trainer,
 )
@@ -99,6 +112,43 @@ def _load_image_tensor(
     return tensor, time_ids
 
 
+def _load_image_tensor_planned(
+    path: str, resolution: int, center_crop: bool, u_left: float, u_top: float, flip: bool
+) -> tuple[Any, tuple[int, int, int, int, int, int]]:
+    """Deterministic variant of ``_load_image_tensor`` for the latent cache: the crop comes
+    as unit fractions (mapped uniformly over the same inclusive integer range ``randint``
+    draws from) and the flip as a bool. Geometry (EXIF transpose, LANCZOS short-side resize,
+    the SDXL ``add_time_ids`` from the original size + actual crop offset) matches
+    ``_load_image_tensor`` exactly; ``center_crop`` reproduces the legacy floor-div center
+    bit-for-bit. The flip does not change time_ids (only the mirrored crop_left does)."""
+    import numpy as np
+    import torch
+    from PIL import Image, ImageOps
+
+    img = ImageOps.exif_transpose(Image.open(path)).convert("RGB")
+    original_w, original_h = img.size
+    scale = resolution / min(original_w, original_h)
+    resized_w = max(resolution, round(original_w * scale))
+    resized_h = max(resolution, round(original_h * scale))
+    img = img.resize((resized_w, resized_h), Image.LANCZOS)
+    if center_crop:
+        left, top = (resized_w - resolution) // 2, (resized_h - resolution) // 2
+    else:
+        left = min(int(u_left * (resized_w - resolution + 1)), max(0, resized_w - resolution))
+        top = min(int(u_top * (resized_h - resolution + 1)), max(0, resized_h - resolution))
+    img = img.crop((left, top, left + resolution, top + resolution))
+    crop_left = left
+    if flip:
+        img = img.transpose(Image.FLIP_LEFT_RIGHT)
+        # Mirror the crop's left origin so the conditioning matches the flipped pixels, the
+        # same mirroring ``_load_image_tensor`` applies on a random flip.
+        crop_left = max(0, resized_w - resolution - left)
+    arr = np.asarray(img, dtype = np.float32) / 255.0
+    tensor = torch.from_numpy(arr).permute(2, 0, 1) * 2.0 - 1.0
+    time_ids = (original_h, original_w, top, crop_left, resolution, resolution)
+    return tensor, time_ids
+
+
 def _encode_sdxl_prompts(
     prompts: list[str], tokenizers: list, text_encoders: list, device: Any
 ) -> tuple:
@@ -126,6 +176,102 @@ def _encode_sdxl_prompts(
     return prompt_embeds, pooled
 
 
+def _build_sdxl_latent_cache(
+    vae, vae_scale, image_paths, cfg, device, weight_dtype, on_event, check_stop
+):
+    """Precompute the per-image latent posterior cache: for each planned crop/flip variant,
+    encode once and store ``(A, B, time_ids)`` on CPU in fp32. ``A`` and ``B`` are the affine
+    posterior parameters (mean/std with the VAE scale folded in) so a per-step sample is
+    ``A + B * randn`` -- distribution-identical to an in-loop ``latent_dist.sample()`` -- and
+    ``time_ids`` is the SDXL micro-conditioning for the crop. The stats stay fp32 so the
+    per-step sample happens in fp32 and only the RESULT is cast to weight_dtype, matching the
+    in-loop path (encode fp32 -> sample fp32 -> scale -> .to(weight_dtype)); fp32 doubles the
+    cache RAM over bf16 but the cache is tiny (a handful of latents per image). Returns None if
+    the build was interrupted by a stop request. ``vae_scale`` is read before the VAE is freed."""
+    import torch
+
+    plan = _plan_cache_variants(
+        len(image_paths), cfg.cache_variants, cfg.center_crop, cfg.random_flip, cfg.seed
+    )
+
+    def _hold(t):
+        t = t.to(torch.float32).cpu()
+        if device == "cuda":
+            try:
+                t = t.pin_memory()
+            except RuntimeError:
+                pass
+        return t
+
+    cache: list[list[tuple]] = []
+    total = len(image_paths)
+    total_variants = sum(len(v) for v in plan)
+    forced = _latent_cache_forced()
+    gated = False
+    for i, path in enumerate(image_paths):
+        variants = []
+        for u_left, u_top, flip in plan[i]:
+            tensor, time_ids = _load_image_tensor_planned(
+                path, cfg.resolution, cfg.center_crop, u_left, u_top, flip
+            )
+            pixel_values = tensor.unsqueeze(0).to(device, dtype = torch.float32)
+            with torch.no_grad():
+                dist = vae.encode(pixel_values).latent_dist
+            a = _hold(dist.mean * vae_scale)
+            b = _hold(dist.std * vae_scale)
+            if not forced and not gated:
+                # Size-gate the automatic cache off the first REAL encoded variant, before
+                # building the rest: thousands of images x variants of two fp32 tensors can
+                # exhaust host/pinned RAM with no fallback. Over budget we bail with the VAE
+                # still resident so the loop encodes latents per step instead.
+                per_variant = a.numel() * a.element_size() + b.numel() * b.element_size()
+                if _latent_cache_over_budget(per_variant, total_variants):
+                    _emit(
+                        on_event,
+                        "warning",
+                        message = (
+                            "Latent cache disabled: estimated "
+                            f"{per_variant * total_variants / 1024 ** 3:.1f} GiB over the "
+                            "budget; encoding latents per step instead. Set "
+                            "UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE=1 to keep it."
+                        ),
+                    )
+                    return LATENT_CACHE_OVER_BUDGET
+                gated = True
+            variants.append((a, b, tuple(time_ids)))
+        cache.append(variants)
+        if (i + 1) % 4 == 0 or i + 1 == total:
+            _emit(on_event, "preparing", stage = "cache_latents", done = i + 1, total = total)
+        if check_stop():
+            return None
+    return cache
+
+
+def _sample_sdxl_cached_latents(cache, idxs, variant_rng, device, weight_dtype):
+    """Draw one latent + its time_ids per index from the cache: pick a variant, then sample
+    the posterior (A + B * randn) with fresh noise per step, exactly like an in-loop
+    ``latent_dist.sample() * vae_scale``. The cached stats are fp32, so the sample is drawn in
+    fp32 and only the RESULT is cast to weight_dtype (matching the in-loop path). Returns
+    ``(latents, batch_time_ids)`` already on ``device`` in the training dtype (scale is folded
+    into the cache)."""
+    import torch
+
+    parts_a, parts_b, tid_rows = [], [], []
+    for i in idxs:
+        variants = cache[i]
+        a, b, time_ids = (
+            variants[variant_rng.randrange(len(variants))] if len(variants) > 1 else variants[0]
+        )
+        parts_a.append(a)
+        parts_b.append(b)
+        tid_rows.append(time_ids)
+    lat_a = torch.cat(parts_a).to(device, non_blocking = True)
+    lat_b = torch.cat(parts_b).to(device, non_blocking = True)
+    latents = (lat_a + lat_b * torch.randn_like(lat_a)).to(dtype = weight_dtype)
+    batch_time_ids = torch.tensor(tid_rows, device = device, dtype = weight_dtype)
+    return latents, batch_time_ids
+
+
 def run_diffusion_lora_training(
     config: DiffusionLoraConfig,
     *,
@@ -174,250 +320,320 @@ def run_diffusion_lora_training(
         precision = "fp16"
     weight_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "no": torch.float32}[precision]
 
-    # Preflight the base model against the same trust gate as inference, before any fetch.
-    _assert_trusted_base_model(cfg.base_model)
+    # TF32 / cudnn.benchmark for the run, restored on the way out (the trainer subprocess is
+    # disposable, but restoring keeps in-process callers -- tests, notebooks -- clean). Wraps
+    # the whole body so every return (early stop and normal) restores the backend flags.
+    snap = _apply_perf_flags(cfg, device)
+    try:
+        # Preflight the base model against the same trust gate as inference, before any fetch.
+        _assert_trusted_base_model(cfg.base_model)
 
-    pairs = discover_image_caption_pairs(
-        cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column
-    )
-    _emit(on_event, "model_load_started", num_images = len(pairs))
+        pairs = discover_image_caption_pairs(
+            cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column
+        )
+        _emit(on_event, "model_load_started", num_images = len(pairs))
 
-    # Honour a stop requested before the (potentially large / slow) base model loads, the
-    # same way the LLM training worker checks its stop thread around model load.
-    if _check_stop():
+        # Honour a stop requested before the (potentially large / slow) base model loads, the
+        # same way the LLM training worker checks its stop thread around model load.
+        if _check_stop():
+            out_dir = Path(cfg.output_dir).expanduser()
+            _emit(
+                on_event,
+                "complete",
+                output_dir = str(out_dir),
+                lora_path = None,
+                stopped = True,
+                steps_run = 0,
+            )
+            return str(out_dir)
+
+        pipe = StableDiffusionXLPipeline.from_pretrained(
+            cfg.base_model, torch_dtype = weight_dtype, token = cfg.hf_token, add_watermarker = False
+        )
+        unet, vae = pipe.unet, pipe.vae
+        tokenizers = [pipe.tokenizer, pipe.tokenizer_2]
+        text_encoders = [pipe.text_encoder, pipe.text_encoder_2]
+        noise_scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
+
+        # Freeze the base; only the LoRA trains. The SDXL VAE overflows fp16, so keep it fp32.
+        for m in (unet, vae, *text_encoders):
+            m.requires_grad_(False)
+        vae.to(device, dtype = torch.float32)
+        for m in (unet, *text_encoders):
+            m.to(device, dtype = weight_dtype)
+
+        unet.add_adapter(
+            LoraConfig(
+                r = cfg.lora_rank,
+                lora_alpha = cfg.lora_alpha,
+                lora_dropout = cfg.lora_dropout,
+                init_lora_weights = "gaussian",
+                target_modules = list(cfg.lora_target_modules),
+            )
+        )
+        if cfg.gradient_checkpointing:
+            unet.enable_gradient_checkpointing()
+        # LoRA params must be fp32 for a stable optimizer under mixed precision.
+        if weight_dtype != torch.float32:
+            cast_training_params(unet, dtype = torch.float32)
+
+        lora_params = [p for p in unet.parameters() if p.requires_grad]
+        optimizer = _make_lora_optimizer(lora_params, cfg.learning_rate)
+        # The scheduler advances once per optimizer update: lr_sched.step() runs a single
+        # time per outer opt_step (after the accumulation inner loop), for cfg.train_steps
+        # total. Count warmup/decay in those optimizer steps -- multiplying by the
+        # accumulation factor would stretch warmup past the run and never reach the decay.
+        lr_sched = get_scheduler(
+            cfg.lr_scheduler,
+            optimizer = optimizer,
+            num_warmup_steps = cfg.lr_warmup_steps,
+            num_training_steps = cfg.train_steps,
+        )
+
+        vae_scale = vae.config.scaling_factor
+        prediction_type = noise_scheduler.config.prediction_type
+
+        # Precompute text embeddings once per unique caption, then free the CLIP text encoders.
+        # SDXL re-encoded captions every step (pure waste: captions are constant) and kept both
+        # text encoders (~1.5 GB) resident. Embeddings are deterministic and this consumes no
+        # torch RNG, so the training math is bit-identical to in-loop encoding -- only faster and
+        # lighter. The env toggle exists purely so the accuracy guard can A/B the two paths.
+        precompute = os.environ.get("UNSLOTH_DIFFUSION_NO_PRECOMPUTE", "") not in ("1", "true")
+        caption_embeds: dict[str, tuple] = {}
+        if precompute:
+            for cap in sorted({c for _, c in pairs}):
+                pe, pooled_c = _encode_sdxl_prompts([cap], tokenizers, text_encoders, device)
+                caption_embeds[cap] = (pe.cpu(), pooled_c.cpu())
+            for te in text_encoders:
+                te.to("cpu")
+            text_encoders = []
+            gc.collect()
+            if device == "cuda":
+                torch.cuda.empty_cache()
+
+        # Precompute the VAE latent cache, then free the VAE: the cache holds the posterior
+        # affine pair (mean/std, scale folded in) so per-step sampling noise is preserved. The
+        # env toggle lets the accuracy guard A/B the cached vs in-loop encode paths.
+        use_cache = cfg.cache_latents and os.environ.get(
+            "UNSLOTH_DIFFUSION_NO_LATENT_CACHE", ""
+        ) not in ("1", "true")
+        latent_cache = None
+        if use_cache:
+            latent_cache = _build_sdxl_latent_cache(
+                vae,
+                vae_scale,
+                [p for p, _ in pairs],
+                cfg,
+                device,
+                weight_dtype,
+                on_event,
+                _check_stop,
+            )
+            if latent_cache is LATENT_CACHE_OVER_BUDGET:
+                # The estimated cache exceeded the host-memory budget; keep the VAE resident
+                # and fall through to the in-loop encode path (latent_cache stays None).
+                latent_cache = None
+            elif latent_cache is None:  # stopped during the cache build; nothing trained yet
+                out_dir = Path(cfg.output_dir).expanduser()
+                _emit(
+                    on_event,
+                    "complete",
+                    output_dir = str(out_dir),
+                    lora_path = None,
+                    stopped = True,
+                    steps_run = 0,
+                )
+                return str(out_dir)
+            else:
+                try:
+                    pipe.vae = None
+                except Exception:  # noqa: BLE001 -- a pipeline without a settable vae keeps it
+                    pass
+                del vae
+                vae = None
+                gc.collect()
+                if device == "cuda":
+                    torch.cuda.empty_cache()
+        # Variant picks use their own stream so the loop's index/noise draws stay on the same
+        # seed-deterministic sequence whether or not the cache is enabled.
+        variant_rng = random.Random(cfg.seed + 1)
+
+        _emit(on_event, "model_load_completed")
+
+        def _next_batch() -> tuple[list[int], list[str], list[str]]:
+            idx = rng.sample(range(len(pairs)), k = min(cfg.train_batch_size, len(pairs)))
+            chosen = [pairs[i] for i in idx]
+            return idx, [c[0] for c in chosen], [c[1] for c in chosen]
+
+        unet.train()
+        stopped = False
+        micro = 0
+        running_loss = 0.0
+        peak_gb = 0.0
+        t_start = time.time()
+        done = 0
+        for opt_step in range(cfg.train_steps):
+            optimizer.zero_grad(set_to_none = True)
+            step_loss = 0.0
+            for _ in range(cfg.gradient_accumulation_steps):
+                idx, img_paths, captions = _next_batch()
+                if latent_cache is not None:
+                    # Scale is folded into the cache; the sampler draws in fp32 and casts the
+                    # result to weight_dtype (matching the in-loop path below).
+                    latents, batch_time_ids = _sample_sdxl_cached_latents(
+                        latent_cache, idx, variant_rng, device, weight_dtype
+                    )
+                else:
+                    loaded = [
+                        _load_image_tensor(p, cfg.resolution, cfg.center_crop, cfg.random_flip, rng)
+                        for p in img_paths
+                    ]
+                    pixel_values = torch.stack([t for t, _ in loaded]).to(
+                        device, dtype = torch.float32
+                    )
+                    # Per-sample SDXL micro-conditioning from the actual crop (original size + offset).
+                    batch_time_ids = torch.tensor(
+                        [tid for _, tid in loaded], device = device, dtype = weight_dtype
+                    )
+
+                    with torch.no_grad():
+                        latents = vae.encode(pixel_values).latent_dist.sample() * vae_scale
+                    latents = latents.to(dtype = weight_dtype)
+
+                noise = torch.randn_like(latents)
+                bsz = latents.shape[0]
+                timesteps = torch.randint(
+                    0, noise_scheduler.config.num_train_timesteps, (bsz,), device = device
+                ).long()
+                noisy = noise_scheduler.add_noise(latents, noise, timesteps)
+
+                if precompute:
+                    prompt_embeds = torch.cat([caption_embeds[c][0] for c in captions]).to(device)
+                    pooled = torch.cat([caption_embeds[c][1] for c in captions]).to(device)
+                else:
+                    prompt_embeds, pooled = _encode_sdxl_prompts(
+                        captions, tokenizers, text_encoders, device
+                    )
+                prompt_embeds = prompt_embeds.to(dtype = weight_dtype)
+                pooled = pooled.to(dtype = weight_dtype)
+                added = {"text_embeds": pooled, "time_ids": batch_time_ids}
+
+                model_pred = unet(
+                    noisy, timesteps, prompt_embeds, added_cond_kwargs = added, return_dict = False
+                )[0]
+
+                if prediction_type == "v_prediction":
+                    target = noise_scheduler.get_velocity(latents, noise, timesteps)
+                else:
+                    target = noise
+
+                if cfg.snr_gamma is not None:
+                    snr = compute_snr(noise_scheduler, timesteps)
+                    w = torch.stack([snr, cfg.snr_gamma * torch.ones_like(timesteps)], dim = 1).min(
+                        dim = 1
+                    )[0]
+                    w = w / snr if prediction_type != "v_prediction" else w / (snr + 1)
+                    loss = F.mse_loss(model_pred.float(), target.float(), reduction = "none")
+                    loss = loss.mean(dim = list(range(1, loss.ndim))) * w
+                    loss = loss.mean()
+                else:
+                    loss = F.mse_loss(model_pred.float(), target.float(), reduction = "mean")
+
+                (loss / cfg.gradient_accumulation_steps).backward()
+                step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps
+                micro += 1
+
+            # max_grad_norm <= 0 means "disable clipping" (the Studio payload sends 0.0 for that);
+            # passing 0.0 to clip_grad_norm_ would scale every gradient to zero (no learning).
+            grad_norm: Optional[float] = None
+            if cfg.max_grad_norm and cfg.max_grad_norm > 0:
+                # clip_grad_norm_ returns the PRE-clip total norm (the grad-norm chart signal).
+                grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm))
+            optimizer.step()
+            lr_sched.step()
+
+            running_loss += step_loss
+            done = opt_step + 1
+            if done % cfg.log_every == 0 or done == cfg.train_steps:
+                # ``learning_rate`` (not ``lr``) is the field the Studio training pump reads, so
+                # these progress events are directly consumable by the existing training
+                # status/SSE machinery when the diffusion trainer is wired into the worker.
+                if device == "cuda":
+                    peak_gb = round(torch.cuda.max_memory_allocated() / 1e9, 2)
+                samples_per_second = round(
+                    (done * cfg.train_batch_size * cfg.gradient_accumulation_steps)
+                    / max(time.time() - t_start, 1e-6),
+                    3,
+                )
+                _emit(
+                    on_event,
+                    "progress",
+                    step = done,
+                    total_steps = cfg.train_steps,
+                    loss = round(step_loss, 5),
+                    avg_loss = round(running_loss / done, 5),
+                    learning_rate = lr_sched.get_last_lr()[0],
+                    grad_norm = round(grad_norm, 5) if grad_norm is not None else None,
+                    samples_per_second = samples_per_second,
+                    peak_memory_gb = peak_gb or None,
+                )
+
+            if _check_stop():
+                stopped = True
+                break
+
+        # Export the trained LoRA in diffusers format (loadable via load_lora_weights), unless
+        # the run was cancelled with save disabled -- then leave no partial adapter behind.
         out_dir = Path(cfg.output_dir).expanduser()
+        lora_path: Optional[str] = None
+        catalog_path: Optional[str] = None
+        if not (stopped and not save_on_stop):
+            out_dir.mkdir(parents = True, exist_ok = True)
+            unet_lora = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))
+            StableDiffusionXLPipeline.save_lora_weights(
+                save_directory = str(out_dir),
+                unet_lora_layers = unet_lora,
+                safe_serialization = True,
+                weight_name = DEFAULT_LORA_FILENAME,
+            )
+            lora_path = str(out_dir / DEFAULT_LORA_FILENAME)
+            # Mirror into the Studio diffusion LoRA directory so the Images picker discovers it
+            # (its scan lists only files directly under loras/diffusion, not subdirectories).
+            catalog_path = _publish_to_lora_catalog(lora_path, cfg)
         _emit(
             on_event,
             "complete",
             output_dir = str(out_dir),
-            lora_path = None,
-            stopped = True,
-            steps_run = 0,
+            lora_path = lora_path,
+            catalog_path = catalog_path,
+            family = cfg.resolved_family,
+            base_model = cfg.base_model,
+            stopped = stopped,
+            steps_run = done if cfg.train_steps else 0,
         )
         return str(out_dir)
-
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        cfg.base_model, torch_dtype = weight_dtype, token = cfg.hf_token, add_watermarker = False
-    )
-    unet, vae = pipe.unet, pipe.vae
-    tokenizers = [pipe.tokenizer, pipe.tokenizer_2]
-    text_encoders = [pipe.text_encoder, pipe.text_encoder_2]
-    noise_scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
-
-    # Freeze the base; only the LoRA trains. The SDXL VAE overflows fp16, so keep it fp32.
-    for m in (unet, vae, *text_encoders):
-        m.requires_grad_(False)
-    vae.to(device, dtype = torch.float32)
-    for m in (unet, *text_encoders):
-        m.to(device, dtype = weight_dtype)
-
-    unet.add_adapter(
-        LoraConfig(
-            r = cfg.lora_rank,
-            lora_alpha = cfg.lora_alpha,
-            lora_dropout = cfg.lora_dropout,
-            init_lora_weights = "gaussian",
-            target_modules = list(cfg.lora_target_modules),
-        )
-    )
-    if cfg.gradient_checkpointing:
-        unet.enable_gradient_checkpointing()
-    # LoRA params must be fp32 for a stable optimizer under mixed precision.
-    if weight_dtype != torch.float32:
-        cast_training_params(unet, dtype = torch.float32)
-
-    lora_params = [p for p in unet.parameters() if p.requires_grad]
-    optimizer = _make_lora_optimizer(lora_params, cfg.learning_rate)
-    # The scheduler advances once per optimizer update: lr_sched.step() runs a single
-    # time per outer opt_step (after the accumulation inner loop), for cfg.train_steps
-    # total. Count warmup/decay in those optimizer steps -- multiplying by the
-    # accumulation factor would stretch warmup past the run and never reach the decay.
-    lr_sched = get_scheduler(
-        cfg.lr_scheduler,
-        optimizer = optimizer,
-        num_warmup_steps = cfg.lr_warmup_steps,
-        num_training_steps = cfg.train_steps,
-    )
-
-    vae_scale = vae.config.scaling_factor
-    prediction_type = noise_scheduler.config.prediction_type
-
-    # Precompute text embeddings once per unique caption, then free the CLIP text encoders.
-    # SDXL re-encoded captions every step (pure waste: captions are constant) and kept both
-    # text encoders (~1.5 GB) resident. Embeddings are deterministic and this consumes no
-    # torch RNG, so the training math is bit-identical to in-loop encoding -- only faster and
-    # lighter. The env toggle exists purely so the accuracy guard can A/B the two paths.
-    precompute = os.environ.get("UNSLOTH_DIFFUSION_NO_PRECOMPUTE", "") not in ("1", "true")
-    caption_embeds: dict[str, tuple] = {}
-    if precompute:
-        for cap in sorted({c for _, c in pairs}):
-            pe, pooled_c = _encode_sdxl_prompts([cap], tokenizers, text_encoders, device)
-            caption_embeds[cap] = (pe.cpu(), pooled_c.cpu())
-        for te in text_encoders:
-            te.to("cpu")
-        text_encoders = []
-        gc.collect()
-        if device == "cuda":
-            torch.cuda.empty_cache()
-
-    _emit(on_event, "model_load_completed")
-
-    def _next_batch() -> tuple[list[str], list[str]]:
-        idx = rng.sample(range(len(pairs)), k = min(cfg.train_batch_size, len(pairs)))
-        chosen = [pairs[i] for i in idx]
-        return [c[0] for c in chosen], [c[1] for c in chosen]
-
-    unet.train()
-    stopped = False
-    micro = 0
-    running_loss = 0.0
-    peak_gb = 0.0
-    t_start = time.time()
-    done = 0
-    for opt_step in range(cfg.train_steps):
-        optimizer.zero_grad(set_to_none = True)
-        step_loss = 0.0
-        for _ in range(cfg.gradient_accumulation_steps):
-            img_paths, captions = _next_batch()
-            loaded = [
-                _load_image_tensor(p, cfg.resolution, cfg.center_crop, cfg.random_flip, rng)
-                for p in img_paths
-            ]
-            pixel_values = torch.stack([t for t, _ in loaded]).to(device, dtype = torch.float32)
-            # Per-sample SDXL micro-conditioning from the actual crop (original size + offset).
-            batch_time_ids = torch.tensor(
-                [tid for _, tid in loaded], device = device, dtype = weight_dtype
-            )
-
-            with torch.no_grad():
-                latents = vae.encode(pixel_values).latent_dist.sample() * vae_scale
-            latents = latents.to(dtype = weight_dtype)
-
-            noise = torch.randn_like(latents)
-            bsz = latents.shape[0]
-            timesteps = torch.randint(
-                0, noise_scheduler.config.num_train_timesteps, (bsz,), device = device
-            ).long()
-            noisy = noise_scheduler.add_noise(latents, noise, timesteps)
-
-            if precompute:
-                prompt_embeds = torch.cat([caption_embeds[c][0] for c in captions]).to(device)
-                pooled = torch.cat([caption_embeds[c][1] for c in captions]).to(device)
-            else:
-                prompt_embeds, pooled = _encode_sdxl_prompts(
-                    captions, tokenizers, text_encoders, device
-                )
-            prompt_embeds = prompt_embeds.to(dtype = weight_dtype)
-            pooled = pooled.to(dtype = weight_dtype)
-            added = {"text_embeds": pooled, "time_ids": batch_time_ids}
-
-            model_pred = unet(
-                noisy, timesteps, prompt_embeds, added_cond_kwargs = added, return_dict = False
-            )[0]
-
-            if prediction_type == "v_prediction":
-                target = noise_scheduler.get_velocity(latents, noise, timesteps)
-            else:
-                target = noise
-
-            if cfg.snr_gamma is not None:
-                snr = compute_snr(noise_scheduler, timesteps)
-                w = torch.stack([snr, cfg.snr_gamma * torch.ones_like(timesteps)], dim = 1).min(
-                    dim = 1
-                )[0]
-                w = w / snr if prediction_type != "v_prediction" else w / (snr + 1)
-                loss = F.mse_loss(model_pred.float(), target.float(), reduction = "none")
-                loss = loss.mean(dim = list(range(1, loss.ndim))) * w
-                loss = loss.mean()
-            else:
-                loss = F.mse_loss(model_pred.float(), target.float(), reduction = "mean")
-
-            (loss / cfg.gradient_accumulation_steps).backward()
-            step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps
-            micro += 1
-
-        # max_grad_norm <= 0 means "disable clipping" (the Studio payload sends 0.0 for that);
-        # passing 0.0 to clip_grad_norm_ would scale every gradient to zero (no learning).
-        grad_norm: Optional[float] = None
-        if cfg.max_grad_norm and cfg.max_grad_norm > 0:
-            # clip_grad_norm_ returns the PRE-clip total norm (the grad-norm chart signal).
-            grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm))
-        optimizer.step()
-        lr_sched.step()
-
-        running_loss += step_loss
-        done = opt_step + 1
-        if done % cfg.log_every == 0 or done == cfg.train_steps:
-            # ``learning_rate`` (not ``lr``) is the field the Studio training pump reads, so
-            # these progress events are directly consumable by the existing training
-            # status/SSE machinery when the diffusion trainer is wired into the worker.
-            if device == "cuda":
-                peak_gb = round(torch.cuda.max_memory_allocated() / 1e9, 2)
-            samples_per_second = round(
-                (done * cfg.train_batch_size * cfg.gradient_accumulation_steps)
-                / max(time.time() - t_start, 1e-6),
-                3,
-            )
-            _emit(
-                on_event,
-                "progress",
-                step = done,
-                total_steps = cfg.train_steps,
-                loss = round(step_loss, 5),
-                avg_loss = round(running_loss / done, 5),
-                learning_rate = lr_sched.get_last_lr()[0],
-                grad_norm = round(grad_norm, 5) if grad_norm is not None else None,
-                samples_per_second = samples_per_second,
-                peak_memory_gb = peak_gb or None,
-            )
-
-        if _check_stop():
-            stopped = True
-            break
-
-    # Export the trained LoRA in diffusers format (loadable via load_lora_weights), unless
-    # the run was cancelled with save disabled -- then leave no partial adapter behind.
-    out_dir = Path(cfg.output_dir).expanduser()
-    lora_path: Optional[str] = None
-    catalog_path: Optional[str] = None
-    if not (stopped and not save_on_stop):
-        out_dir.mkdir(parents = True, exist_ok = True)
-        unet_lora = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))
-        StableDiffusionXLPipeline.save_lora_weights(
-            save_directory = str(out_dir),
-            unet_lora_layers = unet_lora,
-            safe_serialization = True,
-            weight_name = DEFAULT_LORA_FILENAME,
-        )
-        lora_path = str(out_dir / DEFAULT_LORA_FILENAME)
-        # Mirror into the Studio diffusion LoRA directory so the Images picker discovers it
-        # (its scan lists only files directly under loras/diffusion, not subdirectories).
-        catalog_path = _publish_to_lora_catalog(lora_path, cfg)
-    _emit(
-        on_event,
-        "complete",
-        output_dir = str(out_dir),
-        lora_path = lora_path,
-        catalog_path = catalog_path,
-        family = cfg.resolved_family,
-        base_model = cfg.base_model,
-        stopped = stopped,
-        steps_run = done if cfg.train_steps else 0,
-    )
-    return str(out_dir)
+    finally:
+        _restore_perf_flags(snap)
 
 
 def _make_lora_optimizer(params: list, lr: float) -> Any:
     """8-bit AdamW (bitsandbytes) by default -- half the optimizer state, no meaningful
-    quality cost for LoRA -- falling back to fp32 AdamW when unavailable or when
-    UNSLOTH_DIFFUSION_FP32_OPTIM is set (used by the accuracy guard)."""
+    quality cost for LoRA -- falling back to torch AdamW (fused on CUDA) when unavailable.
+    UNSLOTH_DIFFUSION_FP32_OPTIM forces plain (non-fused) AdamW: the accuracy guard wants the
+    reference optimizer, so it must not take the fused path."""
     import torch
 
-    if os.environ.get("UNSLOTH_DIFFUSION_FP32_OPTIM", "") not in ("1", "true"):
+    if os.environ.get("UNSLOTH_DIFFUSION_FP32_OPTIM", "") in ("1", "true"):
+        return torch.optim.AdamW(params, lr = lr)
+    try:
+        import bitsandbytes as bnb
+        return bnb.optim.AdamW8bit(params, lr = lr)
+    except Exception:  # noqa: BLE001 -- bnb missing / no CUDA: fall back to torch AdamW
+        pass
+    if torch.cuda.is_available():
         try:
-            import bitsandbytes as bnb
-            return bnb.optim.AdamW8bit(params, lr = lr)
-        except Exception:  # noqa: BLE001 -- bnb missing / no CUDA: fall back to torch AdamW
+            return torch.optim.AdamW(params, lr = lr, fused = True)
+        except Exception:  # noqa: BLE001 -- fused unsupported on this build/device
             pass
     return torch.optim.AdamW(params, lr = lr)
 

studio/backend/core/training/diffusion_train_common.py

@@ -17,6 +17,7 @@ from __future__ import annotations
 
 import json
 import os
+import random
 import re
 import time
 from dataclasses import dataclass, field, replace
@@ -123,6 +124,82 @@ def resolve_trainable_family(base_model: str, model_family: Optional[str] = None
     return "sdxl"
 
 
+def repo_is_prequantized(base_model: str) -> bool:
+    """Heuristic: a repo whose name marks a bitsandbytes 4-bit build already ships a
+    quantized transformer, so it loads as-is for nf4 and cannot serve the dense
+    (bf16/int8/fp8) base precisions."""
+    name = str(base_model or "").lower()
+    return "bnb-4bit" in name or "-4bit" in name or "int4" in name or "nf4" in name
+
+
+def _module_is_torchao_stub(module: Any) -> bool:
+    """True iff ``module`` is the Unsloth Windows-ROCm torchao import stub rather than the
+    real package. The stub (core/_torchao_stub.py) satisfies find_spec and even lets
+    ``from torchao.quantization import quantize_`` succeed -- but the imported symbols are
+    no-op stub types, so the quantization never happens. Every stub module carries the
+    ``_unsloth_stub`` sentinel, so match on it (comparing against the stub module's own
+    sentinel object, not identity of a re-created one)."""
+    if module is None:
+        return False
+    sentinel = getattr(module, "_unsloth_stub", None)
+    if sentinel is None:
+        return False
+    try:
+        from core._torchao_stub import _STUB_SENTINEL
+    except Exception:  # noqa: BLE001 -- stub module absent -> nothing to compare against
+        return False
+    return sentinel is _STUB_SENTINEL
+
+
+def has_functional_torchao() -> bool:
+    """True iff the real torchao quantization API is importable (not the Windows-ROCm stub).
+
+    ``_int8_quantize_base`` needs ``Int8WeightOnlyConfig`` + ``quantize_`` from
+    ``torchao.quantization`` and has no runtime fallback, so gate both the auto int8 pick
+    and the advertised int8 mode on a FUNCTIONAL import: a plain ``find_spec("torchao")``
+    is satisfied by the stub, whose quantize_ is a no-op that leaves the transformer dense
+    while compile is disabled as if it were int8. Import the exact symbols the int8 path
+    uses and reject the stub module. Never raises."""
+    try:
+        import importlib
+
+        quant = importlib.import_module("torchao.quantization")
+        if _module_is_torchao_stub(quant):
+            return False
+        # The symbols the int8 path actually imports must exist on the real module.
+        return hasattr(quant, "Int8WeightOnlyConfig") and hasattr(quant, "quantize_")
+    except Exception:  # noqa: BLE001 -- torchao absent / broken build -> treat as unavailable
+        return False
+
+
+def train_precision_modes() -> tuple[list[str], str]:
+    """(supported base_precision modes, recommended pick) for the current machine: nf4
+    always works; bf16/auto need a bf16-capable CUDA GPU (Ampere+); int8/fp8 additionally
+    need a FUNCTIONAL torchao (their explicit paths import torchao with no fallback, and the
+    Windows-ROCm stub only looks installed). fp8 also needs an fp8-capable GPU (sm89+). The
+    dense modes all train in bf16 compute, which the DiT trainer requires, so a non-bf16 CUDA
+    GPU (T4/V100/RTX 20xx) is offered only nf4 -- otherwise /info would advertise a start that
+    evicts resident models and then fails the trainer's bf16 guard. Used by the /info endpoint
+    so the UI can gate the precision selector. Never raises."""
+    modes = ["nf4"]
+    recommended = "nf4"
+    try:
+        import torch
+        if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
+            modes.append("bf16")
+            torchao_ok = has_functional_torchao()
+            if torchao_ok:
+                modes.append("int8")
+            major, minor = torch.cuda.get_device_capability()
+            if torchao_ok and (major, minor) >= (8, 9) and hasattr(torch, "float8_e4m3fn"):
+                modes.append("fp8")
+            modes.append("auto")
+            recommended = "auto"
+    except Exception:  # noqa: BLE001 -- no torch / probe failure -> nf4 only
+        pass
+    return modes, recommended
+
+
 def get_trainer(family: str) -> Callable[..., str]:
     """Return the training entrypoint for ``family``. Imports the trainer module lazily so
     this shared module stays free of the heavy trainer imports (and any import cycle)."""
@@ -170,19 +247,96 @@ _FAMILY_VRAM_NOTES = {
     "z-image": "6B model, QLoRA (nf4) by default (~12 GB+). bf16 only.",
 }
 
+# The flow-matching DiT families (run by diffusion_dit_trainer). They expose the
+# base_precision / compile levers and require bf16 compute on CUDA; SDXL is absent because
+# it uses its own mixed_precision path. Kept as a set so the UI gate, the bf16 preflight,
+# and any future dispatch stay in sync.
+_DIT_TRAIN_FAMILIES = frozenset({"flux.1", "qwen-image", "z-image"})
+
+
+def bf16_unsupported_reason(resolved_family: str) -> Optional[str]:
+    """Return a user-facing error string if ``resolved_family`` needs bf16 compute that the
+    live GPU cannot provide, else None. The DiT trainer requires a bf16-capable GPU (Ampere
+    or newer) and otherwise raises deep in model load; the start route uses this to fail fast
+    BEFORE evicting resident GPU workloads. CPU-only hosts (which fall back to fp32 for
+    import/unit tests) and SDXL (its own mixed_precision path) are exempt. Never raises."""
+    if (resolved_family or "").strip().lower() not in _DIT_TRAIN_FAMILIES:
+        return None
+    try:
+        import torch
+        if torch.cuda.is_available() and not torch.cuda.is_bf16_supported():
+            return (
+                "This trainer requires a bfloat16-capable GPU (Ampere or newer); this CUDA "
+                "device does not support bf16. Train the DiT families on a newer GPU."
+            )
+    except Exception:  # noqa: BLE001 -- torch probe failure must not block a start
+        return None
+    return None
+
+
+def training_precision_preflight_error(resolved_family: str, base_precision: str) -> Optional[str]:
+    """Reason the requested DiT precision cannot run on this host, else None -- checked by the
+    start route BEFORE evicting resident GPU workloads (the trainer's own checks fire only in the
+    child, after eviction). Three gates, all mirroring _resolve_base_precision so a doomed run is
+    rejected before teardown: the bf16-GPU requirement (bf16_unsupported_reason); the dense
+    precisions (bf16/int8/fp8) requiring a CUDA GPU; and an explicit int8 needing a FUNCTIONAL
+    torchao (its _int8_quantize_base has no fallback). Never raises."""
+    reason = bf16_unsupported_reason(resolved_family)
+    if reason:
+        return reason
+    fam = (resolved_family or "").strip().lower()
+    mode = (base_precision or "").strip().lower()
+    if fam in _DIT_TRAIN_FAMILIES and mode in ("bf16", "int8", "fp8"):
+        # The DiT trainer's dense precisions all require CUDA (_resolve_base_precision rejects
+        # bf16/int8/fp8 on device != "cuda"). bf16_unsupported_reason exempts a CPU-only host (the
+        # fp32 fallback for import/unit tests), so without this a dense request on a GPU-less host
+        # would pass the preflight, evict resident workloads, then raise only in the child.
+        try:
+            import torch
+            has_cuda = torch.cuda.is_available()
+        except Exception:  # noqa: BLE001 -- no torch / probe failure -> treat as no CUDA
+            has_cuda = False
+        if not has_cuda:
+            return (
+                f"base_precision={mode!r} needs a CUDA GPU; this host has none. "
+                "Use base_precision='nf4' or 'auto'."
+            )
+        if mode == "int8" and not has_functional_torchao():
+            return (
+                "base_precision='int8' needs a functional torchao install; this host's torchao is "
+                "missing or the non-functional Windows-ROCm stub. Use 'nf4', 'bf16', or 'auto'."
+            )
+    return None
+
 
 def family_train_infos() -> list[dict[str, Any]]:
     """Describe every trainable family for the Train UI: name, label, the default + allowed
     base repos, the recommended starting hyperparameters, and a VRAM/access note. Built from
     the family registry so it stays in sync with what the trainers actually support."""
     from core.inference.diffusion_families import detect_family
+    from core.inference.diffusion_transformer_quant import _family_denied
 
+    dit_modes, dit_recommended = train_precision_modes()
     infos: list[dict[str, Any]] = []
     for name in trainable_family_names():
         fam = detect_family("", override = name)
         if fam is None:
             continue
         repos = list(fam.train_base_repos) or [fam.base_repo]
+        # base_precision / compile apply to the DiT trainer only; SDXL keeps its
+        # mixed_precision lever, so the UI hides the selector for it.
+        is_dit = name in _DIT_TRAIN_FAMILIES
+        # On a non-bf16 CUDA GPU the start route's preflight rejects EVERY DiT family (even nf4,
+        # since the DiT trainer requires bf16 unconditionally on CUDA), so advertise no precision
+        # for it -- otherwise /info offers an nf4 DiT option that always 400s. Otherwise drop any
+        # scheme this family's DiT corrupts (fp8 on Qwen-Image: activation outliers exceed fp8's
+        # range; the inference path denies the same set), so the UI never offers a mode
+        # normalized() would then reject.
+        dit_block = bf16_unsupported_reason(name) if is_dit else None
+        if not is_dit or dit_block:
+            fam_modes: list[str] = []
+        else:
+            fam_modes = [m for m in dit_modes if not _family_denied(name, m)]
         infos.append(
             {
                 "name": name,
@@ -190,7 +344,10 @@ def family_train_infos() -> list[dict[str, Any]]:
                 "default_base": repos[0],
                 "base_repos": repos,
                 "defaults": train_defaults(name),
-                "vram_note": _FAMILY_VRAM_NOTES.get(name, ""),
+                "vram_note": dit_block or _FAMILY_VRAM_NOTES.get(name, ""),
+                "precision_modes": fam_modes,
+                "recommended_precision": "nf4" if (not is_dit or dit_block) else dit_recommended,
+                "supports_compile": bool(is_dit and not dit_block),
             }
         )
     return infos
@@ -228,6 +385,23 @@ class DiffusionLoraConfig:
     caption_column: str = "text"  # column in metadata.jsonl
     adapter_name: str = "default"
     hf_token: Optional[str] = None
+    # Precompute the VAE latents once (freeing the VAE for the whole run) instead of
+    # re-encoding every step. ``cache_variants`` crop/flip draws are frozen per image;
+    # the per-step VAE sampling noise itself is preserved (see the DiT trainer docstring).
+    cache_latents: bool = True
+    cache_variants: int = 4
+    # Regional torch.compile of the transformer blocks: "off" | "on" | "auto" (auto turns
+    # it on only for a dense, non-bitsandbytes base where it is a clean win).
+    compile_transformer: str = "auto"
+    # TF32 matmuls + high fp32 matmul precision + cudnn autotuning for the run. Near-lossless;
+    # disable for strict bit-reproducibility A/Bs.
+    enable_tf32: bool = True
+    # DiT base transformer precision: "nf4" (bitsandbytes QLoRA, the memory floor and the
+    # default), "bf16" (dense, fastest eager, compile-friendly), "int8" (torchao
+    # weight-only, half of bf16), "fp8" (torchao float8 training compute on the frozen
+    # linears, Ada/Hopper/Blackwell + compile), or "auto" (pick by free VRAM + GPU class).
+    # Non-nf4 modes need a dense base repo (not a prequant bnb-4bit one). SDXL ignores it.
+    base_precision: str = "nf4"
     # How often to emit a progress event (in optimizer steps).
     log_every: int = 1
     # Optional explicit family override ("sdxl" / "flux.1" / ...); None = detect from
@@ -259,6 +433,43 @@ class DiffusionLoraConfig:
             raise ValueError("resolution must be a multiple of 8 and >= 64")
         if self.mixed_precision not in ("bf16", "fp16", "no"):
             raise ValueError("mixed_precision must be one of bf16 / fp16 / no")
+        if not 1 <= int(self.cache_variants) <= 16:
+            raise ValueError("cache_variants must be between 1 and 16")
+        compile_transformer = str(self.compile_transformer or "auto").strip().lower()
+        if compile_transformer not in ("off", "on", "auto"):
+            raise ValueError("compile_transformer must be one of off / on / auto")
+        base_precision = str(self.base_precision or "nf4").strip().lower()
+        if base_precision not in ("nf4", "bf16", "int8", "fp8", "auto"):
+            raise ValueError("base_precision must be one of nf4 / bf16 / int8 / fp8 / auto")
+        # base_precision is a DiT-only lever (nf4/bf16/int8/fp8/auto for the transformer
+        # load); SDXL uses its own mixed_precision path and ignores base_precision entirely,
+        # so the dense-mode gates (prequant base / non-bf16 compute) apply only to the DiT
+        # families. The mode-name validity check above still runs for every family.
+        if resolved_family != "sdxl" and base_precision in ("bf16", "int8", "fp8"):
+            if repo_is_prequantized(self.base_model):
+                raise ValueError(
+                    f"base_precision={base_precision!r} needs a dense base repo, but "
+                    f"'{self.base_model}' is already bitsandbytes-quantized. Pick the "
+                    f"family's dense (bf16) base repo for this mode, or use nf4/auto."
+                )
+            if self.mixed_precision != "bf16":
+                raise ValueError(
+                    f"base_precision={base_precision!r} trains in bf16 compute; set "
+                    f"mixed_precision to bf16."
+                )
+            # Some DiT families are corrupted by fp8's activation range: outliers exceed even
+            # per-row fp8's dynamic range, so the frozen linears' float8 training compute
+            # learns against a garbage forward pass. The inference path already denies these
+            # schemes; mirror that deny here so the run fails fast instead of silently
+            # producing a broken adapter. int8 (per-token) is unaffected and stays allowed.
+            from core.inference.diffusion_transformer_quant import _family_denied
+
+            if _family_denied(resolved_family, base_precision):
+                raise ValueError(
+                    f"base_precision={base_precision!r} is not supported for "
+                    f"{resolved_family}: its activations exceed fp8's range and corrupt the "
+                    f"trained result. Use 'nf4', 'int8', 'bf16', or 'auto'."
+                )
         # A zero/negative gamma would zero out (or invert) the min-SNR weight and
         # silently train on a degenerate loss; None is the documented disable.
         if self.snr_gamma is not None and float(self.snr_gamma) <= 0:
@@ -283,6 +494,9 @@ class DiffusionLoraConfig:
             lora_target_modules = targets,
             max_grad_norm = float(self.max_grad_norm),
             hf_token = token or None,
+            cache_variants = int(self.cache_variants),
+            compile_transformer = compile_transformer,
+            base_precision = base_precision,
             resolved_family = resolved_family,
         )
 
@@ -367,6 +581,145 @@ def _emit(on_event: Optional[EventCb], type_: str, **kw: Any) -> None:
         on_event({"type": type_, "ts": time.time(), **kw})
 
 
+def _plan_cache_variants(
+    num_images: int, cache_variants: int, center_crop: bool, random_flip: bool, seed: int
+) -> list[list[tuple[float, float, bool]]]:
+    """Seed-deterministic crop/flip plan for the latent cache: per image, up to
+    ``cache_variants`` draws of (u_left, u_top, flip) with the crop as unit fractions the
+    loader maps onto its integer crop range. Uses its own rng stream so the training
+    loop's draws are untouched. Center-crop / no-flip collapse duplicate variants (a
+    center crop without flip is one variant no matter how many draws), so callers encode
+    each distinct variant exactly once. Pure (no torch) for CPU unit tests."""
+    crop_rng = random.Random(seed)
+    plan: list[list[tuple[float, float, bool]]] = []
+    for _ in range(max(0, num_images)):
+        variants: list[tuple[float, float, bool]] = []
+        for _ in range(max(1, cache_variants)):
+            u_left, u_top = crop_rng.random(), crop_rng.random()
+            flip = bool(random_flip and crop_rng.random() < 0.5)
+            if center_crop:
+                u_left = u_top = 0.5  # loader ignores the fractions for a center crop
+            key = (u_left, u_top, flip)
+            if key not in variants:
+                variants.append(key)
+        plan.append(variants)
+    return plan
+
+
+# Host-memory budget for the AUTOMATIC latent cache. The cache holds two fp32 posterior
+# tensors (mean/std, VAE scale folded in) per crop/flip variant per image, pinned on a CUDA
+# host. At 1024px an SDXL variant is ~0.5 MiB and a 16-channel DiT variant several times
+# that, so a few thousand images x cache_variants can exhaust host or pinned RAM with no
+# fallback. Over this budget the default falls back to per-step VAE encoding. A fixed
+# constant (rather than a psutil RAM fraction) keeps the gate dependency-free and identical
+# across hosts; it is deliberately conservative, well under a typical training host's RAM.
+_LATENT_CACHE_BUDGET_BYTES = 4 * 1024**3  # 4 GiB
+
+# Returned by the cache builders when the estimated cache exceeds the budget: the caller
+# keeps the VAE resident and encodes each step's latents in-loop. A distinct sentinel from
+# ``None`` (which means a stop was requested mid-build) so the two are not conflated.
+LATENT_CACHE_OVER_BUDGET: Any = object()
+
+
+def _latent_cache_forced() -> bool:
+    """The user explicitly forced the latent cache on, bypassing the size gate. This is the
+    explicit opt-in counterpart to ``UNSLOTH_DIFFUSION_NO_LATENT_CACHE`` (the explicit
+    opt-out); only the automatic default is size-gated, so an explicit choice is honoured
+    verbatim in either direction."""
+    return os.environ.get("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", "") in ("1", "true")
+
+
+def _latent_cache_over_budget(
+    per_variant_bytes: int,
+    total_variants: int,
+    budget_bytes: Optional[int] = None,
+) -> bool:
+    """True when a cache of ``total_variants`` entries, each two fp32 tensors totalling
+    ``per_variant_bytes``, is estimated to exceed ``budget_bytes``. ``per_variant_bytes`` is
+    measured from a real encoded latent, so the estimate tracks the actual per-family tensor
+    shape (SDXL 4-channel vs. a packed 16-channel DiT latent) rather than a guess. The budget
+    is read from the module constant at call time when not given, so tests can override it."""
+    if budget_bytes is None:
+        budget_bytes = _LATENT_CACHE_BUDGET_BYTES
+    return per_variant_bytes * max(0, total_variants) > budget_bytes
+
+
+def _apply_perf_flags(
+    cfg: "DiffusionLoraConfig",
+    device: str,
+    cudnn_benchmark: bool = False,
+) -> dict:
+    """Set the run-scoped torch backend knobs: TF32 matmuls + high fp32 matmul precision
+    when ``cfg.enable_tf32`` is on, strict fp32 (all TF32 flags cleared) when it is off,
+    plus cudnn autotuning when the caller opts in. Autotune is
+    for the conv-heavy SDXL U-Net only: measured on B200, it DOUBLES peak VRAM (fp32 VAE
+    conv workspaces) while the DiT loop -- pure matmuls once the latent cache is built --
+    gains nothing from it. Returns a snapshot for ``_restore_perf_flags``. Best-effort:
+    missing attributes on a CPU/other-vendor build are skipped."""
+    from core.inference.diffusion_speed import snapshot_backend_flags
+
+    snap: dict[str, Any] = {"flags": snapshot_backend_flags(), "matmul_precision": None}
+    if device != "cuda":
+        return snap
+    try:
+        import torch
+
+        snap["matmul_precision"] = torch.get_float32_matmul_precision()
+        if cfg.enable_tf32:
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.allow_tf32 = True
+            torch.set_float32_matmul_precision("high")
+        else:
+            # The opt-out is a strict-fp32 A/B mode, so actively clear the flags rather
+            # than inherit ambient state (cudnn TF32 defaults to ON in torch).
+            torch.backends.cuda.matmul.allow_tf32 = False
+            torch.backends.cudnn.allow_tf32 = False
+            torch.set_float32_matmul_precision("highest")
+        if cudnn_benchmark:
+            torch.backends.cudnn.benchmark = True
+        # The cuDNN SDPA backend's TRAINING graph is broken for the FLUX attention shapes
+        # on torch 2.10 + cu130 (B200): mha_graph.execute fails, then poisons the context
+        # into illegal memory accesses. Flash / mem-efficient SDPA are mathematically
+        # equivalent, so pin those for the run (restored on exit).
+        cuda_backends = getattr(torch.backends, "cuda", None)
+        if cuda_backends is not None and hasattr(cuda_backends, "enable_cudnn_sdp"):
+            try:
+                snap["cudnn_sdp"] = bool(cuda_backends.cudnn_sdp_enabled())
+            except Exception:  # noqa: BLE001 -- flag unreadable: skip the tweak entirely
+                snap["cudnn_sdp"] = None
+            if snap["cudnn_sdp"]:
+                cuda_backends.enable_cudnn_sdp(False)
+    except Exception:  # noqa: BLE001 -- perf flags are never fatal
+        pass
+    return snap
+
+
+def _restore_perf_flags(snap: Optional[dict]) -> None:
+    """Undo ``_apply_perf_flags`` (the trainer subprocess is disposable, but in-process
+    callers -- tests, notebooks -- must not inherit mutated globals)."""
+    if not snap:
+        return
+    from core.inference.diffusion_speed import restore_backend_flags
+
+    restore_backend_flags(snap.get("flags"))
+    try:
+        import torch
+
+        if snap.get("matmul_precision"):
+            torch.set_float32_matmul_precision(snap["matmul_precision"])
+        # Restore the exact pre-run cudnn SDPA state; None means the flag was unreadable
+        # (or absent) at apply time and was never touched.
+        cuda_backends = getattr(torch.backends, "cuda", None)
+        if (
+            snap.get("cudnn_sdp") is not None
+            and cuda_backends is not None
+            and hasattr(cuda_backends, "enable_cudnn_sdp")
+        ):
+            cuda_backends.enable_cudnn_sdp(bool(snap["cudnn_sdp"]))
+    except Exception:  # noqa: BLE001 -- best-effort restore
+        pass
+
+
 def _assert_trusted_base_model(base_model: str) -> None:
     """Gate the training base model the same way the inference backend gates non-GGUF loads:
     a local path or a trusted repo (``unsloth/*`` or an allowlisted official base). This runs
@@ -446,6 +799,15 @@ def _coerce_gradient_checkpointing(value: Any) -> bool:
     return bool(value)
 
 
+def _coerce_bool(value: Any) -> bool:
+    """Coerce a flag that may arrive as a string through the generic Studio config path
+    (e.g. "false" / "0" / "off"). A non-empty string like "false" is otherwise truthy, so
+    an opt-out would silently no-op. A real bool passes through."""
+    if isinstance(value, str):
+        return value.strip().lower() not in ("", "none", "false", "0", "no", "off")
+    return bool(value)
+
+
 def _config_from_dict(config: dict) -> DiffusionLoraConfig:
     """Build a DiffusionLoraConfig from a plain dict. Unknown keys are ignored so a richer
     request payload (UI form) does not break construction; a small set of generic Studio
@@ -465,4 +827,7 @@ def _config_from_dict(config: dict) -> DiffusionLoraConfig:
         kwargs["gradient_checkpointing"] = _coerce_gradient_checkpointing(
             kwargs["gradient_checkpointing"]
         )
+    for flag in ("cache_latents", "enable_tf32"):
+        if flag in kwargs:
+            kwargs[flag] = _coerce_bool(kwargs[flag])
     return DiffusionLoraConfig(**kwargs)

studio/backend/core/training/diffusion_training_service.py

@@ -17,6 +17,7 @@ runs a scripted target on a thread.
 
 from __future__ import annotations
 
+import math
 import multiprocessing as mp
 import threading
 import time
@@ -31,6 +32,21 @@ _CTX = mp.get_context("spawn")
 _TERMINAL = ("complete", "error")
 
 
+def _finite_or_none(value: Any) -> Optional[float]:
+    """Coerce a numeric progress field to a finite float, or None. A divergent run (or a
+    grad clip that returns inf) can push loss / grad_norm to NaN or +/-Infinity, and those
+    are invalid in strict JSON -- FastAPI's encoder would emit the JS-only NaN/Infinity
+    tokens that break a strict client parse. Nulling them here (the single service ingestion
+    point both trainers feed) keeps every status snapshot and persisted record JSON-safe."""
+    if value is None:
+        return None
+    try:
+        f = float(value)
+    except (TypeError, ValueError):
+        return None
+    return f if math.isfinite(f) else None
+
+
 def _run_diffusion_child(*, event_queue: Any, stop_queue: Any, config: dict) -> None:
     # Imported lazily so this module (and the route layer) stays torch-free at import.
     from .diffusion_lora_trainer import run_diffusion_training_process
@@ -108,21 +124,13 @@ def _append_metric(
         return
     if istep <= 0 or loss is None:
         return
-    try:
-        floss = float(loss)
-    except (TypeError, ValueError):
+    floss = _finite_or_none(loss)
+    if floss is None:  # non-numeric or non-finite (NaN/Inf): skip, keep the curve JSON-safe
         return
-    if floss != floss:  # NaN guard
-        return
-
-    def _opt_float(v: Any) -> Optional[float]:
-        try:
-            return float(v) if v is not None else None
-        except (TypeError, ValueError):
-            return None
-
-    flr = _opt_float(lr)
-    fgn = _opt_float(grad_norm)
+    # lr / grad_norm may be None (sparse series) or non-finite; a non-finite value is
+    # nulled, not dropped, so a bad point never taints the (loss-driven) history.
+    flr = _finite_or_none(lr)
+    fgn = _finite_or_none(grad_norm)
     steps = state["metric_steps"]
     losses = state["metric_loss"]
     lrs = state["metric_lr"]
@@ -228,17 +236,25 @@ class DiffusionTrainingService:
             self._pump.start()
             return job_id
 
-    def stop(self) -> bool:
-        """Request a clean stop (the trainer finishes the current step and saves a partial
-        adapter). Returns True if a stop was signalled, False if nothing was running."""
+    def stop(self, save: bool = True) -> bool:
+        """Request a clean stop: the trainer finishes the current step, then either saves
+        a partial adapter (``save=True``, the default) or discards the run (``save=False``,
+        matching the LLM trainer's cancel). Returns True if a stop was signalled, False if
+        nothing was running."""
         with self._lock:
             if self._proc is None or not self._proc.is_alive() or self._stop_queue is None:
                 return False
             try:
-                self._stop_queue.put(True)
+                # Bare True keeps the wire format older trainers expect; the dict form
+                # carries the no-save cancel flag the trainer's _check_stop understands.
+                self._stop_queue.put(True if save else {"save": False})
             except Exception:  # noqa: BLE001
                 return False
-            self._state["message"] = "Stop requested; finishing the current step..."
+            self._state["message"] = (
+                "Stop requested; finishing the current step and saving a partial adapter..."
+                if save
+                else "Cancel requested; finishing the current step (no adapter will be saved)..."
+            )
             self._state["updated_at"] = time.time()
             return True
 
@@ -302,15 +318,47 @@ class DiffusionTrainingService:
                     s["num_images"] = ev.get("num_images")
             elif etype == "model_load_completed":
                 s.update(in_model_load = False, message = "Training...")
+            elif etype == "preparing":
+                # A long precompute phase (e.g. the VAE latent cache) between model load and
+                # the first step; surfaced so the UI shows visible progress instead of a
+                # silent "Loading base model..." stall.
+                done, total = ev.get("done"), ev.get("total")
+                stage = str(ev.get("stage", "prepare")).replace("_", " ")
+                s.update(
+                    status = "running",
+                    in_model_load = True,
+                    message = (
+                        f"Preparing ({stage} {done}/{total})..."
+                        if done is not None and total is not None
+                        else f"Preparing ({stage})..."
+                    ),
+                )
+            elif etype == "warning":
+                # Non-fatal trainer notes (e.g. torch.compile falling back to eager); keep
+                # training state, surface the text.
+                s["message"] = str(ev.get("message", "warning"))
             elif etype == "progress":
+                # Null any non-finite float (NaN/Inf from a divergent step or an inf grad
+                # norm) so the JSON status stays strict-parseable; a missing key keeps the
+                # last value, a present-but-non-finite one becomes None.
+                loss = _finite_or_none(ev["loss"]) if "loss" in ev else s["loss"]
+                avg_loss = _finite_or_none(ev["avg_loss"]) if "avg_loss" in ev else s["avg_loss"]
+                learning_rate = (
+                    _finite_or_none(ev["learning_rate"])
+                    if "learning_rate" in ev
+                    else s["learning_rate"]
+                )
+                grad_norm = (
+                    _finite_or_none(ev["grad_norm"]) if "grad_norm" in ev else s["grad_norm"]
+                )
                 s.update(
                     status = "running",
                     step = ev.get("step", s["step"]),
                     total_steps = ev.get("total_steps", s["total_steps"]),
-                    loss = ev.get("loss", s["loss"]),
-                    avg_loss = ev.get("avg_loss", s["avg_loss"]),
-                    learning_rate = ev.get("learning_rate", s["learning_rate"]),
-                    grad_norm = ev.get("grad_norm", s["grad_norm"]),
+                    loss = loss,
+                    avg_loss = avg_loss,
+                    learning_rate = learning_rate,
+                    grad_norm = grad_norm,
                     message = "Training...",
                 )
                 # Fold optional perf fields (emitted by the trainers) so the UI can show
@@ -337,7 +385,11 @@ class DiffusionTrainingService:
                     status = "stopped" if ev.get("stopped") else "completed",
                     output_dir = ev.get("output_dir"),
                     lora_path = ev.get("lora_path"),
-                    message = "Stopped (partial adapter saved)."
+                    message = (
+                        "Stopped (partial adapter saved)."
+                        if ev.get("lora_path")
+                        else "Stopped (no adapter saved)."
+                    )
                     if ev.get("stopped")
                     else "Training complete.",
                 )

studio/backend/models/training.py

@@ -719,6 +719,34 @@ class DiffusionTrainingStartRequest(BaseModel):
     random_flip: bool = Field(True)
     caption_column: str = Field("text")
     hf_token: Optional[str] = Field(None)
+    cache_latents: bool = Field(
+        True, description = "Precompute VAE latents once and free the VAE for the run"
+    )
+    cache_variants: int = Field(
+        4, ge = 1, le = 16, description = "Frozen crop/flip variants per image in the latent cache"
+    )
+    compile_transformer: Literal["off", "on", "auto"] = Field(
+        "auto", description = "Regional torch.compile of the transformer blocks"
+    )
+    enable_tf32: bool = Field(
+        True, description = "TF32 matmuls + cudnn autotuning (near-lossless speedup)"
+    )
+    base_precision: Literal["nf4", "bf16", "int8", "fp8", "auto"] = Field(
+        "nf4",
+        description = (
+            "DiT base transformer precision: nf4 QLoRA (memory floor, default), bf16 dense, "
+            "int8 torchao weight-only, fp8 float8 training compute (Ada/Hopper/Blackwell), "
+            "or auto (pick by free VRAM + GPU class). Dense modes need a non-prequant base."
+        ),
+    )
+
+
+class DiffusionTrainingStopRequest(BaseModel):
+    """Optional body for stopping a diffusion training job. ``save`` mirrors the LLM
+    trainer's stop: True (default) exports the partial adapter, False cancels without
+    leaving one behind."""
+
+    save: bool = Field(True)
 
 
 class DiffusionTrainingStartResponse(BaseModel):
@@ -790,6 +818,12 @@ class DiffusionTrainableFamily(BaseModel):
     base_repos: List[str] = Field(default_factory = list)
     defaults: dict = Field(default_factory = dict)
     vram_note: str = ""
+    # base_precision modes this machine supports for the family (empty = the family has no
+    # precision selector, e.g. SDXL), plus the recommended pick and whether regional
+    # torch.compile applies. Defaults keep older backends' payloads valid.
+    precision_modes: List[str] = Field(default_factory = list)
+    recommended_precision: str = "nf4"
+    supports_compile: bool = False
 
 
 class DiffusionTrainingInfoResponse(BaseModel):

studio/backend/routes/training.py

@@ -73,6 +73,7 @@ from models.training import (
     DiffusionTrainingStartRequest,
     DiffusionTrainingStartResponse,
     DiffusionTrainingStatusResponse,
+    DiffusionTrainingStopRequest,
 )
 from models.responses import TrainingStopResponse, TrainingMetricsResponse
 from pydantic import BaseModel as PydanticBaseModel
@@ -1233,10 +1234,23 @@ async def start_diffusion_training(
     from core.training.diffusion_lora_trainer import _config_from_dict
 
     try:
-        _config_from_dict(config).normalized()
+        normalized_cfg = _config_from_dict(config).normalized()
     except ValueError as e:
         raise HTTPException(status_code = 400, detail = str(e))
 
+    # Preflight the requested DiT precision BEFORE freeing GPU residents: the DiT trainer's own
+    # checks (a bf16-capable GPU is required; an explicit int8 needs a functional torchao) fire
+    # only in the child, AFTER _free_gpu_for_diffusion_training() already evicted the user's
+    # chat/Images model. Fail fast (400) so a pre-Ampere GPU (T4 / V100 / RTX 20xx) or a
+    # stub-torchao host never tears down resident models for a run that cannot start.
+    from core.training.diffusion_train_common import training_precision_preflight_error
+
+    _precision_reason = training_precision_preflight_error(
+        normalized_cfg.resolved_family, normalized_cfg.base_precision
+    )
+    if _precision_reason:
+        raise HTTPException(status_code = 400, detail = _precision_reason)
+
     # Run the trainers' trust gate here too (both assert the same predicate before
     # from_pretrained), so an untrusted/typoed base 400s BEFORE freeing GPU residents
     # instead of tearing down the user's chat/Images model and failing in the child.
@@ -1291,11 +1305,17 @@ async def start_diffusion_training(
 
 
 @router.post("/diffusion/stop")
-async def stop_diffusion_training(current_subject: str = Depends(get_current_subject)):
-    """Request a clean stop of the running diffusion training job (partial adapter saved)."""
+async def stop_diffusion_training(
+    body: Optional[DiffusionTrainingStopRequest] = None,
+    current_subject: str = Depends(get_current_subject),
+):
+    """Request a clean stop of the running diffusion training job. The optional body's
+    ``save`` mirrors the LLM /stop: true (default, also for an empty POST) exports the
+    partial adapter, false cancels without saving one."""
     from core.training.diffusion_training_service import get_diffusion_training_service
 
-    stopped = get_diffusion_training_service().stop()
+    save = body.save if body is not None else True
+    stopped = get_diffusion_training_service().stop(save = save)
     return {"status": "stopping" if stopped else "idle"}
 
 

studio/backend/tests/test_diffusion_base_precision.py

@@ -0,0 +1,547 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""CPU-only unit tests for the DiT base_precision work.
+
+Covers the new precision plumbing the precision PR adds: the ``base_precision``
+config validation (dense-vs-prequant + mixed-precision gating), the prequant-repo
+heuristic and its trainer alias, the pure ``auto`` precision policy table, the
+explicit-mode passthrough of ``_resolve_base_precision``, the fp8 module filter, the
+fp8 branch of the compile policy, the ``train_precision_modes`` machine probe, the
+family-info precision fields, and the request-model ``base_precision`` field. No GPU /
+model load: every helper here is pure or name-based, so the config validation runs on
+name matching (``resolve_trainable_family`` is offline) and the torch probe is monkeypatched.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch.nn as nn
+
+import core.training.diffusion_train_common as common
+from core.training import diffusion_dit_trainer as dit
+from core.training.diffusion_train_common import (
+    DiffusionLoraConfig,
+    _config_from_dict,
+    repo_is_prequantized,
+    train_precision_modes,
+)
+from models.training import DiffusionTrainingStartRequest
+
+# A dense (non-prequant) DiT base and a prequant bnb-4bit base. Both resolve a trainer
+# family from their names alone, so normalized() runs without a network call.
+_FLUX_DENSE = "black-forest-labs/FLUX.1-dev"
+_Z_PREQUANT = "unsloth/Z-Image-Turbo-unsloth-bnb-4bit"
+# An SDXL base whose name LOOKS prequant (bnb-4bit): SDXL ignores base_precision, so the
+# dense-mode gates must not fire for it even with a dense mode + fp16 compute.
+_SDXL_PREQUANT_NAME = "some/sdxl-model-bnb-4bit"
+# A dense Qwen-Image base: its DiT is corrupted by fp8 (activation outliers), so fp8 is
+# denied for training the same way the inference path denies it.
+_QWEN_DENSE = "Qwen/Qwen-Image"
+
+
+def _cfg(base_model = _FLUX_DENSE, **kw) -> DiffusionLoraConfig:
+    return DiffusionLoraConfig(base_model = base_model, data_dir = "d", output_dir = "o", **kw)
+
+
+# ── base_precision validation ─────────────────────────────────────────────────
+def test_base_precision_validation():
+    # Default normalizes to the nf4 memory floor.
+    assert _cfg().normalized().base_precision == "nf4"
+
+    # An unknown mode is rejected by name.
+    with pytest.raises(ValueError, match = "base_precision"):
+        _cfg(base_precision = "banana").normalized()
+
+    # A dense mode is case/space-insensitive and stored lowered: " FP8 " on a dense base
+    # with bf16 compute normalizes cleanly to "fp8".
+    norm = _cfg(base_precision = " FP8 ", mixed_precision = "bf16").normalized()
+    assert norm.base_precision == "fp8"
+
+    # A dense mode against a prequant (bnb-4bit) base is refused: the repo already ships a
+    # 4-bit transformer and cannot serve the dense precisions.
+    with pytest.raises(ValueError, match = "dense base repo"):
+        _cfg(base_model = _Z_PREQUANT, base_precision = "bf16").normalized()
+
+    # A dense mode with non-bf16 compute is refused: these modes train in bf16 compute.
+    with pytest.raises(ValueError, match = "bf16 compute"):
+        _cfg(base_precision = "int8", mixed_precision = "fp16").normalized()
+
+    # "auto" is ACCEPTED by normalized() even on a prequant base: the concrete mode is
+    # resolved at runtime against the live GPU, not at config validation.
+    assert _cfg(base_model = _Z_PREQUANT, base_precision = "auto").normalized().base_precision == "auto"
+
+
+def test_base_precision_denies_fp8_for_corrupted_family():
+    # fp8 corrupts the Qwen-Image DiT (activation outliers exceed fp8's range), so a dense
+    # Qwen base with base_precision="fp8" is refused up front -- mirroring the inference deny.
+    with pytest.raises(ValueError, match = "fp8"):
+        _cfg(base_model = _QWEN_DENSE, base_precision = "fp8", mixed_precision = "bf16").normalized()
+
+    # The deny is fp8-specific: int8 (per-token, unaffected) and the other dense modes stay
+    # allowed for the same Qwen base.
+    for mode in ("nf4", "bf16", "int8", "auto"):
+        norm = _cfg(
+            base_model = _QWEN_DENSE, base_precision = mode, mixed_precision = "bf16"
+        ).normalized()
+        assert norm.resolved_family == "qwen-image"
+        assert norm.base_precision == mode
+
+    # A family the deny does not cover (FLUX) still accepts fp8.
+    flux = _cfg(base_model = _FLUX_DENSE, base_precision = "fp8", mixed_precision = "bf16").normalized()
+    assert flux.resolved_family == "flux.1"
+    assert flux.base_precision == "fp8"
+
+
+def test_family_train_infos_drops_denied_fp8_for_qwen(monkeypatch):
+    # /info advertises the machine's DiT modes per family, but a family whose DiT the mode
+    # corrupts must not offer it: with fp8 in the machine list, Qwen-Image drops fp8 while
+    # FLUX keeps it, so the UI never surfaces a mode normalized() would reject.
+    monkeypatch.setattr(
+        common, "train_precision_modes", lambda: (["nf4", "bf16", "int8", "fp8", "auto"], "auto")
+    )
+    # family_train_infos reads the live GPU via bf16_unsupported_reason; pin it to "bf16 OK" so
+    # this positive-path assertion is deterministic across GPU types (a non-bf16 CUDA box would
+    # otherwise empty every DiT family's modes). The empty-on-non-bf16 path is covered separately.
+    monkeypatch.setattr(common, "bf16_unsupported_reason", lambda name: None)
+    infos = {i["name"]: i for i in common.family_train_infos()}
+    assert "fp8" not in infos["qwen-image"]["precision_modes"]
+    assert "int8" in infos["qwen-image"]["precision_modes"]  # int8 is fine on Qwen
+    assert "fp8" in infos["flux.1"]["precision_modes"]
+
+
+def test_resolve_base_precision_explicit_int8_gates_on_torchao(monkeypatch):
+    # Explicit int8 has no runtime fallback, so a missing/stub torchao must fail fast here
+    # rather than load dense with compile disabled. Gate the explicit request the same way
+    # auto + /info already gate it.
+    spec = dit._SPECS["flux.1"]
+    cfg = _cfg(base_precision = "int8")
+
+    monkeypatch.setattr(dit, "has_functional_torchao", lambda: False)  # torchao absent / stub
+    with pytest.raises(ValueError, match = "torchao"):
+        dit._resolve_base_precision(cfg, spec, "cuda")
+
+    # With a functional torchao the explicit int8 passes straight through.
+    monkeypatch.setattr(dit, "has_functional_torchao", lambda: True)
+    assert dit._resolve_base_precision(cfg, spec, "cuda") == "int8"
+
+    # The gate is int8-specific: explicit bf16/fp8 pass through regardless of torchao (fp8 has
+    # its own graceful fallback; bf16 needs no torchao).
+    monkeypatch.setattr(dit, "has_functional_torchao", lambda: False)
+    assert dit._resolve_base_precision(_cfg(base_precision = "bf16"), spec, "cuda") == "bf16"
+    assert dit._resolve_base_precision(_cfg(base_precision = "fp8"), spec, "cuda") == "fp8"
+
+
+def test_bf16_unsupported_reason(monkeypatch):
+    # The route uses this to fail fast on a non-bf16 GPU BEFORE evicting resident workloads.
+    import torch
+
+    from core.training.diffusion_train_common import bf16_unsupported_reason
+
+    # SDXL (own mixed_precision path) and unknown families are always exempt.
+    assert bf16_unsupported_reason("sdxl") is None
+    assert bf16_unsupported_reason("") is None
+
+    # A DiT family on a CUDA GPU without bf16 -> a clear reason.
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: False)
+    assert "bfloat16" in (bf16_unsupported_reason("flux.1") or "")
+
+    # A bf16-capable GPU -> no reason.
+    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True)
+    assert bf16_unsupported_reason("qwen-image") is None
+
+    # A CPU-only host (fp32 fallback for import/unit tests) -> no reason even for a DiT family.
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    assert bf16_unsupported_reason("z-image") is None
+
+
+def test_training_precision_preflight_error(monkeypatch):
+    # The start route calls this BEFORE evicting resident GPU workloads: it folds the bf16-GPU
+    # requirement together with the explicit-int8 torchao requirement, so both fail fast instead
+    # of only surfacing in the trainer child after the GPU has already been freed.
+    import torch
+
+    from core.training.diffusion_train_common import training_precision_preflight_error
+
+    # Present a bf16-capable CUDA GPU so the int8 gate (not the bf16 gate) is what we exercise.
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True)
+
+    # The bf16 gate takes precedence: a non-bf16 GPU rejects any DiT precision first.
+    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: False)
+    assert "bfloat16" in (training_precision_preflight_error("flux.1", "int8") or "")
+    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True)
+
+    # Explicit int8 on a DiT family with a NON-functional torchao -> a clear int8 reason
+    # (its _int8_quantize_base has no fallback, so the child would otherwise raise post-eviction).
+    monkeypatch.setattr(common, "has_functional_torchao", lambda: False)
+    reason = training_precision_preflight_error("qwen-image", "int8")
+    assert reason is not None and "int8" in reason and "torchao" in reason
+
+    # The same int8 request is fine once torchao is functional.
+    monkeypatch.setattr(common, "has_functional_torchao", lambda: True)
+    assert training_precision_preflight_error("qwen-image", "int8") is None
+
+    # With a broken torchao, only EXPLICIT int8 is gated -- nf4/bf16/auto pass, and the int8
+    # gate never applies to a non-DiT (SDXL) or unknown family.
+    monkeypatch.setattr(common, "has_functional_torchao", lambda: False)
+    assert training_precision_preflight_error("flux.1", "nf4") is None
+    assert training_precision_preflight_error("flux.1", "auto") is None
+    assert training_precision_preflight_error("sdxl", "int8") is None
+    assert training_precision_preflight_error("", "int8") is None
+
+    # On a CUDA-ABSENT host, bf16_unsupported_reason exempts CPU-only, but the DiT trainer's dense
+    # precisions still require CUDA (mirroring _resolve_base_precision), so bf16/int8/fp8 for a DiT
+    # family are rejected UP FRONT rather than after eviction. nf4/auto (and SDXL) still pass.
+    monkeypatch.setattr(common, "has_functional_torchao", lambda: True)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    for dense in ("bf16", "int8", "fp8"):
+        reason = training_precision_preflight_error("flux.1", dense)
+        assert reason is not None and "CUDA" in reason
+    assert training_precision_preflight_error("flux.1", "nf4") is None
+    assert training_precision_preflight_error("flux.1", "auto") is None
+    assert training_precision_preflight_error("sdxl", "bf16") is None
+
+
+def test_family_train_infos_empties_dit_modes_on_non_bf16(monkeypatch):
+    # On a non-bf16 GPU the start route rejects EVERY DiT family (even nf4), so /info must not
+    # advertise a DiT precision option that always 400s: the modes empty, the reason surfaces in
+    # vram_note, compile is off, and the recommendation degrades to nf4. SDXL (non-DiT) is exempt.
+    from core.training.diffusion_train_common import _DIT_TRAIN_FAMILIES, family_train_infos
+
+    monkeypatch.setattr(common, "bf16_unsupported_reason", lambda name: "no bfloat16 on this GPU")
+
+    infos = {info["name"]: info for info in family_train_infos()}
+    dit_seen = False
+    for name, info in infos.items():
+        if name in _DIT_TRAIN_FAMILIES:
+            dit_seen = True
+            assert info["precision_modes"] == []
+            assert info["vram_note"] == "no bfloat16 on this GPU"
+            assert info["recommended_precision"] == "nf4"
+            assert info["supports_compile"] is False
+    assert dit_seen  # the registry must still expose at least one DiT family to have covered it
+
+
+def test_base_precision_gates_skip_sdxl():
+    # SDXL ignores base_precision, so the dense-mode gates (prequant base / non-bf16 compute)
+    # must not fire for it: a prequant-looking SDXL name with base_precision="bf16" does not
+    # raise, and the mode is still stored lowered.
+    norm = _cfg(base_model = _SDXL_PREQUANT_NAME, base_precision = "bf16").normalized()
+    assert norm.resolved_family == "sdxl"
+    assert norm.base_precision == "bf16"
+
+    # The non-bf16-compute gate is also skipped for SDXL (fp16 is a valid SDXL mixed
+    # precision), even with a dense base_precision requested.
+    norm2 = _cfg(
+        base_model = "stabilityai/stable-diffusion-xl-base-1.0",
+        base_precision = "int8",
+        mixed_precision = "fp16",
+    ).normalized()
+    assert norm2.resolved_family == "sdxl"
+
+    # The mode-name validity check still runs for SDXL: an unknown mode is rejected.
+    with pytest.raises(ValueError, match = "base_precision"):
+        _cfg(base_model = _SDXL_PREQUANT_NAME, base_precision = "banana").normalized()
+
+    # The gates STILL fire for a DiT family: a prequant DiT base with a dense mode raises.
+    with pytest.raises(ValueError, match = "dense base repo"):
+        _cfg(base_model = _Z_PREQUANT, base_precision = "bf16").normalized()
+
+
+# ── repo_is_prequantized heuristic + trainer alias ────────────────────────────
+@pytest.mark.parametrize(
+    "repo, expected",
+    [
+        ("unsloth/Qwen-Image-2512-unsloth-bnb-4bit", True),
+        ("some/model-4bit", True),
+        ("some/model-int4", True),
+        ("some/model-nf4", True),
+        ("black-forest-labs/FLUX.1-dev", False),
+        ("Tongyi-MAI/Z-Image-Turbo", False),
+    ],
+)
+def test_repo_is_prequantized_cases(repo, expected):
+    assert repo_is_prequantized(repo) is expected
+
+
+def test_repo_is_prequantized_alias_is_same_object():
+    # The trainer keeps a module-level alias for callers/tests; it must be the exact same
+    # function object as the common heuristic (moved there for config validation).
+    assert dit._repo_is_prequantized is repo_is_prequantized
+
+
+# ── _pick_auto_precision policy table (pure) ──────────────────────────────────
+def test_pick_auto_precision_policy_table():
+    p = dit._pick_auto_precision
+
+    # A prequant base always resolves to nf4 (it can only serve 4-bit).
+    assert p(True, "cuda", 140, 23.8, (10, 0), True) == "nf4"
+    # No CUDA -> nf4 (the dense modes need a GPU).
+    assert p(False, "cpu", 140, 23.8, (10, 0), True) == "nf4"
+    # Missing free-VRAM number -> the safe nf4 mode.
+    assert p(False, "cuda", None, 23.8, (10, 0), True) == "nf4"
+
+    # Plenty of free VRAM -> bf16 regardless of fp8 capability: compiled bf16 measured
+    # FASTER than torchao float8 at LoRA-training shapes, so fp8 is opt-in only.
+    assert p(False, "cuda", 140, 23.8, (10, 0), True) == "bf16"
+    assert p(False, "cuda", 140, 23.8, (8, 0), True) == "bf16"
+    assert p(False, "cuda", 140, 23.8, (10, 0), False) == "bf16"
+
+    # Middle band (30 > 23.8 * 1.15 = 27.4, but not > 23.8 * 1.5 = 35.7) -> int8.
+    assert p(False, "cuda", 30, 23.8, (10, 0), True) == "int8"
+    # int8 needs torchao at runtime (no fallback), so the int8 band drops to nf4 when
+    # torchao is not importable while the bf16 band is unaffected.
+    assert p(False, "cuda", 30, 23.8, (10, 0), True, False) == "nf4"
+    assert p(False, "cuda", 140, 23.8, (10, 0), True, False) == "bf16"
+    # int8 still materialises the full bf16 transformer before quantize_ shrinks it, so
+    # free VRAM below the dense-load transient (25 < 27.4) must fall back to nf4 even
+    # though the QUANTIZED weights would have fit.
+    assert p(False, "cuda", 25, 23.8, (10, 0), True) == "nf4"
+    # Too little free VRAM for any dense load -> nf4.
+    assert p(False, "cuda", 10, 23.8, (10, 0), True) == "nf4"
+
+
+# ── _resolve_base_precision passthrough ───────────────────────────────────────
+def test_resolve_base_precision_passes_explicit_through():
+    # An explicit mode passes straight through without probing the GPU (normalized() already
+    # validated it); the spec is only consulted for "auto".
+    spec = dit._SPECS["flux.1"]
+    cfg = _cfg(base_precision = "bf16")
+    assert dit._resolve_base_precision(cfg, spec, "cuda") == "bf16"
+
+    # The dense modes are CUDA-only: an explicit request on a GPU-less host fails fast
+    # (before any model load) instead of silently proceeding; /info never advertised it.
+    with pytest.raises(ValueError, match = "CUDA"):
+        dit._resolve_base_precision(cfg, spec, "cpu")
+    # nf4 stays a passthrough on any device (the bnb load path owns its own errors).
+    assert dit._resolve_base_precision(_cfg(base_precision = "nf4"), spec, "cpu") == "nf4"
+
+
+def test_resolve_auto_requires_bf16_compute():
+    # auto may resolve to bf16/int8 which train in bf16 compute, so a non-bf16
+    # mixed_precision pins auto to the nf4 floor BEFORE any GPU probe (pure, no CUDA
+    # needed here) -- mirroring the normalized() rule for explicit dense modes.
+    spec = dit._SPECS["flux.1"]
+    cfg = _cfg(base_precision = "auto", mixed_precision = "fp16")
+    assert dit._resolve_base_precision(cfg, spec, "cuda") == "nf4"
+
+
+def test_resolve_auto_int8_band_gates_on_torchao(monkeypatch):
+    # The int8 auto band needs a FUNCTIONAL torchao at runtime; when torchao is not
+    # importable _resolve_base_precision must fall to nf4 instead of picking an int8 that
+    # would crash in _int8_quantize_base. Drive the probe into the int8 band and toggle the
+    # functional-torchao probe (shared with train_precision_modes, imported into the trainer).
+    import torch
+
+    spec = dit._SPECS["flux.1"]  # dense_bf16_gb = 23.8
+    cfg = _cfg(base_precision = "auto", mixed_precision = "bf16")
+
+    class _FakeCuda:
+        # Free VRAM in the int8 band (30 > 23.8 * 1.15) but below the bf16 band.
+        @staticmethod
+        def mem_get_info():
+            return (int(30 * 1e9), int(80 * 1e9))
+
+        @staticmethod
+        def get_device_capability():
+            return (10, 0)
+
+    monkeypatch.setattr(torch, "cuda", _FakeCuda)
+
+    monkeypatch.setattr(dit, "has_functional_torchao", lambda: False)  # torchao absent / stub
+    assert dit._resolve_base_precision(cfg, spec, "cuda") == "nf4"
+
+    # With a functional torchao the same band picks int8.
+    monkeypatch.setattr(dit, "has_functional_torchao", lambda: True)
+    assert dit._resolve_base_precision(cfg, spec, "cuda") == "int8"
+
+
+def test_resolve_auto_int8_band_treats_stub_as_absent(monkeypatch):
+    # Simulate the Windows-ROCm torchao STUB: has_functional_torchao returns False (the
+    # stub satisfies find_spec but its quantize_ is a no-op), so the int8 band must fall to
+    # nf4 rather than pick an int8 whose quantization silently does nothing.
+    import torch
+
+    spec = dit._SPECS["flux.1"]
+    cfg = _cfg(base_precision = "auto", mixed_precision = "bf16")
+
+    class _FakeCuda:
+        @staticmethod
+        def mem_get_info():
+            return (int(30 * 1e9), int(80 * 1e9))
+
+        @staticmethod
+        def get_device_capability():
+            return (10, 0)
+
+    monkeypatch.setattr(torch, "cuda", _FakeCuda)
+    # The stub scenario: the probe reports no functional torchao.
+    monkeypatch.setattr(dit, "has_functional_torchao", lambda: False)
+    assert dit._resolve_base_precision(cfg, spec, "cuda") == "nf4"
+
+
+def test_has_functional_torchao_rejects_stub(monkeypatch):
+    # has_functional_torchao must reject the Unsloth import stub: even though
+    # `from torchao.quantization import quantize_` would succeed against the stub, the
+    # symbols are no-op stub types. Simulate a stub torchao.quantization module carrying the
+    # stub sentinel and assert the probe returns False.
+    import importlib
+    import types
+
+    from core._torchao_stub import _STUB_SENTINEL
+
+    real_import_module = importlib.import_module
+
+    stub_quant = types.ModuleType("torchao.quantization")
+    stub_quant._unsloth_stub = _STUB_SENTINEL
+
+    def _fake_import(name, *args, **kwargs):
+        if name == "torchao.quantization":
+            return stub_quant
+        return real_import_module(name, *args, **kwargs)
+
+    monkeypatch.setattr(importlib, "import_module", _fake_import)
+    assert common.has_functional_torchao() is False
+
+    # A real module exposing the int8 symbols (no stub sentinel) probes True.
+    real_like = types.ModuleType("torchao.quantization")
+    real_like.Int8WeightOnlyConfig = object
+    real_like.quantize_ = lambda *a, **k: None
+
+    def _fake_import_real(name, *args, **kwargs):
+        if name == "torchao.quantization":
+            return real_like
+        return real_import_module(name, *args, **kwargs)
+
+    monkeypatch.setattr(importlib, "import_module", _fake_import_real)
+    assert common.has_functional_torchao() is True
+
+
+# ── _fp8_module_filter ────────────────────────────────────────────────────────
+def test_fp8_module_filter():
+    lin = nn.Linear(64, 64)
+    # A plain feed-forward Linear with divisible dims gets float8 training compute.
+    assert dit._fp8_module_filter(lin, "transformer_blocks.0.ff.net.0") is True
+    # A LoRA-owned module is skipped (adapters stay high precision).
+    assert dit._fp8_module_filter(lin, "transformer_blocks.0.attn.to_q.lora_A.default") is False
+    # The output projection is skipped.
+    assert dit._fp8_module_filter(lin, "proj_out") is False
+    # An in_features not divisible by 16 is rejected (float8 kernels reject the shape).
+    assert dit._fp8_module_filter(nn.Linear(30, 64), "transformer_blocks.0.ff.net.0") is False
+    # A non-Linear module is never float8.
+    assert dit._fp8_module_filter(nn.LayerNorm(64), "transformer_blocks.0.norm") is False
+
+
+# ── _should_compile fp8 branch ────────────────────────────────────────────────
+def test_should_compile_fp8_branch():
+    # fp8 is only competitive compiled, so auto arms compile for it on a dense (non-bnb)
+    # cuda base.
+    cfg = _cfg(compile_transformer = "auto")
+    assert dit._should_compile(cfg, False, "cuda", "fp8") is True
+    # fp8 forces compile under auto even when the base is (hypothetically) reported as bnb.
+    assert dit._should_compile(cfg, True, "cuda", "fp8") is True
+    # An explicit "off" still wins over fp8 -- compile stays off.
+    assert dit._should_compile(_cfg(compile_transformer = "off"), False, "cuda", "fp8") is False
+
+
+# ── train_precision_modes machine probe ───────────────────────────────────────
+def test_train_precision_modes_no_cuda(monkeypatch):
+    # Patch the torch module attribute the function imports so it observes a CPU-only box:
+    # no CUDA -> the nf4-only floor with nf4 recommended, and it never raises.
+    import torch
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    assert train_precision_modes() == (["nf4"], "nf4")
+
+
+def test_train_precision_modes_gates_int8_fp8_on_torchao(monkeypatch):
+    # int8/fp8 are only advertised when torchao is FUNCTIONAL: on a CUDA host WITHOUT a real
+    # torchao (or with only the Windows-ROCm stub) /info must not offer int8/fp8, since their
+    # explicit paths import torchao with no fallback. bf16 + auto stay advertised.
+    import torch
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (10, 0))
+
+    # No functional torchao (absent or stub): bf16 + auto only, int8/fp8 dropped.
+    monkeypatch.setattr(common, "has_functional_torchao", lambda: False)
+    modes, recommended = train_precision_modes()
+    assert modes == ["nf4", "bf16", "auto"]
+    assert "int8" not in modes and "fp8" not in modes
+    assert recommended == "auto"
+
+    # With a functional torchao on an fp8-capable GPU, int8 + fp8 are advertised again.
+    monkeypatch.setattr(common, "has_functional_torchao", lambda: True)
+    modes2, _ = train_precision_modes()
+    assert "int8" in modes2 and "fp8" in modes2
+
+
+def test_train_precision_modes_gates_dense_on_bf16_support(monkeypatch):
+    # The dense modes (bf16/int8/fp8/auto) all train in bf16 compute, which the DiT trainer
+    # requires. On a CUDA GPU that cannot do bf16 (T4/V100/RTX 20xx), /info must offer ONLY
+    # nf4 -- otherwise the UI advertises a start that evicts resident models and then fails the
+    # trainer's bf16 guard.
+    import torch
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: False)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (7, 5))  # Turing, no bf16
+    monkeypatch.setattr(common, "has_functional_torchao", lambda: True)
+    modes, recommended = train_precision_modes()
+    assert modes == ["nf4"]
+    assert recommended == "nf4"
+
+
+# ── family_train_infos precision fields ───────────────────────────────────────
+def test_family_train_infos_carries_precision_fields(monkeypatch):
+    # Pin the machine probe so the DiT families carry a deterministic mode list, while SDXL
+    # (no precision selector) stays empty regardless of the probe.
+    monkeypatch.setattr(common, "train_precision_modes", lambda: (["nf4", "bf16"], "auto"))
+    # Also pin bf16_unsupported_reason (family_train_infos reads the live GPU through it): "bf16 OK"
+    # so this positive-path assertion is deterministic across GPU types, not just on CPU-only CI.
+    monkeypatch.setattr(common, "bf16_unsupported_reason", lambda name: None)
+    infos = {i["name"]: i for i in common.family_train_infos()}
+
+    flux = infos["flux.1"]
+    assert flux["precision_modes"] == ["nf4", "bf16"]
+    assert flux["recommended_precision"] == "auto"
+    assert flux["supports_compile"] is True
+
+    sdxl = infos["sdxl"]
+    assert sdxl["precision_modes"] == []
+    assert sdxl["recommended_precision"] == "nf4"
+    assert sdxl["supports_compile"] is False
+
+
+# ── request model base_precision field ────────────────────────────────────────
+def test_request_model_base_precision():
+    # The request defaults to the nf4 memory floor.
+    req = DiffusionTrainingStartRequest(base_model = "x", data_dir = "d", output_dir = "o")
+    assert req.base_precision == "nf4"
+
+    # An allowed dense mode is accepted.
+    assert (
+        DiffusionTrainingStartRequest(
+            base_model = "x", data_dir = "d", output_dir = "o", base_precision = "fp8"
+        ).base_precision
+        == "fp8"
+    )
+
+    # An out-of-Literal value is rejected by pydantic.
+    with pytest.raises(Exception):
+        DiffusionTrainingStartRequest(
+            base_model = "x", data_dir = "d", output_dir = "o", base_precision = "int4"
+        )
+
+    # The generic Studio dict path carries base_precision through onto DiffusionLoraConfig.
+    cfg = _config_from_dict(
+        {
+            "base_model": _FLUX_DENSE,
+            "data_dir": "d",
+            "output_dir": "o",
+            "base_precision": "bf16",
+        }
+    )
+    assert cfg.base_precision == "bf16"

studio/backend/tests/test_diffusion_train_perf.py

@@ -0,0 +1,469 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""CPU-only unit tests for the diffusion training performance work.
+
+Covers the new pure helpers and small policy functions that the perf PR adds:
+the seed-deterministic latent-cache crop/flip plan, the per-family collate fns, the
+index-based sigma gather, the new config validation + request-model fields, the
+torch.compile policy, the stop save/cancel flag, and the ``preparing`` / ``warning``
+service events. No GPU / model load: the collates and gathers run on CPU tensors, the
+scheduler is default-initialised (no ``from_pretrained``), and the route/service tests
+inject in-thread fakes exactly like ``test_diffusion_training.py``.
+"""
+
+from __future__ import annotations
+
+import itertools
+
+import pytest
+import torch
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from auth.authentication import get_current_subject
+from core.training.diffusion_dit_trainer import (
+    _flux_collate,
+    _gather_sigmas,
+    _qwen_collate,
+    _sample_timesteps,
+    _should_compile,
+    _zimage_collate,
+)
+from core.training.diffusion_train_common import (
+    DiffusionLoraConfig,
+    LATENT_CACHE_OVER_BUDGET,
+    _apply_perf_flags,
+    _config_from_dict,
+    _latent_cache_forced,
+    _latent_cache_over_budget,
+    _plan_cache_variants,
+    _restore_perf_flags,
+)
+import core.training.diffusion_lora_trainer as sdxl_trainer
+import core.training.diffusion_train_common as train_common
+from core.training.diffusion_training_service import DiffusionTrainingService
+from models.training import DiffusionTrainingStartRequest, DiffusionTrainingStopRequest
+from routes.training import router as training_router
+
+# A trainable SDXL base so DiffusionLoraConfig.normalized() resolves a family without a
+# network call (resolve_trainable_family is pure name matching for this repo).
+_SDXL = "stabilityai/stable-diffusion-xl-base-1.0"
+
+
+def _cfg(**kw) -> DiffusionLoraConfig:
+    return DiffusionLoraConfig(base_model = _SDXL, data_dir = "d", output_dir = "o", **kw)
+
+
+# ── _plan_cache_variants (pure, seed-deterministic) ───────────────────────────
+def test_plan_cache_variants_deterministic_and_deduped():
+    # Same seed -> byte-identical plan (its own rng stream, so it is fully reproducible).
+    p1 = _plan_cache_variants(3, 4, center_crop = False, random_flip = True, seed = 123)
+    p2 = _plan_cache_variants(3, 4, center_crop = False, random_flip = True, seed = 123)
+    assert p1 == p2
+    assert len(p1) == 3
+
+    # cache_variants=1 -> exactly one variant per image.
+    p_one = _plan_cache_variants(3, 1, center_crop = False, random_flip = True, seed = 7)
+    assert [len(v) for v in p_one] == [1, 1, 1]
+
+    # A center crop with no flip collapses to a single distinct variant no matter how many
+    # draws are requested, and that variant is the fixed (0.5, 0.5, False) center.
+    p_cc = _plan_cache_variants(2, 8, center_crop = True, random_flip = False, seed = 7)
+    assert [len(v) for v in p_cc] == [1, 1]
+    assert p_cc[0][0] == (0.5, 0.5, False)
+
+    # A center crop WITH flip has at most two distinct variants (flip on/off; crop is fixed).
+    p_cf = _plan_cache_variants(2, 8, center_crop = True, random_flip = True, seed = 7)
+    assert all(len(v) <= 2 for v in p_cf)
+
+    # Every crop fraction is a valid unit fraction the loader can map onto its crop range.
+    for u_left, u_top, flip in itertools.chain.from_iterable(p1):
+        assert 0.0 <= u_left < 1.0
+        assert 0.0 <= u_top < 1.0
+        assert isinstance(flip, bool)
+
+
+# ── per-family collate fns ────────────────────────────────────────────────────
+def test_flux_collate_shapes():
+    # FLUX embeds are fixed length: 3 entries batch by a plain cat; text_ids are shared.
+    entries = [(torch.randn(1, 512, 32), torch.randn(1, 16), torch.randn(512, 3)) for _ in range(3)]
+    pe, pooled, text_ids = _flux_collate(entries, "cpu", torch.float32)
+    assert pe.shape == (3, 512, 32)
+    assert pooled.shape == (3, 16)
+    assert text_ids.shape == (512, 3)
+    # Position ids stay float32 regardless of the requested weight dtype.
+    assert pe.dtype == torch.float32
+    assert pooled.dtype == torch.float32
+    assert text_ids.dtype == torch.float32
+
+
+def test_qwen_collate_pads_and_masks():
+    dim = 8
+    # A short (mask=None) and a long (mask=ones) entry -> pad to the batch max and build the
+    # validity mask, with the padded tail of the short sample masked out.
+    short = (torch.randn(1, 5, dim), None)
+    long = (torch.randn(1, 9, dim), torch.ones(1, 9, dtype = torch.int64))
+    pe, mask = _qwen_collate([short, long], "cpu", torch.float32)
+    assert pe.shape == (2, 9, dim)
+    assert mask.shape == (2, 9)
+    assert torch.equal(mask[0, 5:], torch.zeros(4, dtype = mask.dtype))
+
+    # A single unpadded sample with a None mask keeps the legacy None mask (no behaviour delta).
+    pe1, mask1 = _qwen_collate([(torch.randn(1, 5, dim), None)], "cpu", torch.float32)
+    assert pe1.shape == (1, 5, dim)
+    assert mask1 is None
+
+    # A single sample pinned to a compile pad bucket must pad AND expose a mask so the padded
+    # positions are attended to as invalid.
+    pe2, mask2 = _qwen_collate([(torch.randn(1, 5, dim), None)], "cpu", torch.float32, pad_to = 16)
+    assert pe2.shape == (1, 16, dim)
+    assert mask2 is not None
+    assert torch.equal(mask2[0, 5:], torch.zeros(11, dtype = mask2.dtype))
+
+
+def test_zimage_collate_list():
+    # Z-Image uses list I/O: the batch is one tuple carrying a list of per-sample tensors, each
+    # cast to the requested dtype.
+    entries = [(torch.randn(7, 2560),), (torch.randn(9, 2560),)]
+    out = _zimage_collate(entries, "cpu", torch.float32)
+    assert isinstance(out, tuple) and len(out) == 1
+    (caps,) = out
+    assert isinstance(caps, list) and len(caps) == 2
+    assert all(t.dtype == torch.float32 for t in caps)
+
+
+# ── index-based sigma gather ──────────────────────────────────────────────────
+def test_gather_sigmas_matches_search_based_gather():
+    # CI installs the backend test deps without diffusers; the scheduler math is what we
+    # are checking, so skip rather than fail there.
+    pytest.importorskip("diffusers")
+    from diffusers import FlowMatchEulerDiscreteScheduler
+
+    torch.manual_seed(0)
+    sched = FlowMatchEulerDiscreteScheduler()  # default init, no from_pretrained / no network
+    timesteps, indices = _sample_timesteps(sched, 16, "cpu")
+
+    # The index path must return exactly what the old per-item timestep-matching search did.
+    schedule_timesteps = sched.timesteps.to("cpu")
+    step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+    assert step_indices == indices.tolist()
+
+    sigma = _gather_sigmas(sched, indices, "cpu", torch.float32, 4)
+    assert sigma.ndim == 4
+    expected = sched.sigmas[step_indices].flatten()
+    while expected.ndim < 4:
+        expected = expected.unsqueeze(-1)
+    assert torch.equal(sigma, expected)
+
+
+# ── config validation of the new perf fields ──────────────────────────────────
+def test_config_validates_new_fields():
+    # Defaults normalize cleanly and carry the new perf fields through.
+    norm = _cfg().normalized()
+    assert norm.cache_variants == 4
+    assert norm.compile_transformer == "auto"
+    assert norm.enable_tf32 is True
+    assert norm.cache_latents is True
+
+    # cache_variants is bounded to 1..16 inclusive.
+    for bad in (0, 17):
+        with pytest.raises(ValueError):
+            _cfg(cache_variants = bad).normalized()
+
+    # An unknown compile mode is rejected.
+    with pytest.raises(ValueError):
+        _cfg(compile_transformer = "banana").normalized()
+
+    # compile_transformer is case/space-insensitive and stored lowered.
+    assert _cfg(compile_transformer = " ON ").normalized().compile_transformer == "on"
+
+    # The generic Studio dict path preserves the flags without inventing defaults.
+    cfg = _config_from_dict(
+        {
+            "base_model": _SDXL,
+            "data_dir": "d",
+            "output_dir": "o",
+            "enable_tf32": False,
+            "cache_latents": False,
+        }
+    )
+    assert cfg.enable_tf32 is False
+    assert cfg.cache_latents is False
+
+    # String flags from the generic Studio dict path are coerced: "false" is otherwise a
+    # non-empty (truthy) string, so an opt-out would silently no-op.
+    cfg = _config_from_dict(
+        {
+            "base_model": _SDXL,
+            "data_dir": "d",
+            "output_dir": "o",
+            "enable_tf32": "false",
+            "cache_latents": "0",
+        }
+    )
+    assert cfg.enable_tf32 is False
+    assert cfg.cache_latents is False
+
+
+# ── torch.compile policy ──────────────────────────────────────────────────────
+def test_should_compile_policy():
+    # off never compiles, even on cuda.
+    assert _should_compile(_cfg(compile_transformer = "off"), False, "cuda") is False
+    # on always compiles on cuda.
+    assert _should_compile(_cfg(compile_transformer = "on"), False, "cuda") is True
+    # auto stays off over a bitsandbytes base (graph breaks in the dequant path).
+    assert _should_compile(_cfg(compile_transformer = "auto"), True, "cuda") is False
+    # auto turns on for the dense bf16 base precision on cuda.
+    assert (
+        _should_compile(_cfg(compile_transformer = "auto"), False, "cuda", base_precision = "bf16")
+        is True
+    )
+    # Any mode is a no-op on cpu.
+    for mode in ("off", "on", "auto"):
+        assert _should_compile(_cfg(compile_transformer = mode), False, "cpu") is False
+
+
+# ── service stop save/cancel flag ─────────────────────────────────────────────
+class _StopQueue:
+    """Records what stop() puts on the wire (put-only for these tests)."""
+
+    def __init__(self) -> None:
+        self.items: list = []
+
+    def put(self, x) -> None:
+        self.items.append(x)
+
+
+class _AliveProc:
+    def is_alive(self) -> bool:
+        return True
+
+
+def test_service_stop_save_flag():
+    svc = DiffusionTrainingService()
+    # Nothing running -> stop is a no-op and returns False.
+    assert svc.stop() is False
+
+    # Attach a fake live proc + stop queue so stop() has a target.
+    svc._proc = _AliveProc()
+    q = _StopQueue()
+    svc._stop_queue = q
+
+    # save=False is the cancel path: the dict form {"save": False} goes on the queue.
+    assert svc.stop(save = False) is True
+    assert q.items[-1] == {"save": False}
+
+    # The default (save) path keeps the bare-True wire format.
+    assert svc.stop() is True
+    assert q.items[-1] is True
+
+
+# ── preparing / warning events + stopped completion messages ──────────────────
+def test_apply_event_preparing_and_warning():
+    svc = DiffusionTrainingService()
+    svc._apply_event({"type": "preparing", "stage": "cache_latents", "done": 4, "total": 8})
+    st = svc.status()
+    assert st["status"] == "running"
+    assert st["in_model_load"] is True
+    assert "4/8" in st["message"]
+
+    svc._apply_event({"type": "warning", "message": "compile disabled"})
+    assert svc.status()["message"] == "compile disabled"
+
+    # A stop with no saved adapter reports the no-adapter message and the stopped status.
+    svc_no = DiffusionTrainingService()
+    svc_no._apply_event({"type": "complete", "stopped": True, "lora_path": None})
+    st_no = svc_no.status()
+    assert st_no["status"] == "stopped"
+    assert st_no["message"] == "Stopped (no adapter saved)."
+
+    # A stop that DID save a partial adapter reports the partial-adapter message.
+    svc_partial = DiffusionTrainingService()
+    svc_partial._apply_event(
+        {"type": "complete", "stopped": True, "lora_path": "/o/pytorch_lora_weights.safetensors"}
+    )
+    assert svc_partial.status()["message"] == "Stopped (partial adapter saved)."
+
+
+# ── route: stop body forwards the save flag ───────────────────────────────────
+class _FakeService:
+    """Records the save flag the /diffusion/stop route forwards. A local copy of the
+    test_diffusion_training.py pattern so the two suites stay decoupled."""
+
+    def __init__(self) -> None:
+        self._running = True
+        self.stopped_with_save = None
+
+    def stop(self, save = True):
+        self.stopped_with_save = save
+        was = self._running
+        self._running = False
+        return was
+
+
+@pytest.fixture
+def client(monkeypatch):
+    fake = _FakeService()
+    monkeypatch.setattr(
+        "core.training.diffusion_training_service.get_diffusion_training_service", lambda: fake
+    )
+    app = FastAPI()
+    app.include_router(training_router, prefix = "/api/train")
+    app.dependency_overrides[get_current_subject] = lambda: "test-user"
+    c = TestClient(app)
+    c._fake = fake  # type: ignore[attr-defined]
+    return c
+
+
+def test_route_stop_save_body(client):
+    # An explicit {"save": false} body forwards save=False to the service.
+    r = client.post("/api/train/diffusion/stop", json = {"save": False})
+    assert r.status_code == 200, r.text
+    assert client._fake.stopped_with_save is False
+
+    # A body-less POST defaults to save=True.
+    r2 = client.post("/api/train/diffusion/stop")
+    assert r2.status_code == 200, r2.text
+    assert client._fake.stopped_with_save is True
+
+
+# ── request models: new perf fields + stop schema ─────────────────────────────
+def test_request_models_new_fields():
+    req = DiffusionTrainingStartRequest(base_model = "b", data_dir = "d", output_dir = "o")
+    assert req.cache_latents is True
+    assert req.cache_variants == 4
+    assert req.compile_transformer == "auto"
+    assert req.enable_tf32 is True
+
+    # cache_variants is validated against its 1..16 bound by pydantic.
+    with pytest.raises(Exception):
+        DiffusionTrainingStartRequest(
+            base_model = "b", data_dir = "d", output_dir = "o", cache_variants = 32
+        )
+
+    # The stop request defaults to saving a partial adapter.
+    assert DiffusionTrainingStopRequest().save is True
+
+
+# ── perf flags round-trip on cpu ──────────────────────────────────────────────
+def test_perf_flags_cpu_roundtrip():
+    # On a cpu device (or a torch build without cuda), applying the perf flags is a no-op
+    # snapshot path and restoring it must not raise.
+    snap = _apply_perf_flags(_cfg(), "cpu")
+    assert isinstance(snap, dict)
+    _restore_perf_flags(snap)  # no exception
+
+
+def test_perf_flags_tf32_off_clears_flags():
+    # enable_tf32=False is the strict-fp32 A/B mode: it must actively clear the TF32 flags
+    # (cudnn TF32 defaults ON in torch) rather than inherit ambient state, and restore must
+    # put the ambient values back. The flag attributes are plain Python state, present and
+    # settable on CPU-only torch builds, so this runs without a GPU.
+    import torch
+
+    before = (
+        torch.backends.cuda.matmul.allow_tf32,
+        torch.backends.cudnn.allow_tf32,
+        torch.get_float32_matmul_precision(),
+    )
+    snap = _apply_perf_flags(_cfg(enable_tf32 = False), "cuda")
+    try:
+        assert torch.backends.cuda.matmul.allow_tf32 is False
+        assert torch.backends.cudnn.allow_tf32 is False
+        assert torch.get_float32_matmul_precision() == "highest"
+    finally:
+        _restore_perf_flags(snap)
+    after = (
+        torch.backends.cuda.matmul.allow_tf32,
+        torch.backends.cudnn.allow_tf32,
+        torch.get_float32_matmul_precision(),
+    )
+    assert after == before
+
+
+# ── latent cache size gate ────────────────────────────────────────────────────
+class _FakeLatentDist:
+    def __init__(self, shape):
+        self.mean = torch.zeros(shape, dtype = torch.float32)
+        self.std = torch.ones(shape, dtype = torch.float32)
+
+
+class _FakeEncoded:
+    def __init__(self, shape):
+        self.latent_dist = _FakeLatentDist(shape)
+
+
+class _FakeVae:
+    # Minimal VAE stand-in: encode() returns a posterior of the requested latent shape so the
+    # builder measures a real per-variant byte size without a model load or image files.
+    def __init__(self, shape):
+        self._shape = shape
+
+    def encode(self, pixel_values):
+        return _FakeEncoded(self._shape)
+
+
+def _fake_planned_loader(path, resolution, center_crop, u_left, u_top, flip):
+    # The fake VAE ignores pixels; return a valid tensor + square SDXL time_ids.
+    tensor = torch.zeros(3, resolution, resolution, dtype = torch.float32)
+    return tensor, (resolution, resolution, 0, 0, resolution, resolution)
+
+
+def _build_fake_sdxl_cache(monkeypatch, num_images, latent_shape):
+    # center_crop + no flip collapses to one variant per image, so total_variants == num_images.
+    monkeypatch.setattr(sdxl_trainer, "_load_image_tensor_planned", _fake_planned_loader)
+    cfg = _cfg(cache_variants = 1, center_crop = True, random_flip = False).normalized()
+    return sdxl_trainer._build_sdxl_latent_cache(
+        _FakeVae(latent_shape),
+        1.0,
+        [f"img{i}.png" for i in range(num_images)],
+        cfg,
+        "cpu",
+        torch.float32,
+        None,
+        lambda: False,
+    )
+
+
+def test_latent_cache_over_budget_boundary():
+    # 32 bytes per variant x 4 variants = 128 bytes; exactly at budget is not "over".
+    assert _latent_cache_over_budget(32, 4, budget_bytes = 200) is False
+    assert _latent_cache_over_budget(32, 4, budget_bytes = 128) is False
+    assert _latent_cache_over_budget(32, 4, budget_bytes = 127) is True
+    # An empty plan can never overflow.
+    assert _latent_cache_over_budget(1_000_000, 0, budget_bytes = 1) is False
+
+
+def test_latent_cache_forced_env(monkeypatch):
+    monkeypatch.delenv("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", raising = False)
+    assert _latent_cache_forced() is False
+    monkeypatch.setenv("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", "1")
+    assert _latent_cache_forced() is True
+
+
+def test_sdxl_cache_built_under_budget(monkeypatch):
+    # Default (4 GiB) budget: a handful of tiny latents fits, so the full cache is returned.
+    monkeypatch.delenv("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", raising = False)
+    cache = _build_fake_sdxl_cache(monkeypatch, num_images = 3, latent_shape = (1, 4, 8, 8))
+    assert cache is not LATENT_CACHE_OVER_BUDGET and cache is not None
+    assert len(cache) == 3
+    assert all(len(variants) == 1 for variants in cache)
+
+
+def test_sdxl_cache_gated_over_budget(monkeypatch):
+    # A budget below one variant forces the gate on the first encode: the sentinel is returned
+    # so the caller keeps the VAE resident and encodes per step.
+    monkeypatch.delenv("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", raising = False)
+    monkeypatch.setattr(train_common, "_LATENT_CACHE_BUDGET_BYTES", 8)
+    cache = _build_fake_sdxl_cache(monkeypatch, num_images = 3, latent_shape = (1, 4, 8, 8))
+    assert cache is LATENT_CACHE_OVER_BUDGET
+
+
+def test_sdxl_cache_force_bypasses_gate(monkeypatch):
+    # An explicit force-on must be honoured verbatim even when the estimate is over budget.
+    monkeypatch.setenv("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", "1")
+    monkeypatch.setattr(train_common, "_LATENT_CACHE_BUDGET_BYTES", 8)
+    cache = _build_fake_sdxl_cache(monkeypatch, num_images = 3, latent_shape = (1, 4, 8, 8))
+    assert cache is not LATENT_CACHE_OVER_BUDGET and cache is not None
+    assert len(cache) == 3

studio/backend/tests/test_diffusion_training.py

@@ -190,6 +190,45 @@ def test_apply_event_transitions():
     assert svc.status()["status"] == "error" and svc.status()["message"] == "boom"
 
 
+def test_progress_nulls_non_finite_floats_for_strict_json():
+    # A divergent step (or an inf grad norm) can push loss / avg_loss / learning_rate to
+    # NaN or Infinity, which strict JSON forbids. The service must null those so the status
+    # snapshot and the metric history stay strict-JSON serializable.
+    import json
+    import math
+
+    svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
+    svc._apply_event(
+        {
+            "type": "progress",
+            "step": 3,
+            "total_steps": 10,
+            "loss": float("nan"),
+            "avg_loss": float("inf"),
+            "learning_rate": float("-inf"),
+        }
+    )
+    snap = svc.status()
+    assert snap["loss"] is None
+    assert snap["avg_loss"] is None
+    assert snap["learning_rate"] is None
+    # The non-finite point is skipped in the history, so the loss series stays clean.
+    assert snap["metric_loss"] == []
+    assert snap["metric_steps"] == []
+    # strict JSON (allow_nan=False) round-trips without a ValueError from NaN/Infinity.
+    json.dumps(snap, allow_nan = False)
+
+    # A finite point after the bad one is recorded and preserved verbatim.
+    svc._apply_event(
+        {"type": "progress", "step": 4, "total_steps": 10, "loss": 0.5, "learning_rate": 1e-4}
+    )
+    snap2 = svc.status()
+    assert snap2["loss"] == 0.5
+    assert snap2["metric_loss"] == [0.5] and snap2["metric_steps"] == [4]
+    assert math.isfinite(snap2["learning_rate"])
+    json.dumps(snap2, allow_nan = False)
+
+
 def test_terminal_events_clear_model_load_flag():
     # A stop or error during model load emits complete/error WITHOUT a preceding
     # model_load_completed, so the terminal update must reset in_model_load or the
@@ -211,6 +250,7 @@ class _FakeService:
     def __init__(self):
         self._running = False
         self.started_with = None
+        self.stopped_with_save = None
         # Extra keys merged into status() so a test can inject metric history / perf fields.
         self.status_extra: dict = {}
 
@@ -219,7 +259,8 @@ class _FakeService:
         self._running = True
         return "job-123"
 
-    def stop(self):
+    def stop(self, save = True):
+        self.stopped_with_save = save
         was = self._running
         self._running = False
         return was
@@ -564,6 +605,38 @@ def test_route_start_refuses_non_sdxl_base_without_freeing_gpu(client, monkeypat
     assert client._fake.started_with is None
 
 
+def test_route_start_refuses_non_bf16_gpu_without_freeing_gpu(client, monkeypatch):
+    # A DiT precision the host cannot run (no bf16 GPU, or explicit int8 without a functional
+    # torchao) must 400 BEFORE resident GPU workloads are freed: otherwise the host tears down the
+    # user's chat/Images model and the run then dies in the trainer child. The route imports
+    # training_precision_preflight_error locally, so patch it on its home module.
+    import routes.training as tr
+
+    freed = []
+    monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: freed.append(1))
+    monkeypatch.setattr(
+        "core.training.diffusion_train_common.training_precision_preflight_error",
+        lambda fam, prec: (
+            "This trainer requires a bfloat16-capable GPU (Ampere or newer)."
+            if fam != "sdxl"
+            else None
+        ),
+    )
+    r = client.post(
+        "/api/train/diffusion/start",
+        json = {**_BODY, "base_model": "black-forest-labs/FLUX.1-dev"},
+    )
+    assert r.status_code == 400
+    assert "bfloat16" in r.json()["detail"]
+    assert freed == []
+    assert client._fake.started_with is None
+
+    # SDXL (its own mixed_precision path) is exempt: the same probe returns None, so an SDXL
+    # start proceeds normally past the preflight.
+    r2 = client.post("/api/train/diffusion/start", json = _BODY)
+    assert r2.status_code == 200, r2.text
+
+
 # ── metric history + perf/family fields (PR A platform) ──────────────────────
 def test_apply_event_records_metric_history_and_perf():
     svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)