Commit graph

348 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
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
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
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
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
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
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
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
Daniel Han
007a21235c
Generalize transformers tier selection by probing AutoConfig (#6550)
* Resolve the transformers tier by probing AutoConfig instead of guessing

When the only signal is a 5.x tokenizer class, get_transformers_tier guessed the
lowest 5.x sidecar (530). That misroutes models whose built-in config parser needs
a higher tier: dense NemotronH ships a 5.x tokenizer but its '-' (MLP) layer only
transformers 5.10 can parse, so 5.3/5.5 raise KeyError '-'. The config.json
transformers_version field records the saving version, not the minimum to load, so
it cannot drive routing either.

Replace the weak tokenizer->530 guesses (local and remote) with a probe: parse
config.json with the built-in parser (trust_remote_code=False) in each sidecar,
escalating 530->550->510, and pick the first that succeeds. This generalizes to any
architecture without hardcoded lists. Strong signals stay fast paths (no subprocess);
the probe runs only when the tier is otherwise ambiguous and is cached by (model,
commit sha). It never executes repo code, never downloads weights, never raises, and
falls back to the legacy 530 guess on a transient/auth/offline failure or when no
sidecar is available. UNSLOTH_DISABLE_TIER_PROBE restores the old behavior.

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

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

* Address review: tier probe fallbacks and cross-platform robustness

Codex:
- Never escalate to 510 on uncertainty. When every sidecar was probed and none
  parsed with the built-in parser, the model is a remote-code / custom model_type
  that loads via its own code; keep the legacy 530 route instead of jumping to
  510 (which would change the behavior of models that worked on the 5.3 stack).
- Only cache the 530 fallback when the result is conclusive (every tier actually
  probed). If a sidecar was missing/uninstallable the environment is incomplete,
  so return 530 uncached and retry on the next call.
- Do not pin the tier cache under an unknown revision: _resolve_commit_sha no
  longer memoizes a None sha (a transient Hub failure is retried), and _probe_tier
  only caches a tier when the commit sha is known.

Gemini:
- Wrap Path.exists() in the sha resolver in try/except OSError (a remote repo id
  can raise WinError 123 on Windows).
- Probe script writes the error to sys.stderr.buffer as UTF-8 bytes so a non-ASCII
  message cannot itself raise UnicodeEncodeError under cp1252.
- subprocess.run decodes stderr with errors="replace" to avoid UnicodeDecodeError
  on non-UTF-8 consoles.

Tests: 72 passed (added partial-sidecar uncached, sha-unresolved not cached,
all-failed stays 530 + cached, sha resolver retries None / handles OSError).

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

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

* Address review round 2: authenticate tier checks, stop memoizing local sigs

Codex:
- Thread hf_token through _check_config_needs_510/550 and
  _check_tokenizer_config_needs_v5 (and the underlying raw fetches). Previously a
  gated/private model whose only 5.x signal is tokenizer_config.json never reached
  the authenticated probe: the unauthenticated raw fetch failed and cached False,
  so the model fell through to the default 4.x tier. The per-check caches are now
  keyed by (model, token) so an unauthenticated miss cannot poison a later authed
  read, mirroring _load_config_json.
- _resolve_commit_sha no longer memoizes a local directory signature. A local
  signature is mutable (size/mtime of config/tokenizer), so a reused/overwritten
  checkpoint path would otherwise keep selecting the previous tier; it is now
  recomputed every call. Only the immutable remote commit sha is memoized.

Tests: 75 passed (added token-cache isolation + auth header, local signature not
memoized, token threaded into all checks/probe).

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

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

* Address review round 3: reach activation with the token, drop SHA tier cache

Codex round 3:
- Thread hf_token into the activation path that actually selects a sidecar. The
  token-aware tier checks added last round were unreachable:
  activate_transformers_for_subprocess called get_transformers_tier without a
  token, and the inference/training/export workers passed only the model name even
  though they hold a request-scoped hf_token. activate_transformers_for_subprocess
  now takes hf_token and the three workers forward config["hf_token"], so a
  gated/private model whose only 5.x signal is an authenticated config/tokenizer is
  routed to the right sidecar instead of falling to default 4.x.
- Stop importing huggingface_hub during tier detection. _probe_tier no longer
  resolves a commit sha, so it never pulls huggingface_hub into the worker before
  the sidecar venv is prepended to sys.path (activation only prepends, never
  purges), which would otherwise pin the default-env hub over the sidecar's
  pinned huggingface_hub==1.8.0.
- The tier cache is now keyed by model_name for the process lifetime (a model's
  required tier is a property of its architecture; cleared on restart). This drops
  the mutable-SHA memo that masked remote revision changes and the mutable
  local-signature memo, removing _resolve_commit_sha / _local_dir_signature /
  _probe_sha_cache entirely.
- Do not cache a probe success that depended on a skipped lower tier: if a lower
  sidecar was unavailable, the lowest valid tier may change once it installs, so
  the result is returned uncached and re-probed next call.

Tests: 73 passed (probe imports no hub; success uncached when a lower tier is
skipped; activation forwards the token).

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

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

* Trim comments to be more succinct

* Re-probe overwritten local checkpoints and authenticate the probe child

The AutoConfig tier probe cached its result under the bare model_name, so a
local checkpoint overwritten in place (same path, new config.json) kept serving
the stale sidecar. Fold a cheap config.json signature (size + mtime) into the
cache key for local paths; remote ids stay name-keyed so no huggingface_hub
import lands before the sidecar is activated.

The probe relies on the implicit HF_TOKEN env, so an inherited
HF_HUB_DISABLE_IMPLICIT_TOKEN=1 left it unauthenticated and a gated repo 401ed
into the 530 fail-safe. Clear that flag in the child env when a token is set.

* Keep tier probes off the log-only path and probe new 5.x archs default-first

- get_transformers_tier gains probe=True/False. needs_transformers_5 (a coarse
  4-vs-5 boolean used only for a spawn log and a vision-check branch) now passes
  probe=False, so a parent/log-only caller never spawns sidecar probes. The real
  activation path keeps probe=True and resolves the exact tier in the worker.
- A config.json saved by transformers 5.x but matched by no fast path is now probed
  default-first: _probe_tier gains include_default + floor, prepending the ambient
  4.57.x tier to the escalation. A model that still parses on the default is left on
  it (no mis-route onto a sidecar); only a config the default parser cannot read
  escalates to the lowest 5.x tier that parses. The transformers_version field is a
  cheap 'worth probing' hint only, read from the already-fetched config (no extra
  network); ordinary 4.x configs never probe.

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

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

* Separate probe cache by mode and keep version-field 5.x visible to needs_transformers_5

- _probe_tier cache was keyed only by config.json signature, so a default-first probe
  that returned 'default' could be handed back to a later tokenizer/known-5.x caller
  (floor=530), leaving a model with a 5.x-only tokenizer on transformers 4.x. Key the
  cache by probe mode (floor + include_default); the legacy 530 mode keeps the bare key.
- The version-field 5.x detection is a cheap config read, not a probe, so run it even
  when probe=False: a standard-tokenizer model whose only signal is transformers_version
  >= 5 now classifies as 5.x via needs_transformers_5 (returns '530' without spawning a
  probe), so the vision-routing fallback uses the 5.x subprocess instead of failing the
  default parser and marking it non-vision. The real activation path still probes
  default-first and may resolve 'default'.

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

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

* Don't treat local checkpoints as Hub ids, and fix stale activation test double

- _load_config_json / _check_tokenizer_config_needs_v5: a local checkpoint dir whose
  config.json / tokenizer_config.json is not yet present was being fetched from the Hub
  as if the path were a repo id, and the 404 miss was cached. A later call after the
  file is written (in-progress checkpoint) then served the stale miss, so a
  TokenizersBackend checkpoint fell through to the default tier. Skip the Hub fetch for
  local dirs and do not cache the miss, so the file is read once it appears.
- test_activate_transformers_version_or_warn_*: the worker now threads hf_token into
  _activate_transformers_version (model_name, hf_token); update the one-arg test doubles
  to the real two-arg signature so the silent-success path stays silent.

* Tighten comments in the AutoConfig probe and tier-selection paths

* Address review: canonical probe cache key and reuse _token_cache_key

- _probe_cache_key resolves config.json to its absolute realpath before
  keying, so a relative path or a changed cwd can't collide with or miss a
  prior probe result. Remote ids still fall back to the name (stat raises,
  caught).
- _cached_config_json reuses _token_cache_key instead of re-hashing the
  token inline, keeping the (model, token) key derivation in one place.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-22 08:20:06 -07:00
Daniel Han
e2e8e5ab46
Studio: show tool-call progress for large GGUF tool arguments (#6484)
* Studio: show tool-call progress for large GGUF tool arguments

The GGUF agentic tool loop only surfaced an early provisional tool card
for render_html, so any other tool (python, terminal, ...) was invisible
in the UI while its arguments streamed. For a large argument such as a
full HTML or code file this left the chat sitting on "Generating..." with
zero progress for tens of seconds while the model was clearly working.

Generalize the provisional tool_start to any enabled tool once its
streamed arguments grow past a threshold (render_html still surfaces
immediately, small-argument tools are unchanged). The provisional and the
real tool_start share the tool_call_id so the frontend reconciles them
into one card. Close the provisional on no-op, denial, parallel-drop,
post-loop, and on stream errors so a card can never spin forever, surface
each parallel call, and skip the early card while a human confirmation
gate is active. Apply the same confirmation-gate guard to the safetensors
agentic loop.

Additional hardening:
- Only emit a provisional card once a real, non-empty tool_call_id is
  known. llama.cpp can stream a tool call with an empty id, and a card
  keyed by "" cannot reconcile with the real tool_start (the frontend
  mints its own id per event), so it would dangle.
- On a connection drop or other mid-iteration failure, close the dangling
  provisional card with an error result instead of an empty success so the
  UI renders it as failed rather than completed.
- Mirror the provisional cleanup in the safetensors loop: close a
  provisional render_html card if the model generator raises mid-stream or
  the controller turns the call into an internal no-op.

Adds regression tests for the empty-id guard, the error-result on a
dropped connection, and the safetensors mid-stream exception cleanup.

* [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: wasimysaid <wasimysdev@gmail.com>
2026-06-22 05:50:10 -07:00
Daniel Han
ab2717afe0
Studio: persistent per-user trust_remote_code approval cache (#6551)
* Studio: persistent per-user trust_remote_code approval cache

The consent gate pins each approval to a content fingerprint (sha256 over every
repo .py), but nothing was persisted, so the dialog reappeared on every fresh
load of the same unchanged repo. This adds an on-disk, per-user approval cache
that lets the gate skip the dialog when the same user reloads the same code,
while keeping the safety guarantees intact.

Two-tier validation, both must hold or the user is re-prompted:
- Commit SHA (cheap, one HfApi.model_info().sha, no download): a match means a
  byte-identical tree to the approved revision, so the scan/download is skipped.
- Content fingerprint (authoritative): used whenever the SHA is unavailable
  (local path / offline) and always recomputed on a SHA miss. A new or edited
  .py changes both the SHA and the fingerprint, so it is caught in every mode.

Safety:
- Keyed per subject; one user's approval never auto-runs code for another.
- CRITICAL is never stored or honored (guarded on both write and read), so a
  hand-edited store cannot smuggle in an auto-approval.
- The malware (HF unsafe-file) gate stays unconditional.
- Fail-safe: a corrupt store, an unresolvable SHA, or any error degrades to
  "ask again", never to "auto-approve". UNSLOTH_TRC_APPROVAL_CACHE_DISABLE=1
  turns the cache off entirely.

New module utils/security/remote_code_approvals.py holds the store
(studio_root()/security/remote_code_approvals.json, atomic write, 0600, RLock)
plus the SHA resolvers. Recording happens at the single gate chokepoint when the
caller supplies the matching fingerprint, so subject is just threaded through
inference/training/export (orchestrators, routes, workers). The scan endpoint
returns already_approved so the frontend can skip the dialog on a cache hit.

Tests: new tests/test_trc_approval_cache.py covers cache miss, SHA-match skip,
SHA-moved re-scan, new-file re-consent, CRITICAL never cached (write + forged
read), disable flag, subject isolation, combined adapter+base key, corrupt
store, and no-subject bypass. Full security suite: 101 passed.

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

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

* Address review: make the approval cache skip only the prompt, never the scan

Codex found that the SHA "no-scan" fast path could run untrusted code without
re-consent. Removed it; the gate now always re-scans and the cache only seeds the
authoritative fingerprint check, so it can skip the dialog but never the scan.

- CRITICAL is hard-blocked on every load (the scan always runs), so a hand-edited
  store that downgrades a CRITICAL repo's severity can no longer auto-run it
  (P2: do not trust editable severity for SHA approvals).
- The fingerprint covers external auto_map repos, so changed third-party code
  always re-prompts even when the primary commit SHA is unchanged; there is no
  longer a SHA path that bypasses the fingerprint (P1: external auto_map repos).
- resolve_commit_sha is resolved fresh on every call (no memoization), so a repo
  whose default branch moves after approval re-prompts instead of reusing a stale
  cached SHA (P1: revalidate mutable Hub SHAs). The SHA is now only a conservative
  secondary gate: a fresh resolvable SHA must match the approved revision, else the
  seed is withheld; a None (local/offline) falls back to the fingerprint.
- Approvals record the scanner ruleset version (SCAN_RULES_VERSION); the gate
  ignores approvals from an older ruleset so reclassified bytes are re-scanned and
  re-shown instead of silently auto-approved (P2: invalidate on scan-policy change).

Tests: test_trc_approval_cache.py rewritten around the prompt-skip semantics
(unchanged repo still scans; SHA move / changed code / scanner-version bump /
disable flag all re-prompt; forged downgraded severity still blocks CRITICAL).
105 passed with test_consent_gate.py.

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

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

* Trim comments to be more succinct

* Keep run-owner subject out of persisted config; serialize approval writes

Threading subject (the run owner's username / API-key id) into the training
config meant _sanitize_db_config persisted it into config_json, which
training-history GET returns to any authenticated user, leaking who started a run
in multi-user installs. Filter subject alongside the token fields; the worker
still receives it from the live config.

The approval store's RLock only guards one process, but approvals are recorded
from separate inference/export/training subprocesses, so concurrent writers could
clobber each other on os.replace and drop an approval (re-prompt). Hold a
best-effort cross-process file lock around the read-modify-write.

* Fail safe on a malformed approval store

A store with the right version but a non-dict shape (e.g. a hand-edited
"subjects": []) passed _load()'s check, then lookup chained .get() on a list and
raised, breaking every remote-code load until the file was removed. Validate that
subjects is a dict in _load(), and tolerate a non-dict per-subject entry in
lookup/record/forget, so a corrupt store fails safe (re-prompt) instead.

* Keep subject out of the MLX W&B run config

_run_mlx_training uploads the whole training config to W&B minus a sensitive set
that only listed hf_token/wandb_token/s3_config, so the authenticated subject
(username / API-key id) was sent to W&B as run config even though DB history
already strips it. Add subject to the W&B-sensitive filter, mirroring
training._sanitize_db_config.

* Tighten the W&B subject-filter comment

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-22 05:12:49 -07:00
Daniel Han
2a05426adb
Auto-install SSM kernels (causal-conv1d, mamba-ssm) for inference loads (#6535)
* Auto-install SSM kernels (causal-conv1d, mamba-ssm) for inference loads

Mamba/SSM hybrids (Nemotron-H/Nano, Falcon-H1, Granite-4.0-H, ...) lazily import
mamba_ssm / causal_conv1d during from_pretrained, so loading them for chat failed
with 'mamba-ssm is required by the Mamba model but cannot be imported'. The training
worker already wheel-first installs these before a fine-tune; the inference worker
did not. Add utils/ssm_runtime.ensure_ssm_runtime and call it from the inference load
path so the same models load for inference. Training worker is untouched; a drift
test keeps the shared detection and pinned versions in lockstep.

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

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

* ssm_runtime: invalidate import caches, skip MLX, cover LoRA base

- Invalidate importlib finder caches in _is_importable and after a successful
  wheel install, so a kernel installed earlier in this same process is actually
  importable when the modeling code lazy-imports it during from_pretrained.
- Skip the SSM kernel install entirely on the MLX (Apple Silicon) load path:
  these are CUDA/ROCm Torch kernels with no MLX use and no macOS prebuilt wheel,
  so the source build would fail before the MLX backend loads the model.
- For LoRA loads, also run detection over the resolved base model, since an
  adapter id like 'me/my-lora' won't match the SSM heuristics but its SSM base
  (Nemotron-H, ...) is what needs the kernels.

Adds tests for cache invalidation and the MLX-skip / LoRA-base worker wiring.

* Tighten SSM autoinstall comments

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

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

* ssm_runtime: verify wheel imports, HIP-aware source build, build heartbeat

Address review feedback:
- Verify a prebuilt wheel actually imports before trusting it; a CUDA/ABI-mismatched
  wheel now falls back to a source build instead of returning success and failing later
  with the cryptic lazy-import error.
- HIP-aware source build: require hipcc on ROCm, inject clang --gcc-install-dir, and use
  the 1800s timeout, mirroring the training worker (ROCm has no prebuilt wheel).
- Emit a status heartbeat every 60s during the source build so a long (ROCm) build does
  not trip the orchestrator's 300s inactivity timeout.

Tests cover the wheel-not-importable fallback and the missing-hipcc ROCm bail.

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

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

* Make causal-conv1d best-effort and harden the SSM source build

- causal-conv1d is a fast path: models that merely want it (Qwen3-Next, LFM2)
  fall back to torch, so a failed install must not reject an otherwise loadable
  chat model on Windows/CPU/macOS or an ABI without a wheel. Only a true SSM
  model's mamba-ssm requirement stays fatal, matching the training worker which
  treats causal-conv1d as best-effort.
- The source build is reached only when not importable, including a wheel that
  installed but failed to import; add --reinstall/--force-reinstall so it
  replaces the broken install instead of no-opping as already satisfied.
- Add --no-cache to the ROCm uv source build to avoid reusing stale artifacts
  from a partial HIP build, mirroring the training worker.

* Address review: install SSM kernels before transformers, harden import + Windows

Codex:
- Install the SSM kernels before importing transformers. run_inference_process
  imported core.inference.inference (which imports unsloth/transformers) before the
  load, and a sidecar transformers can evaluate its optional-backend gates against
  the import state; installing causal_conv1d/mamba_ssm afterwards left those gates
  unsatisfied and a Nemotron/Falcon/Granite load still failed with "mamba-ssm is
  required". The initial model's kernels are now installed in run_inference_process
  before the ML import, via a shared _ensure_ssm_kernels helper; _handle_load keeps
  calling it (idempotent) for a LoRA's base and for later in-process loads.
- _is_importable now treats any import failure as "not importable", not only
  ImportError. An ABI-incompatible native kernel (undefined symbol after a torch/CUDA
  upgrade) raises OSError/RuntimeError; letting those escape reported
  ssm_runtime_install_failed instead of falling back to reinstall/source build.
- Skip causal-conv1d on Windows (no prebuilt wheel), mirroring the training worker.
  A causal-conv1d-only model (Qwen3-Next/LFM2) no longer drops a chat load into a
  multi-minute untimed source build; it uses the torch fallback. mamba-ssm is still
  attempted for true SSM hybrids.

Tests: test_ssm_runtime.py +5 (broken-kernel exceptions read as not-importable;
causal-conv1d skipped on win32 while mamba-ssm still installs). 36 passed.

* Trim comments to be more succinct

* Run security gates before installing SSM kernels

The SSM kernel auto-install is name-based (model_is_ssm is a substring match, no
config fetch), so a model id merely containing an SSM substring triggered a
native-package install (possibly a slow source build) before the malware and
remote-code consent gates ran. Extract those gates into _run_security_gates and
call it before the kernel install in both the pre-import path of
run_inference_process and in _handle_load, so a blocked or nonexistent model is
refused before any build. The gates are metadata-only and do not import
transformers, so they are safe to run before the pre-import install.

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

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

* Resolve remote LoRA bases before importing transformers

_resolve_base_model only reads a local adapter_config.json, so a remote LoRA
adapter whose own id has no SSM substring but whose base is a Nemotron/Falcon/
Granite model had its base discovered only by ModelConfig in _handle_load, after
transformers was imported and its optional-backend availability snapshotted, so
the SSM kernel install there was too late. Add _remote_lora_base, a metadata-only
adapter_config.json fetch (no huggingface_hub / transformers import), and use it
in the pre-import path so the base is gated and its kernels pre-installed.

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

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

* Gate only loaded roots, tier on the resolved base, read offline LoRA cache

Three follow-ups to the pre-import resolution:

- The security gate reused the SSM target list, which for a local full fine-tune
  includes the config.json-recorded base. That base is never loaded, so scanning
  it could falsely block a safe local checkpoint. Gate only the model plus a
  genuine LoRA base (matching _handle_load's mc.is_lora), separate from the
  broader SSM-install list.

- Tier activation ran on the raw adapter id, so a remote LoRA whose base needs a
  sidecar transformers version imported the default and failed. Resolve the base
  once up front and activate on it.

- _remote_lora_base bailed on offline before checking the hub cache, missing a
  cached adapter's base. Read the cached adapter_config.json when offline or when
  the fetch fails.

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

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

* Keep the pre-import gate transformers-free; harden remote LoRA resolution

The pre-import security gate called security_load_subdirs, which imports
model_config and thus transformers, snapshotting optional-backend availability
before the SSM kernels are installed and defeating the ordering. Add
compute_subdirs to _run_security_gates and pass False in the preflight so it scans
from the root only (transformers-free); _handle_load still runs the authoritative
gate with full subdir scoping after the import.

_remote_lora_base now skips existing local relative paths (is_local_path) so a
checkpoint like outputs/run1 is never treated as a Hub repo, and distinguishes a
definitive 404 (not a LoRA -> None) from transient/offline failures (read the
cache), so a repo that is now a full model no longer resolves a stale cached base.

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

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

* Probe a real model id for SSM kernels; respect HF_ENDPOINT

model_is_ssm is a substring match, so an arbitrary name could false-match and
force a mamba-ssm install that fails the load for a non-SSM model:
- a LoRA adapter id like user/falcon-h1-lora (the SSM-relevant code is the base's);
- a local checkpoint under an SSM-named parent dir, e.g. /runs/falcon-h1/llama-ckpt.

Add ssm_probe_identifier, which resolves the base (or a bare local checkpoint's
basename) and feed that to ensure_ssm_runtime from both the pre-import path and
_handle_load, so detection runs against a real model id, never an adapter id or
parent folders.

_remote_lora_base now honors HF_ENDPOINT so enterprise/mirror deployments resolve
the adapter base instead of always hitting huggingface.co.

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

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

* Tighten comments in the pre-import SSM gate/install path

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <michaelhan2050@gmail.com>
2026-06-22 04:48:29 -07:00
Daniel Han
cba73457df
Studio: self-heal unsloth namespace shadows; clearer failed-load messages (#6532)
* Studio: self-heal unsloth namespace-package shadows in all subprocess workers

A directory named `unsloth` (or `unsloth_zoo`) without an __init__.py on
PYTHONPATH/sys.path, a stray source checkout or a polluted PYTHONPATH, makes
`import unsloth` resolve to an empty namespace package, so a worker's
`from unsloth import FastLanguageModel` dies with a cryptic
"cannot import name ... (unknown location)".

The LLM training path already recovered from this via `_ensure_real_packages`
in trainer.py (PR #6269), but the inference, export, and embedding-training
subprocesses imported Unsloth directly with no guard. Extract that helper into
a shared, dependency-free core/import_guards.py and call it before the Unsloth
import in every subprocess: it drops the offending sys.path entries, imports
the real packages (unsloth before unsloth_zoo so the pre-zoo GPU fixes run),
then restores sys.path. trainer.py now imports the shared helper instead of its
local copy.

Covers both unsloth and unsloth_zoo and both namespace origin forms (None and
"namespace"). The existing PR #6269 test now exercises the shared helper.

* Studio: distinguish a failed model load from no model in the attach gates

A failed load never sets the checkpoint, so the image and audio attach gates
fell through to "Load a model before adding images/audio", which reads as if
the user simply forgot to pick a model rather than that the load errored. Add a
dedicated lastModelLoadError to the chat runtime store, set only when an actual
load attempt fails (not on refresh, list, status, or unload errors, which keep
using modelsError) and cleared when the next load starts. The image gate (all
three call sites) and the audio gate now use it to report a failed load and
point at the server logs, while still blocking in exactly the same cases.

* Tighten namespace-shadow guard and load-error comments
2026-06-21 22:43:31 -07:00
oobabooga
e6b4480832
Studio: simplify the inference backend (#6490) 2026-06-21 20:01:09 -03:00
UmranPros
9e83399f9e
Studio: fix Gemma 4 separate-drafter MTP detection and fallback (#6459)
Recognise the Gemma 4 separate-drafter MTP family, auto-download the drafter with retry, fall back to n-gram with a clear reason when it cannot be resolved, and retry the download on reload. Gemma 3n (ships no drafter) and embedded-MTP models (Qwen) are unaffected.

Fixes #6406
2026-06-20 04:43:37 -07:00
Daniel Han
52877bba05
Studio: gate the MTP target-KV reserve to MTP spec mode, not just MLA (#6449)
_estimate_mtp_overhead_bytes is also reached for the separate-drafter spec modes
(draft-simple / draft-eagle3) through _user_draft_via_extras. Those modes load a
small distinct drafter with its own KV -- already counted in the draft KV +
weights -- and keep no duplicated full target context; only MTP runs a second
context over the target model's own KV geometry (llama.cpp ctx_tgt). Charging the
~main-KV-sized f16 copy there over-reserved by tens of GiB on an MLA model and
needlessly shrank the advertised context, the same under-advertising #6312 set
out to fix.

Thread mtp_keeps_target_ctx through _estimate_mtp_overhead_bytes (True for MTP,
False for separate-drafter modes) and derive _engaged_is_mtp at the fit call site
so the target copy is added only when the engaged mode is actually MTP. MLA + MTP
(GLM-5.2 / DeepSeek / Kimi) is unchanged, so the GLM-5.2 OOM fix is preserved;
non-MLA and the draft-simple / draft-eagle3 paths no longer pay the copy.

test_mtp_mla_target_ctx.py adds a case asserting the separate-drafter reserve
collapses to the draft KV (no target copy) while the default MTP path keeps it.
2026-06-19 05:51:35 -07:00
Daniel Han
76a2b9edf1
Studio: Auto disables MTP for MLA models (GLM-5.2 et al.); UNSLOTH_MLA_MTP_ENABLED to re-enable (#6468)
* Studio: Auto disables MTP for MLA models (GLM-5.2 et al.); UNSLOTH_MLA_MTP_ENABLED to re-enable

Studio's Auto speculative mode promotes any embedded-MTP model >=3B to
--spec-type draft-mtp. For MLA models (GLM-5.2/DeepSeek/Kimi) that is a
regression: llama.cpp's MLA/DSA MTP path keeps a duplicated full target-KV
context and recomputes the sparse-attention indexer every draft step, so it
runs ~2x slower than no speculation (GLM-5.2 UD-IQ1_S bench: 27 vs 45 tok/s,
flat across draft depth 1..6 and 96-100% acceptance, on both prose and code).
vLLM/SGLang get a speedup from the same model, so this is a llama.cpp
implementation gap, not a model property.

Auto now drops embedded MTP for MLA models and falls back to ngram-mod (or
spec-off when the binary lacks ngram-mod), mirroring the existing sub-3B
fallback. The metadata separator is kv_lora_rank: it is present on MLA models
and absent on non-MLA embedded-MTP models (Qwen3.x-MTP), whose MTP module is
structurally identical but fast, so a "full layer" heuristic cannot tell them
apart. Qwen MTP, separate drafters (Gemma, --model-draft), and non-MTP models
are unchanged.

Explicit overrides still engage the slower MTP route: choosing MTP / MTP+Ngram
in Settings, or passing --spec-type in extra args. UNSLOTH_MLA_MTP_ENABLED=1
re-enables Auto promotion for MLA once the upstream path is optimized.

A new spec_fallback_reason value "mla_mtp_disabled" surfaces this as an
Auto-mode policy downgrade (not a binary/update problem), with a settings
banner that points users at the MTP override. It is deliberately kept out of
the "Update llama.cpp" affordance since updating does not help.

Tests: resolver-matrix rows for MLA->ngram-mod / MLA-no-ngram->off /
non-MLA-Qwen->draft-mtp / MLA-separate-drafter->draft-mtp /
non-MTP-MLA->default / forced mtp|mtp+ngram on MLA->draft-mtp / env flag;
kv_lora_rank metadata fixtures; and reload-skip coverage (Auto ngram-mod is
idempotent, forced mtp bounces a reload).

* [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-19 05:40:16 -07:00
oobabooga
b4fa48e1b3
Studio: simplify the inference orchestrator and worker (#6439) 2026-06-18 17:23:16 -03:00
Leo Borcherding
5be8835de5
Studio: skip tensor-parallel for vision models; fix MTP drafter VRAM reserve on Windows (#6416)
---------

Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-06-18 17:21:57 -03:00
Daniel Han
1390e721cb
Studio: reserve the duplicated MTP target KV context for MLA models (GLM-5.2 OOM) (#6447)
* Studio: reserve the duplicated MTP target KV context for MLA models

GLM-5.2 UD-IQ1_S advertised its native 1,048,576-token context, loaded, then
crashed cublasCreate on the first generation with "CUDA error: the resource
allocation failed" on a 2x B200 box. The model loaded fine; the decode OOMed.

Cause: when MTP speculative decoding is engaged, llama.cpp keeps a second full
copy of the target model's KV context for draft verification (ctx_tgt=yes in the
spec log), at f16. On an MLA model that copy is ~the main KV again -- for GLM-5.2
at 1M ctx llama.cpp sized it at ~97.5 GiB -- but the auto-fit reserve only
counted the tiny embedded draft head (~2 GiB), 46x too low. So weights (~202 GiB)
+ main KV (~83 GiB) + a 2 GiB reserve looked like it fit in 2x182 GiB, when the
real footprint with the ~97 GiB MTP copy is ~382 GiB and overruns the cards.
Disabling speculative decoding removed the copy and the same context ran fine.

_estimate_mtp_overhead_bytes now adds the duplicated target context (the main KV
re-estimated at f16) for MLA models, so auto-fit backs the context off (or selects
more GPUs) instead of advertising one that OOMs. It is gated strictly on MLA
(kv_lora_rank present), which is exactly the family that keeps the extra copy
(GLM-5.x, DeepSeek, Kimi-K2); non-MLA MTP (Qwen, Gemma) is byte-for-byte
unchanged. The reserve stays deterministic from GGUF dims, matching #6312.

test_mtp_mla_target_ctx.py covers it: the MLA reserve includes the f16 target
copy and dominates the draft head, the copy is f16 regardless of the main cache
type and scales with context, non-MLA embedded heads keep overhead == draft KV,
and _fit_context_to_vram on the GLM-5.2 / 2x B200 budget now returns a context
below the requested 1M where the old draft-only reserve kept the full 1M.

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

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

* Studio: stub structlog/loggers in MTP-MLA test so it is import-order-independent

The new test imports core.inference.llama_cpp, which pulls in orchestrator ->
structlog. In the lightweight test env structlog is absent, so when this file is
collected before test_mtp_vram_budget.py (it sorts first) or run directly,
collection aborted with ModuleNotFoundError. Install the same loggers/structlog
(+ conditional httpx) stubs the sibling MTP tests use before the import, matching
the established per-file convention.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-18 10:24:33 -07:00
Daniel Han
0c1127cb08
Studio: Bypass Permissions menu fix, decimal GB sizes, and GLM-5.2 high/max/disabled thinking (#6444)
* Studio: fix Bypass Permissions menu freeze and show decimal GB for model sizes

Bypass Permissions freeze: the warning dialog lived inside the composer
"+"/More dropdown and kept the menu mounted via onSelect preventDefault,
so confirming or cancelling the dialog left both popovers frozen open.
Lift the dialog out of the menu into a store-driven
BypassPermissionsConfirmDialog mounted at a stable spot in the composer.
The menu item now closes normally on select and just toggles a new
bypassConfirmOpen store flag, so the popovers dismiss as expected.

Model search sizes: formatBytes divided bytes by 1024 but labelled the
result "GB", so unsloth/GLM-5.2-GGUF:UD-IQ1_S showed 201.8 GB where
Hugging Face reports 217 GB. Switch the search display to decimal
(base-1000) units to match what Hugging Face reports. The GPU-fit math
stays base-1024 since VRAM capacity is binary.

* Studio: address review feedback and add GLM-5.2 high/max/disabled thinking

Review feedback on the Bypass Permissions and size-format changes:

- Mount the Bypass Permissions warning dialog once at the chat-page root
  instead of inside each Composer. It is driven by global store state, so
  the per-composer mount meant Compare mode (multiple composers) rendered
  duplicate dialogs and the shared-composer menu had none. A single root
  mount fixes both.
- Defer opening the dialog past Radix's menu-close focus restoration with
  setTimeout(0), so the dropdown does not steal focus back and break the
  dialog's focus trap.
- Clamp the unit index in formatBytes so units[i] cannot go out of bounds
  past TB (and to absorb log() float error at exact powers of 1000).

GLM-5.2 reasoning levels:

GLM-5.2's template gates thinking with enable_thinking and also reads a
reasoning_effort level ('high' or 'max'), so it needs high / max /
disabled rather than the binary toggle it got before (its style was
detected as enable_thinking, which made 'high' unreachable). Add a new
reasoning style 'enable_thinking_effort' that reuses the effort dropdown
but, unlike gpt-oss, can be fully disabled:

- detect_reasoning_flags classifies a template that has both
  enable_thinking and reasoning_effort, extracting the discrete levels
  from the quoted effort literals it branches on. Templates with only one
  of the two (gpt-oss, Qwen3, DeepSeek, GLM-4.6) are unchanged.
- _request_reasoning_kwargs maps the new style to enable_thinking plus an
  in-range reasoning_effort; disabling sends enable_thinking=false. The
  gpt-oss reasoning_effort path is left untouched.
- The backend reports reasoning_effort_levels on the load/status response;
  the frontend carries them through to the effort dropdown and sends
  enable_thinking + reasoning_effort for this style.

Verified: backend reasoning kwargs render the real GLM-5.2 template to
"Reasoning Effort: High/Max" (thinking) and an empty <think></think>
(disabled); tsc, eslint, i18n parity and the production build all pass.

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

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

* Studio: address review feedback on reasoning effort and formatBytes

- chat-adapter localReasoningEffort: accept 'minimal' so a template that
  branches on it (extracted into reasoning_effort_levels) is sent through
  instead of being coerced to 'low' and then dropped by the backend.
- formatBytes: return '0 B' for non-finite / non-positive sizes (missing
  metadata -> NaN, Infinity, negatives) and clamp the unit index lower
  bound to 0, so sub-1-byte values can't produce a negative index.

* Studio: hybrid reasoning none gate and decimal GB in load progress

- _request_reasoning_kwargs: for enable_thinking_effort models, treat a
  raw reasoning_effort='none' (OpenAI 'no reasoning' sentinel) as the
  enable_thinking=false off gate, so a direct API caller can disable
  thinking even without passing enable_thinking. The frontend already
  sends enable_thinking=false; this only affects raw API callers.
- use-chat-model-runtime: the download / 'X of Y GB in memory' load
  progress divided bytes by 1024**3 but labelled GB, so it disagreed with
  the model picker and Hugging Face. Use decimal GB (1e9) to match.

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

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

* Studio: carry hybrid reasoning levels on all load paths and harden formatBytes

Review follow-ups on the enable_thinking_effort work:

- Every model-load path now copies reasoning_effort_levels and derives
  supportsReasoningOff, via a shared reasoningCapsFromLoad() helper. The
  shared/Compare composer load and the three chat-adapter auto-load paths
  previously set only reasoningStyle, so a GLM-style hybrid model loaded
  through Compare or first-chat auto-load fell back to the default
  low|medium|high and lost its Max / Off controls.
- The local send path clamps the effort to the loaded model's advertised
  levels (clampReasoningEffortToLevels) instead of a hard-coded list. A
  stale "max" carried over from an external provider no longer reaches a
  pure reasoning_effort (gpt-oss) model that only accepts none|low|medium|
  high, where the backend would have dropped it.
- formatBytes divides iteratively instead of via Math.log, which has float
  error at exact powers of 1000 (log(1e12)/log(1000) = 3.9999... would
  label 1 TB as "1000 GB"). Keeps the non-finite/non-positive guard.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-18 10:10:01 -07:00
Daniel Han
6d27160dcc
Studio: graceful recovery ladder when llama-server hard-crashes at startup (#6291)
* Studio: fall back to text-only when a vision projector hard-crashes llama-server

The text-only mmproj fallback (#6075) only fired when llama-server printed a
recognizable projector-format error ("Unknown projector type", exit -6). An
installed llama.cpp that predates a model's projector can instead SIGSEGV
(exit -11) with no parseable output, e.g. unsloth/Qwen3.5-4B-MTP-GGUF +
mmproj-F16 on an older gfx1151 prebuilt: llama-server crashes on load, the
--fit off retry crashes the same way, and load_model gives up with a hard 500
instead of dropping vision.

Generalize the decision: a vision (--mmproj) launch killed by a signal (POSIX
returncode < 0, e.g. -11 SIGSEGV / -6 SIGABRT; Windows 0xC0000000+ access
violation) is treated like a projector incompatibility, so the load retries
once text-only. The retry is skipped if a cancel/unload is pending, mirroring
the MTP guard. Clean non-zero exits (bad GGUF, port bind) and hung processes
keep their own handling; non-vision launches are unaffected.

Reproduced and verified on gfx1151 (Radeon 8060S, ROCm 7.2.1): a current
prebuilt (llama.cpp b9596) loads the exact model + args fine, confirming the
crash is a stale prebuilt. With a wrapper that SIGSEGVs on --mmproj, Studio now
recovers: the load returns 200 (is_vision=false) and serves at ~31 tok/s
text-only instead of failing. New _is_signal_crash helper plus tests pin the
decision.

Also normalize a few em-dashes to ASCII punctuation in existing comments.

* Studio: refine mmproj hard-crash fallback (signal scope + last argv)

- Limit _is_signal_crash to genuine program faults (SIGSEGV, SIGABRT,
  SIGILL, SIGFPE, SIGBUS) and Windows 0xC0000000+ statuses. SIGKILL,
  SIGTERM and SIGINT no longer count, so an OOM-killer, unload or
  supervisor kill is not masked as a projector incompatibility.
- Strip --mmproj from the last attempted argv so the text-only retry
  keeps --fit off / --spec-default instead of resurrecting the original
  spec flags (matters for MTP vision models on an older llama.cpp).
- Drop stray temp files committed by mistake and gitignore the "~" dir
  so they cannot be re-added.

* Studio: tighten comments in mmproj hard-crash fallback

* Studio: retry --flash-attn off before dropping vision on a startup crash

When llama-server hard-crashes at startup, the recovery chain now tries the
least-destructive mitigation first. Flash-attention kernels SIGSEGV at load on
some ROCm/GPU builds (often inside the vision tower's attention); disabling
flash attention keeps BOTH vision and MTP, so a hard program fault with
--flash-attn on now retries once with --flash-attn off before the MTP-drop or
the text-only (mmproj-strip) fallbacks. _is_signal_crash already gates this to
genuine faults (SIGSEGV/SIGABRT/SIGILL/SIGFPE/SIGBUS), so an OOM-kill or unload
(SIGKILL/SIGTERM/SIGINT) does not trigger a retry.

Field context: a gfx1151 user crashes loading a vision GGUF even on the latest
prebuilt, so an update cannot help, and the same model and args load fine on
another gfx1151 box, pointing at a runtime/flash-attn fault. New
_with_flash_attn_off helper plus tests. Verified on hardware with a wrapper
that SIGSEGVs on --flash-attn on: Studio recovers with is_vision=true (vision
and MTP intact) instead of failing or losing vision.

* Studio: name the OOM kill on a too-large model load

When the OS kills llama-server with no diagnostic output (SIGKILL/SIGTERM,
almost always the OOM killer, e.g. a BF16 model too large for the WSL VM's
RAM cap), the recovery ladder correctly does not retry an external kill, so
this is the message the user sees. It fell through to the generic "is the
GGUF valid / out of memory" text. Make it actionable: name the signal and
point at a smaller or more quantized GGUF, a lower context length, or raising
the WSL memory limit. Output-based diagnoses still win and a hard fault keeps
the generic fallback.

* Studio: refuse a model too large for system RAM on a unified-memory APU

On gfx1150/gfx1151 APUs the weights load into shared system RAM (GGML
unified memory). _get_gpu_free_memory reports the full ROCm/APU budget as
free (often ~100 GB), but under WSL the VM's RAM cap is the real ceiling.
Studio trusted the budget, spawned a load larger than RAM, and the OS killed
it mid-flight, taking the Studio process with it (a silent "Terminated" with
no error, the model resident in RAM not VRAM).

Add a pre-flight guard on the APU path: if the weights exceed available
system RAM (psutil, then /proc/meminfo), refuse before spawning with a clear
message (smaller/more-quantized GGUF, lower context, or raise the WSL memory
limit). Weights only so KV/context auto-reduction is not double-counted;
unknown RAM never refuses; non-APU and discrete-GPU paths are untouched.

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

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

* Studio: refine recovery ladder (keep diagnosed errors, flip all flash-attn)

Two review points on the hard-crash recovery ladder:

1. The signal-only text-only fallback stripped --mmproj on any hard fault,
   even when llama-server had already printed a non-projector cause (an OOM
   such as "cudaMalloc failed: out of memory", an unsupported architecture, or
   a tensor-parallel limit). That masked the real error and told the user to
   update llama.cpp for vision. New _output_has_nonprojector_diagnostic gates
   the signal path: it fires only when no such marker is present, so a bare
   SIGSEGV with no output still retries text-only, but a diagnosed OOM surfaces
   the real error instead of silently dropping vision.

2. _with_flash_attn_off only flipped the first --flash-attn. llama.cpp is
   last-wins, so a leftover enable from extra_args (--flash-attn on, -fa on, or
   the = form) could keep flash attention on and re-crash the retry. It now
   flips every occurrence and returns None only when nothing is flippable.

test_llama_cpp_mmproj_fallback.py and the classification/APU suites: 103 passed.

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

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

* Studio: pass the text-only retry's exit code to the failure classifier

When the text-only fallback retry itself fails, read its exit code before
_kill_process() clears it and forward it to _classify_llama_start_failure, so an
OS-killed retry surfaces the actionable out-of-memory message instead of the
generic one (matching the primary failure path).

* Studio: scope APU RAM guard to selected GPUs, count MTP drafter, neutral SIGTERM

Three refinements to the startup recovery work in this PR:

- The unified-memory APU RAM guard fired whenever any visible GPU was a
  gfx1150/gfx1151 APU, so on a mixed APU+dGPU host it could refuse a valid
  load placed on the discrete GPU. Scope _amd_apu_wants_unified_memory to the
  selected gpu_indices (physical ids, mapped via CUDA_VISIBLE_DEVICES like
  _is_datacenter_gpu); None still means every visible GPU. Applied to both the
  RAM guard and the GGML_CUDA_ENABLE_UNIFIED_MEMORY env set.
- The RAM guard counted only the main GGUF plus mmproj, so a separate MTP
  drafter (also resident in unified system RAM, even when offloaded to CPU)
  could push the load past the RAM cap and still get OS-killed mid-load. Add
  the drafter weights to the APU RAM total.
- The startup classifier reported SIGTERM (-15) as 'most likely out of memory',
  but SIGTERM is also how an unload/cancel or a supervisor stops the server.
  Keep the OOM wording for SIGKILL (-9, the OOM killer) and report -15
  neutrally.

Tests updated/added accordingly.

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

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

* Studio: address review on the APU guard and decode-probe ladder

- Map APU physical ids via the active ROCm mask (HIP, then ROCR, then CUDA),
  mirroring _get_gpu_memory, so a HIP_VISIBLE_DEVICES-selected APU is matched.
- Only add the MTP drafter to the APU RAM total when MTP will actually engage,
  so a stale LLAMA_ARG_SPEC_DRAFT_MODEL cannot refuse a non-MTP load.
- After an MTP first-decode hard fault, retry --flash-attn off (keeps MTP)
  before dropping speculative decoding, matching the startup rung.
- Fold the --flash-attn= / -fa= rewrite into one branch.

Tests: tensor-parallel decode-probe assertion updated for the FA-off rung.

* Studio: tighten two comments in the APU guard and RAM preflight

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

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

* Studio: refine flash-attn retry and APU RAM guard per review

- _with_flash_attn_off now decides on the effective last-wins value: it returns
  None when FA is already off (no wasted retry), and neutralizes a bare
  --flash-attn / -fa (which llama.cpp reads as on) so the retry cannot re-enable
  it. Length is preserved so downstream index slices stay valid.
- _amd_apu_wants_unified_memory uses 'gpu_indices is not None' so an empty
  selection is respected (not treated as all-visible).
- The APU RAM refusal now checks the base model only (main + mmproj); an
  optional MTP drafter is dropped by the existing MTP-drop fallback rather than
  causing a hard pre-spawn refusal of an otherwise loadable model.

Tests: bare-flag / effective-off / empty-selection / HIP-mask cases added.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-18 09:07:25 -07:00
Daniel Han
c42c1d56e8
Studio: free chat model VRAM at training start only when the GPU is tight (#6243)
* Studio: free chat model VRAM at training start only when the GPU is tight

The training start route unconditionally tore down the transformers/MLX
inference subprocess before training, and never stopped the llama.cpp GGUF
server at all, so a loaded GGUF chat model kept holding VRAM for the whole
run. Conversely the HF model was always unloaded even when there was plenty
of room to keep it.

Make the unload VRAM aware and cover every inference backend:

- Add routes/training_vram.py with summarize_resident_chat(),
  can_keep_chat_during_training() and free_chat_models_for_training(). The
  keep/unload decision reuses the same estimator and live per device free
  VRAM reader the training GPU selection already uses (auto_select_gpu_ids,
  estimate_required_model_memory_gb, get_visible_gpu_utilization), so the
  probe agrees with the placement computed later in start_training.
- When a chat model is resident and training fits alongside it with a
  conservative margin (required_gb * 1.15 + 4 GB), keep it loaded so the
  user can train and chat at the same time; on a multi GPU box training
  lands on a different GPU and both coexist. Otherwise unload the HF/MLX
  orchestrator and the llama.cpp GGUF server before training starts.
- The export subprocess shutdown stays unconditional and now runs first so
  its freed VRAM is reflected in the decision.

Default deny: non CUDA backends, unestimable models, or any probe error
fall back to the previous always unload behavior.

Adds tests/test_training_vram_coexistence.py and updates two existing route
tests in test_gpu_selection.py.

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

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

* Studio: per-GPU floor for explicit GPU lists + don't unload chat on invalid gpu_ids

Address review feedback on the chat coexistence probe:

- Explicit gpu_ids mode now enforces a per-GPU floor in addition to the
  aggregate free-VRAM check, mirroring auto_select_gpu_ids' min_per_gpu_N.
  Without it, an uneven split such as free [45, 10] for a 40 GB job passed
  the aggregate threshold and kept chat loaded even though the 10 GB GPU
  could not hold its training shard, risking an OOM.
- Invalid explicit gpu_ids (ids outside the visible set, or a UUID/MIG
  mask) make resolve_requested_gpu_ids raise. That request is rejected with
  a 400 before training starts, so leave the resident chat model untouched
  instead of unloading it.
- Tighten the target_modules / gpu_ids type hints to List[str] / List[int].

Adds tests for the per-GPU floor (uneven split unloads, even split keeps)
and for invalid gpu_ids keeping the chat model loaded.

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

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

* Studio: only free chat VRAM once training will start; handle in-flight and CPU-only chat

Address the second review pass on the chat-coexistence path:

- Run the chat/export VRAM teardown as a before_spawn hook inside
  TrainingBackend.start_training, fired only after the start guards pass.
  Previously the route freed chat VRAM before calling start_training, so a
  refused start (e.g. a lingering pump thread) would tear down the resident
  chat model even though no training job began.
- Treat an in-flight HF chat load (loading_models set, no active model yet)
  as not safely sizeable: free it rather than risk both OOMing as the load
  keeps allocating after training starts.
- Do not count or tear down a GGUF llama-server confirmed to run entirely on
  CPU (_gpu_offload_active is False): it holds no VRAM, so killing it cannot
  help training fit.

Adds tests for the before_spawn hook (runs on start, skipped when a
subprocess is alive or a pump thread will not die, survives a hook error),
the in-flight load flag, and the CPU-only GGUF exclusion in both the resident
summary and the unload path.

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

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

* Studio: treat any in-flight chat load (HF swap / mid-start GGUF) as unsafe to keep

Tighten the in-flight detection in summarize_resident_chat so the keep check
never sizes a load that is still allocating:

- Flag loading on ANY non-empty loading_models, not only when active_model_name
  is empty. load_model adds the new model to loading_models before clearing the
  old active_model_name, so a replacement load during a swap was previously
  sized as a normal resident and could OOM as the new model finishes loading.
- Flag a GGUF server that is active but not yet healthy (is_loaded False) as
  in-flight: it is still mmaping/offloading layers, so its final VRAM footprint
  is unknown.

Consolidates the signal into a single resident["loading"] flag; the route frees
the chat model whenever it is set. Adds tests for the replacement HF load and
the mid-start GGUF cases.

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

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

* Studio: tighten comments in chat/training VRAM coexistence (comments only)

* Studio: run before_spawn VRAM hook only after GPU-selection validation

Reviewers found the before_spawn hook fired before prepare_gpu_selection
validated gpu_ids (and before config build), so a refused start (invalid
gpu_ids -> 400, or a bad grad-clip value) could still tear down chat/export
VRAM. Move the hook to immediately before proc.start(), once all synchronous
validation and process construction have passed. This also fixes the route's
in-flight-chat loading branch, since that teardown runs inside the same hook.

Add test_hook_skipped_when_gpu_selection_rejects.

* Studio: recompute GPU auto-selection after the before_spawn VRAM hook

Codex P2: with before_spawn moved after prepare_gpu_selection, placement was
frozen against the pre-teardown VRAM state while the hook freed export/chat
afterward. Auto-selection could pin training onto a GPU the hook then cleared
(or onto a kept chat model). Split validation from placement: explicit gpu_ids
are still validated before the hook (raise -> 400, no teardown; explicit
placement is VRAM-independent), but VRAM-dependent auto-selection now runs
after the hook so it sees the freed memory.

Add test_auto_placement_runs_after_hook and test_explicit_placement_validated_before_hook.

* Studio: allow chatting during training (lift sidebar gate + VRAM-aware load guard) (#6335)

* Studio: allow chatting during training (lift sidebar gate + VRAM-aware load guard)

The sidebar disabled New Chat, project, and home navigation while a training
run was active, so users could not chat during training even though the backend
serves inference fine alongside a run. This removes that gate and adds a backend
guard so the one genuinely risky operation, loading a new local chat model
mid-training, is refused with a clear 409 when it would not fit beside the run.

Frontend (app-sidebar.tsx): drop the chatDisabled = isTrainingRunning gate and
its consumers. Navigation triggers no model load on its own, so chat stays
usable during training.

Backend (routes/training_vram.py, routes/inference.py): add
can_load_chat_during_training plus a load/validate guard that sizes the same
effective load the loader performs (LoRA 4-bit to 16-bit resolved first, HF auto
placement via auto_select_gpu_ids, explicit multi-GPU per-GPU floor, GGUF sized
from on-disk shards and companions or the selected remote variant). It is a
no-op when training is inactive, never blocks external providers or
already-resident models, and default-denies only on a CUDA sizing failure so a
load can never OOM the run. Validate refuses early with the real settings so the
frontend does not unload the resident chat model for a load that would be
rejected.

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

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

* Studio: address review feedback for chat-during-training load guard

- Run the load/validate VRAM guard via asyncio.to_thread so the sync
  nvidia-smi + HF metadata work never blocks the event loop.
- Size the GGUF KV cache at the requested context (_estimate_gguf_kv_gb)
  and add it to the local GGUF estimate so large-context picks are not
  under-counted.
- Keep the requested quantization when adapter_config.json is malformed
  (not a JSON object) instead of raising in _effective_load_in_4bit.

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

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

* Studio: size the training load guard at the launcher's effective GGUF context

The GGUF KV-cache estimate used max_seq_length only, but the llama.cpp
launcher honors a user --ctx-size/-c in llama_extra_args. A load such as
max_seq_length=4096 with --ctx-size 131072 was sized against a 4k cache
while the server allocates 131k, so the guard could approve a long-context
GGUF load that then OOMs training. Size the guard's KV at the larger of
max_seq_length and the parsed --ctx-size (reusing the launcher's own
parse_ctx_override), keeping the conservative f16 cache so the estimate is
never smaller than what the server allocates.

The chat model picker also validated with the raw max_seq_length while
/load sizes with resolveLoadMaxSeqLength, so validate could pass, unload
the current model, then have /load reject the native-context load. Validate
now uses the same effective context; the load path is unchanged.

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

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

* Studio: size the GGUF training guard at the server parallel-slot count

The KV-cache estimate assumed a single slot, but llama-server allocates the
cache across --parallel slots (app.state.llama_parallel_slots). On a Studio
launched with --parallel N>1 the guard under-sized the cache N-fold and could
approve a GGUF chat load that then OOMs training. Thread the same slot count
the loader uses into the guard's KV estimate; default 1 leaves single-slot
setups unchanged.

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

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

* Trim comments for chat-during-training guard

* Studio: keep chat generation alive across navigation; Train spinner + Return to Chat

Hoist the base chat runtime above the routed outlet so navigating to Train (or any tab) no longer aborts an in-flight generation; only an explicit Stop cancels. Add a Train sidebar spinner and swap New Chat to Return to Chat while a run is active, with a lightweight completion watch so the spinner clears from any tab. Also respawn a chat llama-server killed mid-session and guard unreadable HF cache dirs that 500'd the hub model list.

* Studio: show Return to Chat on the Train tab whenever a chat is live

Previously the top sidebar item only swapped to Return to Chat while training was running; on the Train tab with an idle/just-finished run it stayed New Chat, which started a fresh thread and cancelled an in-flight generation. Show Return to Chat (and navigate back, preserving the run) whenever a generation is running or its thread is still active, or training is in progress.

* Studio: keep a running chat alive when starting a New Chat

Starting a New Chat (or switching threads) while a generation was in flight
remounted the single-chat runtime provider, which detached the in-flight run
and cut the previous chat off (it showed up frozen / empty when reopened).

Key the single-chat view by project instead of by thread or new-chat nonce so
the provider stays mounted and assistant-ui switches to a fresh thread in place.
The previous generation keeps streaming in the background and autosaves on
completion, and returning to that thread reattaches the live run instead of
reloading a half-saved one.

Also:
- "Return to Chat" now lands on the thread that is still generating rather than
  the empty new chat that became active after New Chat.
- Skip the explicit /inference/cancel POST when an abort comes from a runtime
  detach (navigation / background switch) rather than an explicit Stop, so a
  backgrounded generation is never cancelled behind the scenes.

* Studio: make model export non-blocking and inline

The Export tab opened a full-screen modal that trapped focus, could not be
closed or cancelled while running, and showed no progress. It also stopped
training and unloaded the chat model before loading, so export could not run
alongside them.

Export now mirrors the training runtime pattern:

- Inline panel embedded where the Export Model button was, with no modal or
  backdrop, so the rest of the UI stays usable during an export.
- Global export runtime store plus an app-root lifecycle hook, so a run keeps
  going and streaming across navigation and is reflected on the Export nav item
  from any tab.
- The worker log stream now stays connected across the load to export phase
  boundary instead of stranding on "Waiting for worker output".
- Progress bar driven by phase and quant index (quant N of M for GGUF), with
  elapsed time and a working Cancel.
- load-checkpoint no longer stops training or unloads inference; export loads in
  its own subprocess in parallel and surfaces out-of-memory as a clear error.
- Add POST /api/export/cancel and is_export_active on /api/export/status.

* Studio: show Return to Chat on the Export tab too

Extend the New Chat to Return to Chat swap to the Export route so leaving a
running chat for Export offers a way back to the live generation, matching the
Train tab.

* Studio: smooth out Export animations and polish the panel

- Drop the height-based reveal animations (source switch, run panel, quant
  picker, hub fields) that caused flashing and reflow; use instant swaps and
  quick opacity fades instead.
- Method and quant cards now transition colors only, with no transition-all or
  hover lift, so selecting a method or quant is crisp instead of jumpy.
- Auto-scroll the export panel into view when it opens and add a scroll-to-bottom
  button when its output is below the fold, like Chat.
- Show Return to Chat on the Export tab while an export is running, matching how
  training drives it on the Train tab.
- Surface the current phase or stage in the live output before the first worker
  line arrives so the panel never looks stuck while progress is advancing.

* Studio: show Return to Chat on every non-chat tab

Generalize the Return to Chat swap from just Train/Export to any non-chat route
(Recipes, Projects, Hub, ...) so a running or active chat is always one click
away, instead of showing New Chat there.

* Studio: stream export logs over the Cloudflare tunnel; drop janky export animations

Exporting over a --secure Cloudflare quick tunnel showed "connecting..." with no
logs while the progress bar advanced. Cloudflare buffers text/event-stream and
only flushes when the stream closes, so the SSE log stream never reached the
browser during the run (direct localhost is unaffected, which is why this only
showed up over the tunnel).

Add a tunnel-safe JSON poll fallback (GET /api/export/logs?since=) that the
runtime lifecycle hook polls while a run is active. Short JSON responses are not
buffered by the proxy, so logs show up in near real time over the tunnel. It
shares the orchestrator's monotonic seq cursor with the SSE stream and the store
de-dupes by seq, so the two transports run together (SSE on localhost, poll over
the tunnel) without double-printing. A successful poll marks the panel
"streaming" instead of leaving it stuck on "connecting...".

Also remove the framer-motion AnimatePresence reveals from the export config and
run panel (quant picker, hub fields, the inline run panel, and the live log
section). The expand/slide animations flashed and felt clunky; the sections now
render in place.

* Studio: recover export over the Cloudflare tunnel when the blocking POST times out (524)

A model export over a --secure Cloudflare quick tunnel showed "Request failed
(524)" even though the export succeeded on the backend (the GGUF was written).
Cloudflare returns 524 when a single request takes longer than ~100s to respond,
and a GGUF conversion routinely runs for minutes, so the blocking per-method
export POST is cut off while the backend keeps going.

Confirm completion via short status polls instead of relying on the long POST
response (the same approach that fixed log streaming):

- The orchestrator records each finished op's outcome (status / output_path /
  error) with a monotonic seq, exposed on GET /api/export/status.
- parseJson now preserves the HTTP status; a 524/520/522/523/502/503 or a
  status-less network drop is classified as a recoverable transport error.
- runExport wraps each phase (load, every export method, each GGUF quant): on a
  recoverable failure it keeps the run alive (logs keep streaming, the panel
  shows "reconnecting...") and polls status until the still-running op finishes,
  then settles from the recorded result, recovering the output path for the
  success banner. A real 4xx still fails immediately; localhost still uses the
  fast POST response. applyBackendStatus also settles a reloaded run from the
  last-op record.

Verified over the tunnel: a 3m14s gemma-4-E4B-it GGUF export now ends on the
success banner with the output path instead of 524.

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

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

* Studio: keep the export method + logs visible after navigating away mid-export

While an export was running, navigating to another tab and back to Export
remounted the page and reset the local form state (exportMethod, quant levels),
so the method card showed unselected and the run panel's log area was hidden
until the card was re-clicked. The run itself lives in the global store and was
unaffected.

Seed exportMethod / quantLevels from the active run's summary via lazy useState
initializers on (re)mount, and gate the panel's log area on the live run
(isExporting / logLines / the run's method) rather than only the local form
selection. The card stays selected and the logs/progress stay visible across
navigation; nothing changes when no run is active.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Studio: address export/training review findings

- Export: guard Start against an empty GGUF quant selection so an inline-panel
  run with no quant can't settle as success with no file produced.
- Export: thread the source HF token into the background load so gated/private
  HF source exports (and gated bases) authenticate, matching the consent path.
- Export: only settle a recovered (non-owned) run as a finished export when the
  last backend op was an export, not a standalone load_checkpoint.
- Training: free the export subprocess whenever an export is active, not only
  once a checkpoint is loaded, so an in-flight export load can't race training
  for VRAM (current_checkpoint is unset during the load phase).

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-18 09:04:01 -07:00
Daniel Han
8d804c9413
Studio: show an actionable message when the GGUF runtime is missing (#6327)
* Studio: show an actionable message when the GGUF runtime is missing

Selecting a GGUF model with no llama-server installed surfaced a generic
"Invalid model" in the UI, because validate_model's catch-all discarded the
real cause. Add LlamaServerNotFoundError (a RuntimeError subclass) raised by the
GGUF preflight in ModelConfig.from_identifier, and catch it in the validate
route so users get an actionable message: run `unsloth studio setup` to
download the prebuilt llama.cpp runtime. Other validation failures keep the safe
generic message. Adds a regression test.

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

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

* Studio: also map missing GGUF runtime to a 400 in load_model

validate_model already surfaces the actionable 'install the runtime'
message for LlamaServerNotFoundError; load_model fell through to the
generic 500 'Failed to load model'. Catch it there too so a GGUF load
without llama-server gives the same install hint instead of a 500.

* Trim comments for PR #6327

* Studio: fix stale validate test after #6398 and surface missing GGUF runtime on /load

- test_other_runtime_errors_do_not_get_gguf_message: after merging #6398,
  validate_model surfaces a RuntimeError's own message, so a plain RuntimeError
  no longer returns "Invalid model". Assert it does not receive the GGUF
  install message instead (the prior assertion was stale after the main merge).
- Raise LlamaServerNotFoundError (not a plain RuntimeError) at the backend
  load-time missing-binary branch, after diffusion routing, so /load returns the
  actionable 400 like remote validation, instead of a generic 500.
- Share LLAMA_SERVER_NOT_FOUND_DETAIL between the from_identifier preflight and
  the load-time raise so the message stays in sync.
- Add a propagation regression test for the non-tensor load path.

* [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-18 06:00:00 -07:00