Commit graph

991 commits

Author SHA1 Message Date
Daniel Han
a7b8f825da Studio diffusion (Phase 4): native stable-diffusion.cpp engine for CPU/Mac
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.
2026-06-25 15:50:10 +00:00
pre-commit-ci[bot]
7ca5573498 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-25 14:43:45 +00:00
Daniel Han
dbb0292561 Studio diffusion (Phase 2C): NVFP4 text-encoder quant (+ generalise fp8 knob)
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.
2026-06-25 14:42:54 +00:00
pre-commit-ci[bot]
54eea13036 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-25 13:56:37 +00:00
Daniel Han
817e14ad19 Studio diffusion (Phase 2B): opt-in fp8 text-encoder layerwise casting
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.
2026-06-25 13:55:08 +00:00
Daniel Han
9de8684a98 Studio diffusion (Phase 3): opt-in speed layer (channels_last / compile / TF32)
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.
2026-06-25 13:55:08 +00:00
Daniel Han
2f00c77e75 Studio diffusion (Phase 2D): streamed block-level offload + functional VAE tiling
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.
2026-06-25 13:55:08 +00:00
Daniel Han
3a0bd55bdf Studio diffusion (Phase 2A): measured-budget memory planner + offload/VAE policy
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.
2026-06-25 13:55:08 +00:00
pre-commit-ci[bot]
7f6e69dd5c [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-25 11:30:41 +00:00
Daniel Han
8ef51d744b Studio diffusion: cross-platform device policy, fp16 guard, lock split, validate-before-evict
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.
2026-06-25 11:14:39 +00:00
pre-commit-ci[bot]
39eb022b52 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-25 04:11:49 +00:00
oobabooga
002b1216ee Address Codex review: FP16 fallback on pre-Ampere, forward HF token, hide edit models, download outside the load lock 2026-06-25 01:08:18 -03:00
pre-commit-ci[bot]
0a2a423b08 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-25 03:21:55 +00:00
oobabooga
2a5de505df Cancel diffusion loads on unload, mask seeds to JS-safe range, skip malformed gallery records 2026-06-24 23:30:29 -03:00
oobabooga
818ddb828f Add FLUX.2-klein-9B to the curated image models 2026-06-24 22:31:57 -03:00
oobabooga
6411f9464a Add FLUX.2-klein-4B via Flux2KleinPipeline 2026-06-24 22:09:59 -03:00
oobabooga
5d80d30168 Support Qwen-Image and FLUX.1 image models alongside Z-Image 2026-06-24 21:42:57 -03:00
oobabooga
5c4d09a1c9 Paginate the image gallery with infinite scroll 2026-06-24 20:48:07 -03:00
oobabooga
2bf9d73b4a Reuse formatEta, compute generation ETA once per step, guard re-renders 2026-06-24 19:50:10 -03:00
oobabooga
54d01a505f Add a live per-step generation progress bar with ETA 2026-06-24 19:41:15 -03:00
oobabooga
e2cd5ab90d Trim verbose comments 2026-06-24 19:21:29 -03:00
oobabooga
a87e092f59 Drop dead PNG-base64 helper and dedupe the image-gen task list 2026-06-24 19:17:49 -03:00
oobabooga
571d506c20 Refine image generation UI: aspect ratio sliders, info hints, action toolbar, and timestamped export filenames 2026-06-24 18:47:33 -03:00
oobabooga
eced3620fa Add batch size and sequential count to image generation with reproducible per-run seeds 2026-06-24 18:02:38 -03:00
oobabooga
caa7efecc0 Persist the image gallery to disk with recipes embedded in each PNG, and match official Z-Image defaults 2026-06-24 17:50:11 -03:00
oobabooga
adebdfc9f4 Fix image load progress, hide diffusion models from chat picker, and gallery polish 2026-06-24 17:17:28 -03:00
oobabooga
453efa7ae7 drop unused diffusion status 'loading' field; cap session gallery 2026-06-24 16:18:14 -03:00
oobabooga
efdc5ad9a2 transparent image rail; hide fine-tuned shortcut; custom folders filtered to diffusion GGUFs 2026-06-24 16:00:17 -03:00
oobabooga
425d61e47a Z-Image-Turbo image generation with chat-style model picker 2026-06-24 15:32:56 -03:00
oobabooga
e8d3cfe2b5 Merge branch 'main' into image-generation 2026-06-24 14:07:59 -03:00
oobabooga
ab6c9ecfee
Studio: honor stream=false on the GGUF agentic tool path (#6570) (#6618)
* 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>
2026-06-24 15:37:08 +01:00
Nilay
e5cf956601
Studio: shareable per-checkpoint preview links (#6486)
* 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>
2026-06-24 06:31:53 -07:00
oobabooga
346d96d7f2
Studio: cap GGUF context to unified memory on Apple Silicon (#6622)
* 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>
2026-06-24 04:39:33 -07:00
Saicharan Ramineni
9d53656614
Make _uv_safe_path space-safe on macOS/Linux (#6503) (#6534)
* 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>
2026-06-24 04:02:24 -07:00
Daniel Han
d1529b1466
Tidy verbose Studio launch messages (#6628)
* Tidy verbose Studio launch messages

The reachability failure hint, the loopback deploy block, and the stop
hint were several lines longer than they needed to be, and the banner
printed a Tip line that just repeated the URL already shown above. Shorten
them while keeping the useful detail (cloud provider names, the SSH
local-forward workaround, the relaunch command, the trusted-network and
macOS Ctrl+C notes). No behavior change, output wording only.

* Refine launch-message wording after review

Apply review feedback so the launch messages read well for both experts
and general users:
- reachability hint: restore the 'from your own computer' cue and put the
  ssh command on its own line so it stops wrapping, name the cloud rules
  precisely (GCP firewall / Azure NSG rule), and restore 'in your browser'
- loopback banner: add the missing colon, drop the Ctrl+C duplication
  (the stop hint already covers it), and explain the exposure in plain
  language instead of 'exposes the API on every interface'
- stop hint: trim so it fits an 80-column terminal without wrapping
- replace two pre-existing em dashes with --
2026-06-24 03:49:19 -07:00
oobabooga
1cc785e5a0
Studio: remove OpenEnv and other unused packages (#6585)
* Studio: drop OpenEnv and unused ExecuTorch/open_spiel install deps

* Studio: drop 8 more unused install deps from extras

* Studio: restore tomli<3.11 for kernels; tidy dep-cleanup comments and tests

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

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

* Studio: refresh scan-packages baseline for scipy _external + unsloth-zoo tests

scipy moved its vendored array_api_compat from scipy/_lib to
scipy/_external, so the four allowlisted array_api_compat __init__.py
entries stopped matching and resurfaced as unsuppressed CRITICAL
"Downloads and executes remote code" findings on all three pip
scan-packages shards (extras, hf-stack, studio). Add the _external
paths next to the existing _lib ones so both scipy layouts stay covered.

Allowlist two unsloth-zoo test-file false positives now present in the
hf-stack shard: tests/test_mlx_save_export_regressions.py (writes to
/tmp dropper) and tests/test_mlx_trainer_internals.py (obfuscation plus
exec/eval).

Drop nine stale entries for packages removed from the Studio
requirements and no longer in any shard closure (evaluate, pytest,
hypothesis, kgb, langid), confirmed absent via with-deps resolution of
all three shards.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-06-23 07:20:47 -07:00
Daniel Han
61eef657e8
Harden MLX self-heal install against supply-chain code execution (#6599)
* 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>
2026-06-23 07:02:14 -07:00
Leo Borcherding
69d8a57ee9
Studio: lazy-import matplotlib so the server starts when the wheel is blocked (#6596)
* Studio: lazy-import matplotlib so the server starts when the wheel is blocked

matplotlib.pyplot was imported at the top of core/training/training.py, on the
server boot path. When matplotlib's native extension fails to load (e.g. an
unsigned wheel blocked by Windows Smart App Control), that import crashed the
whole Studio server at startup instead of just disabling loss plots.

Move it into a lazy _load_pyplot() helper called from _create_loss_plot, using
the headless Agg backend, and return None when matplotlib is unavailable so
plotting degrades gracefully. The plot return was already Optional, so callers
need no changes. Keep the type-only import under TYPE_CHECKING and quote the
annotations.

Fixes #6588

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

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

* Studio: pin matplotlib==3.11.0

Pin matplotlib to the current latest so a new unsigned release does not
reintroduce the Smart App Control block on Windows. Belt-and-suspenders on
top of the lazy import. Pinned in both studio.txt and extras.txt.

* Pin matplotlib to 3.10.9 so Studio still installs on Python 3.10

matplotlib 3.11.0 requires Python >=3.11, so the pin had no installable wheel on
Python 3.10 (still supported) and pip install failed there. 3.10.9 is the latest
3.10.x (requires-python >=3.10) and covers Python 3.10 through 3.13. Also tighten
the lazy-import docstrings.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <danielhanchen@gmail.com>
2026-06-23 06:22:20 -07:00
Wasim Yousef Said
37166efcfc
Fix Gemma 4 GGUF OpenAI API streams (#6476)
* 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>
2026-06-23 06:13:56 -07:00
Daniel Han
6866362da7
studio: report the true reasoning duration and fix Stop for thinking models (#6521)
* 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>
2026-06-23 14:59:56 +02:00
Daniel Han
1ffffc1cf1
Studio: correct anyio<4.14 comments to the real #6483 cause (#6581)
#6579 reworded these comments to attribute the failure to a half-resolved
install and claimed a clean 4.14 is fine on 3.13. That is wrong: #6483 is a
genuine anyio 4.14 + Python 3.13 regression. 4.14 added a per-task cancel
scope in its asyncio backend (TaskHandle/_run_coro) that gets exited in the
wrong task under starlette's collapsing task group, raising the cancel-scope
RuntimeError on streaming; 4.13 has no such code and is unaffected (the
reporter confirmed 4.13.0 fixes it). The TaskHandle ImportError is only the
secondary macOS-arm symptom from the mlx-vs-cap version fight. Comments only.
2026-06-23 05:50:53 -07:00
Daniel Han
935f6c50ef
studio: tighten torchao Windows-ROCm comments and test docstrings (#6610) 2026-06-23 05:49:25 -07:00
Daniel Han
76cbddb859
Studio: allow --secure with --api-only (headless secure API server) and add --api-only to unsloth studio run (#6591)
* 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>
2026-06-23 05:44:56 -07:00
Daniel Han
55c392ff7c
studio: fix sentence-transformers RAG embedder on Windows ROCm (torchao) (#6608)
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.
2026-06-23 05:39:02 -07:00
Daniel Han
21bdc8fa8c
Studio: treat data-center Blackwell (sm_100/sm_103) as Blackwell in llama.cpp prebuilt selection (#6584)
* 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>
2026-06-23 05:15:09 -07:00
Long Yixing
dad11e8c0c
Fix Studio export checkpoint ordering (#6602)
* fix(studio): sort export checkpoints by step

* [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>
2026-06-23 12:25:53 +01:00
Daniel Han
bebc93d8fc
fix(studio): handle multimodal list content in inference text paths (#4383) (#6480)
* 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>
2026-06-23 01:26:11 -07:00
Daniel Han
e226e0ac35
CI: fix import-hoist false positive, vision-cache test cwd, llama.cpp CLI smoke (#6598)
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.
2026-06-23 01:16:47 -07:00
oobabooga
a1423e05d2 Studio: local diffusion image generation page 2026-06-23 00:39:36 -03:00
Saicharan Ramineni
7ecbf5a770
Use UTF-8 for Python code-execution subprocess I/O (#6489 class) (#6548)
* 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>
2026-06-22 09:06:03 -07:00