Commit graph

359 commits

Author SHA1 Message Date
Daniel Han
2d09508951
Studio diffusion (Phase 6): img2img / inpaint / edit / LoRA / upscale on the native engine (#6680)
* 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.

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

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

* 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.

* 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.

* Studio diffusion (Phase 5): image quality-vs-quant accuracy harness

Add scripts/diffusion_quality.py, the accuracy analogue of the KLD workflow: hold
prompt + seed fixed, render a grid with a reference quant (default BF16), then render
each candidate quant and measure drift from the reference. Records mean PSNR + SSIM
(pure-numpy, no skimage/scipy) and optional CLIP text-alignment + image-similarity
(transformers, --clip), plus file size, latency, and peak VRAM, then prints a
quality-vs-cost table and recommends the smallest quant within a quality budget.
--selftest validates the metrics on synthetic images with no GPU or model.

Verified on Z-Image (B200): the table degrades monotonically with quant size
(Q8 -> Q4 -> Q2: PSNR 21.7 -> 15.5, SSIM 0.82 -> 0.61), while CLIP-text stays flat
(~0.34) -- quantization erodes fine detail far more than prompt adherence.

* 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.

* 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.

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

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

* 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.

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

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

* 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.

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

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

* Studio diffusion (Phase 6): img2img / inpaint / edit / LoRA / upscale on the native engine

Builds on Phase 4's native stable-diffusion.cpp engine, extending it from
text-to-image to the wider feature surface, since sd.cpp supports all of these
through the binary already. Pure command-builder additions plus one engine
method, so the txt2img path is unchanged.

- sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img +
  strength make a run img2img, adding mask makes it inpaint, ref_images drives
  FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and
  lora_dir + the <lora:name:weight> prompt syntax select LoRAs. New
  SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run
  mode (input image + esrgan model, no prompt / text encoders).
- sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so
  generate() (now carrying the conditioning flags) and a new upscale() reuse
  the same streaming / error / output-check path.
- scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img /
  --strength / --upscale-model / --upscale-repeats.

Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the
upscale builder and its validation, and the engine's img2img + upscale paths.
Full diffusion suite 176 passing.

Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the
init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale
(512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing
coherent images. Video and the diffusers-path feature wiring are deferred.

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

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

* Studio diffusion (Phase 4): enforce the sd-cli timeout while reading output

Iterating proc.stdout directly blocks until the stream closes, so a sd-cli that hangs
without producing output (or without closing stdout) would never reach proc.wait and the
wall-clock timeout was silently bypassed. Drain stdout on a daemon thread and wait on the
PROCESS, so the main thread always enforces the timeout and kills a hung process (which
closes the pipe and ends the reader). Add a test that times out even when stdout blocks,
and make the no-binary test hermetic so a host-installed sd-cli can't leak in.

* Studio diffusion (Phase 4) review fixes: sd.cpp installer + engine hardening

- install_sd_cpp_prebuilt: download the release archive with urlopen + an explicit
  timeout + copyfileobj (urlretrieve has no timeout and hangs on a stalled socket);
  extract through a per-member containment check (Zip-Slip guard); expanduser the
  --install-dir so a tilde path is not taken literally; and on Windows CUDA also fetch
  the separately-published cudart runtime DLL archive so sd-cli.exe can start.
- sd_cpp_engine: find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the
  installer, so a custom-root install is discovered without UNSLOTH_SD_CPP_PATH; start
  sd-cli with the parent-death child_popen_kwargs so it is not orphaned on a backend
  crash; reap the SIGKILLed child (proc.wait) so a cancel/timeout does not leave a zombie.
- tests: Zip-Slip rejection, normal extraction, studio-home discovery.

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

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

* Studio diffusion (Phase 4) review round 2: collect sd-cli batch outputs

Codex review: when batch_count > 1, stable-diffusion.cpp's save_results() writes
the numbered files <stem>_<idx><suffix> (base_0.png, base_1.png, ...) instead of
the literal --output path. SdCppEngine.generate checked only the literal path, so
a batch generation would exit 0 and then raise 'no image' (or return a stale
file). generate now returns the literal path when present and otherwise falls
back to the numbered siblings; single-image behavior is unchanged.

Test: a fake sd-cli that writes img_0.png/img_1.png (not img.png) is collected
without error.

* Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats

Codex review on the native engine arg builder:

- build_sd_cpp_command emitted --width/--height unconditionally, so an
  img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of
  the input. width/height are now Optional (None = unset): an image-conditioned
  run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives
  the size from the input image (set_width_and_height_if_unset); a plain txt2img
  run with unset dims keeps the prior 1024x1024 default; explicit dims are always
  honored. width/height are read only by the builder, so the type change is local.

- build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...)
  that silently swallowed repeats=0 into sd-cli's default of one pass, turning an
  explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError
  and emits the flag for any explicit value != 1.

Tests: img2img unset dims omit width/height (init_img and ref_images), explicit
dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at
the default. (Two pre-existing binary-discovery tests fail only because a real
sd-cli is installed in this dev environment; unrelated to this change.)

* [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: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-07-01 15:18:38 -03:00
Daniel Han
96940d87b7
Studio diffusion (Phase 4): native stable-diffusion.cpp engine for CPU/Mac (#6679)
* 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.

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

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

* 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.

* 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.

* Studio diffusion (Phase 5): image quality-vs-quant accuracy harness

Add scripts/diffusion_quality.py, the accuracy analogue of the KLD workflow: hold
prompt + seed fixed, render a grid with a reference quant (default BF16), then render
each candidate quant and measure drift from the reference. Records mean PSNR + SSIM
(pure-numpy, no skimage/scipy) and optional CLIP text-alignment + image-similarity
(transformers, --clip), plus file size, latency, and peak VRAM, then prints a
quality-vs-cost table and recommends the smallest quant within a quality budget.
--selftest validates the metrics on synthetic images with no GPU or model.

Verified on Z-Image (B200): the table degrades monotonically with quant size
(Q8 -> Q4 -> Q2: PSNR 21.7 -> 15.5, SSIM 0.82 -> 0.61), while CLIP-text stays flat
(~0.34) -- quantization erodes fine detail far more than prompt adherence.

* 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.

* 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.

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

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

* 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.

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

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

* 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.

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

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

* Studio diffusion (Phase 4): enforce the sd-cli timeout while reading output

Iterating proc.stdout directly blocks until the stream closes, so a sd-cli that hangs
without producing output (or without closing stdout) would never reach proc.wait and the
wall-clock timeout was silently bypassed. Drain stdout on a daemon thread and wait on the
PROCESS, so the main thread always enforces the timeout and kills a hung process (which
closes the pipe and ends the reader). Add a test that times out even when stdout blocks,
and make the no-binary test hermetic so a host-installed sd-cli can't leak in.

* Studio diffusion (Phase 4) review fixes: sd.cpp installer + engine hardening

- install_sd_cpp_prebuilt: download the release archive with urlopen + an explicit
  timeout + copyfileobj (urlretrieve has no timeout and hangs on a stalled socket);
  extract through a per-member containment check (Zip-Slip guard); expanduser the
  --install-dir so a tilde path is not taken literally; and on Windows CUDA also fetch
  the separately-published cudart runtime DLL archive so sd-cli.exe can start.
- sd_cpp_engine: find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the
  installer, so a custom-root install is discovered without UNSLOTH_SD_CPP_PATH; start
  sd-cli with the parent-death child_popen_kwargs so it is not orphaned on a backend
  crash; reap the SIGKILLed child (proc.wait) so a cancel/timeout does not leave a zombie.
- tests: Zip-Slip rejection, normal extraction, studio-home discovery.

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

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

* Studio diffusion (Phase 4) review round 2: collect sd-cli batch outputs

Codex review: when batch_count > 1, stable-diffusion.cpp's save_results() writes
the numbered files <stem>_<idx><suffix> (base_0.png, base_1.png, ...) instead of
the literal --output path. SdCppEngine.generate checked only the literal path, so
a batch generation would exit 0 and then raise 'no image' (or return a stale
file). generate now returns the literal path when present and otherwise falls
back to the numbered siblings; single-image behavior is unchanged.

Test: a fake sd-cli that writes img_0.png/img_1.png (not img.png) is collected
without error.

---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-07-01 15:03:53 -03:00
pre-commit-ci[bot]
387cfd268b [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-01 10:47:53 +00:00
Daniel Han
2cae1c1e7c Merge remote-tracking branch 'origin/main' into image-generation
# Conflicts:
#	studio/backend/routes/models.py
#	studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
2026-07-01 10:47:21 +00:00
Daniel Han
8cc05ac89c
Reduce comments across recent fixes (#6776)
Condense the verbose comments and docstrings added by the recent
chat template, GPT-OSS detection, PEFT tensor-parallel, and Studio
inference proxy fixes. Comments and whitespace only; no code changes.
2026-06-30 23:13:36 -07:00
oobabooga
8b69f22809 Fix fp8 text-encoder quant crashing generation on T5 and tied-embedding encoders 2026-06-30 23:22:13 -03:00
Anish Umale
d0f8d40c36
studio: allow updating HF models through UI (#5388)
* add models for /update endpoint

* add logic for identifying out of date hf models

* add endpoint for updating hf models

* add relevant field to GgufVariantDetail

* make exception handling better

* add update_available flag for cached_models, and moved /update endpoint from inference -> models

* hook up /update endpoint on the frontend

* implement update scenarios for the model picker

* fix bug where downloaded flag for an older revision was being wrongly set to false

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

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

* fix import and make hf calls async

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

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

* remove has_vision from UpdateRequest

* fix ci

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

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

* clear cancel event before updating gguf variant

* set _cancel_event back if it was set initially

* add hf_token to get_paths_info

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

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

* studio: harden model update endpoint and update checks

- update_hf_model: pass snapshot_download local_dir (local_path is not a
  valid kwarg and 500s when updating bicodec audio models)
- get_gguf_variants: wrap the remote update check so a network, rate-limit,
  gated, or offline failure degrades to "no update info" instead of failing
  the whole variant listing, matching list_cached_models
- add regression tests for both paths

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

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

* Studio: HF model update detection and Update action for cached models

Surface an "Update available" cue and a managed Update action for cached
on-device models. /api/hub/update-status compares each cached main GGUF
file's local blobs against the remote main revision using set membership
across all cached revisions, so a repo that was already updated (and still
holds the old snapshot alongside the new one) is not falsely flagged.

The Update action re-downloads through the download manager so it shows in
the Downloads panel with progress and cancel. The frontend wires the Update
button into the GGUF, on-device, and model-selector cards and keeps the
quant label fully visible when the action buttons crowd the row.

Adds regression tests for the multi-revision update check.

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

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

* Studio: accept force_download kwarg in hf_xet_fallback test double

The download seam now passes force_download to the attempt callable; the _FakeAttempt mock did not accept it, failing 6 tests with TypeError. Add the keyword (default False) so the scripted-results double matches the seam.

* Fix Studio model update regressions

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

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

* Address Studio update review feedback

* Address Studio update edge cases

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

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

* Share GGUF update status helper

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

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

* Fix GGUF update detection and cache cleanup

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

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

* Fix cached GGUF update badges

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: shimmyshimmer <107991372+shimmyshimmer@users.noreply.github.com>
Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
2026-07-01 01:54:57 +03:00
pre-commit-ci[bot]
53077b5ae3 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-30 22:34:29 +00:00
Daniel Han
6ae6fb1c45
Studio diffusion (Phase 2): memory planner, streamed offload, fp8 TE, speed layer, quality harness (#6675)
---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-06-30 19:30:57 -03:00
Daniel Han
ecf028780d
Studio diffusion (Phase 1): cross-platform device policy, fp16 guard, lock split, validate-before-evict (#6670)
---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-06-30 16:33:47 -03:00
Tai An
7337729e57
fix(studio/llama_cpp): disable trust_env on the loopback health probe (#6750) (#6752)
* fix(studio/llama_cpp): disable trust_env on the loopback health probe

_wait_for_health() polls http://127.0.0.1:<port>/health with the default
httpx trust_env=True, so an ambient HTTP(S)_PROXY in the environment is
applied to the loopback request. A proxy that returns 503 for 127.0.0.1
makes every probe fail, so the loop runs until timeout and Studio load
hangs (trust_env=False returns 200 immediately).

Pass trust_env=False so the local readiness probe never goes through a
proxy. This mirrors the existing trust_env=False handling in the sibling
llama_http / external_provider HTTP clients.

* test(offline_gguf_cache): accept trust_env kwarg in fake_get mock

_wait_for_health now calls httpx.get(..., trust_env=False); update the retry test's fake_get to accept the kwarg so it doesn't raise TypeError.

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

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

* fix(studio/llama_cpp): bypass proxies for loopback clients

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

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

* fix(studio/routes): bypass proxies for llama streams

* [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-30 19:09:26 +02:00
Daniel Han
e8945cab46
Whole-document context for RAG chat attachments (#6693)
* Add whole-document context mode to RAG chat attachments

Thread-attached files are injected in full when they fit a token budget,
instead of only top-K retrieved chunks, so the model reads the entire file
for summarize/reason-over-document requests. Oversized files fall back to
top-K retrieval so the context window is never blown. KB and project
corpora are unchanged (still retrieval).

- core/rag/store.py: all_chunks_for_scope returns every completed-document
  chunk for a scope, ordered document-then-index, joined with filename.
- core/rag/tool.py: whole_document_context renders the chunks as the same
  <chunk> blocks + citation source-map retrieval produces, returns None
  when empty or over budget.
- core/inference/tools.py: build_rag_autoinject tries whole-document first
  for thread scopes, falls through to search_for_autoinject otherwise.
- core/rag/config.py: THREAD_WHOLE_DOC + WHOLE_DOC_MAX_TOKENS (env-tunable).
- tests/test_rag_whole_document.py: store ordering, whole-doc render +
  budget cutoff, auto-inject whole-doc vs top-K fallback, KB never whole-doc.

* Add scanned-PDF OCR fallback to RAG ingestion

A PDF page with no extractable text layer (a scanned or image-only page)
previously ingested as empty, so image PDFs were invisible to retrieval and
whole-document context. Such pages are now rendered and transcribed by the
loaded vision model during ingestion, so they become searchable and readable
like any other page. This restores OCR for the RAG document flow without a
separate extraction pipeline.

- core/rag/parsers.py: render_pdf_pages renders whole pages (1-based) to PNG.
- core/rag/captioner.py: factor the shared vision call into _vision_complete;
  add _ocr_one + ocr_pages (transcribe rendered pages, OCR_MAX_PAGES bound).
- core/rag/ingestion.py: _ocr_scanned_pages runs right after parse, replacing
  text on near-empty PDF pages. No-op when OCR is off, no page is scanned, or
  no vision model is loaded (degrades like figure captioning).
- core/rag/config.py: OCR_SCANNED, OCR_MIN_CHARS, OCR_MAX_PAGES, OCR_DPI,
  OCR_TIMEOUT_S, OCR_MAX_TOKENS (env-tunable).
- tests/test_rag_ocr_fallback.py: page render, ocr_pages gating + cap, scanned
  PDF end-to-end OCR into chunks + whole-doc, born-digital skips OCR, disabled
  leaves the page empty.

* Broaden OCR prompt to figures/tables and guard against repetition runaway

The OCR prompt now asks the vision model to also transcribe text inside figures,
diagrams, charts and tables, so labels and table cells on scanned pages are
indexed rather than skipped. Verified on real documents that this does not
regress plain-text transcription.

Some vision models loop on sparse images (e.g. a title-only cover) and emit the
same line hundreds of times. _collapse_runaway caps any run of identical
consecutive lines so a pathological page cannot flood the index; legitimate
short repeats (a label appearing a few times) survive. Applied in ocr_pages.

* Restrict whole-document injection to thread attachments only

whole_document_context resolved the combined project+thread scope, so a project
chat (the frontend sends both thread_id and project_id) injected the entire
project corpus in full, contradicting the design that project and KB corpora stay
retrieval-only. A large project corpus could also push the total over budget and
drop a small thread attachment back to top-K.

Resolve the thread scope alone in whole_document_context, and in
build_rag_autoinject only enter whole-doc mode when a thread attachment is present
and no KB is selected (a KB pick is exclusive: search that corpus). Project
sources and KBs keep top-K retrieval. Adds regression tests for the mixed
project+thread payload, the budget isolation, and KB precedence.

* Address review: keep project retrieval, harden budget + OCR guards

Follow-up to the 8-reviewer pass on the whole-document + OCR work.

- Preserve project grounding in project chats. The thread-scope-only fix made
  whole-doc exclusive of retrieval, so a thread attachment silently dropped the
  project corpus for that turn. build_rag_autoinject now whole-docs the thread
  attachment AND retrieves the project sources top-K, merged under one citation
  numbering via tool.render_sources. KB selection stays exclusive.
- Budget: a NULL/zero token_count no longer bypasses the cap (length-based
  fallback in _row_token_count), so a malformed huge doc can't inject in full.
- OCR runaway guard: _collapse_runaway now also caps each distinct line at a
  generous total across the page (not just consecutive), bounding the
  interleaved/alternating loops weak models emit; blank-line floods collapse too.
- OCR: warn when a scanned PDF exceeds OCR_MAX_PAGES (pages past the cap stay
  untranscribed) instead of silently dropping them.
- Document the known limits: OCR'd pages have no PDF highlight regions; vision
  models need a micro-batch >= image tokens (Gemma-family) or the server aborts.
- Tests for project-retrieval composition, NULL-token budget, and interleaved
  runaway; drop the now-superseded exclude-project test.

* Add OCR toggle to RAG retrieval settings

Make scanned-PDF OCR user-controllable per upload instead of only via the
RAG_OCR_SCANNED config default. The retrieval settings panel gains an OCR
scanned pages switch (persisted in localStorage, on by default); the chosen
value is read fresh at upload time and sent with each document upload.

Backend: the three upload routes accept an optional ocr form field and pass it
through start_ingestion to _ocr_scanned_pages, which now treats None as use the
config default and an explicit bool as an override. The on/off policy lives only
in _ocr_scanned_pages now, so ocr_pages no longer re-checks the config (that
double gate would have blocked a per-upload ocr=True while the default was off).

Tests cover both override directions (force on while config off, force off while
config on).

* Add "Describe figures & charts" toggle with chart-aware captions

Surface RAG figure captioning as a user control and make it actually useful for
graphs and plots. The figure detection already clustered vector drawings and
raster images into regions and rendered them, but captioning was off by default,
had no UI, and used a thin generic prompt.

Accuracy: the caption prompt now asks for chart type, axis titles and units,
legend or series, salient trends and readable values, and table columns, while
forbidding invented numbers. The token budget is configurable (CAPTION_MAX_TOKENS)
and captions pass through the same runaway guard as OCR so a looping vision model
cannot flood the index.

Control: a per-upload caption override threads from the three upload routes through
start_ingestion and _run, with the on/off policy single-sourced in _run (caption
self-gating removed from caption_images, mirroring the OCR change) so a force-on
override works when the config default is off. The frontend adds a "Describe
figures & charts" switch in the retrieval settings, persisted in localStorage and
sent with each upload. Default on; it is a no-op without a vision model and bounded
to CAPTION_MAX_IMAGES figures per document.

Tests cover the new caption_images contract, the runaway guard on captions, the
chart-aware prompt and token budget (and that OCR keeps its own prompt and budget),
and both override directions end to end through ingestion.

* Generalize figure understanding: transcribe-first prompt + high-DPI tiling

Make figure/chart description work across any visual and any model strength, not
just a strong VLM on simple figures. Two changes, validated by a recall benchmark
on authoritative documents (ResNet/Attention papers, USDA, UN UDHR).

1. Transcribe-first caption prompt. The caption now asks the model to transcribe
   every visible label verbatim (titles, axis labels and units, legends, every
   box/node/arrow label, table cells, equations) and then add a one-line summary,
   instead of only describing the figure. Transcription is the most model-robust
   visual task, so weak models that cannot reason about a chart still recover its
   labels.

2. High-DPI tiling of figure pages. Figure-bearing pages are rendered as an
   overlapping grid of high-DPI tiles (plus a full-page pass for context); each
   tile is transcribed, then merged and de-duplicated. This keeps small diagram
   labels legible and covers every sub-figure without relying on exact region
   detection, which previously missed sub-figures and small labels.

Supporting changes: figure render DPI 130 -> 200 with a clip margin so edge labels
are not lost; vision calls are deterministic (temperature 0) so transcription does
not randomly drop labels; the repetition guard now applies to captions too. New
config knobs: FIGURE_DPI, FIGURE_MARGIN_FRAC, FIGURE_TILE_ROWS/COLS, FIGURE_TILE_
OVERLAP, FIGURE_FULLPAGE, CAPTION_MAX_PAGES, larger CAPTION_MAX_TOKENS, and
CAPTION_MAX_IMAGES as a per-document tile budget.

Measured figure context recall (per-label, dense academic figures):
  Qwen2.5-VL: 0.50 -> 0.83 (overall 0.81 -> 0.94)
  Gemma-4-E2B (weak): ~0 with loops -> 0.83 (overall 0.91)
Born-digital text and scanned-page recall are unchanged (no regression).

parsers gains _figure_boxes (shared detection), pages_with_figures, and
render_pdf_figure_tiles; captioner gains merge_page_captions and a temperature
parameter; ingestion routes figure captioning through the tiled path.

* Fix RAG review issues: whole-doc budget pre-check, figure gating, empty re-ingest, vision auth

Whole-document context now runs a cheap token-sum pre-check (store.scope_token_estimate)
before hydrating every chunk's text, so an attachment that cannot fit the budget is
rejected without loading the whole corpus into memory. The estimate mirrors
all_chunks_for_scope's filter and the per-row token-count fallback exactly.

Ingestion skips all figure work (PDF rasterization and detection, not just the caption
call) unless a vision model is loaded, so a text-only deployment pays nothing. When OCR
is enabled, scanned/image-only pages are excluded from figure tiling since OCR already
transcribes them whole, avoiding double vision work and overlapping index entries; a
scanned figure page is still tiled when OCR is off.

start_ingestion no longer dedupes forever to a prior ingest that produced zero chunks
(e.g. a scanned PDF uploaded before a vision model was loaded): the empty record is
dropped and the content is re-ingested.

Vision OCR and caption requests now send the backend Authorization header, so they
match the chat endpoint and do not 401 under direct-stream (--api-key) mode.

Adds tests for the budget estimate, scanned-page exclusion, the vision-model gate, the
empty re-ingest path, and the auth-header passthrough.

* Trim RAG vision-ingestion comments and docstrings

Tighten the verbose multi-line docstrings and comments added across the RAG vision
ingestion work (captioner, config, parsers, ingestion, store, tool, build_rag_autoinject,
the RAG tests, and the chat-store/upload-hook frontend toggles) to one or two lines while
keeping their intent. No code changed: verified comment/docstring-only against the prior
commit, and the RAG test suite still passes.

* Fix figure-tiling exclusion and client dedupe for re-ingestable docs

Figure tiling now excludes only the pages OCR actually transcribed, not every
text-less page. _ocr_scanned_pages returns the set of pages it OCR'd, and _run passes
that to pages_with_figures as exclude_pages (replacing the ocr_on-keyed min_text_chars
heuristic). A scanned page that OCR skipped (past OCR_MAX_PAGES, or whose OCR returned
empty) is no longer dropped from captioning, so a chart on such a page still gets a
caption.

The document panel's upload dedupe no longer skips re-selecting a file whose only
matching doc completed with zero chunks. Such a doc is re-ingestable (e.g. a scan
attached before a vision model loaded), and the backend re-ingests on the same content
hash, so the client must let it reach the backend; healthy or still-indexing docs are
still skipped. The SSE complete frame's chunk count is recorded on the doc so the
check is exact.

Adds a regression test for the un-OCR'd scanned figure page and updates the
pages_with_figures test to the exclude_pages interface.

* Address review findings: whole-doc budget guard, job numChunks, dead code, upload cap

whole_document_context now treats a non-positive max_tokens as "never inject" instead
of injecting the whole corpus unbounded, so RAG_WHOLE_DOC_MAX_TOKENS=0 tightens rather
than disables the budget (the real off switch stays RAG_THREAD_WHOLE_DOC=0).

The job-status endpoint and get_job_status now expose num_chunks (joined from the
document), and the upload hook threads it through the SSE-fallback completion paths
(reconcile + poll). Previously a document that finished via the connection-cap fallback
had no chunk count client-side, so the re-ingest dedupe wrongly treated it as empty and
re-uploaded it. IndexJob/JobEvent gain the field and the untyped cast is dropped.

Removes the dead render_pdf_figures function (superseded by the tiling path), its test,
and the unused FIGURE_MARGIN_FRAC config knob.

Adds an upload size cap (RAG_MAX_UPLOAD_BYTES, default 200 MB; 413 on exceed with the
partial file cleaned up) so a pathological file can't drive unbounded parse + vision
work. render_pdf_figure_tiles clamps rows/cols to >= 1 (no ZeroDivisionError on a
misconfigured grid). Captioning progress is reported after OCR so the bar is monotonic.
sqlite connections set busy_timeout=5000 so a long figure/scan ingest holding its
connection doesn't make a concurrent ingest/read fail with "database is locked".

Adds tests for the non-positive budget, the zero-grid clamp, job-status num_chunks, and
the oversize-upload rejection.

* Extract PDF text as layout-aware Markdown via pymupdf4llm

parsers._pdf now extracts each PDF page as Markdown with pymupdf4llm.to_markdown
(page_chunks=True) instead of flat page.get_text("text"), so tables, headings and lists
keep their structure in the indexed chunks and retrieve far better (a table's cells stay
associated with their row instead of flattening into a token stream). Gated by
RAG_PDF_MARKDOWN (default on); falls back to plain PyMuPDF text when the toggle is off,
pymupdf4llm is missing, extraction fails, or a page yields no Markdown. The scanned-page
OCR and figure-tiling passes operate on rendered pixels and are unaffected; docx/html/txt
keep their existing extractors.

The preview-highlight locator already strips Markdown punctuation when building anchors;
it now also splits anchor tokens on pipes so a Markdown table row still anchors to the
raw PDF word stream.

Declares pymupdf4llm as a studio/RAG dependency (was only transitively present via the
data-designer plugin). Adds parser tests (Markdown table reaches the page text, the
plain-text fallback, the missing-lib fallback) and a locator test for table-pipe anchoring.

* Pin pymupdf4llm to 0.3.4 so the package scan does not pull onnxruntime

The lockstep pymupdf4llm 1.27.x line makes pymupdf-layout a hard dependency,
which in turn pulls onnxruntime (plus numpy/networkx/protobuf). The security-audit
pip scan-packages job resolves requirements --with-deps, so adding pymupdf4llm to
no-torch-runtime.txt and studio.txt surfaced onnxruntime's un-baselined CRITICAL
finding and flipped the hf-stack shard from pass to fail.

pymupdf4llm 0.3.x keeps pymupdf-layout behind an optional [layout] extra, so a plain
install resolves to pymupdf + tabulate only and never touches onnxruntime. 0.3.4
requires pymupdf>=1.27.1, satisfied by our pinned pymupdf==1.27.2.3, and to_markdown
(page_chunks=True) produces equivalent layout-aware Markdown on real PDFs (verified on
the Attention, ResNet and USDA documents). Production already installs these files
--no-deps, so onnxruntime was never shipped at runtime; this only fixes the scanner.

The parser test now asserts Markdown markup (heading or table pipes) rather than table
pipes specifically, since 0.3.4 emits a heading but not a pipe table on the tiny
borderless synthetic fixture; both markers are absent from the plain-text fallback.

* Fix RAG whole-doc review findings

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

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

* Address RAG whole-doc review follow-ups

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

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

* Address RAG review follow-up edge cases

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

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

* Reserve image budget for whole-document RAG

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

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

---------

Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-30 15:55:23 +02:00
Daniel Han
98a01e70cd
Studio: restore tensor parallelism for vision/mmproj GGUFs (#6659)
* Studio: restore tensor parallelism for vision/mmproj GGUFs

#6416 disabled --split-mode tensor for any GGUF that ships an mmproj projector to
dodge a GGML_ASSERT crash (#6415) seen on an older llama.cpp build with consumer
Blackwell (sm_120). The blanket skip silently dropped tensor_parallel=true for
every multimodal/MTP GGUF (e.g. Qwen3.6-35B-A3B-MTP); on hardware where the model
fits on one GPU the load then collapsed to a single GPU. mmproj + --split-mode
tensor works on current builds (verified end to end on B200/sm_100), so the skip
was disabling a working configuration.

Make the vision skip self-healing per binary:
- attempt tensor for vision models by default
- skip upfront only on a binary already seen to abort on tensor + mmproj this
  session (_vision_tensor_split_aborts), recorded when such a launch crashes at
  startup (_record_vision_tensor_split_abort). Process scoped, so a studio update
  re-probes the new build. The route-level layer-split fallback stays the net.
- add _select_gpus(min_gpus=...) so a downgraded tensor request can keep multiple
  GPUs instead of collapsing to one (default 1, no behavior change).

Add tests/test_tp_vision_regression.py: an AST allowlist guard over the
tensor_parallel drop sites (which would have flagged #6416), plus cache and
_select_gpus coverage. No GPU required.

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

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

* Studio: address review on vision tensor-parallel self-healing

Three fixes from the PR review:

- Record a vision-tensor abort only after every startup retry fails. The first
  version cached the binary on the first spawn crash, which on every build
  (including capable ones) is the benign --fit step abort that the existing
  --fit off retry resolves. That poisoned the cache so the next vision load in
  the same process skipped tensor. Recording now happens at the post-retry
  failure block (after fit-off, flash-attn-off and MTP-drop), so a binary that
  actually works is never cached.

- Gate the record on the tensor/mmproj crash signature: a hard signal fault
  (_is_signal_crash) with no non-tensor cause (_output_has_nonprojector_diagnostic
  excludes OOM and unknown-arch), so an OOM, bad extra args, or MTP/flash-attn
  crash no longer marks an otherwise capable binary incompatible.

- Preserve the multi-GPU request on the cached downgrade. The vision gate now
  raises _layer_min_gpus to the visible GPU count and threads it through the
  layer-split GPU selection (_select_gpus min_gpus and the subset loops), so a
  downgraded tensor request still spreads across GPUs instead of collapsing to a
  single card the model happens to fit.

Verified two vision+tensor loads in one backend process both tensor-split across
4 GPUs (the benign fit abort no longer poisons the cache). Tests updated.

* Studio: harden vision tensor-parallel self-healing (review round 2)

Address the second review round on the vision/mmproj tensor-parallel fix:

- Preserve vision on the first load: a --split-mode tensor + --mmproj
  GGML_ASSERT now raises so the route-level tensor->layer fallback retries
  layer split with the projector intact, instead of stripping --mmproj and
  silently loading text-only (which returned success and skipped the fallback,
  losing vision on the first load until the next cached load).

- Symmetric multi-GPU preservation: the pooled-VRAM tensor downgrade now raises
  _layer_min_gpus from the usable tensor GPUs like the vision downgrade, so it
  no longer collapses a multi-GPU request to a single card.

- Base the layer fallback minimum on usable GPUs: _select_gpus caps min_gpus to
  the count of cards with usable VRAM, so a downgrade never forces a nearly-full
  card in (or trips --fit) just to hit the count.

- Re-probe after in-app updates: key the per-binary abort cache on (path, mtime)
  like _capability_cache, so POST /api/llama/update swapping the binary in place
  (no backend restart) re-probes the new build instead of inheriting the old
  build's abort.

- Bump _layer_min_gpus for a known-bad vision binary independent of the tensor
  drop, so the route fallback's layer retry (tensor already off) still spreads
  across GPUs.

Adds deterministic non-GPU regression tests for each.

* Studio: gate cached-vision layer minimum on the current tensor request

The cached-vision _layer_min_gpus bump fired for every later vision load on a
binary recorded as tensor+mmproj-incompatible, including loads that did not
request tensor parallelism. A plain non-tensor vision load that fits on one card
would then grab every GPU just because an earlier TP attempt aborted in the same
backend process.

Re-tie the bump to the current tensor request (back inside the tensor-drop
guard), so only a downgraded tensor request preserves the multi-GPU spread; a
non-tensor vision load minimizes device count as before.

* Studio: preserve GPU count + confirm assert on vision tensor fallback

Third review round on the vision/mmproj tensor-parallel fix:

- Preserve multi-GPU on the first tensor->layer fallback. The route-level retry
  runs tensor-off, so the in-function downgrades can't see the original tensor
  request and a fits-on-one-card model loaded the first successful fallback on a
  single GPU. The GGUF load closure now passes preserve_multi_gpu_on_layer (the
  toggle asked for tensor, this attempt is layer) and load_model raises
  _layer_min_gpus for it, so the downgrade still spreads across GPUs.

- Cap the auto-context layer loops to usable GPUs. They bypass _select_gpus, so a
  raised _layer_min_gpus could force a nearly-full card into the subset (or trip
  --fit). They now start from _auto_min_gpus, capped to the GPUs with usable VRAM.

- Confirm the tensor/mmproj assert before caching. Recording (and the layer-retry
  raise) now require the ggml assert marker via _is_tensor_split_assert, not the
  bare-signal predicate shared with the projector-incompat branch, so a corrupt
  or too-new projector that SIGSEGVs independent of split mode is no longer cached
  as tensor/mmproj-incompatible.

Adds deterministic non-GPU regression tests for each.

* Studio: extend multi-GPU fallback to extra/env tensor + overhead-aware cap

Fourth review round on the vision/mmproj tensor-parallel fix:

- Preserve multi-GPU fallback for all tensor requests, not just the UI toggle.
  Tensor can also be requested via --split-mode tensor in extra args or an
  inherited LLAMA_ARG_SPLIT_MODE=tensor env; the fallback retries those too, so
  the preserve_multi_gpu_on_layer hint now keys off _effective_tensor_parallel
  (the same check the fallback uses), comparing the overall request against the
  current attempt instead of only request.tensor_parallel.

- Cap the auto-context layer fallback to GPUs that can pay the per-device layer
  overhead. The cap counted any card with positive usable VRAM, so a nearly-full
  GPU with a few MiB free stayed eligible and could be exposed to llama.cpp and
  OOM. It now mirrors _select_gpus: a card counts only if usable VRAM exceeds the
  per-device pipeline overhead.

Adds deterministic non-GPU regression tests for both.

* Studio: match the #6415 split-axis assert + replay layer-preserve hint

Fifth review round on the vision/mmproj tensor-parallel fix:

- Narrow the tensor/mmproj crash signature. _is_tensor_split_assert matched any
  GGML_ASSERT/GGML_ABORT, so an unrelated invariant a corrupt GGUF or projector
  trips with --mmproj present could be cached as tensor/mmproj-incompatible. It
  now matches the specific #6415 warmup assertion
  (GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0) in ggml-backend-meta),
  whose split-axis signature is inherent to tensor splitting. A reworded future
  assert just re-crashes-then-falls-back (vision preserved via layer split)
  instead of poisoning the cache for other models.

- Persist the layer-preserve hint for respawns. A successful tensor->layer
  fallback committed _last_load_kwargs without preserve_multi_gpu_on_layer, so
  _respawn_if_dead replayed only --split-mode layer + tensor_parallel=False and a
  mid-session respawn of a fits-on-one-card model came back single-GPU. The hint
  is now in the replay snapshot, so recovery keeps the multi-GPU placement.

Adds deterministic non-GPU regression tests for both.

* Studio: tighten comments on the vision tensor-parallel fix

Make the comments and docstrings added by this PR succinct: collapse the
multi-line block comments in llama_cpp.py / inference.py to one or two lines,
trim the verbose test docstrings (the names and assert messages already carry the
intent), and shorten the module docstring. No code changes; verified comment-only
with scripts/comment_tools.py check --strip-docstrings.

* Studio: cache vision tensor abort only on the split-axis token

_is_tensor_split_assert also accepted any GGML_ASSERT/GGML_ABORT from
ggml-backend-meta, but that file holds many asserts, so an unrelated
scheduler/projector/model invariant on an --mmproj launch could cache the binary
as tensor/mmproj-incompatible and make later compatible vision models skip tensor
parallelism. Match the GGML_BACKEND_SPLIT_AXIS_* token itself (unique to the
#6415 warmup assert), not the source file name.

* Studio: don't leak the httpx test stub into later tests

The regression module stubbed httpx via sys.modules.setdefault, which installs
the lightweight stub even when real httpx is present but not yet imported. The
stub then persists for the whole pytest process, so provider/HF tests collected
later (importing httpx or huggingface_hub.errors) got a module missing
HTTPError/Response. Mirror the neighboring llama_cpp helper tests: import real
httpx first and only fall back to a stub on ImportError.

* Studio: latch the #6415 tensor-split abort on the first spawn, key it per model

The self-heal recorded the --split-mode tensor abort only in the post-retry
failure block, after the flash-attn-off retry. But SPLIT_MODE_TENSOR requires
flash_attn, so the flash-off retry can't run tensor and its output no longer
carries the warmup split-axis assert (ggml-backend-meta :541). The record
therefore never fired on the real reproducer and the crash loop repeated on
every load (reported by oobabooga on #6659).

Latch instead on the first spawn that shows the signal crash + split-axis
marker: record it, kill the process, and raise straight to the route's layer
fallback, skipping the futile flash-attn/MTP retry ladder for this crash.

The crash is a tensor-split geometry limit (e.g. MQA n_head_kv=1 splitting to
GGML_BACKEND_SPLIT_AXIS_0), not a vision/mmproj property: it reproduces without
--mmproj and even single-GPU tensor. So drop the vision/mmproj scoping, rename
_vision_tensor_* -> _tensor_split_*, and key the session cache on
(binary, mtime, model) rather than (binary, mtime) so one model's abort no
longer skips tensor for every other model on the same build.

Regression tests updated to pin the early-spawn record, the per-model cache,
and that an unrelated ggml-backend-meta assert is not treated as the marker.

* Studio: reload on explicit tensor-off after a multi-GPU layer fallback

When a tensor load is downgraded to layer but kept multi-GPU to honor the
tensor request (preserve_multi_gpu_on_layer, the geometry-cache gate, or the
budget downgrade), the server reports tensor_parallel=False with --split-mode
layer stored. A later Apply that explicitly turns the tensor toggle off then
matched the loaded state and deduped to already_loaded, so Studio kept the
fallback's all-GPU CUDA_VISIBLE_DEVICES placement instead of re-selecting
normal placement (a single GPU for a model that fits on one card).

Latch a _layer_preserves_tensor_intent flag in load_model whenever a tensor
request is downgraded to layer with the multi-GPU floor raised
(_layer_min_gpus > 1), clear it when tensor stays on or on unload, and force a
reload in _request_matches_loaded_settings when the user explicitly turns the
tensor toggle off while that flag is set. An Apply that does not touch the
toggle still dedupes, so a working multi-GPU layer server is not churned.

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

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

* Studio: address reviewer.py findings on the tensor-split self-heal

P1 (dedup): tensor intent can be dropped via extras, not only the toggle. An
explicit llama_extra_args=["--split-mode", "layer"] matches the stored fallback
extras, so _request_matches_loaded_settings deduped to the preserved all-GPU
placement instead of reloading. Now reload when layer_preserves_tensor_intent
and the user explicitly drops tensor via the toggle OR via extras
(_effective_tensor_parallel of the explicit extras is false).

P1 (downgrade symmetry): the len(tp_gpus) < 2 compute-buffer downgrade cleared
tensor_parallel without raising _layer_min_gpus, unlike the budget and geometry
downgrades. GPUs below tensor's replicated compute-buffer reserve can still take
layer split's lower overhead, so keep the multi-GPU request (len(gpus) >= 2) and
let _select_gpus cap unusable cards.

P2 (cache key): key the tensor-split abort cache on st_mtime_ns, so a binary
replaced in place within the same second after an abort is re-probed instead of
inheriting the stale entry.

P2 (test hygiene): load routes/inference.py via importlib in the regression
tests instead of importing the routes package, which runs routes/__init__.py and
pulls in every router (e.g. python-multipart). Added regression coverage for the
extras-off reload, the compute-buffer multi-GPU preservation, and the same-second
nanosecond cache invalidation.

* Studio: record the tensor-split abort on the Windows CRT abort exit too

The first-spawn split-axis latch only recorded when _is_signal_crash matched
(POSIX signal or 0xC0000000+ NTSTATUS). On MSVC builds GGML_ASSERT terminates
through the CRT abort() path with exit code 3, which is neither, so the cache
never filled on Windows and every later load of the same bad binary/model
repeated the tensor crash before falling back to layer.

The split-axis marker is definitive, so accept either a signal crash or the
Windows abort() exit (3) when the marker is present. Add _is_abort_exit and a
unit test, and assert the early latch honors it.

* Studio: fix UnboundLocalError on --fit-on fallback, reload backend fast path

Two follow-ups from review on the tensor-split self-heal:

UnboundLocalError: _layer_min_gpus was initialized inside the GPU-selection try.
If NVML probing or GGUF/mmproj sizing raised, the except path logged "using
--fit on" and fell through to the command builder, where the new
self._layer_preserves_tensor_intent = _layer_min_gpus > 1 then raised, turning a
safe --fit-on layer fallback into a hard load failure. Bind _layer_min_gpus
before the try so the except path always has it.

Backend fast path: _request_matches_loaded_settings forces a reload when a
preserved tensor->layer fallback gets an explicit tensor-off request, but
load_model's own _already_in_target_state still matched the tensor-off/layer
settings and short-circuited, so the placement re-selection never ran. Mirror
the guard there: reload when layer_preserves_tensor_intent and the request drops
tensor intent. The flag clears on that reload, so there's no loop.

Added regression coverage for both.

* Studio: testable tensor-split record decision; skip futile fit-off retry

Follow-ups from a deeper review of the tensor-split self-heal:

Extract the record decision into _should_record_tensor_split_abort(rc, output)
(marker AND (signal crash OR Windows abort)) and call it from the early latch.
The combined boolean was only covered by source-inspection substring checks, so
an or->and typo would silently stop recording on Windows (CRT abort exit 3 is
not a signal) with every test still green. Add a behavioral test over the
POSIX / Windows / NTSTATUS / clean-exit / SIGKILL / no-marker matrix.

Skip the --fit off retry inside _spawn_and_wait when the crash already shows the
split-axis marker: that abort is fit-independent, so the retry just warms up and
crashes a second time before the latch records it. Skipping it lets the caller
latch immediately and corrects the latch comment.

Also clarify the dedup-guard comments (toggle read from model_fields_set vs
extras via _effective_tensor_parallel without env; the backend fast path is
intentionally broader and only ever forces a reload).

* Studio: don't reload-loop tensor-off requests under env tensor

The preserved-fallback reload guard fired on the raw tensor toggle, ignoring
LLAMA_ARG_SPLIT_MODE=tensor. For an env-driven tensor user, an explicit
tensor_parallel=false request then forced a reload that re-engaged tensor via
the env and re-created the same preserved layer fallback, so every /load
reloaded -- bypassing the env-downgrade matching that exists to avoid exactly
this loop.

Gate the guard on the env-aware effective tensor state: reload only when an
explicit toggle/extras change leaves _effective_tensor_parallel (which consults
the env) off. If the env still forces tensor, fall through to the existing
env-downgrade match, which dedupes instead of looping. Added a regression test
with LLAMA_ARG_SPLIT_MODE=tensor set.

* Studio: tighten comments and test docstrings on the TP self-heal

Condense the verbose comments and test docstrings added across the review rounds
into fewer, succinct lines without changing their intent: the early-latch and
downgrade-site rationale, the cache/key and helper docstrings, the dedup-guard
comments, and the per-test docstrings. No code changes (AST-verified comments
and docstrings only); tests and lint unchanged.

* Studio: clear preserved tensor flag on diffusion; carry it across non-drop reloads

Two follow-ups on the preserved-fallback machinery:

Diffusion: the DiffusionGemma path early-returns from load_model before the
command builder that sets/clears _layer_preserves_tensor_intent, so the flag
from a prior tensor->layer fallback leaked onto a later diffusion load and
forced needless reloads of the diffusion server on tensor-off/extra Applies.
Clear it when starting diffusion.

Settings reload: the preserve hint was recomputed only from the new request, so
a reload for an unrelated setting (e.g. max_seq_length) with the tensor toggle
omitted dropped a preserved multi-GPU layer placement back to one GPU. Carry
llama_backend.layer_preserves_tensor_intent into the hint when the request is
not an explicit tensor-off/extras-off drop, so a fitting model stays multi-GPU.

Added regression tests for the diffusion clear, the carry-forward, and the
updated tensor-intent computation.

* Studio: gate the preserve carry-forward on the same model being loaded

The tensor-intent carry-forward read llama_backend.layer_preserves_tensor_intent
without checking it belonged to the model being loaded. On a direct model switch
(load B without an explicit /unload of A), the flag is still set from A's
downgrade (it isn't reset until B's load_model reaches the command builder, after
the route reads it), so a plain load of B got preserve_multi_gpu_on_layer=True
and was spread across all GPUs even though it fits on one and the user never
requested tensor for it. The backend dedup doesn't have this leak (it checks
model_identifier first); the leak was only in the route hint.

Extract the decision into _carry_preserved_tensor_intent(preserved, same_model,
explicit_drop) and gate it on the backend still holding the same model. Add a
behavioral truth-table test (catches a `not` inversion and a missing same-model
guard) and tighten the compute-buffer downgrade test to bound its source window.

* Studio: match the HF quant too when carrying preserved tensor intent

The same-model guard on the preserve carry-forward compared only model_identifier,
which is variant-agnostic for HF repos. A later load of the same repo with a
different gguf_variant (which already bypassed dedupe on the variant mismatch)
was treated as the same model, so a request that omits tensor settings inherited
the prior variant's preserved intent and forced multi-GPU layer placement for a
quant that never requested tensor. Also require the loaded hf_variant to match for
HF repos (local direct-file loads already differ by model_identifier path). Added
a regression test for the variant guard.

* Studio: match the loaded GGUF by path too when carrying preserved tensor intent

A local directory holding multiple GGUF variants keeps one variant-agnostic
model_identifier (the directory) while config.gguf_file selects the file, so the
same-model guard let variant B inherit variant A's preserved tensor->layer
fallback and forced B onto multi-GPU. Mirror _already_in_target_state's identity
logic: match by resolved path when both sides have a local file, else by HF
variant. #6659

* Studio: let implicit same-settings reloads dedupe after a preserved fallback

The backend _already_in_target_state mirror forced a reload on ANY effective
tensor-off request once a tensor->layer fallback was preserved. In the HF
auto-pick / local-directory flows the route-level dedup is skipped, so an
identical /load with tensor omitted reached this guard and reloaded every time
even without an explicit drop. Thread the route's preserve_multi_gpu_on_layer
decision in so only an explicit drop reloads; implicit carry-forward dedupes. #6659

* Studio: only an explicit tensor/split-mode change drops preserved intent

The explicit-drop test treated request.llama_extra_args is not None as a drop,
so a same-model reload that merely added an unrelated pass-through arg (e.g.
--top-k 20) without touching the tensor field or --split-mode disabled the
carry-forward and collapsed a fitting model back to one GPU. A drop now requires
an explicit tensor_parallel field change or a non-tensor --split-mode override,
via a shared _is_explicit_tensor_drop helper used by both the already-loaded
dedup and the load carry-forward so the two readers agree. #6659

* Studio: treat an explicit clear of extras as a tensor drop

When tensor intent was extras-driven (--split-mode tensor) and fell back to a
preserved layer split, a later request that explicitly clears extras
(llama_extra_args=[]) but omits tensor_parallel left the empty list with no
split-mode override, so the carry-forward kept the model pinned multi-GPU instead
of returning to normal layer selection. _is_explicit_tensor_drop now also counts
an explicit empty-list clear as a drop, while an unrelated extra (--top-k) or
inherit (None) still carries the preserved intent. #6659

* Studio: don't treat the UI's tensor_parallel echo as a tensor drop

The Studio frontend always sends tensor_parallel and copies the /load response's
resolved value back into its state, so after a tensor->layer fallback every
ctx/settings reload carries tensor_parallel=false even though the user never
changed it. Keying the drop on the field (or on an empty extras clear) collapsed
the preserved multi-GPU placement on the next reload. A fallback also always
stores --split-mode layer, never a tensor split mode, so a clear never wipes
tensor intent. _is_explicit_tensor_drop now drops only on an explicit non-tensor
--split-mode override; the bare field echo, an empty clear, an unrelated extra,
and inherit all keep the preserved placement, and --split-mode tensor /
tensor_parallel=true re-engage tensor. #6659

* Studio: match the resolved config.identifier when carrying tensor intent

The same-model guard for the carry-forward compared the raw request id, but
ModelConfig.from_identifier normalizes it (adds the unsloth/ prefix for a
shorthand, fixes repo-id case) before load_model stores config.identifier. So a
ctx/settings reload using the shorthand id missed the match, dropped
_carry_preserved_tensor_intent, and could collapse a preserved multi-GPU layer
placement to one GPU. Compare against config.identifier (what the backend stores),
keeping it symmetric with _already_in_target_state. #6659

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-27 01:52:18 -07:00
oobabooga
0607351a81 Merge remote-tracking branch 'origin/main' into image-generation
# Conflicts:
#	studio/backend/routes/models.py
2026-06-26 23:23:02 -03:00
oobabooga
9c6f852a43 Fix image-generation GPU eviction, load lifecycle, and Images picker gating 2026-06-26 23:18:36 -03:00
Daniel Han
b11966b2db
studio: list the full local model catalog from /v1/models (#6519)
---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-06-26 20:42:06 -03:00
Daniel Han
9451aef51e
studio: return a clean model id from the OpenAI API instead of the local .gguf path (#6518)
---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-06-26 16:07:53 -03:00
Lee Jackson
1d4c4d4a76
Merge branch 'main' into image-generation 2026-06-26 18:22:14 +01:00
Daniel Han
2ef394137a
Studio: harden background consumer loops and streaming paths against silent UI freezes (#6653)
* Studio: harden the data-recipe and inference consumer loops against pump death

Follow-up to #6643. The same single-unsupervised-consumer pattern the training
pump had lives in two sibling loops, with the same failure mode: one bad event
kills the only thread that updates the in-memory state every UI surface reads,
while the worker subprocess keeps running.

- data_recipe JobManager._pump_loop: a malformed worker log line that makes
  parse_log_message raise no longer kills the pump. Guard _handle_event, the
  queue read, and the worker-exit finalize, and broaden _drain_queue so a drain
  error still finalizes the job instead of leaving it wedged "active" (which also
  leaked the workflow-scoped API key until its 24h expiry).
- inference InferenceOrchestrator._dispatcher_loop: guard the routing body so a
  malformed response or a mailbox put error can't kill the dispatcher and hang
  every in-flight generation (callers key liveness on the subprocess, not on
  this thread).

Adds regression tests for both.

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

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

* Studio: extend consumer-loop hardening to RAG, hub, auth, and stream-reader paths

Continuation of the data-recipe and inference pump hardening: the same
"background producer updates in-memory state that a single unsupervised
consumer surfaces to the UI" pattern shows up in several more Studio paths,
each able to silently freeze a UI surface while the worker keeps running.

RAG ingestion SSE (core/rag/ingestion.py):
- job_events polled the queue with a blocking get and never noticed client
  disconnect or a dead worker, so a closed tab or a producer that died
  without emitting a terminal event left the stream hanging. It now polls
  with a timeout, emits heartbeats, ends on terminal job status, caps idle
  time, and always pops the job registry in finally.
- Added _reap_finished_jobs() and call it from start_ingestion so finished
  job state does not accumulate.

Startup reconcile (storage/rag_db.py, main.py):
- reconcile_orphaned_ingestion_jobs() marks ingestion jobs (and their
  documents) that were left non-terminal by a previous crash as failed, so
  the UI does not show jobs stuck "running" forever after a restart. Wired
  in at startup next to cleanup_orphaned_runs().

Hub download watcher (hub/services/download_lifecycle.py):
- _watch() could leave a job pinned "running" if finalize raised. Body is
  now guarded: on failure it logs and sets the job to error, and always
  invalidates the hf cache scan in finally.

External provider stream (core/inference/external_provider.py):
- read timeout was None (no stall ceiling); set to 300s so a wedged
  upstream surfaces as an error instead of an indefinitely hung stream.

Auth store (auth/storage.py):
- Enable WAL + busy_timeout on the auth DB so token validation (read on
  every request) and login writes stop serialising on the rollback journal.
  Matches studio_db / rag_db / providers_db.

Login rate limiter (routes/auth.py):
- _LOGIN_IP_BUCKETS could grow unbounded under spoofed-IP traffic; cap it
  and prune stale buckets, mirroring the per-account bucket handling.

Training progress SSE (routes/training.py):
- Break promptly on client disconnect instead of waiting for the next
  yield to fail on a closed socket, matching the export / data-recipe SSE
  routes.

llama-server stdout drain (core/inference/llama_cpp.py):
- Broaden the drain guard so an unexpected decode/read error logs at debug
  and stops the drainer cleanly instead of escaping the thread.

Frontend stream readers (chat-api.ts, rag-api.ts):
- Wrap the SSE read loops in try/finally + reader.cancel() so early return
  ([DONE]), thrown errors, and consumer aborts release the reader lock
  instead of holding it until GC.

Tests:
- test_training_progress_stream_nan: fake request now implements the async
  is_disconnected() the route polls, matching the other SSE route fakes.

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

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

* Studio: address Codex review feedback on the consumer-loop hardening

Four follow-ups from the automated review, all on code this PR introduced:

- Data-recipe pump (manager.py): a queue read that keeps raising an error
  outside the read's narrow catch set (e.g. a broken queue pipe after the
  child died) hit the `continue` guard and skipped the dead-worker finalize
  below, spinning forever and leaving the job wedged "active" with its
  workflow key unretired. On a read failure, fall through to finalize when
  the worker is no longer alive. Added a regression test.

- RAG ingestion SSE (ingestion.py): the 5-minute idle cap could end the
  stream while the job was still pending/running (a large document spends
  minutes in embedding/storing with no per-batch progress event). The route
  then sends [DONE], and the client treats a no-terminal-frame end as
  completion, marking the document indexed mid-ingestion. Drop the idle cap:
  while the worker is alive and non-terminal we keep heartbeating; the stream
  ends only on terminal DB status, the None sentinel, or client disconnect.

- Login rate limiter (auth.py): the per-IP path pruned but then added the
  new IP unconditionally, so a spoofed-source-IP spray kept _LOGIN_IP_BUCKETS
  unbounded and made every new IP pay a full-dict prune scan. Gate the add on
  the cap, mirroring the account path.

- Hub download watcher (download_lifecycle.py): if finalize raised before it
  reaped (proc.wait) and dropped the worker (e.g. an I/O error draining
  stderr), the crash path published a terminal state while the live Popen
  stayed registered and kept writing the cache, and the terminal set_job let
  claim() admit a retry on the same repo. Terminate + drop the worker before
  setting the terminal state.

* Studio: keep login throttling working when the per-IP bucket dict saturates

Review follow-up. The previous cap fix skipped creating a bucket for a new IP
once _LOGIN_IP_BUCKETS was full, returning ip_fails=0. Under a sustained spray
that also fills the account dict, every failure from such an IP then looked
first-seen and _login_blocked had no bucket to enforce, so the cap effectively
disabled throttling once saturated.

Bound the dict with a FIFO eviction instead: if the IP is new and the dict is
full, reclaim expired buckets (rate-limited so a burst of distinct IPs can't
make each failure an O(n) sweep) and, if still full, evict the oldest-inserted
IP. The new IP always gets a real bucket, so a saturating (e.g. spoofed
X-Forwarded-For) spray stays throttled while memory stays bounded. Added a
regression test that saturates the dict and asserts a later IP is still blocked.

* Studio: address Codex review (RAG queue lifecycle, stream error, orphan chunks)

Three follow-ups on the Phase 6 changes:

- RAG ingestion SSE (ingestion.py): job_events removed the per-job queue in its
  finally on ANY exit, including an early client disconnect while the worker is
  still running. That dropped the worker's later events (the queue is the only
  one _emit writes to) and made a reconnect find no queue and receive only
  [DONE], which the client treats as completion. Only drop the queue on a
  terminal exit (None sentinel / terminal DB status); leftover terminal queues
  are still swept by _reap_finished_jobs. Added queue-lifecycle tests.

- External provider stream (routes/inference.py): once the 300s read timeout can
  fire, the stream's except path failed the monitor but ended without an error
  frame or [DONE], so the chat client saw a bare EOF and saved the timed-out
  answer as a successful partial with no error. Emit an SSE error frame (and
  [DONE]) on stream failure so the client surfaces it.

- RAG startup reconcile (storage/rag_db.py): marking a half-ingested document
  failed left its chunks/fts/vec rows intact, and retrieval filters by scope not
  status, so a failed document could still be retrieved and cited. Purge the
  document's chunks when reconciling it to failed (the doc row stays for
  re-ingest).

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

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

* Studio: release the remaining SSE stream readers (training, data-recipe, export)

reviewer.py follow-up. The chat and RAG SSE readers were wrapped in
try/finally + reader.cancel(), but the other three readers built on the same
response.body.getReader() pattern were left without it: streamTrainingProgress,
streamRecipeJobEvents, and streamExportLogs leak the ReadableStreamDefaultReader
lock (held until GC) when the consumer aborts, returns early, or a parse/callback
throws. Wrap each in try/finally + reader.cancel() (export already had a
try/catch, so it only needed the finally). All five frontend SSE readers now
release the reader symmetrically.

* Tighten resilience comments and docstrings

Condense the verbose explanatory comments and internal-helper docstrings added
in this branch to shorter, clearer forms. Comment/whitespace only; verified no
code changed via AST diff. No behaviour change.

* Studio: keep chunks for completed docs during ingestion reconcile

Startup reconciliation flips orphaned (non-terminal) ingestion jobs to failed and
purges the document's chunks so a failed source can't be retrieved. But it dropped
the chunks unconditionally, so a document the worker had already committed as
'completed' before the crash (only its job row left non-terminal) lost every chunk
while still reporting 'completed'. That leaves an empty source that retrieval can't
return and dedup (status != 'failed') blocks from re-ingest.

Only purge chunks when the document UPDATE actually transitions it to failed; an
already-completed document keeps its chunks. Adds reconcile regression tests for
both the completed-doc and genuine in-flight-orphan cases.

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

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

* Studio: drop a finished RAG job's queue when the client disconnects

job_events kept the per-job queue until it consumed the None sentinel, so a UI
that stops on the terminal event (its reader.cancel aborts the stream before
[DONE]) left the queue registered until the next _reap_finished_jobs sweep; a
batch of uploads followed by idling retained them all.

_run writes the terminal DB status before emitting the terminal event, so on
generator exit, drop the queue when the job's DB row is already terminal (worker
done, nothing to resume) and keep it only while the worker is still running. Adds
a disconnect-after-terminal-event regression test.

* Remove stray async task output files committed by mistake

* Studio: harden login IP throttle and end progress stream on disconnect

Two Codex review items:

Login per-IP throttle: when the per-IP bucket dict saturated, FIFO eviction could
drop a still-hot (blocked) bucket, so an IP could flood the dict with distinct
(or spoofed) source IPs to push out its own bucket and retry as first-seen. Stop
evicting hot buckets; a new IP that can't fit now shares a bounded overflow
counter that still trips the per-IP threshold, so a saturating spray stays
throttled and no live counter is reset.

Progress SSE: on client disconnect the polling loop only broke and fell through
to the unconditional final 'complete' frame, so a buffered or proxying consumer
could read a still-active run as completed. Return from the generator instead.

Adds regression tests for both (spray cannot reset a hot bucket; disconnect while
active emits no complete frame).

* Studio: shard the login overflow counter and stop cancelling chat stream after [DONE]

Two Codex review items:

Login throttle overflow: the single shared overflow counter meant that once a
saturating spray pushed it past the per-IP threshold, _login_blocked returned 429
for every new unbucketed source IP, before credentials were checked -- a global
login denial. Shard the overflow into a fixed array of counters keyed by hash(ip),
so a hot shard only throttles the IPs that map to it while a single source's
repeated failures still concentrate in one shard and stay throttled. Memory stays
bounded and no live bucket is evicted. Adds a regression test that a hot overflow
shard does not block an unrelated IP.

Chat stream: the reader.cancel() in the SSE finally fired even after a natural
[DONE]/EOF. The backend finalizes its api-monitor entry right after yielding the
sentinel (the local pass-through finishes after the last yield), so a client
cancel there can be observed as a disconnect and mark a completed request as
cancelled. Track natural completion and only cancel on an early/abnormal exit.
(No frontend unit test: the Studio frontend has no test harness.)

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

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

* Studio: give prep-timeout test fakes an is_disconnected method

The progress stream now ends on client disconnect (await request.is_disconnected()
before falling through to the terminal frame). After merging that into the
prep-timeout tests added later on main, their _FakeRequest/_ReconnectRequest must
provide is_disconnected or the generator raises AttributeError under CI.

* Studio: keep the login overflow throttle when bucket capacity frees up

_login_blocked only consulted the per-IP overflow shard while the bucket dict was
still at capacity. If a slot freed before the 60s window expired (e.g. another
IP's successful login calls _clear_login_bucket), a source counted in a hot shard
stopped being blocked and its next failure got a fresh per-IP bucket, resetting
the throttle the overflow path exists to preserve. Always max in the IP's shard
(shards are empty outside saturation, so it is a no-op in the common case). Adds a
regression test that a hot source stays throttled after a bucket frees.

* Studio: clear a login IP's overflow throttle on successful login

_clear_login_bucket reset the per-IP and per-account buckets on a successful
login but not the overflow shard, so after the dict saturated and an IP was
counted in overflow, a later successful login left those entries behind and the
next failed attempt could immediately return 429.

Store overflow entries as (timestamp, ip) so a source is throttled by its own
count within the shard (also removing cross-IP collateral within a shard), and
drop just that IP's entries in _clear_login_bucket. Adds a regression test that a
successful login clears the overflow throttle.

* Studio: bound the login overflow shard memory under high-cardinality spray

The per-IP overflow tracked failures in a time-pruned deque of (timestamp, ip)
tuples, so a spoofed-X-Forwarded-For spray of distinct one-off IPs grew memory and
the per-check scan with request cardinality for the whole window -- undermining
the bucket cap that exists to bound memory. Replace each shard with a fixed-
capacity dict (ip -> [count, window_start]): O(1) lookups, and when a shard is
full a one-off IP evicts the lowest-count entry (Space-Saving) so memory is hard-
bounded while a persistent attacker keeps a high count and is never evicted. Adds
a regression test that shards stay within the per-shard cap under a 5000-IP spray.

* Studio: purge chunks for already-failed docs during ingestion reconcile

The reconcile chunk-purge was gated on the documents UPDATE actually flipping a
non-terminal doc to failed. A doc the worker had already marked 'failed' before
the crash (job row left non-terminal) was not re-flipped, so its committed chunks
were kept and stayed retrievable/citable, since retrieval filters by scope not
status. Purge chunks whenever the document is not 'completed' (failed, in-flight,
or gone), preserving the completed-doc carve-out. Adds a regression test.

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

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

* Studio: don't inherit an evicted IP's count onto a new overflow source

When a full overflow shard evicted the lowest-count entry, the new source
inherited that count (Space-Saving base + 1). If a shard was saturated with hot
entries, an unrelated new IP could land at/over the threshold and be 429'd after a
single attempt -- cross-IP collateral despite the per-source-isolation intent.
New entries now start clean at count 1; the only cost is that a heavy hitter that
is the lowest-count entry in a fully saturated shard can briefly reset, which is
preferable to blocking a bystander. Adds a regression test.

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

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

* Studio: carry overflow failures into a new IP bucket on transition

_login_blocked took max(per-IP bucket, overflow shard) rather than combining them,
so a source could log (threshold-1) failures in overflow during saturation and,
once a bucket slot freed, another (threshold-1) in a fresh bucket within the same
window -- roughly doubling the per-IP limit. When a saturated-era IP first gets a
real bucket, migrate its windowed overflow count into that bucket (and drop the
overflow entry) so the combined failures throttle at the intended limit. Adds a
regression test.

* Studio: reconcile a completed doc's orphaned job to completed, not failed

When a crash left an ingestion job non-terminal after its document was already
committed as completed, reconcile marked the job failed. After restart the upload
UI has no in-memory SSE queue and falls back to getJob(), which treats a failed
job as an indexing failure and removes/toasts a document that is actually
searchable. Mark the job completed (keeping its chunks) when its document is
completed. Extends the completed-doc reconcile test to assert the job status.

* Studio: clamp the overflow failure count migrated into a login bucket

A saturated source could accrue an unbounded overflow count, then materialize
one deque entry per recorded failure when a bucket slot freed, allocating an
arbitrarily large deque under the login lock. Only at-or-above the per-IP
threshold matters for blocking, so cap the count there at the record and take
sites; the migration is now bounded without weakening the limit.

* Studio: keep the RAG job stream alive on a transient status read

The heartbeat poll read the job row unguarded; a momentarily-locked DB would
raise out of job_events, which the SSE route turns into a terminal error frame,
and the UI drops a document whose worker is still running. Treat a failed status
read as non-terminal: heartbeat and retry, and keep the queue so a reconnect can
resume.

* Studio: set busy_timeout before journal_mode on the auth DB

Switching journal_mode needs a lock, so if a refresh-token write already holds
one, journal_mode=WAL raises SQLITE_BUSY and the shared try leaves the
connection on SQLite's default zero lock wait. Set busy_timeout first so the
switch waits instead of failing.

* [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-26 03:31:33 -07:00
oobabooga
a0472d2bff Match edit keywords by id segment and sanitize generate-route errors 2026-06-26 01:44:59 -03:00
oobabooga
7699ce763d Guard superseded loads, redact load errors, drop chat models from the image picker, clamp runs 2026-06-25 12:59:34 -03: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