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.
generate() assigned self._gen only at the pipe() call, after deferred
compile, LoRA resolution/application, and ControlNet download/build had
run. Across that setup window generate_progress() reported inactive even
though _generate_lock was held, so a reloaded page's mount probe showed
idle and let a second generate queue behind the first.
Publish an active step-0 _GenState the moment the generation lock is
acquired, before the setup work, and clear it in the outer finally so a
setup-time error cannot leave the UI stuck active. Mirrors the video
backend's queued phase and the training start guard.
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.
Register the dims the forward actually compiled with: image-conditioned
workflows (img2img, inpaint, upscale, edit) run at the input image's size,
not the slider's, so recording the slider values marked never-compiled
shapes as covered and warm restarts kept paying compile for the real one.
Validate a request-supplied transformer_prequant_path (existence plus the
UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH allowlist) before treating prequant as
available at the resident-fit re-check: an unusable path skipped the dense
fit check up front and then fell back to materializing dense bf16 after
the previous pipeline was evicted, recreating the post-eviction OOM path.
Shared as usable_prequant_source, also used by the auto-policy planner.
* 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>
diffusion_compile_cache: auto mode now saves the Mega-cache bundle after the
first compiled generation (UNSLOTH_DIFFUSION_COMPILE_CACHE_SAVE=0 opts out), so
users get warm restarts without the distributor env; a bundle hit starts clean
(no pointless rewrite of the just-loaded artifacts) and explicit mode 1/on keeps
the distributor-style re-save. New register_shape + manifest shape coverage: a
STATIC compile produces new artifacts per (width, height, batch), so the
generate path registers each generation's shape and an uncovered shape
re-dirties the context, growing the bundle to cover every shape the session
used. Measured (B200, real backend): Qwen-Image deferred gen-3 hitch 29.1 ->
22.2 s warm with bit-identical output (7.9 MB bundle, ~0.5 s save); SDXL gen-3
115.7 -> 24.7 s and a mid-session 768px recompile 65.8 -> 12.6 s (bundle 63.6 ->
98.7 MB after the 768 re-save).
diffusion_speed: U-Net denoisers (UNet2DConditionModel; no _repeated_blocks, so
the regional compile never reached them) now get a whole-module STATIC
torch.compile on the default tier, plus fused QKV projections and a compiled VAE
decode. Measured on SDXL (30 steps / 7.0 / 1024px, 4 prompts, LPIPS vs the
bit-exact reference): 6.16 -> 3.14 s end to end (1.96x) at LPIPS 0.035, steady
state 0.70-0.88 s/image through the real backend. Rejected on measurement:
dynamic=True whole-module (366 s compile for 39.3 ms/step vs static's 73 s for
26.9), regional BasicTransformerBlock only (45.0 ms/step; ResNet convs stay
eager), max-autotune + inductor flags (25.9 ms/step for a 445 s warmup),
channels-last UNet alone (neutral). DiT tiers unchanged: fused QKV measured
exactly neutral under the regional compile (Qwen-Image 6.53 vs 6.52 s), so it
stays max-only there, and the DiT VAE decode stays eager (a few % of a DiT
generation). compiled_shapes_are_static tells the cache layer which loads are
per-shape (max tier, U-Net whole-module).
diffusion: register each generation's shape with the compile cache before the
save, pass pipe.unet to the cache fingerprint when the pipe has no transformer,
and correct the transformer_quant resolved reason on dense loads (it claimed a
GGUF transformer was loaded on every non-quantized pipeline load).
Tests: 333 passing across the related suites (speed 42, compile_cache 27, cache
40, precision 20, backend, base_precision, transformer_quant, memory); ruff
clean. Full measurement record: outputs/image_optim_round2_audit.md.
Applies the video round-2 accuracy findings to the image diffusion stack and fixes
two real image-path bugs found while measuring. All numbers B200, production
settings (family default steps/guidance, 1024px, seed 42, 4 fixed prompts), LPIPS
(AlexNet) via the new scripts/image_speedmem_bench.py, which drives the production
lever functions in the loader's own order.
- inductor precision parity: emulate_precision_casts=True on the regional-compile
path (fused pointwise kernels keep fp32 intermediates where eager rounds to bf16
between ops). Pairwise LPIPS of the compiled tier vs the same-stack eager tier:
Qwen-Image 0.019 to 0.006 at identical speed (72.4 vs 72.5 ms/step), FLUX.1-dev
0.046 to 0.029 at +2% step time (69.8 vs 68.3, reproduced), FLUX.2-klein-4B
0.018 to 0.017 at identical speed. Snapshot/restored with the other process-wide
backend flags so an off load never inherits it.
- cache x compile composition: re-point each cache hook's fn_ref.original_forward
at a torch.compile'd wrapper of the same bound method (armed only where the
speed layer compiled the block; restored before every disable_cache and before
the partial-hook cleanup). Qwen-Image FBCache computed steps 91.8 to 71.2 ms
(back at the uncached compiled rate), 1.21x end to end (7.36 to 6.06 s per 4
images); FLUX.1-dev already traced through its FBCache hook and is measured
neutral (same-process armed vs unarmed latents bit-identical). Skip counts
within noise (13 vs 11 of 76; pairwise LPIPS 0.005).
- FBCache mid-session toggle crash: diffusers 0.39 caches the HookRegistry child
list on first cache_context use, so an uncached generation followed by a
20+-step generation (the auto toggle path) enabled hooks the context never
reached and crashed with "No context is set" (reproduced live on FLUX.1-dev).
Invalidate the stale child cache after every enable_cache.
- TE fp8_dynamic zero-row guard: torchao per-row fp8 derives a per-output-channel
scale from the row amax, so an all-zero weight row is 0/0 = NaN. SDXL's
text_encoder_2 (OpenCLIP bigG) ships exactly such a row, and every explicit
fp8_dynamic SDXL render came out black; keep zero-row Linears dense (LPIPS
0.976 black to 0.096 working). Other families' encoders have no such rows and
are byte-identical.
- No AUTO TE quant exists on the image branch (text_encoder_quant defaults dense,
explicit-only), so the video round's auto-dense retune has no image analogue;
the explicit lever's cost is now measured (TE fp8_dynamic alone, LPIPS vs
bit-exact: Qwen-Image 0.038, FLUX.1-dev 0.084, SDXL 0.096; no speed win, VRAM
-6.5 GB on Qwen-Image) for the docs.
Tests: 96 passing across the cache/speed/precision suites (11 new arming, 2
child-registry, 2 zero-row, 4 inductor-flag); ruff clean.
POST /video/generate previously held the response open for the whole
generation (multi-minute for 720p), so in --secure mode the Cloudflare
quick tunnel's ~100s origin-response cap returned a 524 while the server
kept generating, and the frontend treated the run as failed.
Generation now follows the same return-at-once pattern as /video/load:
begin_generate validates synchronously (409 on no model or on a second
concurrent generate via a new busy sentinel) and runs the existing
generate + gallery-persist pipeline, with the route's exact error
mapping, on a daemon thread. GET /video/generate-progress gains optional
terminal fields: phase completed carries the saved gallery record, phase
failed a client-safe error; active only drops together with a terminal
phase. The cancel event is registered before the worker starts so
/video/generate/cancel keeps working across the whole job.
VideoGenerateResponse becomes an accepted acknowledgement (status
started, video kept as an always-null compat field). The video page
fires the POST, then drives completion off the progress poll it already
runs (completed prepends the clip, failed surfaces the error, the
cancelled sentinel stays toast-free). The API-key training-start guards
now also probe the video backend for an in-flight background clip, since
it is no longer visible as an in-flight HTTP request to the keep-warm
counter.
Route tests keep the fake backend for load/generate/status but inherit
the real job machinery, covering immediate accept, concurrent 409, the
terminal completed record, sanitized/ValueError/cancelled failures, and
cancel of a running job.
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.
* Studio: resolve the repo-root MTP drafter after the MTP/ GGUF rename
The Gemma 4 QAT GGUF repos renamed the higher-precision MTP/ subdir
copies from gemma-4-...-<quant>-MTP.gguf to mtp-gemma-4-...-<quant>.gguf,
so their basenames now start with the same mtp- prefix as the small
repo-root drafter (mtp-gemma-4-E4B-it.gguf).
The drafter selectors filtered candidates by a mtp- basename prefix and
took the first in sort order. With the new names the MTP/ copies also
match, and because MTP/ (uppercase) sorts before the lowercase root file,
selection flipped to the large BF16 copy under MTP/ instead of the root
drafter both functions document they should pick.
Restrict both selectors, and the companion byte estimate, to root-level
mtp-*.gguf so the MTP/ copies stay explicit-selection only:
- core/inference/llama_cpp.py _pick_mtp (loader auto-download)
- hub/utils/gguf_plan.py preferred_mtp_sibling (Hub variant plans)
- routes/inference.py _remote_gguf_companion_bytes (VRAM headroom)
Also reuse a drafter already in the local cache before downloading, so a
device that already holds a copy on disk does not re-fetch it.
Old-scheme names keep working (they have no root-level mtp- sibling to
mis-select). Adds regression tests for the new naming, both selection
paths, and the on-disk reuse.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: gate MTP drafter cache reuse to offline mode
Reuse the cached drafter only when HF is offline. Online, route back
through _download_companion_gguf/hf_hub_download so the current revision
is checked (etag) and a changed drafter is refetched, matching the
offline-only cross-snapshot reuse already used for the main GGUF. This
avoids pairing freshly downloaded weights with a stale cached draft.
Make the reuse tests offline and add an online-skips-reuse test.
* Studio: prefer a root MTP drafter across all cached snapshots
Offline reuse scanned snapshots one at a time and returned the first
snapshot that held any drafter, only preferring root within it. A newer
partial snapshot with just the MTP/ copy could shadow the small root
drafter in an older snapshot. Collect drafters across all snapshots and
prefer any repo-root file before an MTP/ copy.
* Studio: keep newest-first snapshot order when reusing cached drafters
Collecting root candidates and sorting by absolute snapshot path could
pick a drafter from an older snapshot. _iter_hf_cache_snapshots yields
newest first and the main GGUF is resolved in that order, so preserve it
(root still preferred over MTP/ copies) to avoid pairing a fresh main
weight with a stale drafter revision.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Studio: add Vulkan llama.cpp support
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address gemini's feedback
* Studio: move the Vulkan VRAM probe into a standalone script
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Improve Vulkan probe error reporting
* Resolve llama-server symlink so Vulkan build is detected
* Drop unreachable Vulkan fallback in GPU free-memory dispatcher
* Skip the Intel GPU probe when NVIDIA or ROCm is present
* Reserve host RAM headroom for Vulkan integrated GPUs
* Add a `UNSLOTH_FORCE_VULKAN` environment variable
* [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
* Honor GGML_VK_VISIBLE_DEVICES, reserve discrete Vulkan VRAM headroom, and clear Intel GPU on --cpu-fallback
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Route Intel and forced-Vulkan hosts to the upstream Vulkan prebuilt, add arm64 Vulkan, keep Vulkan out of RAG auto-detect
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Clear the fork release pin when routing a Vulkan host to the upstream repo
* Gate auto-Vulkan routing on no physical NVIDIA so hidden CUDA devices aren't used
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Pin Vulkan launches with --device Vulkan<i> instead of the raw GGML_VK_VISIBLE_DEVICES index space
* Let user --device override the Vulkan pin, and gate direct Vulkan asset picks on no physical NVIDIA
* Update RAG auto-backend test mocks for the _resolve_auto binary and Vulkan probes
* Keep the add_dll_directory handle alive through the Vulkan probe DLL loads
* Revert RAG auto Vulkan guard, guard multi-backend Vulkan detection, and preserve forced Vulkan across updates
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Use getattr for RTLD_GLOBAL in the Vulkan probe CDLL mode
* Skip CUDA/ROCm APU and datacenter GPU tuning on Vulkan builds
On a Vulkan llama.cpp build gpu_indices are ggml compact ordinals, not
CUDA/ROCm physical ids, so _amd_apu_wants_unified_memory and
_apply_datacenter_env were reading the wrong device. On a mixed AMD APU
plus discrete GPU host that could raise a spurious system-RAM shortfall
and block a valid discrete-GPU load. Gate all three call sites on
not is_vulkan_backend; the Vulkan path already reserves iGPU host
headroom and the backend ignores GGML_CUDA_* anyway.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Tighten Vulkan-guard comment in load_model
* Reduce comments in Vulkan support to be more succinct
* Resolve shell-wrapper llama-server entrypoint to the real lib dir
create_exec_entrypoint falls back to a #!/bin/sh wrapper at the install
root when it cannot symlink into build/bin. _find_llama_server_binary
returns that root entrypoint, but Path.resolve() does not follow a shell
wrapper, so _llama_lib_dir returned the install root and _is_vulkan_backend
missed libggml-vulkan.so -- silently skipping the Vulkan probe and --device
pin on an otherwise valid Vulkan install. Follow the wrapper's exec target
to build/bin. Regression test: test_shell_wrapper_entrypoint_resolves_to_real_lib_dir.
* [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: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: danielhanchen <danielhanchen@gmail.com>
- diffusion_cache: do not engage FBCache when the selected pipeline opens no cache_context.
A CacheMixin transformer is necessary but not sufficient -- Flux Kontext / img2img /
inpaint / controlnet reuse the CacheMixin FluxTransformer2DModel yet their __call__ never
opens a cache_context, so the First-Block-Cache hook raised 'No context is set' on the
first forward, crashing every default FLUX.1-Kontext edit (28 steps, above the FBCache
threshold). Detect it from the pipeline __call__ source, resolved off the instance so the
per-expert proxy view delegates to the real pipe.
- diffusion_attention: honor an explicit aiter backend on ROCm/AMD targets instead of
dropping it via the NVIDIA-only guard (aiter is the AMD ROCm kernel; it only works there).
- video: clear the CUDA cache on a failed load so a partially built pipeline's reserved VRAM
does not OOM the next load (mirrors the image backend), and re-check cancellation after the
export/mux so a clip cancelled during the blocking encode is discarded, not persisted.
- diffusion_auto_policy / diffusion_prequant: validate a request-supplied prequant path
override (present AND allowlisted) before budgeting the small prequant plan, so the loader
does not skip the dense shards and then rebuild dense after evicting the resident pipeline.
- diffusion_controlnet: family-gate a curated ControlNet addressed by its full repo id, not
only its short catalog id, so a cross-family repo id 400s up front instead of downloading
and loading through the wrong ControlNet class.
Five studio diffusion source files added by earlier sub-PRs (Krea 2 Turbo,
Images LoRA, ControlNet) were missing the AGPL-3.0 SPDX header the rest of
studio/backend carries. Add it so the whole backend is consistently licensed.
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.
* Studio: render thinking blocks for safetensors inference with prefilled <think> templates
Reasoning templates like Qwen3.6 end the generation prompt with an open
<think> tag. skip_prompt streaming drops it, so the frontend never sees
the opening tag and shows reasoning as plain text. Detect the prefill
and re-emit it at the start of the stream on the transformers and MLX
paths. Also stop stripping think tags in _clean_generated_text when a
tokenizer marks them special.
* Studio: guard think re-emit for special close tags, yield prefill early
Address review feedback:
- Guard: skip re-emitting the open <think> when the tokenizer marks </think>
as a special token, since skip_special_tokens would strip the model's close
tag and leave an unclosed block that swallows the answer. Falls back to
plain text (pre-fix behaviour) for those tokenizers.
- Yield the prefilled <think> before the first token so the thinking block
renders during prompt prefill instead of after the first generated token.
- Drop the now-unnecessary _clean_generated_text think-tag exemption; the
guard handles the special-token case at the source.
No mainstream reasoning model (Qwen3.6, Qwen3, DeepSeek-R1, QwQ, GLM-4.6)
marks think tags special, so behaviour is unchanged for them.
* [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: Lyxot <longyixing331@gmail.com>
* Studio: allow CPU-only DiffusionGemma by granting the diffusion runner the CPU device
* Studio: mark CPU-only DiffusionGemma as non-GPU-resident for training VRAM preflight
* Studio: keep the CPU DiffusionGemma change minimal (revert VRAM-flag tweak; Metal hosts still hold unified memory)
* Studio: keep CPU DiffusionGemma fallback fully CPU-masked so a masked GPU host does not re-expose GPU 0
* feat(cli): detect MLX distributed launch context
* feat(mlx): wire distributed inference backend
* feat(cli): broadcast MLX distributed chat turns
* fix(cli): wait indefinitely for distributed chat turns
* fix(cli): report MLX distributed load errors cleanly
* fix(mlx): route distributed vlm through loader
* fix(cli): detect inline MLX host JSON
* fix(studio): harden distributed object sharing
* fix(studio): select JACCL distributed backend
* fix(cli): abort distributed error paths
* Distinguish real stream errors from model text via GenStreamError in distributed CLI
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fail loud when MLX distributed init returns a singleton group
The worker only reaches this block when distributed was explicitly
requested. A singleton (size 1) group means the launch failed to form a
real group (MLX built without distributed support, or an invalid launch
env/hostfile); silently continuing leaves nonzero ranks looping forever
on share_distributed_object. Raise instead so the surrounding handler
returns a clear load error.
* Tighten MLX distributed inference comments
---------
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: danielhanchen <unslothai@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* 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>
* fix: handle case-variant GGUF cache hits for unsloth start
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* gguf cache: keep split shards co-located and isolate cache tests properly
When a cached main shard was reused from an older snapshot, the extra shards
were resolved independently and could come from a different snapshot dir (or a
fresh download into the current ref), leaving llama.cpp unable to load a
multi-shard GGUF whose pieces are split across directories. Only reuse a cached
main shard when every sibling shard sits in the same snapshot; otherwise fetch
the whole set together so they stay co-located.
Also patch huggingface_hub.constants.HF_HUB_CACHE (not just the HF_HUB_CACHE env
var) in the two cache tests that seeded a temp cache: the snapshot lookup reads
the module constant, so the env-only override let the real cache leak in and
skip an asserted download.
* Do not let a companion-only cache snapshot shadow real GGUF variants
When listing GGUF variants from the local HF cache, a newer snapshot may
contain only a companion file (for example a vision projector fetched on
demand) while the actual quant files live in an older snapshot. The prior
scan returned the first snapshot whose vision flag was set, yielding an
empty variant list and hiding the real quants. Keep scanning older
snapshots for actual variants and carry the vision flag across snapshots.
Also record the disk-space fallback variant's size in expected_sizes so
the later cache-reuse probe can size-verify the fallback main shard
instead of only checking for its existence.
* Propagate cached repo casing to companions and preflight split co-location
Two fixes to the case-variant GGUF cache reuse:
- Resolve the requested repo id to its cached canonical casing once in
load_model, up front, and pass it to the main GGUF and its companions
(mmproj / MTP drafter). Previously only _download_gguf resolved the
casing internally, so a case-variant request loaded the main file from
the canonical cache dir while the companions kept the requested casing
and missed the cached vision projector / drafter offline. Extracted the
resolution into a shared _resolve_repo_id_casing helper.
- Apply the split-shard co-location check in the disk-space preflight. When
a split GGUF's shards are cached across different snapshots the whole set
is refetched later, so counting them as cached made the preflight read 0
bytes to download, skip the smaller-variant fallback, and then fail the
full download on a low-disk machine.
* Reuse a co-located split GGUF snapshot and fix split fallback size probe
- When reusing a cached split GGUF, scan snapshots for one that holds the
whole set co-located instead of taking the newest snapshot's first shard.
A newer snapshot with only the first shard no longer shadows an older
complete snapshot, so an already-cached split model is reused rather than
refetched (which would fail offline).
- The disk-space fallback records its size in expected_sizes only for a
single-file fallback. _find_smallest_fitting_variant returns the whole
variant size, so using it as the first shard's expected size rejected a
valid cached first shard of a split fallback and forced a re-download.
* Scan for a complete split snapshot in the preflight; require a loaded catalog hit
- The disk-space preflight now uses the same co-located snapshot scan as the
download path (_cached_colocated_split_main) instead of the newest-snapshot
probe, so a newer snapshot holding only the first shard no longer masks an
older complete one and trips the smaller-variant fallback for a fully cached
split model.
- _resolve_model only attaches to a /v1/models entry that is actually loaded
(loaded != False). /v1/models also lists cached-but-unloaded catalog entries,
and matching one by case skipped /api/inference/load and left the agent
pointed at a model that is not resident.
* Restrict cross-snapshot GGUF cache reuse to offline
Reusing a same-name blob from an older or case-variant snapshot bypasses the
Hub revision/etag check, so a repo that updates a GGUF in place could serve
stale weights online. Gate the cross-snapshot and case-variant reuse (both the
disk-space preflight accounting and the download path) on HF_HUB_OFFLINE.
Online, hf_hub_download fetches the current revision and resumes a partial
download, so the reuse is unnecessary there; offline it remains the resilience
fallback. Marked the two reuse regression tests as the offline scenarios they
represent and added an online test asserting a fresh fetch.
* Harden offline cache reuse and hub-id detection
Three follow-ups on the case-variant GGUF cache path:
- Honor every truthy HF_HUB_OFFLINE spelling (1/true/yes/on), not just "1", when
gating the cross-snapshot and case-variant cache reuse. With HF_HUB_OFFLINE=true
the Hub calls are already offline, so the reuse must trigger or the cached GGUF
fails to load; route both the preflight accounting and the download path through
the same offline parse the rest of the backend uses.
- Resolve mmproj/MTP companions from the actual cached snapshot when offline.
resolve_cached_repo_id_case can keep a partial lower-case spelling when any dir
exists under the requested casing, so an hf_hub_download on that casing misses the
canonical companion; scan every case-variant snapshot and return the cached path.
- Restrict the case-insensitive model-id match to syntactically valid hub ids
(a single namespace/name over the HF charset). A server-side relative path such
as models/Llama/Foo.gguf is no longer treated as a hub id, so it cannot
casefold-match a differently cased path on a case-sensitive filesystem. This is
host independent, unlike the local-existence probe which cannot see a server path.
* Only casefold-match model ids against a loopback Studio
A two-segment string like Models/Foo is indistinguishable from a hub id, and the
local Path.exists() probe in _is_hub_model_id cannot see a path that exists only
on a remote Studio host. So against a remote server, casefolding could attach to
a distinct server-side path (Models/Foo vs models/foo) on a case-sensitive
filesystem. Gate the case-insensitive match on is_loopback_url(base): only a
local Studio, where the existence probe is authoritative, casefolds. For a remote
Studio the match is exact and a case-mismatched request falls through to
/api/inference/load, whose already-loaded dedup resolves it correctly.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: Wasim Yousef Said <wasimysdev@gmail.com>
* Studio: heal DiffusionGemma tool calls into structured tool_calls
* Fall back to supports_tools for backends without the passthrough capability
* Route DiffusionGemma client tools through passthrough when enable_tools is on
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Drop orphaned strip_tool_call_markup import after syncing with main
* Tighten supports_tool_passthrough comment
* Re-run CI on current main
---------
Co-authored-by: danielhanchen <unslothai@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.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).
img2img and inpaint take their output size from the uploaded image and only
snap it to a multiple of 16, so an ordinary phone photo (up to the 4096/side
decode cap, 4x the txt2img 2048 ceiling and ~16x the area) drove an OOM-scale
latent and an opaque 500 on a normal card, while txt2img, upscale, edit, and
FLUX.2-klein inpaint are all already megapixel-bounded. Clamp the init longest
side to 2048 (the txt2img ceiling) before deriving width/height; edit is exempt
since its pipeline resizes to ~1MP internally.
_cast_nvfp4 quantized every nn.Linear with no filter, unlike the int8 and fp8
torchao text-encoder modes which exclude the VLM vision tower / lm_head / T5 wo.
On qwen-image / qwen-image-edit that 4-bit quantized the Qwen2.5-VL image tower,
degrading the edit/image conditioning the sibling schemes protect. Apply the
same make_filter_fn exclusion (require_bf16, mirroring _cast_fp8_dynamic).
union_control_mode() only matched the short catalog id, so a client naming a
curated union ControlNet by its bare HF repo id (the form resolve_controlnet
documents and accepts) got None and the generate path omitted control_mode,
which makes FluxControlNetModel.forward raise controlnet_mode cannot be None.
Fall back to matching repo_id against the curated entries so the bare repo id
resolves to the same union mode; non-union bare repos still return None.
The images page also wired Reapply for a resident single_file model with no
checkpoint filename (status carries none), but the backend rejects a
single_file/gguf load without a filename, so clicking Reapply 400'd. Narrow the
resident-Reapply wiring to pipeline (the one kind that needs no filename),
matching the existing GGUF handling, so single_file stays a no-op instead of
erroring.
Resolve the app-sidebar.tsx conflict: main refactored the chat-export dropdown to a
format-based CHAT_EXPORT_OPTIONS + dynamic-import exportConversationByFormat dispatcher,
which the merged body already uses. Keep main's dispatcher and drop the branch's static
export imports; keep TestTubeOutlineIcon imported from the shared @/lib/hugeicons-derived
module (also used by images-page) rather than main's duplicate inline definition. The
branch's Images and Video nav items are preserved.
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.
* Speed up Studio startup path
* Studio: recheck managed binary executability on preflight cache hit and ignore stale unauthenticated platform fetches
Preflight: a matching capability cache fingerprint no longer skips the
runnability check when the managed binary's executable bit was cleared
(size and mtime unchanged, since chmod bumps ctime not mtime). The cache
fast path now confirms the binary is still executable, otherwise it falls
back to the CLI help probe so preflight reports Stale and can repair,
instead of returning Ready and failing later at backend start. Adds a
regression test.
Frontend: now that first render is no longer gated on fetchDeviceType,
the initial unauthenticated health call can resolve after an
authenticated platform fetch. Guard the store so a late unauthenticated
or failed non-forced response cannot overwrite an already authoritative
device type, tunnel URL, or secure flag. Forced refreshes and the first
unauthenticated load are unaffected.
* Studio: use access(X_OK) for the preflight cache executability guard
A mode bitmask treats any execute bit as launchable, but the executable
bits can be set only for another owner or group, or be denied by an ACL,
so the current user could still hit PermissionDenied at launch and the
cached fast path would wrongly return Ready. access(X_OK) checks real
executability for the calling user, so an ownership or permission change
correctly falls back to the CLI help probe and the Stale repair path.
* Studio: ignore any stale non-forced platform fetch once authoritative
Extend the platform store guard so a non-forced health response never
overwrites an already authoritative result, not only unauthenticated
ones. With a saved token the post-render non-forced request can be
authenticated but older than a later forced refresh that already picked
up the tunnel URL and secure flag; if that earlier request resolves last
it would null those fields. Now any non-forced response is dropped once
the store holds a server-reported platform. Forced refreshes and the
first authoritative write are unaffected.
* Studio: run the managed CLI help probe before trusting the preflight cache
Restore running the managed CLI help probe before returning Ready from
the desktop capability cache, so a managed install whose venv interpreter
or a runtime dependency is broken (while path, size, mtime, and markers
are unchanged) is reported Stale for repair rather than proceeding to a
backend start that cannot spawn. The capability cache still skips the
heavier desktop-capabilities probe on a hit, so a warm cache runs one
probe instead of two. Removes the executable-access shortcut, which the
help probe now subsumes.
---------
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
The native sd.cpp one-shot engine re-reads its companion VAE and text-encoder
files from the HF cache on every generation, but the delete-cached guard only
compared the loaded main diffusion repo_id. Deleting a companion repo such as
comfyanonymous/flux_text_encoders while a FLUX GGUF was loaded therefore
succeeded and broke the next generation. Add a loaded_repo_ids() accessor to the
native backend (mirroring loading_repo_ids(), reconstructed from the committed
family's VAE + text-encoder repos) and have the guard refuse those companions
too. Server mode and the diffusers engine hold companions in-process/VRAM, so
they are unaffected.
Video: a local FILE picked for a gguf/single_file load is handed straight to the
loader (the resolver returns the file itself, ignoring gguf_filename), so its own
suffix must match the kind. A .gguf picked as single_file (or a .safetensors picked
as gguf) slipped past the gguf_filename checks, evicting the resident GPU owner
before failing in from_single_file. Reject the mismatch in validate, before the handoff.
Engine router: publish the newly selected engine only after the old one finishes
unloading. The arbiter's diffusion evictor unloads the active engine, so flipping the
active name to the new (empty) engine first let a concurrent chat/video acquire evict
that empty engine and take the GPU while the old model was still freeing VRAM.
Two evict/OOM fixes on the diffusion load paths:
- The video load moved a pipeline onto the GPU (apply_memory_plan) and
committed it while holding no lock, so an unload / GPU-arbiter eviction --
which bumps the load token and then barriers on _generate_lock before
freeing -- could hand VIDEO to chat/images and let the new owner allocate
concurrently with the in-flight placement, OOMing. Hold _generate_lock
across placement + the locked commit, mirroring the image backend, so an
evicting owner waits until this worker's placement is torn down or
committed. Lock order stays _generate_lock -> _lock (unload takes _lock
then releases it before the barrier), so there is no deadlock.
- resolve_local_single_file reinterpreted an On-Device folder as a base
single_file load whenever it held exactly one .safetensors, so a PEFT LoRA
adapter folder (adapter_config.json + adapter_model.safetensors) with a
family-token name was picked as a base checkpoint, evicting the resident
model before from_single_file failed on the adapter weights. Skip adapter
folders (adapter_config.json) and the adapter_model basename so the pick
stays a pipeline load and 400s in validation, before the GPU handoff.
Adds regression tests for both.
The video load-request preflight gated its local-pipeline check on
root.is_dir(), so a bare local file (e.g. /models/ltx-2.safetensors) sent
with model_kind=pipeline skipped it, passed validation, and the route then
evicted the resident GPU model before from_pretrained failed on the
non-directory path. Gate on root.exists() instead (mirroring the image
loader's diffusion.validate_load_request), so a local file is rejected up
front. Add a regression test.
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).
union_control_mode fell back to control_mode=0 (the canny head) for ANY unmapped control
type, so a typo'd or unsupported value like 'detph' silently conditioned the map as canny
instead of failing. preprocess_control passes non-canny maps through unchanged, so that map
would be interpreted under the wrong mode with no error. Now only 'passthrough' (or an empty
type) keeps the deliberate mode-0 default; any other unknown type raises ValueError, which
the generate route maps to a 400. Known modes are unchanged.
* Studio: account for DeepSeek-V4 compute buffer in context auto-fit
DeepSeek-V4-Flash's lightning indexer plus compressed sparse attention reserve a
large context-scaling compute buffer that _compute_buffer_ctx_bytes did not model
(the KQ-mask and dequant-scratch rates both miss it, even with an f16 cache).
Measured on UD-Q4_K_XL at ub 512 it is about 65.5 GiB at 1M context, which the
mask estimate puts near 1.5 GiB, so the auto-fit kept the full 1M train context
and llama-server OOM'd allocating the ~70 GB buffer, then spilled to CPU (~4
tok/s). Add a deepseek4-gated flat plus per-token term so the fit caps the context
(about 256k on a B200) and the model stays fully on GPU.
* [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>
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.
* Add DeepSeek-V4-Flash-GGUF to Studio with none/high/max reasoning
Adds unsloth/DeepSeek-V4-Flash-GGUF as a default selectable model with the
recommended decoding defaults (temperature 1.0, top_p 1.0 from the official
generation_config.json) and its three tier reasoning control. The high/max
ladder is surfaced for deepseek-v4 model ids and flows through the existing
enable_thinking_effort reasoning style via chat_template_kwargs, so no
frontend changes are needed.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio DeepSeek-V4: segment-scope high, enable thinking for lone effort, render tests
Match deepseek-v4 on whole repo-name segments so a future deepseek-v40 or
deepseek40 cannot false-match the synthetic 'high'. In _request_reasoning_kwargs,
emit enable_thinking when a named effort level is sent without it, so the
newly exposed High mode renders thinking-on over the API (the UI already sent
it explicitly). Add a none/high/max render-path test file (jinja behind
importorskip) with a lone-high regression.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
