FLUX.1-dev, Qwen-Image, and Z-Image LoRA training (#6820)

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

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

* Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override

Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want
fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and
prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match
on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40;
GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it.

Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s),
no overflow, quality below the quant noise floor. So the default leans to accuracy on
data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force)
lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off).

187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override
threading).

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

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

* Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9

scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer.
Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with
use_triton_kernel=False (the default triton path needs the missing MSLK library), but
only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than
using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's
CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto
on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast.

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

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

* Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder

Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4
tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS,
2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the
diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both
slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166
vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is
correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add
scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end).

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

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

* Studio diffusion (Phase 9): pre-quantized transformer loading

The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on
the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs
13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline
(scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton
on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True)
the quantized weights, so the dense bf16 never touches the GPU.

Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching
GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same
torchao config + min_features filter the runtime path uses, applied ahead of time.

New core/inference/diffusion_prequant.py (resolve_prequant_source +
load_prequantized_transformer, best-effort, lazy imports). diffusion.py
_load_dense_quant_pipeline tries the pre-quant source first and falls back to the
dense materialise+quantise path, then to GGUF, so the default is unchanged.
DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty
prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for
the resolver, the meta-init+assign loader, and the backend branch selection +
fallbacks; GPU verification via scripts/verify_prequant_backend.py.

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

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

* Studio diffusion (Phase 10): attention-backend selection

Add a selectable attention kernel via the diffusers set_attention_backend
dispatcher. Attention is memory-bandwidth bound, so a better kernel is an
end-to-end win orthogonal to the linear-weight quantisation (it speeds the QK/PV
matmuls torchao never touches) and composes with torch.compile.

auto picks the best exact backend for the device: cuDNN fused attention
(_native_cudnn) on NVIDIA when a speed profile is active, measured ~1.18x
end-to-end on a B200 (Z-Image 1024px/8 steps) with LPIPS ~0.004 vs the default
(below the compile/quant noise floor); native SDPA elsewhere and when speed=off
(so off stays bit-identical). Explicit native/cudnn/flash/flash3/flash4/sage/
xformers/aiter are honored, and an unavailable kernel falls back to the default
rather than failing the load.

New core/inference/diffusion_attention.py (normalize + per-device select + apply,
best-effort, lazy imports). Set on pipe.transformer BEFORE compile in load_pipeline;
attention_backend threads through begin_load / load_pipeline / status like the other
load knobs. New request field attention_backend + status field. Hermetic CPU tests
for normalize / select policy / apply fallback, plus route threading + 422. Measured
via scripts/perf_levers_probe.py.

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

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

* Studio diffusion (Phase 11): prefer int8 on consumer GPUs in the auto ladder

Consumer / workstation GPUs halve fp8 (and fp16/bf16) FP32-accumulate tensor-core
throughput, while int8 runs at full rate (int32 accumulate is not nerfed). Public
benchmarks (SDNQ across RTX 3090/4090/5090, AMD, Intel) confirm int8 via torch._int_mm
is as fast or faster than fp8 on every consumer part, and the only path on pre-Ada
consumer cards without fp8 tensor cores. So when transformer_quant=auto, reorder the
arch tier to put int8 first on a consumer/workstation GPU (detected by the existing
_is_consumer_gpu name heuristic), while data-center HBM parts keep fp8 first.

Pure ladder reorder via _prefer_consumer_scheme; no new flags. Verified non-regression
on a B200 (still picks fp8). Hermetic tests for consumer Blackwell/Ada/workstation
(-> int8) and data-center Ada/Hopper/Blackwell (-> fp8).

* Studio diffusion (Phase 12): First-Block-Cache step caching for many-step DiT

Add opt-in step caching (First-Block-Cache) for the diffusion transformer. Across
denoise steps a DiT's output settles, so once the first block's residual barely
changes the remaining blocks are skipped and their cached output reused. diffusers
ships it natively (FirstBlockCacheConfig + transformer.enable_cache, with the
standalone apply_first_block_cache hook as a fallback).

Measured on Flux.1-dev (28 steps, 1024px): ~1.4x on top of torch.compile (2.83 ->
2.03s) at LPIPS ~0.08 vs the no-cache output, well inside the quality bar.

OFF by default and a per-load opt-in: the win scales with step count, so it is for
many-step models (Flux / Qwen-Image) and pointless for few-step distilled models
(e.g. Z-Image-Turbo at ~8 steps), where a single skipped step is a large fraction
of the trajectory. It composes with regional compile only with fullgraph=False (the
cache's per-step decision is a torch.compiler.disable graph break), which the speed
layer now switches to automatically when a cache is engaged. Best-effort: a model
whose block signature the hook does not recognise is caught and the load proceeds
uncached.

- new core/inference/diffusion_cache.py: normalize_transformer_cache + apply_step_cache
  (enable_cache / apply_first_block_cache fallback; threshold auto-raised for a
  quantised transformer per ParaAttention's fp8 guidance; lazy diffusers import).
- diffusion_speed.py: apply_speed_optims takes cache_active; compile drops fullgraph
  when a cache is engaged.
- diffusion.py: apply_step_cache before compile; thread transformer_cache /
  transformer_cache_threshold through begin_load -> load_pipeline and report the
  engaged mode in status().
- models/inference.py + routes/inference.py: transformer_cache (off | fbcache) and
  transformer_cache_threshold request fields, engaged mode in the status response.
- hermetic tests for normalisation, the enable_cache / hook-fallback paths, threshold
  selection, and best-effort failure handling, plus route threading + validation.
- scripts/fbcache_flux_probe.py: the Flux validation probe (latency / speedup / VRAM /
  LPIPS vs the compiled no-cache baseline).

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

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

* Studio diffusion (Phase 14): fix int8 dense quant on Flux / Qwen (skip M=1 modulation linears)

The opt-in dense int8 transformer path crashed on Flux.1 and Qwen-Image with
'torch._int_mm: self.size(0) needs to be greater than 16, but got 1'. int8 dynamic quant
goes through torch._int_mm, which requires the activation row count M > 16. A DiT's AdaLN
modulation projections (Flux norm1.linear 3072->18432, Qwen img_mod.1 / txt_mod.1, Flux.2
*_modulation.linear) and its timestep / guidance / pooled-text conditioning embedders are
computed once from the [batch, dim] conditioning vector (M = batch = 1), not per token, so
they hit _int_mm at M=1 and crash. Their feature dims are large, so the existing
min_features filter did not exclude them.

Fix: the int8 filter now also skips any Linear whose fully-qualified name matches a
modulation / conditioning-embedder token (norm, _mod, modulation, timestep_embed,
guidance_embed, time_text_embed, pooled). These layers run at M=1 once per block and are a
negligible share of the FLOPs, so int8 keeps the full speedup on the attention / FFN layers
(M = sequence length). fp8 / nvfp4 / mxfp8 use scaled_mm, which has no M>16 limit and
quantises these layers fine, so the exclusion is int8-only. Sequence embedders
(context_embedder / x_embedder / txt_in, M = seq) are deliberately not excluded -- note
'context_embedder' contains the substring 'text_embed', which is why the token is the
specific 'time_text_embed', not 'text_embed'.

Measured on a B200 (1024px, transformer_quant=int8 + speed=default), int8 now runs on every
supported model and is the fastest dense path on Flux/Qwen (int8 runs full-rate vs fp8's
FP32-accumulate): FLUX.1-dev 9.62s eager -> 1.98s (4.86x, vs fp8 2.15s), Qwen-Image -> 1.87s
(5.57x, vs fp8 2.09s), FLUX.1-schnell -> 0.41s (3.59x). Z-Image and Flux.2-klein (already
working) are unchanged.

- diffusion_transformer_quant.py: add _INT8_EXCLUDE_NAME_TOKENS; make_filter_fn takes
  exclude_name_tokens; quantize_transformer passes it for int8 only.
- hermetic test that the int8 filter excludes the modulation / embedder linears (and keeps
  attention / FFN / sequence-embedder linears), while fp8 keeps them.
- scripts/int8_linear_probe.py: the meta-device probe used to enumerate each transformer's
  Linear layers and derive the exclusion list.

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

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

* Studio diffusion (Phase 15): build int8 pre-quantized checkpoints (skip M=1 modulation linears)

The prequant-checkpoint builder applied the dense quant filter without the int8-only
M=1 modulation / conditioning-embedder exclusion the runtime path uses, so a built int8
checkpoint baked those projections as int8 and crashed (torch._int_mm needs M>16) at the
first denoise step on Flux / Qwen. Factor the scheme->exclusion decision into a shared
exclude_tokens_for_scheme() used by both the runtime quantise path and the offline builder
so they can never drift, and apply it in build_prequant_checkpoint.py. int8 prequant now
produces a working checkpoint on every supported model, giving int8 (the consumer-preferred
scheme) the same ~2x load-VRAM and download reduction fp8 already had.

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

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

* Studio diffusion (Phase 16): route no-GPU loads to the native sd.cpp engine

When no CUDA/ROCm/XPU GPU is available, route diffusion load/generate to the
native stable-diffusion.cpp engine instead of diffusers, with diffusers as the
guaranteed fallback. On CPU sd.cpp is 1.4-2.8x faster and uses 1.5-2.2x less RAM.

- diffusion_engine_router: centralised engine selection (built on the existing
  select_diffusion_engine), env opt-outs, MPS gating, recorded fallback reason.
- sd_cpp_backend (SdCppDiffusionBackend): the diffusers backend method surface
  backed by sd-cli, with lazy binary install, registry-driven asset fetch,
  step-progress parsing, and cancellation.
- diffusion_families: per-family single-file VAE + text-encoder asset mapping.
- sd_cpp_engine: cancellation support (process-group kill + SdCppCancelled).
- routes/inference + gpu_arbiter: drive the active engine via the router; the
  API now reports the active engine and any fallback reason.
- tests for the backend, router, route selection, and cancellation.

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

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

* Phase 16 review fixes: engine-switch unload, sd.cpp error mapping, per-image seeds, Qwen sampler

Address review feedback on #6724:
- engine router: unload the engine being deactivated on a switch, so the old
  model is not left resident-but-unreachable (the evictor only targets the active
  engine).
- generate route: sd.cpp execution errors (nonzero exit / timeout / missing
  output) now map to 500, not 409 (which only means not-loaded / cancelled).
- native batch: return per-image seeds and persist the actual seed for each image
  so every batch image is reproducible.
- Qwen-Image native path: apply --sampling-method euler --flow-shift 3 per the
  stable-diffusion.cpp docs; other families keep sd-cli defaults.
- honor speed_mode (native --diffusion-fa) and, off-CPU, memory_mode/cpu_offload
  offload flags on the native load instead of hardcoding them off.
- fail the load when the sd-cli binary is present but not runnable (version()
  now returns None on exec error / nonzero exit).
- size estimate: only treat the transformer asset as a possible local path.

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

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

* Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in

load_prequantized_transformer ends in torch.load(weights_only=False), which executes
arbitrary code from the pickle. The transformer_prequant_path load-request field reached
that unpickle for any local file an authenticated caller named, so a request could trigger
remote code execution. Refuse the source.kind=='path' branch unless the operator sets
UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted
and unaffected. Document the requirement on the API field and add gate tests.

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

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

* Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa

- apply_attention_backend now restores the native default when no backend is requested or a
  kernel fails. diffusers keeps a process-wide active attention backend that
  set_attention_backend updates, and a fresh transformer's processors follow it, so a load
  that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile
  load pinned, breaking the bit-identical/off guarantee.
- select_attention_backend drops flash3/flash4 up front when the CUDA capability is below
  Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit
  request on the wrong card set fine then crashed mid-generation; it now falls back to native.
- Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a
  valid alias of native) is accepted instead of 422-rejected by Pydantic.
- Drop the dead replace('-','_') normalization (no alias uses dashes/underscores).
- perf_levers_probe.py output dir is now relative to the script, not a hardcoded path.

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

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

* Studio diffusion (Phase 12): only engage FBCache on context-aware transformers; quantized threshold for GGUF

- apply_step_cache now engages only via the transformer's native enable_cache (the diffusers
  CacheMixin path), which exists exactly when the pipeline wraps the transformer call in a
  cache_context. The standalone apply_first_block_cache fallback installed on non-CacheMixin
  transformers too (e.g. Z-Image), whose pipeline opens no cache_context, so the load reported
  transformer_cache=fbcache and then the first generation crashed inside the hook. Such a model
  now runs uncached per the best-effort contract.
- GGUF transformers are quantized (the default Studio load path), so they now use the higher
  quantized FBCache threshold when the caller leaves it unset, instead of the dense default
  that could keep the cache from triggering.
- fbcache_flux_probe.py: compile cached runs with fullgraph=False (FBCache is a graph break, so
  fullgraph=True failed warmup and silently measured an eager cached run); output dir is now
  relative to the script, not a hardcoded path.

* Studio diffusion (Phase 11): keep professional RTX cards on the fp8 ladder

_is_consumer_gpu treated professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) as
consumer because their names carry no datacenter token, so the auto ladder moved int8 ahead
of fp8 and the fp8 path chose fast accumulate for them. The rest of the backend already
classifies these as datacenter/professional (llama_cpp.py _DATACENTER_GPU_RE), so detect the
same RTX PRO 6000 / RTX 6000 Ada markers here and keep fp8 first with precise accumulate.

Also fix the consumer-Blackwell test to use compute capability (10, 0) instead of (12, 0).

* Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config

Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just
TypeError on older torchao), which would break the mxfp8 config helper instead of falling
back to the default. Catch both so the fallback is robust.

quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds
CUDA memory during the per-row VRAM probe; output dir relative to the script.

* Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load

- snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform
  missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the
  whole snapshot. restore_backend_flags restores each flag independently so one failure can't
  leave the others leaked process-wide.
- load_pipeline restores the flags (and clears the GPU cache) when the build fails after
  apply_speed_optims mutated the process-wide flags but before _state captured them for unload
  to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and
  contaminated later off generations.

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

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

* Studio diffusion (Phase 4): enforce the sd-cli timeout while reading output

Iterating proc.stdout directly blocks until the stream closes, so a sd-cli that hangs
without producing output (or without closing stdout) would never reach proc.wait and the
wall-clock timeout was silently bypassed. Drain stdout on a daemon thread and wait on the
PROCESS, so the main thread always enforces the timeout and kills a hung process (which
closes the pipe and ends the reader). Add a test that times out even when stdout blocks,
and make the no-binary test hermetic so a host-installed sd-cli can't leak in.

* Studio diffusion (Phase 14): guard the int8 exclusion filter against a None fqn

The filter callback can be invoked without a module name, so fqn.lower() would raise
AttributeError on None. Fall back to an empty name (nothing matches the exclusion tokens,
so the linear is kept) instead of crashing the quantise pass.

* Studio diffusion (Phase 16) review fixes: native engine robustness

- sd_cpp_backend: stop truncating explicit seeds to 53 bits (mask to int64);
  a large requested seed was silently collapsed (2**53 -> 0) and distinct seeds
  aliased to the same image. Random seeds stay 53-bit (JS-safe).
- sd_cpp_backend: sanitize empty/whitespace hf_token to None so HfApi/hf_hub
  fall back to anonymous instead of failing auth on a blank token.
- sd_cpp_backend: a superseding load now cancels the in-flight generation, so the
  old sd-cli can no longer return/persist an image from the previous model.
- diffusion_engine_router: run the previous engine's unload() OUTSIDE the lock so a
  slow 10+ GB free / CUDA sync does not block engine selection.
- diffusion_engine_router: probe sd-cli runnability (version()) before committing to
  native, so a present-but-unrunnable binary falls back to diffusers at selection.
- diffusion_device: resolve a torch-free CPU target when torch is unavailable, so a
  CPU-only install can still reach the native sd.cpp engine instead of failing load.
- tests updated for the runnability probe + a not-runnable fallback case.

* Studio diffusion (Phase 9) review fixes: prequant safety + validation

- SECURITY: a request-supplied local pre-quant path is now unpickled only when it
  resolves inside an operator-configured ALLOWLIST of directories
  (UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in,
  once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on
  any path a load request named (arbitrary code execution). realpath() blocks symlink
  escapes; a bare on/off toggle is no longer a wildcard.
- Validate the checkpoint's min_features against the runtime Linear filter, so a
  checkpoint that quantised a different layer set is rejected instead of silently
  loading a model that mismatches the dense path while reporting the same scheme.
- Tolerant base_model_id compare (exact or same final path/repo segment), so a local
  path or fork of the canonical base is accepted instead of falling back to dense.
- _has_meta_tensors uses any(chain(...)) (no intermediate lists).
- prequant verify/probe scripts use repo-relative paths (+ env overrides), not the
  author's absolute /mnt paths.
- tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail.

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

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

* Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts

- diffusion_memory: when group offload is unavailable and the plan falls back to
  whole-module offload, enable VAE tiling (the group plan left it off, but the fallback
  is the low-VRAM path where the decode spike can OOM). Covers both the group and
  sequential fallback branches.
- perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a
  balanced bit-identity regression actually fails the check instead of exiting 0.
- compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and
  the load-progress poll has a 30 min deadline instead of looping forever on a hang.
- test for the group->model fallback enabling VAE tiling.

* Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path

- diffusion: a torchao-quantized transformer is committed only compiled. A dense model
  resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF
  it replaced), so when transformer_quant engaged and speed resolved to off, promote to
  default (regional compile); warn loudly if compile still does not engage.
- diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the
  CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK);
  otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF.
- nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path).
- test asserts the eager-quant -> default-compile promotion.

* Studio diffusion (Phase 10) review fixes: attention gating + probe isolation

- diffusion_attention: gate the auto cuDNN-attention upgrade on SM80+; on pre-Ampere
  NVIDIA (T4/V100) cuDNN fused SDPA is accepted at set time but fails at first generation,
  so auto now stays on native SDPA there.
- diffusion_attention: _active_attention_backend handles get_active_backend() returning an
  enum/None (not a tuple); the old  unpack always raised and was swallowed, so
  the native-restore short-circuit never fired.
- perf_levers_probe: free the resident pipe on a skipped (attn/fbcache) variant; run LPIPS
  on CPU so it isn't charged to every variant's peak VRAM; reset force_fuse_int_mm_with_mul
  so the inductor_flags variant doesn't leak into later compiled rows.
- tests for the SM80 cuDNN gate.

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

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

* Studio diffusion (Phase 4) review fixes: sd.cpp installer + engine hardening

- install_sd_cpp_prebuilt: download the release archive with urlopen + an explicit
  timeout + copyfileobj (urlretrieve has no timeout and hangs on a stalled socket);
  extract through a per-member containment check (Zip-Slip guard); expanduser the
  --install-dir so a tilde path is not taken literally; and on Windows CUDA also fetch
  the separately-published cudart runtime DLL archive so sd-cli.exe can start.
- sd_cpp_engine: find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the
  installer, so a custom-root install is discovered without UNSLOTH_SD_CPP_PATH; start
  sd-cli with the parent-death child_popen_kwargs so it is not orphaned on a backend
  crash; reap the SIGKILLed child (proc.wait) so a cancel/timeout does not leave a zombie.
- tests: Zip-Slip rejection, normal extraction, studio-home discovery.

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

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

* Studio diffusion (Phase 4) review round 2: collect sd-cli batch outputs

Codex review: when batch_count > 1, stable-diffusion.cpp's save_results() writes
the numbered files <stem>_<idx><suffix> (base_0.png, base_1.png, ...) instead of
the literal --output path. SdCppEngine.generate checked only the literal path, so
a batch generation would exit 0 and then raise 'no image' (or return a stale
file). generate now returns the literal path when present and otherwise falls
back to the numbered siblings; single-image behavior is unchanged.

Test: a fake sd-cli that writes img_0.png/img_1.png (not img.png) is collected
without error.

* Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats

Codex review on the native engine arg builder:

- build_sd_cpp_command emitted --width/--height unconditionally, so an
  img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of
  the input. width/height are now Optional (None = unset): an image-conditioned
  run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives
  the size from the input image (set_width_and_height_if_unset); a plain txt2img
  run with unset dims keeps the prior 1024x1024 default; explicit dims are always
  honored. width/height are read only by the builder, so the type change is local.

- build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...)
  that silently swallowed repeats=0 into sd-cli's default of one pass, turning an
  explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError
  and emits the flag for any explicit value != 1.

Tests: img2img unset dims omit width/height (init_img and ref_images), explicit
dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at
the default. (Two pre-existing binary-discovery tests fail only because a real
sd-cli is installed in this dev environment; unrelated to this change.)

* Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc

Codex review: the transformer_prequant_path field description still told operators
to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the
prior security fix made that variable a directory allowlist -- _allowed_prequant_roots
deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator
following the documented =1 would have every transformer_prequant_path request
silently refused. The description now states it must name one or more allowlisted
directories and that a bare on/off value is not accepted.

Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does
not say =1, and describes an allowlist/directory (guards against doc drift).

* Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset

Codex review on attention-backend selection:

- Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on
  pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first
  generation with no fallback. select_attention_backend now applies
  _cudnn_attention_supported() to an explicit cuDNN request too.

- flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a
  Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3
  is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a
  (min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+.

- apply_attention_backend's success path left diffusers' process-wide active
  backend pinned to the kernel it set; a later component whose processors are
  unconfigured (backend None) would inherit it. It now resets the global registry
  to native after a successful per-transformer set (the transformer keeps its own
  backend), best-effort. Also fixed _active_attention_backend: get_active_backend()
  returns a (name, fn) tuple, so the prior code stringified the tuple and never
  matched a name, defeating the native-restore short-circuit.

Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on
SM90; global registry reset after a successful set; _active_attention_backend
reads the tuple return.

* Studio diffusion (Phase 11) review round 2: keep GH200/B300 on the fp8 ladder

Codex review: _DATACENTER_GPU_TOKENS omitted GH200 (Grace-Hopper) and B300
(Blackwell Ultra), though it has the distinct GB200/GB300 superchip tokens. So
_is_consumer_gpu returned True for 'NVIDIA GH200 480GB' / 'NVIDIA B300', and the
auto ladder moved int8 ahead of fp8 on those data-center parts -- contradicting
llama_cpp.py's datacenter regex, which lists both. Added GH200 and B300 so they
are treated as data-center class and keep the intended fp8-first behavior.

Test: extends the datacenter parametrize with 'NVIDIA B300' and
'NVIDIA GH200 480GB' (now _is_consumer_gpu False).

* Studio diffusion (Phase 14) review round 2: apply int8 M=1 exclusion in the builder

Codex review: the M=1 modulation/embedder exclusion was wired only into the dense
runtime quantiser; the offline builder scripts/build_prequant_checkpoint.py called
make_filter_fn(min_features) with no exclusion. So an int8 prequant checkpoint
quantised the AdaLN modulation and conditioning-embedder linears, and loading it
via transformer_prequant_path (the load path only loads already-quantised tensors,
it can't re-skip them) reintroduced the torch._int_mm M=1 crash this phase fixes
for the runtime path.

Extracted int8_exclude_name_tokens(scheme) as the single source of truth (int8 ->
the M=1 exclusion, every other scheme -> none) and use it in both the runtime
quantiser and the builder, so a prequant artifact's quantised-layer set always
matches the runtime. fp8/fp4/mx artifacts are byte-identical (empty exclusion).

Test: int8_exclude_name_tokens returns the exclusion for int8 and () for
fp8/nvfp4/mxfp8.

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

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

* Studio diffusion (Phase 16) review round 2: native CPU arbiter, status offload, load race

Codex review on the native-engine routing:

- The /images/load route took the GPU arbiter (acquire_for(DIFFUSION) -> evict chat)
  unconditionally after engine selection. A native sd.cpp load on a pure-CPU host
  never touches the GPU, so that needlessly tore down the resident chat model. The
  handoff is now gated: diffusers always takes it, a force-native sd.cpp load on a
  CUDA/XPU/MPS box still takes it, but a native sd.cpp load on a CPU host skips it.

- sd_cpp status() hardcoded offload_policy 'none' / cpu_offload False even when
  _run_load computed real offload flags (balanced/low_vram/cpu_offload off-CPU), so
  the setting was unverifiable. status now derives them from state.offload_flags
  (still 'none' on CPU, where the flags are empty).

- _run_load committed the new state without cancelling/waiting on a generation that
  started during the (slow) asset download, so a stale sd-cli run against the OLD
  model could finish afterward and persist an image from the previous model once the
  new load reported ready. The commit now signals the in-flight cancel and waits on
  _generate_lock before swapping _state (taken only at commit, so the download never
  serialises against generation), mirroring the diffusers load path.

Tests: CPU native load skips the arbiter while a GPU native load takes it; status
reports offload active when flags are set; _run_load cancels and waits for an
in-flight generation before committing.

* Studio diffusion (Phase 14) review round 2: align helper name with the stack

Rename the int8 exclusion helper to exclude_tokens_for_scheme, matching the
identical helper already present higher in the diffusion stack (Phase 16). The
helper definition, the runtime quantiser call, and the offline builder are now
byte-identical to that version, so the two branches no longer introduce a
divergent name for the same single-source-of-truth and the stack merges without
a conflict on this fix. No behavior change.

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

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

* Studio diffusion: eager patches + torch.compile cache speed phase

Adds the opt-in speed path for the GGUF diffusion transformer behind a
selectable speed mode (default off, so output is unchanged until a profile
is chosen):

- diffusion_eager_patches.py: shared eager fast-paths (channels_last,
  attention/backend selection, fused norms and QKV) installed at load and
  rolled back on unload or failed load.
- diffusion_compile_cache.py / diffusion_gguf_compile.py: a persistent
  torch.compile cache and the GGUF-transformer compile wiring.
- diffusion_arch_patches.py: architecture-specific patches.
- diffusion_patch_backend.py: shared install/restore plumbing.
- diffusion_speed.py: speed-profile planning.

Tests for each module plus the benchmarking and probe scripts used to
measure speed, memory, and accuracy of the path.

* Studio diffusion: image workflows (safetensors, image-conditioned, editing) + Images UI

Backend:
- Load non-GGUF safetensors models: full bnb-4bit pipelines and single-file
  fp8 transformers, gated to the unsloth org plus a curated allowlist.
- Image-conditioned workflows built with Pipeline.from_pipe so they reuse the
  loaded transformer/VAE/text-encoder with no extra VRAM: img2img, inpaint,
  outpaint, and a hires-fix upscale pass.
- Instruction editing as its own family kind (Qwen-Image-Edit-2511,
  FLUX.1-Kontext-dev) and FLUX.2-klein reference conditioning (single and
  multi-reference) plus klein inpaint.
- Auto-resize odd-sized inputs to a multiple of 16 (and resize the matched
  mask) so img2img/inpaint/edit no longer reject non-/16 uploads. Bound the
  decoded image size and cap upscale output to avoid OOM on large inputs.
- Fixes: from_pipe defaulting to a float32 recast that crashed torchao
  quantized transformers; image-conditioned calls forcing the slider size
  onto the input image. Native sd.cpp engine rejects image-conditioned and
  reference requests it cannot serve.

Frontend:
- Redesigned Images page with capability-gated workflow tabs (Create,
  Transform, Inpaint, Extend, Upscale, Reference, Edit), a brush mask editor,
  client-side outpaint, and a multi-reference picker.
- Advanced options moved to a right-docked panel mirroring Chat: closed by
  default, toggled by a single fixed top-bar button that stays in place.

sd.cpp installer: pin the release, verify each download's sha256, add a
download timeout, and make the source repo configurable for a future mirror.

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

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

* Studio Images: correct the Advanced panel comment (closed by default, fixed toggle)

* Studio diffusion: LoRA adapters for the Images workflow

Add community LoRA support across both diffusion backends, the single
biggest step toward broad image-workflow coverage.

Backend
- New shared module core/inference/diffusion_lora.py: adapter discovery
  (local scan + curated catalog + owner/name[:file] Hub refs), download
  via hf_hub_download_with_xet_fallback, alias sanitization, native
  managed-dir materialization with collision-broken aliases, prompt-tag
  injection (deduped against user-typed tags), and a supports_lora gate.
- Native sd-cli: resolve + materialize selected LoRAs into a per-run
  managed dir, inject <lora:ALIAS:w> tags, pass --lora-model-dir with
  --lora-apply-mode auto. The arg builder already emitted these flags.
- Diffusers: non-fused load_lora_weights + set_adapters manager, tracked
  on the pipe so an unchanged selection is a no-op and a model swap
  resets; cleared on unload. Never fuses (breaks quantized transformers
  and blocks live weight tweaks).
- Gated off where unsupported: torchao fp8/int8 dense, GGUF-via-diffusers,
  and native Qwen-Image (no LoRA name-conversion branch upstream).
- Request contract: optional loras on DiffusionGenerateRequest; empty or
  omitted is identical to today. supports_lora surfaced in status; chosen
  LoRAs persisted in gallery recipe metadata.
- New GET /api/models/diffusion-loras for the picker (family-filtered).

Frontend
- Repeatable multi-LoRA picker (adapter select + weight slider 0..2 +
  remove), gated by the loaded model's supports_lora and family, max 8.

Tests
- New test_diffusion_lora.py (14): helpers, request validation, native
  tag/dir wiring, diffusers set_adapters manager, supports_lora matrix.

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

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

* Studio diffusion: ControlNet for the Images workflow (diffusers)

Add ControlNet conditioning, the #2 most-used diffusion workflow after
LoRA, on the diffusers backend for the families with ControlNet pipelines
(FLUX.1 and Qwen-Image), with Union models as the default picks.

Backend
- New core/inference/diffusion_controlnet.py: family-gated discovery
  (curated Union models + local dirs + bare owner/name repos), resolution
  to a loadable repo/dir, control-image preprocessing (passthrough +
  a dependency-free canny edge map), and a supports_controlnet gate.
- diffusion.py: a ControlNet manager parallel to the LoRA one. Loads the
  (small) ControlNet model once via from_pretrained (cached by id) and
  builds the family's ControlNet pipeline via Pipeline.from_pipe(base,
  controlnet=model), reusing the resident base modules at their loaded
  dtype (no reload, no recast). Passes the control image + conditioning
  scale + guidance start/end at generate time; cleared on unload.
- Families: FLUX.1 -> FluxControlNetPipeline/Model, Qwen-Image ->
  QwenImageControlNetPipeline/Model. Others declare none (gated off).
- Gated off for the native engine, GGUF-via-diffusers, and torchao
  fp8/int8 dense (same rule as LoRA). v1 conditions txt2img only.
- Request contract: optional controlnet on DiffusionGenerateRequest;
  supports_controlnet in status; the choice persisted in gallery meta.
- New GET /api/models/diffusion-controlnets for the picker.

Frontend
- A ControlNet control in the Images rail (model select + control-image
  upload + control-type select + strength slider), gated by the loaded
  model's supports_controlnet + family, shown for text-to-image.

Tests
- New test_diffusion_controlnet.py (10): discovery/resolve/preprocess/gate
  helpers, request validation, family wiring, and the diffusers pipe
  manager (loads once, caches, from_pipe with controlnet, rejects
  unsupported families).

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

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

* Studio ControlNet: show the picker on the Create tab (workflow id is 'create', not 'txt2img')

The ControlNet control gated on workflow === 'txt2img', but the Images workflow tab ids are create/transform/inpaint/extend/upscale/reference/edit -- there is no 'txt2img'. So the picker never rendered even with a ControlNet-capable model loaded. Gate on 'create' (the text-to-image tab) for both the picker and the request wiring. Found via a live Playwright capture of the running Studio.

* Studio: do not force diffusers pipelines cross-tagged gguf into the GGUF variant expander

Some diffusers image repos (e.g. unsloth/Qwen-Image-2512-unsloth-bnb-4bit) carry a
stray "gguf" tag on the Hub but ship no .gguf files. The model search classified
them as GGUF from the bare tag, so the picker rendered the GGUF variant expander,
which then dead-ended at "No GGUF variants found." Trust the bare gguf tag only when
the repo is not a diffusers pipeline; the -GGUF name suffix and real gguf metadata
(populated via expand=gguf) remain authoritative, so genuine GGUF repos are unaffected.

* Studio Images: load non-curated unsloth/on-device diffusers repos instead of no-op

handleModelSelect only loaded curated safetensors ids and GGUF variant picks; any other
non-GGUF pick (an on-device diffusers folder, or a future unsloth diffusers image repo
surfaced by search) silently did nothing. Treat such a pick as a full diffusers pipeline
load when the id is unsloth-hosted or on-device (the backend infers the family + base repo
and gates loads to unsloth/* or local paths), and show a clear message otherwise instead
of silently ignoring the click. Curated and GGUF paths are unchanged.

* Studio Images: keep curated safetensors models in Recommended after download

The curated bnb-4bit / fp8 diffusion rows were filtered out of the Images picker's
Recommended list once cached (curatedSafetensorsRows dropped anything in downloadedSet),
so they vanished from the picker after the first load and could only be found by typing an
exact search. The row already renders a downloaded badge, matching how GGUF Recommended
rows stay visible when cached. Drop the exclusion so the curated safetensors always list.

* Studio diffusion LoRA: sanitize dots out of adapter aliases

The LoRA alias is used as the diffusers PEFT adapter name, and PEFT rejects names
containing "." (module name can't contain "."). sanitize_alias kept dots, so a LoRA whose
filename carries a version tag (e.g. Qwen-Image-2512-Lightning-8steps-V1.0-bf16) failed to
apply with a 400. Replace dots too; the alias stays a valid native <lora:NAME:w> filename
stem. Adds regression coverage for internal dots.

* Studio Images: clarify the GGUF transformer-quant Advanced control

Renamed the confusing "Transformer quant / GGUF default" control to "GGUF speed mode"
with an "Off (run the GGUF)" default, and reworded the hint to state plainly that FP8/INT8/
FP4 load the FULL base model (larger download + more VRAM) rather than re-packing the GGUF,
falling back to the GGUF if it can't fit. Behavior unchanged; labels/hint only.

* Studio Images: list on-device unsloth diffusion models in the picker

The Images picker's On Device tab hid every non-GGUF cached repo whenever a
task filter was active, so downloaded unsloth diffusion pipelines (bnb-4bit
and FP8 safetensors) never showed up there. List cached repos that pass the
task gate, limited under a filter to unsloth-hosted ones so base repos (which
fail the diffusion load trust gate) don't appear only to dead-end on click.
Chat behavior is unchanged: the task gate still drops image repos there.

* Studio: hide single-file image checkpoints from the chat model picker

The chat picker treats a cached repo as an image model, and hides it, only
when it ships a diffusers model_index.json. Single-file, ComfyUI, and
ControlNet image checkpoints (an FP8 Qwen-Image, a z-image safetensors, a
Qwen-Image ControlNet) carry none, so they surfaced as loadable chat models.
Fall back to resolving the repo id against the known diffusion families, the
same resolver the Images backend loads from, so these checkpoints are tagged
text-to-image and stay in the Images picker only.

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

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

* Studio Images: add the FLUX.2-dev model family

Loading unsloth/FLUX.2-dev-GGUF failed because detect_family knew only the
Qwen3-based FLUX.2-klein, so FLUX.2-dev (the full, Mistral-based Flux2Pipeline)
resolved to nothing and the load errored. Add a flux.2-dev family: Flux2Pipeline
+ Flux2Transformer2DModel over the black-forest-labs/FLUX.2-dev base repo (gated,
reachable with an HF token), with its FLUX.2 32-channel VAE and Mistral text
encoder wired for the sd-cli path from the open Comfy-Org/flux2-dev mirror.
text-to-image only: diffusers 0.38 ships no Flux2 img2img / inpaint pipeline for
dev. Frontend gets sensible dev defaults (28 steps, guidance 4), distinct from
klein's turbo defaults. Verified live: GGUF load resolves the family + gated base
repo and generates a real 1024x1024 image on GPU.

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

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

* Studio Images: clearer error for an unsupported diffusion model

When a repo id resolves to no diffusion family the load raised 'Could not infer a
diffusion family... Pass family_override (z-image)', which points at an unrelated
family and doesn't say what is supported. Replace it with a message that lists the
supported families (from a new supported_family_names helper) and notes that video
models and image models whose diffusers transformer has no single-file loader are
not supported. Applies to both the diffusers and native sd.cpp load paths. Also
refreshes two stale family-registry comments that still called FLUX.2-dev omitted.

* [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

* [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

* [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 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.

* 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

---------

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>

studio/backend/core/inference/diffusion_families.py

@@ -141,6 +141,10 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
         transformer_class = "FluxTransformer2DModel",
         base_repo = "black-forest-labs/FLUX.1-schnell",
         aliases = ("flux1", "flux-1"),
+        # LoRA training targets the guidance-distilled FLUX.1-dev via the DiT trainer
+        # (QLoRA nf4). The dev repo is gated on the Hub, so a user HF token is required.
+        trainable = True,
+        train_base_repos = ("black-forest-labs/FLUX.1-dev",),
         img2img_pipeline_class = "FluxImg2ImgPipeline",
         inpaint_pipeline_class = "FluxInpaintPipeline",
         controlnet_pipeline_class = "FluxControlNetPipeline",
@@ -238,6 +242,9 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
         base_repo = "Qwen/Qwen-Image",
         cfg_kwarg = "true_cfg_scale",
         aliases = ("qwen_image", "qwenimage"),
+        # LoRA training via the DiT trainer, defaulting to the prequant nf4 repo (QLoRA).
+        trainable = True,
+        train_base_repos = ("unsloth/Qwen-Image-2512-unsloth-bnb-4bit", "Qwen/Qwen-Image"),
         img2img_pipeline_class = "QwenImageImg2ImgPipeline",
         inpaint_pipeline_class = "QwenImageInpaintPipeline",
         controlnet_pipeline_class = "QwenImageControlNetPipeline",
@@ -262,6 +269,10 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
         transformer_class = "ZImageTransformer2DModel",
         base_repo = "Tongyi-MAI/Z-Image-Turbo",
         aliases = ("zimage", "z_image"),
+        # LoRA training via the DiT trainer (bf16 only). Defaults to the prequant nf4 repo
+        # for QLoRA; the bf16 Tongyi-MAI base is the alternative.
+        trainable = True,
+        train_base_repos = ("unsloth/Z-Image-Turbo-unsloth-bnb-4bit", "Tongyi-MAI/Z-Image-Turbo"),
         img2img_pipeline_class = "ZImageImg2ImgPipeline",
         inpaint_pipeline_class = "ZImageInpaintPipeline",
         # Z-Image's MLP down-projections peak near 9e5, which overflows float16.
@@ -309,7 +320,6 @@ def trainable_family_names() -> tuple[str, ...]:
     """Names of families Studio can train a LoRA on, in registry order."""
     return tuple(fam.name for fam in _FAMILIES if fam.trainable)
 
-
 # Editing / inpaint checkpoints share an arch keyword but need a different
 # pipeline and an input image, which this text-to-image backend doesn't drive.
 # "layered" rejects Qwen-Image-Layered: its transformer sets additional_t_cond=True

studio/backend/core/training/diffusion_dit_trainer.py

@@ -0,0 +1,706 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Flow-matching LoRA training for the DiT image families (FLUX.1-dev, Qwen-Image, Z-Image).
+
+These are rectified-flow transformers, not the SDXL U-Net, so they share only the plumbing
+in ``diffusion_train_common`` (config, dataset discovery, events, stop, publishing). The
+training math here is flow matching: sample a sigma with the logit-normal density used by
+the diffusers dreambooth scripts, form ``noisy = (1 - sigma) * latents + sigma * noise``,
+predict the velocity, and regress it onto ``target = noise - latents``.
+
+The per-family differences (latent normalisation + packing, the transformer forward
+signature, and the LoRA save entrypoint) live in small ``_FamilySpec`` objects; the loop
+itself is family-agnostic. Verified against diffusers 0.38.0.
+
+Memory: the text encoder(s) are the largest module (T5-XXL ~9 GB for FLUX, Qwen2.5-VL ~7 GB
+for Qwen-Image, Qwen3 for Z-Image), so captions are encoded ONCE up front and the encoders
+are freed before the loop. The transformer trains as a QLoRA (nf4) adapter by default with
+gradient checkpointing and 8-bit AdamW, so only the (small) LoRA params + optimizer state
+and the frozen 4-bit base sit in VRAM during the loop.
+"""
+
+from __future__ import annotations
+
+import gc
+import random
+import time
+from contextlib import nullcontext
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable, Optional
+
+from core.training.diffusion_train_common import (
+    DEFAULT_LORA_FILENAME,
+    DiffusionLoraConfig,
+    EventCb,
+    StopCb,
+    _assert_trusted_base_model,
+    _emit,
+    _publish_to_lora_catalog,
+    discover_image_caption_pairs,
+)
+
+# Per-family LoRA target modules (attention projections). FLUX / Qwen double-stream blocks
+# also carry added-kv projections; Z-Image is single-stream. Kept here (not in the generic
+# DEFAULT_LORA_TARGETS) because they are architecture-specific.
+_FLUX_TARGETS = (
+    "to_q",
+    "to_k",
+    "to_v",
+    "to_out.0",
+    "add_q_proj",
+    "add_k_proj",
+    "add_v_proj",
+    "to_add_out",
+)
+_QWEN_TARGETS = _FLUX_TARGETS
+_ZIMAGE_TARGETS = ("to_q", "to_k", "to_v", "to_out.0")
+
+
+@dataclass
+class _FamilySpec:
+    """Everything the shared loop needs that differs by family."""
+
+    family: str
+    lora_targets: tuple[str, ...]
+    # bf16 only (Z-Image overflows fp16 and its RoPE/embedder run in fp32).
+    force_bf16: bool
+    # Builds (pipe, transformer, vae) with the transformer loaded as a trainable nf4 QLoRA
+    # when qlora=True. Returns the pipeline (for save_lora_weights + encode_prompt), the
+    # transformer to attach LoRA to, and the VAE (kept resident for latent encoding).
+    load: Callable[..., tuple[Any, Any, Any]]
+    # Encode a list of captions -> a per-caption tuple of CPU tensors (the family's embeds).
+    encode_prompts: Callable[..., list[tuple]]
+    # Encode a pixel tensor [B,3,H,W] in [-1,1] -> latents (family-normalised, on device).
+    encode_latents: Callable[..., Any]
+    # One transformer forward: (transformer, noisy, timesteps, sigmas, embeds_batch, cfg,
+    # device, weight_dtype) -> model_pred aligned with target = noise - latents.
+    forward: Callable[..., Any]
+    # Save the LoRA in diffusers format via the family pipeline's save_lora_weights.
+    save: Callable[..., None]
+
+
+# ── shared flow-matching helpers ──────────────────────────────────────────────
+def _get_sigmas(scheduler, timesteps, device, dtype, n_dim):
+    """Gather per-sample sigmas for ``timesteps`` and broadcast to ``n_dim`` (matches the
+    diffusers dreambooth get_sigmas helper)."""
+    import torch
+
+    sigmas = scheduler.sigmas.to(device = device, dtype = dtype)
+    schedule_timesteps = scheduler.timesteps.to(device)
+    step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+    sigma = sigmas[step_indices].flatten()
+    while sigma.ndim < n_dim:
+        sigma = sigma.unsqueeze(-1)
+    return sigma
+
+
+def _sample_timesteps(scheduler, batch_size, device):
+    """Logit-normal density timestep sampling (weighting_scheme='logit_normal'), returning
+    (timesteps, indices) into the scheduler's schedule."""
+    import torch
+    from diffusers.training_utils import compute_density_for_timestep_sampling
+
+    u = compute_density_for_timestep_sampling(
+        weighting_scheme = "logit_normal",
+        batch_size = batch_size,
+        logit_mean = 0.0,
+        logit_std = 1.0,
+        mode_scale = 1.29,
+    )
+    num_train = scheduler.config.num_train_timesteps
+    indices = (u * num_train).long().clamp(0, num_train - 1)
+    timesteps = scheduler.timesteps.to(device)[indices].to(device)
+    return timesteps
+
+
+def _encoders_to_device(pipe, device) -> None:
+    """Move the pipeline's (non-quantized) text encoders to ``device`` before encoding.
+
+    A QLoRA FLUX load places the nf4 transformer on GPU but leaves the text encoders on
+    CPU, so encode_prompt would mix devices. Best-effort per encoder: a 4-bit encoder that
+    is already placed raises on .to() and is left as-is."""
+    for attr in ("text_encoder", "text_encoder_2", "text_encoder_3"):
+        enc = getattr(pipe, attr, None)
+        if enc is None:
+            continue
+        try:
+            enc.to(device)
+        except (ValueError, RuntimeError, NotImplementedError):
+            pass  # already-placed 4-bit encoder / non-movable module
+
+
+def _bnb_4bit_config():
+    from diffusers import BitsAndBytesConfig as DiffusersBnb
+    import torch
+    return DiffusersBnb(
+        load_in_4bit = True,
+        bnb_4bit_quant_type = "nf4",
+        bnb_4bit_compute_dtype = torch.bfloat16,
+    )
+
+
+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 we load it as-is rather than re-quantizing on the fly. A
+    dense (bf16) base instead gets on-the-fly nf4 quantization for QLoRA."""
+    name = str(base_model or "").lower()
+    return "bnb-4bit" in name or "-4bit" in name or "int4" in name or "nf4" in name
+
+
+def _load_quantized_transformer(transformer_cls, cfg):
+    """Load ``cfg.base_model``'s transformer subfolder as a trainable nf4 QLoRA module."""
+    import torch
+    return transformer_cls.from_pretrained(
+        cfg.base_model,
+        subfolder = "transformer",
+        quantization_config = _bnb_4bit_config(),
+        torch_dtype = torch.bfloat16,
+        token = cfg.hf_token,
+    )
+
+
+# ── FLUX.1-dev ────────────────────────────────────────────────────────────────
+def _flux_load(cfg, device, weight_dtype, qlora):
+    import torch
+    from diffusers import FluxPipeline, FluxTransformer2DModel
+
+    if qlora:
+        transformer = FluxTransformer2DModel.from_pretrained(
+            cfg.base_model,
+            subfolder = "transformer",
+            quantization_config = _bnb_4bit_config(),
+            torch_dtype = torch.bfloat16,
+            token = cfg.hf_token,
+        )
+        pipe = FluxPipeline.from_pretrained(
+            cfg.base_model,
+            transformer = transformer,
+            torch_dtype = torch.bfloat16,
+            token = cfg.hf_token,
+        )
+    else:
+        pipe = FluxPipeline.from_pretrained(
+            cfg.base_model, torch_dtype = weight_dtype, token = cfg.hf_token
+        )
+        transformer = pipe.transformer
+    pipe.vae.to(device, dtype = torch.float32)
+    return pipe, transformer, pipe.vae
+
+
+def _flux_encode_prompts(pipe, captions, device):
+    import torch
+
+    _encoders_to_device(pipe, device)
+    out = []
+    with torch.no_grad():
+        for cap in captions:
+            pe, pooled, text_ids = pipe.encode_prompt(
+                prompt = cap,
+                prompt_2 = cap,
+                device = device,
+                num_images_per_prompt = 1,
+                max_sequence_length = 512,
+            )
+            out.append((pe.cpu(), pooled.cpu(), text_ids.cpu()))
+    return out
+
+
+def _flux_encode_latents(vae, pixel_values):
+    import torch
+
+    with torch.no_grad():
+        lat = vae.encode(pixel_values.to(torch.float32)).latent_dist.sample()
+    lat = (lat - vae.config.shift_factor) * vae.config.scaling_factor
+    return lat
+
+
+def _flux_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device, weight_dtype):
+    import torch
+    from diffusers import FluxPipeline
+
+    pe, pooled, text_ids = embeds_batch
+    bsz, c, h, w = noisy.shape
+    packed = FluxPipeline._pack_latents(noisy, bsz, c, h, w)
+    # Position ids drive RoPE and are indices, not activations -- keep them float32 (the
+    # dtype diffusers' own pipeline builds) regardless of the bf16 training dtype.
+    img_ids = FluxPipeline._prepare_latent_image_ids(bsz, h // 2, w // 2, device, torch.float32)
+    guidance = torch.full((bsz,), 1.0, device = device, dtype = torch.float32)
+    model_pred = transformer(
+        hidden_states = packed,
+        timestep = timesteps / 1000,
+        guidance = guidance,
+        pooled_projections = pooled.to(weight_dtype),
+        encoder_hidden_states = pe.to(weight_dtype),
+        txt_ids = text_ids.to(torch.float32),
+        img_ids = img_ids,
+        return_dict = False,
+    )[0]
+    return FluxPipeline._unpack_latents(model_pred, h * 8, w * 8, 8)
+
+
+def _flux_save(pipe_cls, out_dir, transformer_lora_layers):
+    from diffusers import FluxPipeline
+    FluxPipeline.save_lora_weights(
+        save_directory = out_dir,
+        transformer_lora_layers = transformer_lora_layers,
+        weight_name = DEFAULT_LORA_FILENAME,
+    )
+
+
+# ── Qwen-Image ────────────────────────────────────────────────────────────────
+def _qwen_load(cfg, device, weight_dtype, qlora):
+    import torch
+    from diffusers import QwenImagePipeline, QwenImageTransformer2DModel
+
+    # The prequant default (unsloth/Qwen-Image-2512-unsloth-bnb-4bit) ships the transformer
+    # 4-bit, so from_pretrained loads it trainable as-is. A dense (bf16) base -- the 20B
+    # Qwen/Qwen-Image -- is quantized to nf4 on the fly so QLoRA still fits.
+    kwargs = {"torch_dtype": torch.bfloat16, "token": cfg.hf_token}
+    if qlora and not _repo_is_prequantized(cfg.base_model):
+        kwargs["transformer"] = _load_quantized_transformer(QwenImageTransformer2DModel, cfg)
+    pipe = QwenImagePipeline.from_pretrained(cfg.base_model, **kwargs)
+    pipe.vae.to(device, dtype = torch.float32)
+    return pipe, pipe.transformer, pipe.vae
+
+
+def _qwen_encode_prompts(pipe, captions, device):
+    import torch
+
+    _encoders_to_device(pipe, device)
+    out = []
+    with torch.no_grad():
+        for cap in captions:
+            pe, mask = pipe.encode_prompt(
+                prompt = cap,
+                device = device,
+                num_images_per_prompt = 1,
+                max_sequence_length = 1024,
+            )
+            out.append((pe.cpu(), mask.cpu() if mask is not None else None))
+    return out
+
+
+def _qwen_encode_latents(vae, pixel_values):
+    import torch
+
+    # AutoencoderKLQwenImage is a 3D (video) VAE: add a temporal dim, encode, drop it back
+    # into a [B,16,1,H,W] latent normalised by the per-channel latents_mean / latents_std.
+    px = pixel_values.to(torch.float32).unsqueeze(2)  # [B,3,1,H,W]
+    with torch.no_grad():
+        lat = vae.encode(px).latent_dist.sample()  # [B,16,1,h,w]
+    z = vae.config.z_dim
+    mean = torch.tensor(vae.config.latents_mean, device = lat.device, dtype = lat.dtype)
+    std = torch.tensor(vae.config.latents_std, device = lat.device, dtype = lat.dtype)
+    mean = mean.view(1, z, 1, 1, 1)
+    std = std.view(1, z, 1, 1, 1)
+    return (lat - mean) / std
+
+
+def _qwen_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device, weight_dtype):
+    import torch
+    from diffusers import QwenImagePipeline
+
+    pe, mask = embeds_batch
+    bsz, c, f, h, w = noisy.shape
+    packed = QwenImagePipeline._pack_latents(noisy, bsz, c, h, w)
+    # Each batch entry is a LIST of one (frame, h/2, w/2) tuple: the transformer indexes
+    # sample[0] / sample[1:] per entry (transformer_qwenimage.py), so a flat list breaks it.
+    img_shapes = [[(1, h // 2, w // 2)]] * bsz
+    pred = transformer(
+        hidden_states = packed,
+        encoder_hidden_states = pe.to(weight_dtype),
+        encoder_hidden_states_mask = mask.to(device) if mask is not None else None,
+        timestep = timesteps / 1000,
+        img_shapes = img_shapes,
+        return_dict = False,
+    )[0]
+    return QwenImagePipeline._unpack_latents(pred, h * 8, w * 8, 8)
+
+
+def _qwen_save(pipe_cls, out_dir, transformer_lora_layers):
+    from diffusers import QwenImagePipeline
+    QwenImagePipeline.save_lora_weights(
+        save_directory = out_dir,
+        transformer_lora_layers = transformer_lora_layers,
+        weight_name = DEFAULT_LORA_FILENAME,
+    )
+
+
+# ── Z-Image ───────────────────────────────────────────────────────────────────
+def _zimage_load(cfg, device, weight_dtype, qlora):
+    import torch
+    from diffusers import ZImagePipeline, ZImageTransformer2DModel
+
+    # Prequant default loads 4-bit as-is; the dense bf16 Tongyi-MAI base is quantized to nf4
+    # on the fly. Z-Image is bf16 only (its RoPE/embedder run fp32; fp16 overflows).
+    kwargs = {"torch_dtype": torch.bfloat16, "token": cfg.hf_token}
+    if qlora and not _repo_is_prequantized(cfg.base_model):
+        kwargs["transformer"] = _load_quantized_transformer(ZImageTransformer2DModel, cfg)
+    pipe = ZImagePipeline.from_pretrained(cfg.base_model, **kwargs)
+    pipe.vae.to(device, dtype = torch.float32)
+    return pipe, pipe.transformer, pipe.vae
+
+
+def _zimage_encode_prompts(pipe, captions, device):
+    import torch
+
+    _encoders_to_device(pipe, device)
+    out = []
+    with torch.no_grad():
+        for cap in captions:
+            pe, _neg = pipe.encode_prompt(
+                prompt = cap,
+                device = device,
+                do_classifier_free_guidance = False,
+                max_sequence_length = 512,
+            )
+            # pe is a list of one variable-length [seq, 2560] tensor per prompt.
+            emb = pe[0] if isinstance(pe, (list, tuple)) else pe
+            out.append((emb.cpu(),))
+    return out
+
+
+def _zimage_encode_latents(vae, pixel_values):
+    import torch
+    with torch.no_grad():
+        lat = vae.encode(pixel_values.to(torch.float32)).latent_dist.mode()
+    return (lat - vae.config.shift_factor) * vae.config.scaling_factor
+
+
+def _zimage_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device, weight_dtype):
+    import torch
+
+    (emb,) = embeds_batch
+    # List I/O: one [C,1,H,W] latent + one [seq,2560] caption per sample. The timestep
+    # convention is REVERSED ((1000 - t) / 1000) and the prediction is NEGATED.
+    x_list = list(noisy.unsqueeze(2).unbind(dim = 0))
+    cap_list = [emb.to(device = device, dtype = weight_dtype)]
+    t_norm = (1000 - timesteps) / 1000
+    out = transformer(x_list, t_norm, cap_list, return_dict = False)[0]
+    return -torch.stack(out, dim = 0).squeeze(2)
+
+
+def _zimage_save(pipe_cls, out_dir, transformer_lora_layers):
+    from diffusers import ZImagePipeline
+    ZImagePipeline.save_lora_weights(
+        save_directory = out_dir,
+        transformer_lora_layers = transformer_lora_layers,
+        weight_name = DEFAULT_LORA_FILENAME,
+    )
+
+
+_SPECS: dict[str, _FamilySpec] = {
+    "flux.1": _FamilySpec(
+        family = "flux.1",
+        lora_targets = _FLUX_TARGETS,
+        force_bf16 = False,
+        load = _flux_load,
+        encode_prompts = _flux_encode_prompts,
+        encode_latents = _flux_encode_latents,
+        forward = _flux_forward,
+        save = _flux_save,
+    ),
+    "qwen-image": _FamilySpec(
+        family = "qwen-image",
+        lora_targets = _QWEN_TARGETS,
+        force_bf16 = True,
+        load = _qwen_load,
+        encode_prompts = _qwen_encode_prompts,
+        encode_latents = _qwen_encode_latents,
+        forward = _qwen_forward,
+        save = _qwen_save,
+    ),
+    "z-image": _FamilySpec(
+        family = "z-image",
+        lora_targets = _ZIMAGE_TARGETS,
+        force_bf16 = True,
+        load = _zimage_load,
+        encode_prompts = _zimage_encode_prompts,
+        encode_latents = _zimage_encode_latents,
+        forward = _zimage_forward,
+        save = _zimage_save,
+    ),
+}
+
+
+# HF repos that gate access behind a license acceptance: training needs a token whose
+# account has accepted the license. Checked by name (no network) so a missing token fails
+# fast with an actionable message instead of a confusing 401 mid-load.
+_GATED_TRAIN_REPOS = frozenset({"black-forest-labs/flux.1-dev"})
+
+
+def _assert_gated_access(base_model: str, hf_token: Optional[str]) -> None:
+    """Raise a clear error before loading a gated base without a token."""
+    name = str(base_model or "").strip().lower()
+    if name in _GATED_TRAIN_REPOS and not (hf_token and str(hf_token).strip()):
+        raise ValueError(
+            f"'{base_model}' is a gated Hugging Face repo. Accept its license on the Hub "
+            f"and add your HF token in Studio settings before training from it."
+        )
+
+
+def _load_pixel_tensor(path, resolution, center_crop, random_flip, rng):
+    """Load an image -> a normalised [3,H,W] tensor in [-1,1]. Same geometry as the SDXL
+    loader but without the SDXL time-ids (DiT families don't use them)."""
+    import numpy as np
+    import torch
+    from PIL import Image, ImageOps
+
+    img = ImageOps.exif_transpose(Image.open(path)).convert("RGB")
+    w0, h0 = img.size
+    scale = resolution / min(w0, h0)
+    rw, rh = max(resolution, round(w0 * scale)), max(resolution, round(h0 * scale))
+    img = img.resize((rw, rh), Image.LANCZOS)
+    if center_crop:
+        left, top = (rw - resolution) // 2, (rh - resolution) // 2
+    else:
+        left = rng.randint(0, max(0, rw - resolution))
+        top = rng.randint(0, max(0, rh - resolution))
+    img = img.crop((left, top, left + resolution, top + resolution))
+    if random_flip and rng.random() < 0.5:
+        img = img.transpose(Image.FLIP_LEFT_RIGHT)
+    arr = np.asarray(img, dtype = np.float32) / 255.0
+    return torch.from_numpy(arr).permute(2, 0, 1) * 2.0 - 1.0
+
+
+def run_dit_lora_training(
+    config: DiffusionLoraConfig,
+    *,
+    on_event: Optional[EventCb] = None,
+    should_stop: Optional[StopCb] = None,
+) -> str:
+    """Train a flow-matching DiT LoRA (FLUX.1-dev / Qwen-Image / Z-Image) and export it."""
+    import torch
+    import torch.nn.functional as F
+    from diffusers import FlowMatchEulerDiscreteScheduler
+    from diffusers.training_utils import cast_training_params
+    from peft import LoraConfig
+    from peft.utils import get_peft_model_state_dict
+
+    cfg = config.normalized()
+    spec = _SPECS.get(cfg.resolved_family)
+    if spec is None:
+        raise ValueError(f"No DiT trainer for family {cfg.resolved_family!r}")
+
+    rng = random.Random(cfg.seed)
+    torch.manual_seed(cfg.seed)
+
+    save_on_stop = True
+
+    def _check_stop() -> bool:
+        nonlocal save_on_stop
+        if should_stop is None:
+            return False
+        sig = should_stop()
+        if not sig:
+            return False
+        if isinstance(sig, dict) and sig.get("save") is False:
+            save_on_stop = False
+        return True
+
+    # DiT families train in bf16 (Z-Image/Qwen require it; FLUX prefers it). A caller that
+    # explicitly asks for fp16 on a bf16-only family is refused rather than silently
+    # upgraded, so the choice is never misrepresented.
+    if cfg.mixed_precision == "fp16" and spec.force_bf16:
+        raise ValueError(
+            f"{spec.family} LoRA training requires bf16: fp16 overflows its fp32 RoPE / "
+            f"embedder internals. Set mixed precision to bf16."
+        )
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    # The flow-matching + 4-bit path is bf16 throughout (fp32 on a CPU-only box, which is
+    # unsupported for real runs but keeps import/unit tests architecture-agnostic).
+    weight_dtype = torch.bfloat16 if device == "cuda" else torch.float32
+    use_lora_targets = tuple(cfg.lora_target_modules) or spec.lora_targets
+
+    _assert_trusted_base_model(cfg.base_model)
+    _assert_gated_access(cfg.base_model, cfg.hf_token)
+    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))
+    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)
+
+    # QLoRA by default for the big DiTs (nf4 transformer). The prequant Qwen/Z-Image repos
+    # are already 4-bit; FLUX quantizes its transformer on the fly.
+    pipe, transformer, vae = spec.load(cfg, device, weight_dtype, qlora = True)
+
+    # Precompute all caption embeddings, then free the (large) text encoder(s): captions are
+    # constant and the encoders are frozen, so this is exact and the biggest memory win.
+    image_paths = [p for p, _ in pairs]
+    captions = [c for _, c in pairs]
+    uniq = sorted(set(captions))
+    encoded = spec.encode_prompts(pipe, uniq, device)
+    caption_embeds = {cap: emb for cap, emb in zip(uniq, encoded)}
+    _free_text_encoders(pipe)
+    gc.collect()
+    if device == "cuda":
+        torch.cuda.empty_cache()
+
+    # Freeze the base; attach the trainable LoRA to the transformer.
+    transformer.requires_grad_(False)
+    transformer.add_adapter(
+        LoraConfig(
+            r = cfg.lora_rank,
+            lora_alpha = cfg.lora_alpha,
+            lora_dropout = cfg.lora_dropout,
+            init_lora_weights = "gaussian",
+            target_modules = list(use_lora_targets),
+        )
+    )
+    if cfg.gradient_checkpointing:
+        # Non-reentrant checkpointing: reentrant recompute of a bnb 4-bit LoRA linear can
+        # trip an illegal memory access on the larger FLUX transformer, and non-reentrant
+        # is the recommended mode anyway (it also handles a checkpointed segment whose
+        # inputs do not require grad, which happens with a frozen 4-bit base).
+        import functools
+        import torch.utils.checkpoint as _ckpt
+        transformer.enable_gradient_checkpointing(
+            gradient_checkpointing_func = functools.partial(_ckpt.checkpoint, use_reentrant = False)
+        )
+    cast_training_params(transformer, dtype = torch.float32)
+    lora_params = [p for p in transformer.parameters() if p.requires_grad]
+
+    optimizer = _make_optimizer(lora_params, cfg.learning_rate)
+    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+        cfg.base_model, subfolder = "scheduler", token = cfg.hf_token
+    )
+
+    _emit(on_event, "model_load_completed")
+
+    transformer.train()
+    stopped = False
+    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):
+            i = rng.randrange(len(image_paths))
+            px = (
+                _load_pixel_tensor(
+                    image_paths[i], cfg.resolution, cfg.center_crop, cfg.random_flip, rng
+                )
+                .unsqueeze(0)
+                .to(device)
+            )
+            latents = spec.encode_latents(vae, px).to(weight_dtype)
+
+            noise = torch.randn_like(latents)
+            timesteps = _sample_timesteps(scheduler, latents.shape[0], device)
+            sigmas = _get_sigmas(scheduler, timesteps, device, weight_dtype, latents.ndim)
+            noisy = (1.0 - sigmas) * latents + sigmas * noise
+
+            emb = caption_embeds[captions[i]]
+            emb_dev = tuple(
+                t.to(device = device, dtype = weight_dtype)
+                if (t is not None and t.is_floating_point())
+                else (t.to(device) if t is not None else None)
+                for t in emb
+            )
+            # bf16 autocast around the forward + loss, matching the diffusers dreambooth
+            # scripts' accelerator.autocast: it reconciles the fp32 LoRA params with the
+            # bnb 4-bit base matmuls in one compute dtype. Without it the 4-bit backward
+            # on FLUX dies with an illegal-address / CUBLAS failure.
+            autocast = (
+                torch.autocast(device_type = "cuda", dtype = torch.bfloat16)
+                if device == "cuda"
+                else nullcontext()
+            )
+            with autocast:
+                model_pred = spec.forward(
+                    transformer, noisy, timesteps, sigmas, emb_dev, cfg, device, weight_dtype
+                )
+                target = noise - latents
+                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
+
+        if cfg.max_grad_norm and cfg.max_grad_norm > 0:
+            torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm)
+        optimizer.step()
+
+        running_loss += step_loss
+        done = opt_step + 1
+        if done % cfg.log_every == 0 or done == cfg.train_steps:
+            if device == "cuda":
+                peak_gb = round(torch.cuda.max_memory_allocated() / 1e9, 2)
+            sps = 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 = cfg.learning_rate,
+                samples_per_second = sps,
+                peak_memory_gb = peak_gb or None,
+            )
+        if _check_stop():
+            stopped = True
+            break
+
+    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)
+        layers = get_peft_model_state_dict(transformer)
+        spec.save(pipe, str(out_dir), layers)
+        lora_path = str(out_dir / DEFAULT_LORA_FILENAME)
+        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)
+
+
+def _make_optimizer(params, lr):
+    """8-bit AdamW (bitsandbytes) when available -- half the optimizer state, no accuracy
+    regression for LoRA -- else the torch AdamW fallback."""
+    import torch
+    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)
+
+
+def _free_text_encoders(pipe) -> None:
+    """Drop every text-encoder / tokenizer the pipeline holds, so the (large) encoders do
+    not sit in VRAM during training. The embeddings are already precomputed."""
+    for attr in (
+        "text_encoder",
+        "text_encoder_2",
+        "text_encoder_3",
+        "tokenizer",
+        "tokenizer_2",
+        "tokenizer_3",
+    ):
+        if getattr(pipe, attr, None) is not None:
+            try:
+                setattr(pipe, attr, None)
+            except Exception:  # noqa: BLE001
+                pass

studio/backend/core/training/diffusion_lora_trainer.py

@@ -25,6 +25,7 @@ Design:
 from __future__ import annotations
 
 import argparse
+import gc
 import os
 import random
 import time
@@ -226,7 +227,7 @@ def run_diffusion_lora_training(
         cast_training_params(unet, dtype = torch.float32)
 
     lora_params = [p for p in unet.parameters() if p.requires_grad]
-    optimizer = torch.optim.AdamW(lora_params, lr = cfg.learning_rate)
+    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
@@ -241,6 +242,24 @@ def run_diffusion_lora_training(
     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]]:
@@ -281,9 +300,13 @@ def run_diffusion_lora_training(
             ).long()
             noisy = noise_scheduler.add_noise(latents, noise, timesteps)
 
-            prompt_embeds, pooled = _encode_sdxl_prompts(
-                captions, tokenizers, text_encoders, device
-            )
+            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}
@@ -381,6 +404,21 @@ def run_diffusion_lora_training(
     return str(out_dir)
 
 
+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)."""
+    import torch
+
+    if os.environ.get("UNSLOTH_DIFFUSION_FP32_OPTIM", "") not in ("1", "true"):
+        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
+    return torch.optim.AdamW(params, lr = lr)
+
+
 def run_diffusion_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> None:
     """mp.Queue subprocess adapter: run a training job, translating the ``on_event``
     callback to ``event_queue`` and a ``stop_queue`` poll to ``should_stop``. Dispatches to

studio/backend/core/training/diffusion_train_common.py

@@ -80,9 +80,13 @@ def resolve_trainable_family(base_model: str, model_family: Optional[str] = None
       compatible: a genuinely wrong pick still fails cleanly later in from_pretrained).
     """
     name = str(base_model or "").strip().lower()
-    if name.endswith(".gguf"):
+    # GGUF weights (a ``.gguf`` file or a ``*-GGUF`` repo) are inference-only: training needs
+    # the full diffusers pipeline (transformer + VAE + text encoders), which a GGUF repo does
+    # not provide. Reject by name even when the family itself is trainable.
+    if name.endswith(".gguf") or "gguf" in name:
         raise ValueError(
-            f"'{base_model}' is a GGUF checkpoint, which can't be trained. {_trainable_hint()}"
+            f"'{base_model}' is a GGUF checkpoint/repo, which can't be a training base "
+            f"(training needs the full diffusers model). {_trainable_hint()}"
         )
     if model_family and str(model_family).strip():
         key = str(model_family).strip().lower()
@@ -91,7 +95,9 @@ def resolve_trainable_family(base_model: str, model_family: Optional[str] = None
             known = ", ".join(supported_family_names())
             raise ValueError(f"Unknown model_family {model_family!r}. Known families: {known}.")
         if not fam.trainable:
-            raise ValueError(f"'{fam.name}' models can't be trained yet. {_trainable_hint()}")
+            raise ValueError(
+                f"'{fam.name}' models can't be trained yet. {_trainable_hint()}"
+            )
         return fam.name
 
     fam = detect_family_for_pick(base_model)
@@ -126,9 +132,72 @@ def get_trainer(family: str) -> Callable[..., str]:
     if key == "sdxl":
         from core.training.diffusion_lora_trainer import run_diffusion_lora_training
         return run_diffusion_lora_training
+    if key in ("flux.1", "qwen-image", "z-image"):
+        from core.training.diffusion_dit_trainer import run_dit_lora_training
+        return run_dit_lora_training
     raise ValueError(f"No trainer is registered for family {family!r}.")
 
 
+# Per-family training defaults surfaced by the Train UI. Distilled/turbo bases and the big
+# DiTs want different rank / learning rate / resolution; these are starting points, not
+# hard limits. Families absent here fall back to the DiffusionLoraConfig defaults.
+FAMILY_TRAIN_DEFAULTS: dict[str, dict[str, Any]] = {
+    "sdxl": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 1024},
+    "flux.1": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 512},
+    "qwen-image": {"lora_rank": 16, "learning_rate": 5e-5, "resolution": 512},
+    "z-image": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 768},
+}
+
+
+def train_defaults(family: str) -> dict[str, Any]:
+    """Recommended starting hyperparameters for ``family`` (empty if unknown)."""
+    return dict(FAMILY_TRAIN_DEFAULTS.get((family or "").strip().lower(), {}))
+
+
+# Display labels + a short VRAM/access note per trainable family, surfaced by the Train UI
+# so users pick a base with realistic expectations. Kept next to the defaults they pair with.
+_FAMILY_LABELS = {
+    "sdxl": "SDXL",
+    "flux.1": "FLUX.1-dev",
+    "qwen-image": "Qwen-Image",
+    "z-image": "Z-Image",
+}
+_FAMILY_VRAM_NOTES = {
+    "sdxl": "Trains on ~12 GB+ (bf16 LoRA). The lightest, fastest option.",
+    "flux.1": (
+        "12B model, QLoRA (nf4) by default (~16 GB+). Gated on Hugging Face: accept the "
+        "FLUX.1-dev license and add your HF token before training."
+    ),
+    "qwen-image": "20B model, QLoRA (nf4) by default (~24 GB+). The heaviest option.",
+    "z-image": "6B model, QLoRA (nf4) by default (~12 GB+). bf16 only.",
+}
+
+
+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
+
+    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]
+        infos.append(
+            {
+                "name": name,
+                "label": _FAMILY_LABELS.get(name, name),
+                "default_base": repos[0],
+                "base_repos": repos,
+                "defaults": train_defaults(name),
+                "vram_note": _FAMILY_VRAM_NOTES.get(name, ""),
+            }
+        )
+    return infos
+
+
 @dataclass
 class DiffusionLoraConfig:
     """Everything a diffusion LoRA training run needs. Only ``base_model`` /

studio/backend/models/training.py

@@ -772,15 +772,25 @@ class DiffusionDatasetSummary(BaseModel):
     caption_count: int
 
 
-class DiffusionTrainingInfoResponse(BaseModel):
-    """Where diffusion training reads/writes on this Studio, plus usable datasets.
+class DiffusionTrainableFamily(BaseModel):
+    """A base-model family the diffusion trainer supports, with UI-facing metadata."""
 
-    Lets the UI show real on-disk locations and offer existing dataset folders,
-    instead of asking users to know the Studio home layout."""
+    name: str
+    label: str
+    default_base: str
+    base_repos: List[str] = Field(default_factory = list)
+    defaults: dict = Field(default_factory = dict)
+    vram_note: str = ""
+
+
+class DiffusionTrainingInfoResponse(BaseModel):
+    """Where diffusion training reads/writes on this Studio, plus usable datasets and the
+    trainable model families (so the UI can offer a base picker with realistic guidance)."""
 
     datasets_root: str
     outputs_root: str
     datasets: List[DiffusionDatasetSummary]
+    families: List[DiffusionTrainableFamily] = Field(default_factory = list)
 
 
 class DiffusionDatasetUploadResponse(BaseModel):

studio/backend/routes/training.py

@@ -61,6 +61,7 @@ from models.training import (
     DiffusionDatasetSummary,
     DiffusionDatasetUploadResponse,
     DiffusionMetricHistory,
+    DiffusionTrainableFamily,
     DiffusionTrainingInfoResponse,
     DiffusionTrainingStartRequest,
     DiffusionTrainingStartResponse,
@@ -1125,6 +1126,44 @@ def _free_gpu_for_diffusion_training() -> None:
         logger.warning("Could not free chat models for diffusion training: %s", e)
 
 
+def _preflight_gated_base(base_model: str, hf_token: Optional[str]) -> None:
+    """HEAD a remote base repo's model_index.json with the caller's token; raise HTTP 400 on
+    401/403 (gated / unauthorized) with an actionable message. Best-effort: a local path,
+    a non-repo string, or a network hiccup passes through so the trainer can surface any real
+    load error itself. Runs before GPU teardown so a doomed start never evicts a loaded model."""
+    import urllib.error
+    import urllib.request
+
+    repo = (base_model or "").strip()
+    # Only remote 'org/name' repos are gated; skip local paths and single-file names.
+    if (
+        not repo
+        or repo.count("/") != 1
+        or repo.startswith((".", "/", "~"))
+        or repo.endswith(".gguf")
+    ):
+        return
+    url = f"https://huggingface.co/{repo}/resolve/main/model_index.json"
+    headers = {"Authorization": f"Bearer {hf_token}"} if hf_token else {}
+    req = urllib.request.Request(url, method = "HEAD", headers = headers)
+    try:
+        urllib.request.urlopen(req, timeout = 5)
+    except urllib.error.HTTPError as e:
+        if e.code in (401, 403):
+            raise HTTPException(
+                status_code = 400,
+                detail = (
+                    f"Access to '{repo}' is gated or unauthorized. Accept the model's license "
+                    f"on its Hugging Face page and add your HF token in Studio settings, then "
+                    f"try again."
+                ),
+            )
+        # 404 (e.g. a repo without a root model_index.json) and other codes are not an
+        # access problem -- let the trainer surface any genuine load error.
+    except Exception:  # noqa: BLE001 -- network/DNS hiccup must not block a start
+        return
+
+
 @router.post("/diffusion/start", response_model = DiffusionTrainingStartResponse)
 async def start_diffusion_training(
     body: DiffusionTrainingStartRequest, current_subject: str = Depends(get_current_subject)
@@ -1169,8 +1208,13 @@ async def start_diffusion_training(
     except ValueError as e:
         raise HTTPException(status_code = 400, detail = str(e))
 
+    # Preflight access to a gated base repo with the user's token BEFORE freeing GPU
+    # residents, so a missing/insufficient token fails fast (400) without tearing down the
+    # user's loaded chat/Images model, and never surfaces as a confusing mid-load 401.
+    _preflight_gated_base(config.get("base_model", ""), config.get("hf_token"))
+
     # Free resident GPU workloads (export / Images pipeline / chat) before the trainer
-    # loads its own SDXL pipeline.
+    # loads its own pipeline.
     _free_gpu_for_diffusion_training()
 
     service = get_diffusion_training_service()
@@ -1259,8 +1303,14 @@ async def diffusion_training_info(current_subject: str = Depends(get_current_sub
                 continue
             if summary.image_count > 0:
                 found.append(summary)
+        from core.training.diffusion_train_common import family_train_infos
+
+        families = [DiffusionTrainableFamily(**info) for info in family_train_infos()]
         return DiffusionTrainingInfoResponse(
-            datasets_root = str(root), outputs_root = str(outputs_root()), datasets = found
+            datasets_root = str(root),
+            outputs_root = str(outputs_root()),
+            datasets = found,
+            families = families,
         )
 
     return await asyncio.to_thread(scan)

studio/backend/tests/test_diffusion_dit_trainer.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Unit tests for the flow-matching DiT LoRA trainer (FLUX.1 / Qwen-Image / Z-Image).
+
+CPU-only: cover family resolution, the per-family spec table, the QLoRA prequant
+heuristic, the bf16-only guard, and the gated-repo name check. The full training loop is
+exercised by the live GPU smokes, not here."""
+
+from __future__ import annotations
+
+import pytest
+
+from core.training.diffusion_dit_trainer import (
+    _GATED_TRAIN_REPOS,
+    _SPECS,
+    _assert_gated_access,
+    _repo_is_prequantized,
+    run_dit_lora_training,
+)
+from core.training.diffusion_train_common import DiffusionLoraConfig, family_train_infos
+
+
+def test_specs_cover_the_three_dit_families():
+    assert set(_SPECS) == {"flux.1", "qwen-image", "z-image"}
+    # FLUX / Qwen share the added-kv attention target set; Z-Image is single-stream.
+    assert "add_q_proj" in _SPECS["flux.1"].lora_targets
+    assert "add_q_proj" in _SPECS["qwen-image"].lora_targets
+    assert "add_q_proj" not in _SPECS["z-image"].lora_targets
+    # Z-Image and Qwen are bf16-only.
+    assert _SPECS["z-image"].force_bf16 is True
+    assert _SPECS["qwen-image"].force_bf16 is True
+
+
+@pytest.mark.parametrize(
+    "repo, expected",
+    [
+        ("unsloth/Qwen-Image-2512-unsloth-bnb-4bit", True),
+        ("unsloth/Z-Image-Turbo-unsloth-bnb-4bit", True),
+        ("some/model-int4", True),
+        ("black-forest-labs/FLUX.1-dev", False),
+        ("Tongyi-MAI/Z-Image-Turbo", False),
+    ],
+)
+def test_prequant_heuristic(repo, expected):
+    assert _repo_is_prequantized(repo) is expected
+
+
+def test_zimage_rejects_fp16_before_loading():
+    # bf16-only families must refuse an explicit fp16 request up front (no model load).
+    cfg = DiffusionLoraConfig(
+        base_model = "Tongyi-MAI/Z-Image-Turbo",
+        data_dir = "does-not-exist",
+        output_dir = "o",
+        mixed_precision = "fp16",
+    )
+    with pytest.raises(ValueError, match = "bf16"):
+        run_dit_lora_training(cfg)
+
+
+def test_gated_access_requires_token():
+    assert "black-forest-labs/flux.1-dev" in _GATED_TRAIN_REPOS
+    # No token -> clear, actionable error before any download.
+    with pytest.raises(ValueError, match = "gated"):
+        _assert_gated_access("black-forest-labs/FLUX.1-dev", None)
+    with pytest.raises(ValueError, match = "gated"):
+        _assert_gated_access("black-forest-labs/FLUX.1-dev", "   ")
+    # With a token, or for a non-gated repo, it is a no-op.
+    _assert_gated_access("black-forest-labs/FLUX.1-dev", "hf_realtoken")
+    _assert_gated_access("Tongyi-MAI/Z-Image-Turbo", None)
+
+
+def test_family_train_infos_lists_dit_families():
+    infos = {i["name"]: i for i in family_train_infos()}
+    for fam in ("sdxl", "flux.1", "qwen-image", "z-image"):
+        assert fam in infos, f"{fam} missing from family_train_infos"
+        assert infos[fam]["default_base"]
+        assert infos[fam]["base_repos"]
+        assert "resolution" in infos[fam]["defaults"]
+    # FLUX default base is the gated dev repo; its note flags the license requirement.
+    assert infos["flux.1"]["default_base"] == "black-forest-labs/FLUX.1-dev"
+    assert "gated" in infos["flux.1"]["vram_note"].lower()
+    # Z-Image defaults to the prequant nf4 repo for QLoRA.
+    assert "4bit" in infos["z-image"]["default_base"].lower()

studio/backend/tests/test_diffusion_lora_trainer.py

@@ -195,22 +195,31 @@ def test_config_rejects_nonpositive_learning_rate():
         ).normalized()
 
 
-def test_config_rejects_known_non_sdxl_base_models():
-    # Known DiT families and GGUF checkpoints must fail at normalise time (an instant
-    # 400 via the API) instead of minutes later inside from_pretrained.
+def test_config_rejects_untrainable_base_models():
+    # GGUF checkpoints and families without a trainer (Kontext editing, SD3) must fail at
+    # normalise time (an instant 400 via the API), not minutes later inside from_pretrained.
     for bad in (
         "unsloth/FLUX.1-dev-GGUF",
-        "black-forest-labs/FLUX.1-schnell",
-        "unsloth/Qwen-Image-2512-unsloth-bnb-4bit",
-        "Tongyi-MAI/Z-Image-Turbo",
+        "z-image-turbo-Q4_K_M.gguf",
         "stabilityai/stable-diffusion-3-medium",
         "unsloth/FLUX.1-Kontext-dev",
-        "z-image-turbo-Q4_K_M.gguf",
     ):
-        with pytest.raises(ValueError, match = "SDXL"):
+        with pytest.raises(ValueError):
             DiffusionLoraConfig(base_model = bad, data_dir = "d", output_dir = "o").normalized()
 
 
+def test_config_resolves_dit_families():
+    # FLUX.1 / Qwen-Image / Z-Image bases now resolve to their DiT trainer families.
+    for base, fam in (
+        ("black-forest-labs/FLUX.1-dev", "flux.1"),
+        ("black-forest-labs/FLUX.1-schnell", "flux.1"),
+        ("unsloth/Qwen-Image-2512-unsloth-bnb-4bit", "qwen-image"),
+        ("Tongyi-MAI/Z-Image-Turbo", "z-image"),
+    ):
+        cfg = DiffusionLoraConfig(base_model = base, data_dir = "d", output_dir = "o").normalized()
+        assert cfg.resolved_family == fam
+
+
 def test_config_accepts_sdxl_and_unknown_base_models():
     # SDXL names and unclassifiable custom names/paths must pass the guard (a wrong
     # custom pick still fails cleanly in from_pretrained).
@@ -227,14 +236,23 @@ def test_config_accepts_sdxl_and_unknown_base_models():
 # ── trainer registry + family resolution + metadata sidecar (PR A platform) ──
 def test_get_trainer_resolves_sdxl():
     from core.training.diffusion_lora_trainer import get_trainer, run_diffusion_lora_training
+
     assert get_trainer("sdxl") is run_diffusion_lora_training
     assert get_trainer("SDXL") is run_diffusion_lora_training  # case-insensitive
 
 
 def test_get_trainer_unknown_family_raises():
     from core.training.diffusion_lora_trainer import get_trainer
+
     with pytest.raises(ValueError, match = "No trainer"):
-        get_trainer("flux.1")  # not registered until the DiT trainers ship
+        get_trainer("flux.2-dev")  # a real family with no registered trainer
+
+
+def test_get_trainer_resolves_dit_families():
+    from core.training.diffusion_dit_trainer import run_dit_lora_training
+    from core.training.diffusion_lora_trainer import get_trainer
+    for fam in ("flux.1", "qwen-image", "z-image"):
+        assert get_trainer(fam) is run_dit_lora_training
 
 
 def test_normalized_sets_resolved_family():
@@ -242,9 +260,7 @@ def test_normalized_sets_resolved_family():
         base_model = "stabilityai/stable-diffusion-xl-base-1.0", data_dir = "d", output_dir = "o"
     ).normalized()
     assert cfg.resolved_family == "sdxl"
-    cfg2 = DiffusionLoraConfig(
-        base_model = "my-custom-thing", data_dir = "d", output_dir = "o"
-    ).normalized()
+    cfg2 = DiffusionLoraConfig(base_model = "my-custom-thing", data_dir = "d", output_dir = "o").normalized()
     assert cfg2.resolved_family == "sdxl"  # unknown -> default SDXL trainer
 
 
@@ -254,9 +270,16 @@ def test_explicit_model_family_validated():
     # A bogus explicit family is rejected up front.
     with pytest.raises(ValueError, match = "Unknown model_family"):
         C(base_model = "b", data_dir = "d", output_dir = "o", model_family = "not-a-family").normalized()
-    # A known-but-not-yet-trainable family is rejected with the SDXL hint.
-    with pytest.raises(ValueError, match = "SDXL"):
-        C(base_model = "b", data_dir = "d", output_dir = "o", model_family = "flux.1").normalized()
+    # A known-but-not-trainable family (Kontext editing) is rejected with a helpful hint.
+    with pytest.raises(ValueError):
+        C(base_model = "b", data_dir = "d", output_dir = "o", model_family = "flux.1-kontext").normalized()
+    # A DiT family that IS trainable resolves to itself.
+    assert (
+        C(base_model = "b", data_dir = "d", output_dir = "o", model_family = "flux.1")
+        .normalized()
+        .resolved_family
+        == "flux.1"
+    )
     # SDXL explicit passes.
     assert (
         C(base_model = "b", data_dir = "d", output_dir = "o", model_family = "sdxl")

studio/backend/tests/test_diffusion_training.py

@@ -484,15 +484,8 @@ def test_route_start_refuses_non_sdxl_base_without_freeing_gpu(client, monkeypat
 def test_apply_event_records_metric_history_and_perf():
     svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
     svc._apply_event(
-        {
-            "type": "progress",
-            "step": 1,
-            "total_steps": 10,
-            "loss": 0.5,
-            "learning_rate": 1e-4,
-            "samples_per_second": 3.2,
-            "peak_memory_gb": 7.1,
-        }
+        {"type": "progress", "step": 1, "total_steps": 10, "loss": 0.5,
+         "learning_rate": 1e-4, "samples_per_second": 3.2, "peak_memory_gb": 7.1}
     )
     svc._apply_event(
         {"type": "progress", "step": 2, "total_steps": 10, "loss": 0.4, "learning_rate": 9e-5}
@@ -538,14 +531,8 @@ def test_metric_history_decimates_at_cap():
 def test_complete_event_records_family_and_catalog():
     svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
     svc._apply_event(
-        {
-            "type": "complete",
-            "output_dir": "/o",
-            "lora_path": "/o/w.safetensors",
-            "catalog_path": "/loras/w.safetensors",
-            "family": "sdxl",
-            "base_model": "b",
-        }
+        {"type": "complete", "output_dir": "/o", "lora_path": "/o/w.safetensors",
+         "catalog_path": "/loras/w.safetensors", "family": "sdxl", "base_model": "b"}
     )
     st = svc.status()
     assert st["status"] == "completed"
@@ -557,12 +544,8 @@ def test_complete_event_records_family_and_catalog():
 def test_status_route_nests_metric_history(client):
     # The status route folds the service's flat arrays into a nested metric_history object.
     client._fake.status_extra = {
-        "metric_steps": [1, 2],
-        "metric_loss": [0.5, 0.4],
-        "metric_lr": [1e-4, 9e-5],
-        "family": "sdxl",
-        "samples_per_second": 2.0,
-        "peak_memory_gb": 6.0,
+        "metric_steps": [1, 2], "metric_loss": [0.5, 0.4], "metric_lr": [1e-4, 9e-5],
+        "family": "sdxl", "samples_per_second": 2.0, "peak_memory_gb": 6.0,
     }
     r = client.get("/api/train/diffusion/status")
     assert r.status_code == 200, r.text
@@ -571,3 +554,55 @@ def test_status_route_nests_metric_history(client):
     assert body["metric_history"]["loss"] == [0.5, 0.4]
     assert body["family"] == "sdxl"
     assert body["samples_per_second"] == 2.0
+
+
+# ── /diffusion/info families + gated-repo preflight (PR B) ──────────────────────
+def test_info_lists_trainable_families(client):
+    r = client.get("/api/train/diffusion/info")
+    assert r.status_code == 200, r.text
+    families = {f["name"]: f for f in r.json()["families"]}
+    for fam in ("sdxl", "flux.1", "qwen-image", "z-image"):
+        assert fam in families
+    assert families["flux.1"]["default_base"] == "black-forest-labs/FLUX.1-dev"
+    assert families["z-image"]["defaults"]["resolution"] == 768
+
+
+def test_start_gated_base_without_access_is_400_and_keeps_gpu(client, monkeypatch):
+    # A gated FLUX base with no valid token must 400 from the HEAD preflight BEFORE the GPU
+    # residents are freed, so a doomed start never evicts the user's loaded model.
+    import urllib.error
+    import urllib.request
+
+    import routes.training as tr
+
+    freed: list[int] = []
+    monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: freed.append(1))
+
+    def _fake_urlopen(req, timeout = None):
+        raise urllib.error.HTTPError(req.full_url, 403, "Forbidden", {}, None)
+
+    monkeypatch.setattr(urllib.request, "urlopen", _fake_urlopen)
+    r = client.post(
+        "/api/train/diffusion/start",
+        json = {**_BODY, "base_model": "black-forest-labs/FLUX.1-dev"},
+    )
+    assert r.status_code == 400
+    assert "gated" in r.json()["detail"].lower()
+    assert freed == []
+    assert client._fake.started_with is None
+
+
+def test_start_ungated_base_preflight_is_noop(client, monkeypatch):
+    # A reachable base (HEAD 200) proceeds to start normally.
+    import urllib.request
+
+    import routes.training as tr
+
+    monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: None)
+    monkeypatch.setattr(urllib.request, "urlopen", lambda req, timeout = None: object())
+    r = client.post(
+        "/api/train/diffusion/start",
+        json = {**_BODY, "base_model": "black-forest-labs/FLUX.1-dev"},
+    )
+    assert r.status_code == 200, r.text
+    assert client._fake.started_with["base_model"] == "black-forest-labs/FLUX.1-dev"