install_sd_cpp_prebuilt: only write the .unsloth-studio-owned marker when the
install created the target directory or it was empty. Adopting a pre-existing,
unowned, non-empty directory (a user's own stable-diffusion.cpp checkout) made
it eligible for the uninstaller's recursive delete.
routes/training upload: make the multi-file promotion transactional. Back up
each displaced original and roll every destination back on any failure, so a
mid-loop rename error can no longer partially overwrite the live dataset.
routes/training _resolve_dataset_folder: reject a symlinked dataset directory
and prove the resolved folder stays under the datasets root, so image
read/caption/delete cannot escape the root through a link.
routes/training delete: escape glob metacharacters in the thumbnail filename so
deleting an image named like [ab].png removes only its own thumbnails.
image_gallery / video_gallery listing: filter records against the response
schema inside the pager via a valid callback, so offset/limit/has_more all count
over accepted records. A leading schema-invalid record no longer returns an
empty page with has_more=true and stalls infinite scroll at offset 0.
image_gallery / video_gallery save: publish via a temp file plus atomic rename
(the sidecar is the video pair's commit marker) and clean up on failure, so a
partial write never surfaces a truncated PNG or strands an orphan MP4.
diffusion_train_common discovery: treat an empty caption sidecar as a metadata
tombstone that still falls through to the dreambooth instance prompt, so
clearing every metadata caption no longer fails with no captioned images found.
diffusion backend unload: wait for an in-flight denoise to exit before tearing
down process-wide patches and state, mirroring the load path.
diffusion_engine_router: serialize the whole check/unload/publish transition so
a concurrent selection cannot return the engine being unloaded.
uninstall.ps1: gate the default sd.cpp process stop on the owner marker so a
user's own sd-server is not terminated for a directory we then keep.
Route on-device single-checkpoint video folders through the single_file loader:
a bare local .safetensors directory (no model_index.json) is advertised as a
pipeline with no filename, so validation rejected it before it could load.
Reinterpret the pick as a single_file load of the sole checkpoint, mirroring the
image load route.
Treat a reserved-but-not-yet-spawned LLM training start as active in
is_training_active() so /images/load, /video/load, and /diffusion/start cannot
race the reserved run for VRAM during the pre-spawn free window. Mirrors the
diffusion training service reservation.
Resume an in-flight image generation on the Images page mount: probe
generate-progress, re-enter the poll loop, and refresh the gallery on completion
so a run started elsewhere is reflected and its saved image appears without a
manual refresh. Seed resident image defaults from the resolved base_repo rather
than a possibly path-shaped repo_id so the first resident generation uses the
right recipe.
* Auto-detect completion masking markers with template table fallback
Studio's train_on_completions previously relied only on the hardcoded
MODEL_TO_TEMPLATE_MAPPER / TEMPLATE_TO_RESPONSES_MAPPER tables and
silently disabled masking when a model was not in the table, so unmapped
models (LFM2-8B-A1B, DeepSeek, and others) trained on full sequences
without telling the user. Several mapped templates (glm, mistral, llama,
starling, zephyr, qwen3-thinking) also carried markers that mask every
assistant token, which made every row drop in the post-masking filter.
Both training callsites (CUDA trainer.py and MLX worker.py) now share
utils.datasets.completion_masking.apply_completion_masking:
- Try unsloth_zoo chat template auto-detection first; it raises loudly
when the template cannot be parsed and never masks the EOS token.
- gpt-oss models keep their manual markers so non-final assistant
<|end|> tokens stay trained, matching current behavior.
- If auto-detection raises, fall back to the template table exactly as
before.
- If the table also misses, emit an explicit user-visible warning that
completion masking could not be applied and full-sequence training
will occur, instead of a quiet log line.
The >30 percent dropped-rows safety net in trainer.py now guards the
auto path as well. Table consumers for inference and chat templates are
unchanged. Validated against one representative tokenizer for every
template in TEMPLATE_TO_RESPONSES_MAPPER plus the unmapped models:
no template regresses; unit tests cover the four decision paths.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Restrict masking fallback to marker detection failures
The auto branch wrapped the whole train_on_responses_only call, so a real
failure while applying the masking (dataset map, tokenization) was treated
as a detection miss and training silently proceeded on full sequences.
Detect markers separately via get_chat_template_parts (test seam via
detect_fn), then apply them with errors propagating, matching the manual
path. Tokenizers with preset unsloth marker attrs skip detection and call
bare so zoo reuses the stored parts.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fail the run when applying completion masking raises
The helper already falls back internally on detection failures and returns
applied=False on a double miss, so an exception reaching the callsites is a
real failure applying the masking. Remove the callsite catches that
downgraded it to full-sequence training; the run now fails visibly instead.
Also use the explicit re-export alias form in utils/datasets/__init__.py for
the two new names, satisfying the import-hoist source lint.
* Import completion masking from its submodule
The import-hoist source lint counts only real name loads, so package-level
re-exports of the two new names cannot satisfy it. Import
apply_completion_masking from utils.datasets.completion_masking directly at
both callsites and leave utils/datasets/__init__.py untouched.
* Completion masking: gpt-oss renames and MLX raw/alpaca parity
Renamed or private gpt-oss checkpoints are name-detected as gpt-oss but miss
the exact-name table; default them to the gpt-oss template markers instead of
falling through to full-sequence training.
Gate the MLX masking call on not raw_text_mode and format_type != alpaca,
mirroring the CUDA path: raw/CPT text has no chat turns to mask and
Alpaca-rendered text lacks the tokenizer's chat markers.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Define raw_text_mode outside the MLX feature-detect block
With an older zoo lacking the append_eos config field, the masking
gate referenced raw_text_mode before assignment. Hoist the assignment
above the feature detection so both consumers see it.
* Gate MLX masking on the formatter's resolved format
format_type auto can resolve to alpaca or raw text; the masking skip
checked only the requested value, so auto-detected Alpaca data got
chat-template markers applied to rendered prompt text. Track the
final_format returned by format_and_template_dataset and gate on it,
matching the CUDA path.
* Unwrap the mlx-lm TokenizerWrapper before marker checks
The wrapper delegates plain reads to the wrapped HF tokenizer but hides
underscore attrs, so preset unsloth markers were invisible and detection
relied on the loader's call patch. Unwrap to the real tokenizer first,
as the zoo MLX resolver does.
* Tighten masking comments
* gpt-oss: auto-detect markers first like every other template
The quantized and BF16 gpt-oss checkpoints ship a chat template without
the channel final header, so the pinned manual markers match nothing
there and masking trained zero tokens. Auto-detection derives markers
from whichever template the checkpoint ships and keeps the final
terminator trained; the manual gpt-oss markers remain the detection
failure fallback, including for renamed checkpoints.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Tighten comments
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Run backend.start_training off the event loop with asyncio.to_thread so the
synchronous diffusion/video unload calls (which wait on engine generation
locks) cannot freeze concurrent requests; guard against overlapping starts
with a _start_in_progress compare-and-set under the service lock.
Resolve bare diffusion dataset names directly under datasets_root() before
falling back to the generic resolver, so an unrelated LLM upload file or
recipe folder sharing the name cannot shadow the image dataset.
Reject exact duplicate filenames within one multipart upload batch: two
parts staged to the same destination would let the later tmp.replace
silently discard the earlier file. Case variants stay exempt per the
existing stem-guard contract.
Require an instance prompt in the train panel when only some images have
captions, since backend discovery silently skips uncaptioned images.
torch.cuda.is_bf16_supported() reports True on pre-Ampere GPUs that only
emulate bf16, so the SDXL LoRA trainer would keep bf16 there and fail at
load/forward. Use native_bf16_supported() (the same compute-capability
probe the DiT trainer already uses) so T4 / V100 / RTX 20xx fall back to
fp16 instead.
* feat(studio): route CLI trainer to MLX backend
* fix(studio): harden MLX trainer routing
* fix(studio): harden MLX trainer adapter routing
* test(studio): assert MLX CLI activation order
* fix(studio): address MLX CLI review feedback
* feat(cli): support MLX in legacy script
* fix(cli): adapt MLX tokenizer for raw text
* fix(cli): omit unsupported MLX eval batch arg
* fix(cli): feed raw text to MLX trainer
* Fix CLI MLX routing and Python 3.9 annotations
Route the MLX backend through create_mlx_trainer_adapter so the torch-free
Apple Silicon path never imports trainer.py (torch/unsloth/trl). Replace
from __future__ import annotations with typing.Optional/Union so the CLI
annotations stay Python 3.9 compatible without the unused-import lint hit.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Strip return_tensors from MLX raw-text tokenizer proxy
On a torch-free MLX install, RawTextDataLoader calls the tokenizer with
return_tensors='pt'; the callable proxy forwarded that to the HF
tokenizer, which tried to build torch tensors and failed before
training. Drop return_tensors so the MLX path returns plain token ids.
* Tighten CLI MLX-backend comments
---------
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
torch.cuda.is_bf16_supported() defaults to counting pre-Ampere bf16 EMULATION as
supported, so on a T4/V100/RTX 20xx the DiT-training bf16 gates all passed even
though the trainer requires native Ampere-or-newer bf16: /diffusion/info advertised
the DiT precision modes, /diffusion/start's preflight let the run through and freed
resident GPU models, then the trainer child hit the real unsupported bf16 path. The
inference device resolver already fixed this (issue #6658) by gating NVIDIA on
capability major >= 8; the training path never got it. Add a shared
native_bf16_supported() helper (NVIDIA cap major >= 8; ROCm keeps the trustworthy
is_bf16_supported()) and use it in the three DiT bf16 sites -- train_precision_modes,
bf16_unsupported_reason, and the trainer guard -- so a pre-Ampere card is offered
nf4 only and never advertises/evicts-then-fails. Tests now exercise the emulation
case (is_bf16_supported True but capability < 8).
The SDXL trainer drew min(train_batch_size, len(pairs)) indices, so a dataset with
fewer images than the batch trained at a smaller effective batch than configured
while the scheduler and samples-per-second still assumed the full batch. The shared
PermutationBatchSampler already refills across permutation cycles to return exactly k
indices, and the DiT trainer calls it with the full batch size, so drop the clamp and
pass train_batch_size through for parity and to honor the configured batch.
Two evict/corruption fixes surfaced by review of the diffusion training path:
- krea-2 sets force_bf16 in the DiT trainer spec, but the route-level
_FORCE_BF16_FAMILIES preflight listed only qwen-image and z-image, so a
krea-2 start with mixed_precision=fp16 passed the route check, reserved
training and evicted resident GPU models, and only the child trainer then
raised. Add krea-2 to the set and a drift-guard test asserting it equals
the trainer specs whose force_bf16 is set.
- The dataset-upload same-stem duplicate check compared stems
case-sensitively, so on case-insensitive filesystems (Windows / default
macOS) sample.png and Sample.jpg both passed even though their caption
sidecars sample.txt / Sample.txt resolve to the same file, silently
sharing and corrupting one caption. Compare stems and the same-name guard
with casefold at both the on-disk and in-batch sites.
Fold PR #6872's image-generation fixes into the branch, deduped against the
round-12 dataset-upload and gallery integrity work already on image-generation.
Fixes carried forward from #6872:
- fp8 single-file transformer memory estimate: an fp8 checkpoint loads with no
quantization_config and diffusers upcasts it to bf16 (~2x resident), so budget
it accordingly in _plan_memory and estimate_safetensors_dense_mib.
- dense-quant OOM-evict preflight: when the GGUF fits resident but the dense bf16
transformer this path materializes does not, skip the fast path up front rather
than evict the current pipeline and OOM in finalization. Combined with the
existing offload->resident candidate re-plan so both the family-table estimate
and the on-disk shard measurement gate engagement (unified on the
transformer_resident_override_mib plan override).
- ControlNet: evict the previous module and its from_pipe wrapper before loading a
new one so swapping ControlNets within a base-model load cannot accumulate to OOM.
- ControlNet union_control_mode: raise on an unknown control type instead of
silently defaulting to canny.
- edit-family mask rejection: raise instead of silently dropping a mask on an
image-editing model that has no inpaint pipeline.
- companion cache: walk the snapshot dir and exclude transformer/ so the
dense-quant prefetch's cached shards do not inflate the companion total and
wrongly force offload.
- training: drop piecewise_constant from the LR scheduler enum and force bf16 for
fp16-incompatible families.
- dataset upload: batch-atomic staging with the same-stem duplicate guard.
- images page: guard negative-prompt restore on guidance>0, clear stale ControlNet
selection on restore, and revert an optimistic quant label when a pipeline load
never starts.
- uninstall (sh + ps1): keep the owner-marker guard on sd.cpp removal.
Conflicts resolved in favour of image-generation's evolved memory system,
loadSpecFor catalog, and stop-and-save (lora_path) run detection; #6872's fp8 and
dense-preflight fixes carried forward on top. All affected backend tests pass
(test_diffusion_backend, test_diffusion_training, test_diffusion_lora_trainer,
test_video_gallery, test_diffusion_controlnet).
piecewise_constant is the only diffusers scheduler that needs a step_rules string, and
neither diffusion trainer passes one (get_scheduler is called with only warmup/training
steps, and there is no config field for it). Accepting it let /diffusion/start pass
normalized(), free the resident GPU workloads, spawn the trainer, and only then crash in
the subprocess (get_piecewise_constant_schedule does step_rules.split(",") on None) -- the
exact evict-then-fail the up-front validation exists to prevent. Reject it now with a clear
400. The remaining six schedulers all run with only warmup/training steps.
Two overlapping /diffusion/start requests can interleave between the is_active()
check and the reservation, so reserve() itself must reject a second reservation
atomically. Otherwise both callers reserve, both free the GPU's resident chat or
image model, and the loser only 409s after the eviction -- the evict-then-fail the
reservation exists to prevent. reserve() now raises under the lock if a start is
already reserved or a job is already running.
start_diffusion_training freed resident GPU models and only then called
service.start(config), which is where is_active() first flips true. During that
free-then-spawn window a concurrent /images/load or /video/load saw training as
inactive, passed its training guard, acquired the GPU, and began a background load,
so the trainer and that pipeline both allocated VRAM. Add reserve()/unreserve() to
the training service (is_active() also reports the reservation) and reserve BEFORE
the free, in a try/finally so a failed start rolls the reservation back. An
overlapping load's guard now refuses during the window. Regression tests: the route
reserves before the free (and the free sees an active service), and the service
reservation marks active then rolls back.
Several image/video/training preflights ran before the route acquires the GPU or
frees resident models, but let a doomed local pick through and only failed deep in
the background load, after the user's chat/Images/Video model was already evicted.
- Local base_repo / base_model: _is_trusted_diffusion_repo accepts any existing
local path, but the base loads via from_pretrained (needs model_index.json). A
local dir that is not a diffusers pipeline passed the trust gate, evicted the
resident model, then failed. Add a shared _assert_local_base_is_pipeline check
and call it in the image, video, and training preflights.
- Dataset images: discover_image_caption_pairs only checked filenames, so a
corrupt or zero-byte upload passed the start-route preflight, freed the GPU, then
crashed the spawned trainer in PIL. Add an opt-in verify_images decode probe
(cheap PIL header check) that the start route enables; the trainers leave it off
since they decode every image anyway.
- Local single-file safetensors: the On-Device scanner advertises a bare
.safetensors directory (no model_index.json) as a text-to-image model, but the
picker starts it as a pipeline with no filename, so every click 400s. Reinterpret
such a pick as a single_file load of the sole checkpoint (resolve_local_single_file)
so the advertised model is actually loadable.
Regression tests for each: local non-pipeline base (image/video/training), the
verify_images decode gate, and resolve_local_single_file.
The start route's precision preflight folded bf16/int8/fp8 into the CUDA
requirement but omitted mxfp8, so an mxfp8 request on a GPU-less host (or an
older CUDA GPU without Blackwell) passed the preflight, evicted resident image
and chat models, then raised only in the spawned trainer child. Mirror
_resolve_base_precision: require CUDA for mxfp8 and re-check the Blackwell
(sm100+) capability up front, so a doomed run is rejected before teardown.
The start-route preflight caught the bf16-GPU and int8-torchao requirements but not the dense
precisions' CUDA requirement: on a GPU-less host bf16_unsupported_reason exempts CPU-only, so a
bf16/fp8 (or int8-with-torchao) DiT request passed the preflight, evicted resident workloads, then
raised only in the trainer child. Add the dense-mode CUDA gate mirroring _resolve_base_precision so
the doomed run is rejected up front. Also pin bf16_unsupported_reason in the two positive-path
family-info tests so they are deterministic across GPU types (a non-bf16 CUDA box would otherwise
empty every DiT family's advertised modes).
The start route preflight only rejected non-bf16 GPUs; an explicit int8 request on
a host with a missing or stub torchao passed the preflight, evicted resident GPU
workloads, then died in the trainer child (its int8 base quantizer has no fallback).
Fold both gates into training_precision_preflight_error so int8-without-torchao fails
fast before eviction. Also empty the advertised DiT precision_modes (and surface the
reason in vram_note, drop compile) whenever the bf16 preflight would reject the family,
so /info never offers an nf4 DiT option the route always 400s.
- DiffusionFamily gains deploy_base_repo (krea/Krea-2-Turbo): deploying a LoRA
trained on Raw now previews it on Turbo, not the non-distilled Raw checkpoint.
Scoped to a same-precision override so it never turns an nf4 train base into a
larger bf16 deploy load; exposed through family_train_infos -> the Train UI's
onDeployClick / historical-run deploy resolve the deploy base.
- _GENERATION_DEFAULTS gains a Krea entry (8 steps, 0 CFG) so the OpenAI
/v1/images/generations route matches the Create UI's documented distilled recipe
instead of falling through to the generic (9, 0.0).
- normalized() + family_train_infos() mirror the inference fp8 deny for
Qwen-Image (activation outliers exceed fp8's range and corrupt the trained
result); int8 stays allowed and the UI no longer advertises fp8 for it.
- _resolve_base_precision() gates an explicit int8 on a FUNCTIONAL torchao, the
same gate auto and /info already apply, so a missing/stub torchao fails fast
instead of silently loading dense with compile disabled.
- train_precision_modes() gates the dense modes (bf16/int8/fp8/auto) on
torch.cuda.is_bf16_supported(), so a non-bf16 CUDA GPU (T4/V100/RTX 20xx) is
offered only nf4 instead of a start that evicts resident models and then fails.
- start_diffusion_training preflights bf16 support for the DiT families BEFORE
_free_gpu_for_diffusion_training(), so any DiT start (nf4 included, since the
trainer requires bf16 unconditionally on CUDA) fails fast without eviction.
- The torchao 0.17 MX training path swaps a matched frozen Linear's weight for a wrapper tensor
whose linear override computes input @ weight_t and drops the bias, so mxfp8'ing a biased frozen
linear silently loses its bias and corrupts the base output the LoRA regresses against (verified
on Blackwell: the bias term is fully dropped). Skip biased linears in _mx_module_filter.
- _resolve_base_precision re-checked explicit dense modes against the live device but only rejected
CPU, so an explicit mxfp8 request on a non-Blackwell CUDA GPU passed and then crashed at the first
MX GEMM after a full dense-transformer load. /info only advertises mxfp8 on sm100+; mirror that
gate here and fail fast for a stale or direct client below Blackwell.
The latent cache holds two fp32 posterior tensors per crop/flip variant per
image, pinned on CUDA hosts, so datasets with thousands of images can exhaust
host or pinned memory with no fallback. Estimate the cache size from the first
real encoded latent and fall back to per-step VAE encoding when it exceeds a
4 GiB budget. UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE bypasses the gate; the
existing UNSLOTH_DIFFUSION_NO_LATENT_CACHE opt-out is unchanged.
Replace the with-replacement per-batch index draw in the SDXL and DiT LoRA
trainers with a shared PermutationBatchSampler that visits every image once per
cycle before repeating, so short runs cover the whole dataset. The sampler
reshuffles from the run's rng so the index stream stays seed-deterministic.
Guard the diffusion run detail route against a valid-JSON non-object record,
which previously raised TypeError and returned a 500; it now 404s like the list
path's shape check.
Add regression tests for both.
The perf rewrite dropped the bf16 capability guard, so a pre-Ampere CUDA
device (T4/V100/RTX 20xx) would die deep in model load with an opaque dtype
error instead of a clear message. Restores parity with the SDXL trainer.
