The opt-in `max` speed tier now compiles the repeated block with
mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode:
Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer
cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are
deliberately avoided -- both crash on the regionally-compiled block (its static
output buffer is overwritten across denoise steps), measured.
Adds two reproducible benchmarks used to validate the optimization research:
- scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head.
- scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes.
Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen;
coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune);
FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo).
Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy
bug and a dead-on-arrival speed path; this fixes both and adds the lossless /
near-lossless wins, all measured on a B200.
Correctness:
- TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32
process-wide and never restored them, so a later `off` load silently inherited
TF32 and was no longer bit-identical. Added snapshot_backend_flags /
restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer
runs and restored on unload. Verified: load max -> unload -> load off is now
byte-identical (PSNR inf) to a fresh off.
- sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and
only checked the timeout after EOF, so a child stuck in model load / GPU init
with no output ignored the timeout. Drained stdout on a reader thread with a
wall-clock deadline. Added a silent-hang regression test.
Speed (diffusers path), near-lossless, opt-in tiers:
- Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and
Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks
compiles and runs ~2.2x faster on the GGUF Z-Image transformer on
torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the
block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB
vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move
output quality. Gate relaxed; default tier delivers it.
- cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs).
- torch.inference_mode() around the pipeline call (lossless, strictly faster than
the no_grad diffusers uses internally).
Memory path:
- VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers;
the balanced (group) tier keeps exact slicing only, so it is now bit-identical to
the resident image (verified PSNR inf) and slightly faster.
- Group offload adds non_blocking + record_stream on the CUDA stream path to
overlap each block's H2D copy with compute (lossless; gated on the installed
diffusers signature so older versions still work).
Native (sd.cpp) path:
- native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a
near-lossless CUDA win that was previously only added on offload tiers; max also
-> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so
it is never auto-on. Engine generate() merges it, de-duped against offload flags.
Default profile: a GGUF model with no explicit speed_mode now resolves to the
`default` profile (resolve_speed_mode), since compile's perturbation sits below the
quantisation noise floor and so does not reduce quality versus the dense reference;
out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models
stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is
always honored, so the byte-identical path remains one flag away and is the
regression reference.
Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/
perf_verify.py (the B200 verification above), and diffusion_bench.py gains
--speed-mode so the speed tiers are benchmarkable.
Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore,
GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags +
the engine de-dup, and the sd-cli silent-hang timeout.
Builds on Phase 4's native stable-diffusion.cpp engine, extending it from
text-to-image to the wider feature surface, since sd.cpp supports all of these
through the binary already. Pure command-builder additions plus one engine
method, so the txt2img path is unchanged.
- sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img +
strength make a run img2img, adding mask makes it inpaint, ref_images drives
FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and
lora_dir + the <lora:name:weight> prompt syntax select LoRAs. New
SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run
mode (input image + esrgan model, no prompt / text encoders).
- sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so
generate() (now carrying the conditioning flags) and a new upscale() reuse
the same streaming / error / output-check path.
- scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img /
--strength / --upscale-model / --upscale-repeats.
Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the
upscale builder and its validation, and the engine's img2img + upscale paths.
Full diffusion suite 176 passing.
Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the
init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale
(512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing
coherent images. Video and the diffusers-path feature wiring are deferred.
Adds the CPU / Apple-Silicon tier of the two-engine strategy, mirroring the
chat backend's llama.cpp shell-out. Diffusers stays the default on CUDA / ROCm
/ XPU; this covers the hardware diffusers serves poorly, consuming the same
split GGUF assets Studio already curates.
- sd_cpp_args.py: pure sd-cli command builder. Maps the family to its
text-encoder flag (Z-Image Qwen3 to --llm, Qwen-Image to --qwen2vl, FLUX.1
CLIP-L + T5), and the diffusers memory policy (none/group/model/sequential)
to sd.cpp's offload flags (--offload-to-cpu / --clip-on-cpu / --vae-on-cpu /
--vae-tiling / --diffusion-fa), so one user knob drives both engines.
- sd_cpp_engine.py: SdCppEngine over a located sd-cli. find_sd_cpp_binary()
with the same precedence as the llama finder (env override, then the Studio
install root, then in-tree, then PATH), an is_available/version probe, and a
one-shot subprocess generate that streams progress and returns the PNG.
runtime_env() prepends the binary's directory to the platform library path
so a prebuilt's bundled libstable-diffusion.so resolves.
select_diffusion_engine() is the pure routing decision (GPU backends to
diffusers, CPU/MPS to native when present).
- install_sd_cpp_prebuilt.py: resolve + download the per-host prebuilt
(macOS-arm64/Metal, Linux x86_64 CPU, Vulkan/ROCm/Windows variants) into the
Studio install root. resolve_release_asset() is a pure, unit-tested
host-to-asset matrix.
- scripts/sd_cpp_smoke.py: end-to-end native generation harness.
Tests (CPU-only, subprocess/filesystem stubbed): 49 new across args, engine,
routing, runtime env, and the installer resolver. Full diffusion suite 166
passing.
Verified on a B200 box: built sd-cli (CUDA) and the prebuilt (CPU) both
generate Z-Image-Turbo Q4_K end to end through SdCppEngine: balanced (group
offload, 5.0s gen), low_vram (full CPU offload + VAE tiling, 13.4s), and the
dynamically-linked CPU prebuilt (50.4s on CPU), all producing coherent images.
Generalise the text-encoder precision knob from a fp8 bool to text_encoder_quant
(fp8 | nvfp4). nvfp4 quantises the companion text encoder to 4-bit via torchao
NVFP4 weight-only (two-level microscaling) on Blackwell's FP4 tensor cores; fp8
stays the broader-hardware path (cc>=8.9). Both are gated, best-effort, and run
before placement; status reports the mode actually engaged. This is the lean
realisation of GGUF-native text-encoder quant: 4-bit on the encoder without the
3045-line port.
Verified on Z-Image (B200, balanced/group where the encoder stays resident), vs the
bf16 encoder: nvfp4 cut generation peak VRAM 48% (10840 -> 5593 MB, the lowest TE
option, below whole-model offload) at near-fp8 quality (16.4 vs 17.1 dB PSNR), and
both quants ran faster than bf16. A memory-vs-quality tradeoff (off by default);
size it per model with the Phase 5 quality harness. diffusion_bench gains
--text-encoder-quant.
129 CPU tests pass.
Add a text_encoder_fp8 knob that casts the companion text encoder(s) to fp8 (e4m3)
storage via diffusers apply_layerwise_casting, upcasting per layer to the bf16
compute dtype while normalisations and embeddings stay full precision. Applied
before placement, gated to CUDA + bf16, best-effort (a failure leaves the encoder
dense). status reports which encoders were cast.
Verified on Z-Image (B200, balanced/group mode where the encoder stays resident):
generation peak VRAM dropped 37% (10840 -> 6791 MB, below the lowest-VRAM offload)
at near-resident speed. It is a memory-vs-quality tradeoff, not free -- ~20 dB PSNR
vs the bf16 encoder, a larger shift than one transformer quant step -- so it is off
by default and documented as such, with the Phase 5 harness to size the cost.
127 CPU tests pass.
Add a speed_mode knob (off by default, so the render path stays bit-identical):
default applies channels_last VAE + regional torch.compile of the denoiser's
repeated block where eligible; max also enables TF32 matmul and fused QKV. Regional
compile is gated off for the GGUF transformer (dequantises per-op) and for families
flagged not compile-friendly (a new supports_torch_compile flag, False for Z-Image),
so it activates automatically only once a non-GGUF bf16 transformer is loaded. Speed
optims run before placement/offload, per the diffusers composition order. status now
reports speed_mode + the optims actually engaged.
Verified on Z-Image (B200): default -> ['channels_last'], max -> ['channels_last',
'tf32'], compile correctly skipped for GGUF; generation works in every mode.
121 CPU tests pass.
Add a streamed 'group' offload tier (diffusers apply_group_offloading, block_level,
use_stream) that keeps the transformer flowing through the GPU a few blocks at a
time while the text encoder / VAE stay resident, and fix VAE tiling to drive the
VAE submodule (pipelines like Z-Image expose enable_tiling on pipe.vae, not the
pipeline). apply_memory_plan now returns the (policy, tiling) actually engaged so
status never overstates either, and group falls back to whole-module offload when
the transformer can't be streamed.
Measured on Z-Image (B200), all lossless (PSNR inf vs resident): balanced/group
cuts generation peak VRAM 32% (15951 -> 10840 MB) at near-resident speed (2.07 ->
2.99s); low_vram/model cuts it 48% (-> 8318 MB) but is slower (7.99s). Mode names
now match that tradeoff: balanced = stream the transformer, low_vram = offload
every component. auto picks group when the companions fit resident, else model.
112 CPU tests pass.
Add a lean, backend-agnostic memory policy that picks a CPU-offload policy and
VAE tiling/slicing from measured free device memory vs the model's estimated
resident footprint, then applies it to the built pipeline. auto stays resident
when the model fits (byte-identical to the prior resident path), and falls to
whole-module offload when tight; fast/balanced/low_vram are explicit overrides.
Sequential submodule offload is unreliable for GGUF transformers on diffusers
0.38, so it falls back to whole-module offload and status reports the policy
actually engaged.
Verified on Z-Image-Turbo Q4_K_M (B200): auto reproduces the resident image with
no VRAM/latency regression (PSNR inf); balanced/low_vram cut generation peak VRAM
47.9% (15951 -> 8318 MB) with byte-identical output, at the expected latency cost.
73 prior + 35 new CPU tests pass.
Phase 1 of porting the richer diffusion stack onto the image-generation backend.
- Add a compartmentalized device/dtype policy module (diffusion_device.py)
resolving CUDA/ROCm/XPU/MPS/CPU with capability flags. Keeps the NVIDIA
capability-based bf16 choice; ROCm and XPU are isolated; MPS uses bf16 or
fp32, never a silent fp16 that renders a black image.
- Add a per-family fp16_incompatible flag (Z-Image) and promote a resolved
float16 to float32 for those families so they do not produce black images.
- Split the backend locks: a generation holds only _generate_lock, so status,
unload, and a new load are never blocked by a long denoise. Add per-generation
cancellation via callback_on_step_end so an eviction or a superseding load
preempts a running generation; a replacement load waits for it to stop before
allocating, so two pipelines never sit in VRAM at once.
- Validate a load request before the GPU handoff so an unloadable pick never
evicts a working chat model, and reject missing local paths up front.
- Add CPU-only tests for the device policy, dtype guard, lock split and
cancellation, and validate-before-evict, plus a GPU benchmark/regression
script (scripts/diffusion_bench.py) measuring latency, peak VRAM, and PSNR
against a saved reference.
* Studio: honor stream=false on the GGUF agentic tool path (#6570)
* Studio: dedup the #6570 non-streaming tool tests and cover cached_tokens
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: cover the cached_tokens metadata fix and clarify the drain comment (#6570)
* Studio: align the GGUF tool drain naming and tighten its comment (#6570)
---------
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>
* checkpoint preview endpoint
* harden new preview endpoints
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address review
* Studio preview: pin adapter, guard streaming submit, robust copy-link
Harden the public per-checkpoint preview surface:
- Pin use_adapter=True in the preview payload sanitizer. Otherwise an
unauthenticated /p caller can POST use_adapter=false, which calls
disable_adapter_layers() on the shared in-memory model without restoring
it; since load_model skips reloads for the same checkpoint, every later
visitor (the page never sends the field) keeps getting base-model output
instead of the fine-tuned checkpoint. Forcing it on also re-enables a
previously disabled adapter and no-ops on merged checkpoints.
- Ignore preview-page submits while a response is streaming. The send
button was disabled but the Enter handler still called requestSubmit(),
so a second request could start before the first reply landed in msgs and
reorder the chat history. Both the keydown and submit handlers now honor
the disabled button.
- Keep the cloudflare-URL polling loop alive across transient startup fetch
errors instead of letting one rejection halt it.
- Build the copy-link from a backend preview_ref (output dir relative to
outputs_root, gated on previewability and the two-segment /p route limit)
so a nested output dir no longer copies a basename-only link that 404s.
Expose preview_ref on training run summaries.
Add route-level security tests (path traversal, payload sanitization,
asset containment, CSP header, HTML title escaping, streaming lock held
until drained) and preview_ref unit tests.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio preview: Safari-safe submit and adapter pin only for LoRA
Follow-ups from cross-browser and route simulations:
- Preview page: send the message from a shared send() helper called by both
the form submit and the Enter key, instead of form.requestSubmit(). The
latter throws on Safari < 16 and older iOS, which broke Enter-to-send there.
Verified across Chromium, Firefox and WebKit with Playwright.
- Only pin use_adapter=True when the resolved checkpoint is a LoRA adapter
(adapter_config.json present); for a merged checkpoint strip it to None.
A merged model has no adapter to toggle, so forcing it on only produced a
per-request "not a PeftModel" warning. The cross-request base-model
contamination fix still holds for LoRA previews.
Add a merged-checkpoint test asserting use_adapter is stripped to None.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio preview: trim verbose comments
Tighten comments across the preview routes, page, checkpoint helpers, and tests
to short single-line notes; drop ones that just restate the code. No behavior
change (verified comment/docstring-only with comment_tools.py check).
* Harden preview routes for PR #6486
- Return a generic 400 detail on a rejected preview path so the public /p
route never echoes the absolute install path (the real reason is logged
server-side instead).
- Strip confirm_tool_calls, session_id and rag_scope in the preview payload
sanitizer so the public surface stays inert regardless of the tool gate.
- Use Path.is_relative_to for the asset containment check, matching the rest
of the codebase.
- Add img-src 'self' and font-src 'self' to the preview page CSP.
- Preview page: on a mid-stream error keep the streamed text, flag the break,
and restore the prompt so the user can retry; drop the unused --font-sans var.
---------
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: Daniel Han <danielhanchen@gmail.com>
* Studio: cap GGUF context to unified memory on Apple Silicon
* Studio: tighten Apple ctx-cap comments and drop the overstated MLX-sync claim
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: reserve flat MTP fraction and floor sparse-KV ctx in the Apple unified-memory cap
The Apple Silicon GGUF context cap mirrored the discrete-GPU auto-fit branch but
missed two protections the discrete path already applies:
- It passed the full unified-memory budget with budget_frac=1.0 without first
reserving the flat MTP fraction the discrete path takes off via _pin_fraction.
With an MTP draft whose KV cannot be byte-sized (e.g. Qwen3.6-MTP, #6529), the
cap filled the whole budget and left nothing for the draft, so unified memory
could still over-commit. Reserve _flat_mtp_reserve up front; this is a no-op
when MTP is not engaged.
- It required _can_estimate_kv(), so a GGUF with sparse KV metadata skipped the
cap entirely and launched at full native context. Mirror the discrete
file-size-only fallback and floor the auto context to 4096 when the cache
cannot be sized.
Adds regression tests for both paths.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: tighten comments in the Apple unified-memory context cap
Condense the verbose comment blocks in the Apple budget helper, the no-GPU
Metal branch, and the context-fit tests. Comments only, no code change
(verified with ast-based comment_tools check); suite still green.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <danielhanchen@gmail.com>
* Copy uv `-c`/`-r` paths to a space-free temp dir on macOS/Linux
uv 0.11.x truncates a constraints/requirements path passed via `-c`/`-r`
at the first space, so `unsloth studio` setup from a repo cloned under a
path containing a space (e.g. `/Users/me/Open Source/unsloth`) fails with:
error: File not found: `/Users/me/Open`
_uv_safe_path() already worked around this on Windows via the 8.3 short
path but returned the space-containing path unchanged on macOS/Linux,
which have no 8.3 equivalent. Extend it to copy the (small, flat)
constraints/requirements file into a space-free temp dir and hand uv the
copy; the temp dirs are removed at process exit. Falls back to the
original path on any error, so it is never worse than before.
Refs unslothai/unsloth#6503
* Route UV_OVERRIDE through _uv_safe_path and fix temp-dir leak (#6503)
The -c/-r fix did not cover UV_OVERRIDE, which uv also truncates at the first
space. On Apple Silicon the overrides file is handed to uv via UV_OVERRIDE at
install time (install_python_stack.py) and during the MLX self-heal
(utils.mlx_repair), so a repo under a path containing a space still broke every
uv call there. Move _uv_safe_path into backend.utils.uv_path_safety so both
sites share it, and route UV_OVERRIDE through it.
Also stop leaking the temp dir when shutil.copyfile fails after mkdtemp, and add
tests for the UV_OVERRIDE channel, the TMPDIR-with-space fallback, the atexit
cleanup, and the no-leak path.
---------
Co-authored-by: danielhanchen <danielhanchen@gmail.com>
* Harden MLX self-heal install against supply-chain execution
The Apple Silicon MLX self-heal runs uv pip install on a daemon thread
during Studio startup, default-on with only an env opt-out, before the
post-install stack check. Two things widened the supply-chain surface:
- it accepted source distributions, whose PEP 517 build backends run
arbitrary code at install time; and
- it forwarded the full process environment, exposing Studio secrets to
that code and letting a poisoned env (UV_FIND_LINKS / UV_DEFAULT_INDEX)
repoint the install at a hostile source.
Require pre-built wheels (--only-binary=:all:) and forward only an
allowlist of variables uv needs (PATH/HOME, proxy + CA settings, cache
dir), setting UV_OVERRIDE ourselves. mlx/mlx-metal ship wheels only and
mlx-lm/mlx-vlm publish py3-none-any wheels, so a healthy self-heal is
unaffected; an unavailable wheel just leaves Studio chat-only as before.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Drop cache-dir env vars from the self-heal allowlist
Address review: a poisoned process env could set UV_CACHE_DIR / XDG_CACHE_HOME
to redirect uv at an attacker-staged cache (cache poisoning, symlink writes),
which partly undercut the index-redirect protection. Drop them from the
allowlist; uv falls back to its safe user-owned default cache, still reused
across runs, so there is no normal-path cost. Test now asserts both are
excluded from the install env.
* [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>
* Fix Gemma 4 GGUF OpenAI API streams
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Avoid duplicate Responses stream disconnect watcher
* Keep reasoning-only Responses output hidden
* Address Gemma stream review comments
* Avoid Responses stream task-group cleanup
* Harden OpenAI chat completion streams
* Address OpenAI stream review issues
* Clean up Studio OpenAI stream helpers
* Fix Studio passthrough cold stream timeout
* Fix tool parser compatibility exports lint
* Preserve audio stream disconnect cancellation
* Avoid synthetic finish after passthrough errors
* Address stream cleanup and Gemma parser reviews
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Gemma 4: parse bare-string tool args and keep safetensors tools for native <|tool_call>
- Quote bare unquoted string values in Gemma native tool-call args (e.g.
{location:Tokyo,unit:celsius}) so they parse; JSON scalars stay typed.
- Stop _detect_safetensors_features from suppressing supports_tools for
templates that emit Gemma native <|tool_call>, which the shared parser
now reads.
- Add 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
* Harden Gemma tool-call parsing and stream-error detection
Address three issues in the Gemma-native tool-call path:
- _quote_gemma_object_keys stopped a bare (unquoted) string value at the
first comma, so an argument like `location:New York, NY` was split
mid-value and the synthesized JSON failed to parse, dropping the whole
tool call. A bare value now ends only at `}` or a comma that begins the
next `key:` pair.
- parse_tool_calls_from_text scanned the entire response for Gemma markers
even inside a tool call already parsed from a `<tool_call>{...}` JSON
block, so a marker-like string inside an argument (data) was promoted to
a second, unintended tool call. Matches inside an already-consumed call
span are now skipped.
- _openai_passthrough_stream relied on _monitor_openai_sse_line to flag a
stream error, which returns early when monitor_id is None
(skip_api_monitor), so an upstream error chunk left saw_stream_error
unset and the synthetic-finish guard emitted a successful finish_reason
after a failed stream. Error chunks are now detected independently of API
monitoring.
Adds tests/test_gemma_tool_parse_edge_cases.py covering the comma and
marker-injection cases.
* Emit the terminal finish_reason chunk in GGUF streams
The OpenAI chat-completions GGUF tool stream and plain stream both built a
final ChatCompletionChunk carrying finish_reason but never yielded it, so
clients received the optional usage chunk and [DONE] with no chunk carrying
finish_reason. OpenAI-compatible consumers rely on that terminal choice to
distinguish stop/length/tool_calls. Yield it before the usage chunk and
[DONE], matching the other streaming paths.
* Parse tool calls in document order and skip nested markers both ways
Unify the JSON- and Gemma-format tool-call passes into a single
position-ordered scan:
- Calls are now emitted in byte order across both formats, so a mixed
output like `<|tool_call>call:create{...}<tool_call|> ... <tool_call>
{"name":"read",...}</tool_call>` executes create before read, matching
the order they appear in (tools run in returned order).
- A candidate that starts inside an already-accepted call's span is
skipped, in both directions: a JSON marker inside a Gemma argument and a
Gemma marker inside a JSON argument are treated as data, not promoted to
a second executable tool call.
Extends tests/test_gemma_tool_parse_edge_cases.py with the ordering and
JSON-in-Gemma nesting cases.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Quote bare Gemma array elements; order finish before trailing usage
- _quote_gemma_object_keys skipped array values, so a Gemma call with a
bare-string array argument like labels:[bug,ui] produced invalid JSON and
the whole tool call was dropped. Array values are now scanned and bare
string elements quoted, while numbers, quoted strings, and JSON literals
are preserved.
- In the OpenAI passthrough stream, a trailing usage-only chunk
(stream_options.include_usage) that arrived before any finish chunk was
relayed before the synthetic finish, producing usage -> finish -> [DONE].
Emit the synthetic finish before that usage chunk so the order matches the
other streams (finish -> usage -> [DONE]).
Extends tests/test_gemma_tool_parse_edge_cases.py with the bare-array cases.
* Harden Gemma array parsing, XML-parameter guard, and stream teardown
Address five review findings on the Gemma tool-call and OpenAI passthrough
streaming paths:
- parse_tool_calls_from_text collected JSON and Gemma markers without the
_inside_open_parameter guard, so a marker embedded in an existing
<function=...><parameter=...> value was promoted to a separate tool call.
Candidates that start inside an open XML parameter are now skipped, matching
the guard the XML-style parser already applies.
- _quote_gemma_array_elements preserved array elements starting with { or [
verbatim, so an array of objects (items:[{path:a}]) or a nested array failed
json.loads and the whole call was dropped. Object and nested-array elements
are now normalised recursively.
- _openai_passthrough_stream synthesized a finish chunk before a trailing
usage-only chunk and set saw_finish_reason, which made the EOF guard skip the
[DONE] sentinel. The EOF path now emits [DONE] whenever the upstream omitted
it, even after a finish chunk was already synthesized.
- /generate/stream drove generation through asyncio.to_thread with no
disconnect watcher, so a client disconnect during a long generation went
unnoticed until the next send. It now runs _await_disconnect_then_cancel
against the request, matching the other local streaming endpoints.
- _SameTaskStreamingResponse closed the body iterator with aclose() on a
send-side disconnect, raising GeneratorExit so the generators' cancellation
handlers (which finish the api_monitor entry) never ran. It now throws
CancelledError, falling back to aclose() when athrow is unavailable.
Extends tests/test_gemma_tool_parse_edge_cases.py with array-of-objects,
nested-array, and marker-inside-XML-parameter cases.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Watch disconnects on Anthropic streams; keep timestamps in Gemma values
Two follow-ups on the streaming and tool-parse paths:
- _anthropic_tool_stream and _anthropic_plain_stream drove generation through
asyncio.to_thread(next, gen, ...) and only polled is_disconnected() between
events, so a client disconnect during prefill or a long generation/tool step
held the decode slot until the next event or a failed send. Both now run the
_await_disconnect_then_cancel watcher used by the other local streams, stop it
in finally, and break promptly when cancel_event is set.
- _GEMMA_NEXT_KEY_RE treated any comma followed by word-chars-then-colon as the
next key, so a bare value such as "meet at 10:00, 11:00 tomorrow" was split
into bogus keys. The next-key token must now be identifier-shaped (start with
a letter or underscore), so a comma before a timestamp, ratio, or other
numeric-then-colon text stays part of the value.
Adds a timestamp-in-bare-value regression test.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Guard nested markers, reset on disconnect, clean unstarted streams
Three follow-ups on the tool-parse and streaming paths:
- parse_tool_calls_from_text only skipped markers that fell inside a span it
had already parsed successfully, so when an unquoted Gemma argument contained
a literal marker (code:<|tool_call>call:terminal{...}<tool_call|>) the outer
object failed to normalize, its span was never recorded, and the inner marker
was promoted to a standalone terminal call. Candidates nested inside any other
candidate's brace span are now skipped regardless of whether the enclosing
candidate parsed, so a marker in malformed outer data is never executed.
- /generate/stream skipped backend.reset_generation_state() when the disconnect
watcher set cancel_event between chunks: the loop broke and the finally's reset
is guarded on cancel_event being unset. A subprocess backend kept decoding
after the client left. The cancel-break path now resets the backend.
- _SameTaskStreamingResponse threw CancelledError / called aclose() on the body
iterator on a send-side disconnect, but neither runs the try/finally of a
generator that never started (early disconnect on http.response.start), so the
passthrough's eagerly-opened upstream httpx stream and cancel-registry entry
leaked. It now tracks whether the body started and, when it did not, runs an
optional unstarted_cleanup hook; the OpenAI passthrough wires it to close the
upstream resp/client and exit the cancel tracker.
Adds a nested-unquoted-marker regression test.
* [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: Daniel Han <danielhanchen@gmail.com>
* studio: report the true reasoning duration and fix the Stop button for thinking models
For a local GGUF the "Thought for N" label was timed entirely on the client by a
brittle edge-detector, so an always-think model (Qwen3 MTP) that buffers its whole
reasoning and flushes it in one chunk showed "1 second" instead of the real
minute-plus. The client cannot time reasoning it receives atomically, so make the
timing backend-authoritative.
Backend: generate_chat_completion_with_tools measures wall-clock reasoning and
emits a Studio reasoning_summary event (duration_ms) at the moment reasoning ends
-- the first answer token, or end-of-stream for a reasoning-only reply -- for both
the tool-detection pass and the final-answer pass. Timing resets per tool
iteration so the final answer's thinking time wins on the client (which takes the
latest reasoning_summary). routes/inference.py forwards the event in the GGUF tool
stream.
Frontend: parse the reasoning_summary SSE into a _reasoningDurationMs chunk and
use it as the authoritative reasoning duration (last write wins), clamped to >= 0
and guarded to a finite number so a malformed or proxied chunk cannot produce a
NaN label; the persisted value wins for the final "Thought for N" label, with the
previous live timer kept only as a fallback when no metadata arrives.
* [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>
* Studio: start the Cloudflare tunnel for --secure even in --api-only, and add --api-only to `unsloth studio run`
--secure exposes ONLY the Cloudflare link (it forces a loopback bind), but
_cloudflare_tunnel_should_start gated the tunnel on `not api_only`, so
`run.py --secure --api-only` started no tunnel and then fail-closed with
"A secure Cloudflare link is not allowed". That blocked the natural headless
use: serve just the API (no web UI) over the authenticated tunnel.
Make --secure start the tunnel regardless of api_only (the non-secure path is
unchanged: tunnel only a 0.0.0.0 bind, never api-only Tauri or Colab). Then
expose --api-only on `unsloth studio run` and forward it through both the
re-exec args and the in-venv run_server call, so
`unsloth studio run --secure --api-only --model ...` is a one-liner secure API
server.
Verified end to end: `run.py --secure --api-only` now brings up the tunnel and
serves /api/health over it (200), with / returning 404 (no UI).
Tests: update the tunnel-gate truth table (secure+api-only now tunnels;
secure+colab still does not) and add --api-only registration + re-exec/in-venv
forwarding coverage to the run CLI tests.
* Trim comments to be succinct (no behavior change)
* studio: address review on parent --api-only and secure api-only CORS
- Reject --api-only on the parent `unsloth studio` group when a subcommand
is invoked, with the same redirect guidance used for --parallel/--secure;
otherwise the flag was silently dropped and the UI served anyway.
- Keep CORS any-origin for secure api-only serving: that mode publishes the
API over Cloudflare for remote browser clients, so the Tauri-only lockdown
(still applied to plain local api-only) would break preflight. Factored the
decision into cors_origins_for_mode() and gate it on api_only and not secure;
run_server exports UNSLOTH_SECURE before importing main.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: suppress TAURI_PORT and de-dup test for headless run --api-only
- run_server gains emit_tauri_port (default True, unchanged for the Tauri/
desktop path). The new headless `run --api-only` path passes False so the
Tauri-only TAURI_PORT= line no longer prepends the documented URL/API key
banner (it ran even under --silent and could break one-liner parsers).
- Remove a duplicate test_reexec_forwards_api_only that shadowed the
parametrized one; fold the --secure --api-only case into it so the secure
headless path is actually collected.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
torchao has no working Windows ROCm build. transformers.quantizers imports it,
and it loads torch's c10d distributed backend at module level, which the AMD
Windows wheels omit (no RCCL). The import aborts, transformers can no longer
expose PreTrainedModel, and the sentence-transformers embedder silently falls
back to the llama-server GGUF embedder. Linux ROCm and NVIDIA are unaffected
(the c10d ops are present / torchao is real there).
The training and export workers already install the shared torchao stub before
importing transformers, but the RAG embedder runs in the main backend process,
which never did. Two fixes, both no-ops off Windows ROCm:
- embeddings.py: install_torchao_windows_rocm_stub() before the first
sentence-transformers import, so an already-installed torchao is neutralized
(fixes existing venvs).
- install_python_stack.py: stop installing torchao on Windows ROCm; it can only
crash on import there, so new venvs never ship it.
Add tests covering the embedder stub call and the install skip.
* Studio: treat data-center Blackwell (sm_100/sm_103) as Blackwell in llama.cpp prebuilt selection
_host_is_blackwell gated on _BLACKWELL_MIN_SM = 120, but data-center Blackwell
parts report a lower compute capability than consumer Blackwell: B100/B200 are
sm_100 and B300/GB300 are sm_103, while RTX 50 is sm_120 and DGX Spark is
sm_121. Because 100 and 103 are both < 120, every data-center Blackwell host was
classified as non-Blackwell, so two GPU-targeting paths never fired for a
B200/B300:
- the Linux blackwell_runtime_override that prefers the highest CUDA-major
runtime line shipping a bundle covering the host SMs (so a cu12x torch could
pin a cuda12 bundle over a native cuda13 one), and
- _drop_blackwell_incapable_windows_cuda, which removes cuda-12.4 builds that
load and validate but run Blackwell on a slow PTX-JIT path.
The result is a B200/B300 being handed a prebuilt that does not natively offload
its SM, i.e. the llama.cpp prebuilt is not really for the GPU. The Blackwell
floor is sm_100, so set _BLACKWELL_MIN_SM = 100. The toolkit floor (12.8) is
unchanged and already correct for sm_100/sm_103.
Surfaced loading unsloth/GLM-5.2-GGUF UD-IQ1_S on 8x B200.
Adds tests covering the sm_100/sm_103 classification, the Linux cuda13
preference for a data-center host, and the Windows cuda-12.4 drop.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Trim comments to be succinct (no behavior change)
* studio: require CUDA 12.9 for sm_103/sm_121 Blackwell prebuilts
sm_103 (B300/GB300) and sm_121 (DGX Spark) have no native compiler
target before CUDA 12.9; the family floor of 12.8 only covers
sm_100/101/120. Make the Windows-CUDA Blackwell filter SM-aware so a
legacy win-cuda-12.8 bundle is dropped on an sm_103/sm_121 host while
sm_100/sm_120 hosts keep the 12.8 floor.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* fix(studio): handle multimodal list content in inference text paths
Studio receives chat message content in two shapes: the legacy string
form, and the OpenAI multimodal list form
([{"type": "text", "text": ...}, {"type": "image_url", ...}]).
Several string-only paths called .strip()/re.sub()/f-string interpolation
on content directly, raising "'list' object has no attribute 'replace'"
for vision models (issue #4383), or rendering the list repr into the
prompt for the manual chat-template formatters.
Add core/inference/message_content.py with content_to_text(), a pure
helper (no heavy imports) that returns strings unchanged and joins the
text parts of a list while dropping image/audio parts. Apply it at every
string-only content site: _generate_vision_response, the audio user-text
extraction, format_chat_prompt, and the llama3/mistral/chatml/alpaca/
generic template formatters. The plain-string path is a no-op, so
existing behavior is unchanged.
Adds tests/test_message_content.py covering str/None/list/tuple,
multimodal drop, multi-part join and empty-part skipping.
* Tighten code comments (no logic change)
* studio: join multimodal text parts with newline for llama.cpp parity
llama.cpp joins multiple text content parts with a newline (common/chat.cpp),
so match that in content_to_text instead of a single space.
---------
Co-authored-by: Daniel Han <michaelhan2050@gmail.com>
Three independent upstream CI fixes that currently fail on every open PR:
verify_import_hoist.py: TARGET-CHANGED only flags a genuine swap (a BEFORE
target no longer reachable in AFTER). A pure superset growth such as adding
import urllib.error next to import urllib.request binds the same top-level
package and loses nothing, so it is no longer a blocker (transformers_version.py).
test_vision_cache.py: run each test from a fresh empty cwd. is_vision_model
calls is_local_path first, and a relative model id that happens to exist on
disk short-circuits before the mocked detection runs; the CI cwd and HF cache
can contain dirs colliding with the synthetic ids, causing 'called 0 times'.
Production code is correct; only the test needed cwd isolation.
consolidated-tests-ci.yml: the llama.cpp smoke probes the first of
llama-cli / llama-mtmd-cli / llama-server that exists instead of hard-requiring
llama-cli, which upstream no longer always builds. llama-cli stays first so it
is preferred when present. Adds Windows .exe + build/bin/Release handling.
* Use UTF-8 for Python code-execution subprocess I/O
Studio's code-execution tool already tells the child to emit UTF-8
(PYTHONIOENCODING=utf-8 in _build_safe_env), but _python_exec writes the
temp script and decodes the subprocess pipe with the OS default codec.
On Windows (cp1252), non-ASCII in model-written code or its output --
arrows, CJK, emoji -- raises UnicodeEncodeError / UnicodeDecodeError and
breaks execution.
Complete the UTF-8 wiring in core/inference/tools.py:
- write the temp script with encoding="utf-8"
- decode _python_exec stdout as utf-8, errors="replace"
- set PYTHONIOENCODING=utf-8 in _build_bypass_env too (matches
_build_safe_env, so the bypass path's child also emits utf-8)
The child is python with PYTHONIOENCODING=utf-8, so it emits UTF-8
regardless of the console code page and the decode is always correct.
Shell execution via cmd.exe has a separate console-code-page story and
is left to a follow-up.
Refs unslothai/unsloth#6489
* Scope Python exec UTF-8 env to Python tool
* Make bash bypass test robust to a host-set PYTHONIOENCODING for PR #6548
Bypass mode preserves benign host env vars, so a host-set PYTHONIOENCODING was
inherited into the bash bypass env and tripped the new assertion even though
_bash_exec never adds it. Clear it in the test so the assertion checks _bash_exec,
not the runner environment.
---------
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
* Restore sys.modules in test_pre_import_gate_is_transformers_free
The test pops transformers and utils.models.model_config from sys.modules to
assert the pre-import security gate does not re-import them, but never put them
back. A later importer then rebound a fresh utils.models.model_config, so tests
that had captured the original instance missed their patches and hit the real
path: test_vision_cache patches _is_vision_model_uncached on the original
module, but is_vision_model (still bound to that original) ran the real network
lookup instead. This produced 17 spurious failures whenever test_ssm_runtime
ran before test_vision_cache in the same process.
Snapshot the removed modules and restore the original objects in a finally, so
the assertions still run against a clean slate while later tests see the same
module instances they captured at import time.
* [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>
* Studio: honor custom HF_HOME for model download and load
_setup_cache_env always derived HF_HUB_CACHE and HF_XET_CACHE from
XDG_CACHE_HOME / ~/.cache, ignoring a user-set HF_HOME. Because it sets
HF_HUB_CACHE explicitly and that variable takes precedence over HF_HOME
in huggingface_hub, the hub cache was pinned to the standard location: a
model already present under a custom HF_HOME was detected but then
re-downloaded from scratch on load.
Seed HF_HUB_CACHE and HF_XET_CACHE from HF_HOME when the user set it
(HF's own default is $HF_HOME/hub and $HF_HOME/xet), and honor the legacy
HUGGINGFACE_HUB_CACHE alias. The hub download workers call
snapshot_download without a cache_dir for both the Xet and HTTP-fallback
paths, so they follow HF_HUB_CACHE; fixing it here unifies detection and
both transports on one root. Explicit HF_HUB_CACHE / HF_XET_CACHE stay
untouched. Adds tests for the custom-HF_HOME, default, explicit-override,
and legacy-alias cases. Fixes#5182.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: do not crash startup when a custom HF_HOME is not writable
Seeding HF_HUB_CACHE/HF_XET_CACHE from HF_HOME means _setup_cache_env now
mkdir's under a user-controlled path. A non-writable or not-yet-mounted
HF_HOME (typo, offline drive) would raise and crash startup, where the old
code silently fell back. Make the mkdir best-effort; the env var is still
set, so HF reports a clear error at download time. Adds a regression test.
* Studio: strip blank HF_HOME and isolate cache-env tests
Address review: a whitespace-only HF_HOME no longer derives " /hub";
strip it and fall back to the default (matches studio_root). Tests set
UNSLOTH_STUDIO_HOME to a tmp dir so _setup_cache_env's UV/VLLM mkdirs do
not touch the real ~/.unsloth/studio. Adds a whitespace regression test.
* [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>