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.
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.
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.
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).
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.
- 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.
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.
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.
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.
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.
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.
Backend:
- validate_load_request rejects a non-.gguf single-file name before the GPU
handoff, so a family-looking repo paired with README.md no longer evicts the
chat model and only fails in the background load.
- detect_family scopes the edit/kontext/inpaint keyword check to the model id
or filename basename, not arbitrary parent directories, so a valid
text-to-image file under a folder named edit is no longer rejected.
- the images gallery listing skips records that fail schema validation, so one
corrupt or hand-dropped PNG can no longer 500 the whole endpoint.
- _terminate reaps the killed sd-cli child so cancellation and timeout paths do
not leak zombie process-table entries.
- the images load route gates the chat-eviction handoff on the resolved device
being non-CPU, so a CPU-only diffusers fallback no longer evicts a resident
chat model for a load that cannot use the GPU.
Frontend:
- treat Images as a chat-like full-height route (no outer padding or scroll) so
its picker is not pushed down and the gallery is not clipped.
- allow /images under the chat-only guard so the native CPU/MPS image path is
reachable on the no-GPU hosts it was built for.
- roll the optimistic quant label back when a same-repo swap fails after the
load started, so the selector never advertises a quant that is not loaded.
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.
plan_diffusion_memory only applies the legacy cpu_offload override when no
memory_mode was supplied, matching the documented API contract that
memory_mode overrides cpu_offload when set; an explicit fast request now
stays resident even if the old flag is also enabled.
The transformer-quant dense path fetches the base repo's transformer/
shards inside the locked finalize phase, where unload and cancellation
cannot preempt the multi-GB download. The load worker now widens the
preemptible prefetch to include those shards when that path can actually
run: quant requested and supported for the device, scheme resolvable, and
no pre-quantized checkpoint shortcutting the dense build.
Memory planning and dense-quant path: size a local diffusers base's
resident companions from its on-disk VAE and text-encoder weights instead
of folding them to zero, feed the distilled variant hint into the runtime
headroom estimate so turbo and schnell models are not over-reserved, place
group-offload companions resident before attaching the transformer hooks
so a failed placement falls back to whole-module offload instead of
crashing, and bail out of the dense transformer download before it starts
when the requested quant scheme is unsupported so the load falls back to
GGUF cleanly.
sd.cpp stack: scrub the native path lease secret from sd-cli child env,
redact native load-progress errors, forward the resolved accelerator when
auto-installing a forced-native binary, release stale diffusion GPU
ownership on CPU-native loads, and remove the sd.cpp install tree on
uninstall.
Prequant and scripts: reject prequant artifacts missing base_model_id
when a base is requested, expanduser before checkpoint existence checks,
record and validate the int8 exclusion filter and fp8 fast-accum in
checkpoint metadata, make verify_prequant_backend allowlist its local
checkpoint and fail on missing or bad LPIPS and on load-peak regressions,
average only finite PSNR values in diffusion_quality, and reset the
process-wide attention backend between perf probe variants.
API and UI: normalize attention_backend casing before Literal validation,
close hidden popovers when leaving the Images page, and clear the stale
quant label when loading a direct local GGUF file.
