ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio diffusion: persistent sd-server for the native engine (load once, serve many) (#6768)

Browse source 
* 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 7): accuracy-preserving speed pass

Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy
bug and a dead-on-arrival speed path; this fixes both and adds the lossless /
near-lossless wins, all measured on a B200.

Correctness:
- TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32
  process-wide and never restored them, so a later `off` load silently inherited
  TF32 and was no longer bit-identical. Added snapshot_backend_flags /
  restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer
  runs and restored on unload. Verified: load max -> unload -> load off is now
  byte-identical (PSNR inf) to a fresh off.
- sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and
  only checked the timeout after EOF, so a child stuck in model load / GPU init
  with no output ignored the timeout. Drained stdout on a reader thread with a
  wall-clock deadline. Added a silent-hang regression test.

Speed (diffusers path), near-lossless, opt-in tiers:
- Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and
  Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks
  compiles and runs ~2.2x faster on the GGUF Z-Image transformer on
  torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the
  block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB
  vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move
  output quality. Gate relaxed; default tier delivers it.
- cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs).
- torch.inference_mode() around the pipeline call (lossless, strictly faster than
  the no_grad diffusers uses internally).

Memory path:
- VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers;
  the balanced (group) tier keeps exact slicing only, so it is now bit-identical to
  the resident image (verified PSNR inf) and slightly faster.
- Group offload adds non_blocking + record_stream on the CUDA stream path to
  overlap each block's H2D copy with compute (lossless; gated on the installed
  diffusers signature so older versions still work).

Native (sd.cpp) path:
- native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a
  near-lossless CUDA win that was previously only added on offload tiers; max also
  -> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so
  it is never auto-on. Engine generate() merges it, de-duped against offload flags.

Default profile: a GGUF model with no explicit speed_mode now resolves to the
`default` profile (resolve_speed_mode), since compile's perturbation sits below the
quantisation noise floor and so does not reduce quality versus the dense reference;
out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models
stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is
always honored, so the byte-identical path remains one flag away and is the
regression reference.

Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/
perf_verify.py (the B200 verification above), and diffusion_bench.py gains
--speed-mode so the speed tiers are benchmarkable.

Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore,
GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags +
the engine de-dup, and the sd-cli silent-hang timeout.

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

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

* Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks

The opt-in `max` speed tier now compiles the repeated block with
mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode:
Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer
cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are
deliberately avoided -- both crash on the regionally-compiled block (its static
output buffer is overwritten across denoise steps), measured.

Adds two reproducible benchmarks used to validate the optimization research:
- scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head.
- scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes.

Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen;
coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune);
FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo).

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

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

* Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source)

Add an opt-in transformer_quant mode that loads the dense bf16 transformer and
torchao-quantises it onto the low-precision tensor cores, instead of the GGUF
transformer (which dequantises to bf16 per matmul and so runs at bf16 rate). On a
B200 (Z-Image-Turbo, 1024px/8 steps): auto picks fp8 at 0.614s vs GGUF+compile's
0.823s (1.34x), int8 0.626s (1.32x), both at lower LPIPS than GGUF's own 4-bit floor.

GGUF+compile stays the low-memory default and the fallback. The mode is gated on
CUDA + bf16 + resident VRAM headroom (the dense load peaks ~21GB vs GGUF's 13GB);
any unsupported arch/scheme, OOM, or quant failure falls back to GGUF with a logged
reason. auto picks the best scheme per GPU via a real quantise+matmul smoke probe
(Blackwell nvfp4/fp8/mxfp8, Ada/Hopper fp8, Ampere int8); a min-features filter skips
the tiny projections that crash int8's torch._int_mm. New module mirrors
diffusion_precision.py; quant runs before compile before placement.

184 -> tests pass; new test_diffusion_transformer_quant.py plus backend/route
coverage. scripts/diffusion_bench.py gains --transformer-quant; scripts/quant_probe.py
is the standalone torchao lever probe.

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

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

* Studio diffusion (Phase 8): consumer-GPU tuning - lock fp8 fast accumulate, prefer fp8 over mxfp8, reject 2:4 sparsity

Consumer Blackwell halves tensor-core throughput on FP32 accumulate (fp8 419 vs 838
TFLOPS with FP16 accumulate; bf16 209), so:
- fp8 config locks use_fast_accum=True (Float8MMConfig). torchao already defaults it on;
  pinning it guards consumer cards against a default change. On B200 it is identical
  speed and slightly better quality (LPIPS 0.050 vs 0.091).
- the Blackwell auto ladder prefers fp8 over mxfp8 (measured faster + more accurate).

2:4 semi-structured sparsity evaluated and rejected (scripts/sparse_accum_probe.py):
2:4 magnitude-prune + fp8 gives LPIPS 0.858 (broken image) with no fine-tune, the
cuSPARSELt kernel errors on torch 2.9, and it does not compose with torch.compile
(our main ~2x). Documented as a dead end, not shipped.

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

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

* Studio diffusion (Phase 8): add fp8 fast-accum overflow verification probe

scripts/fp8_overflow_check.py hooks every quantised linear during a real Z-Image
generation and reports max-abs + non-finite counts for use_fast_accum True vs False.
Confirms fast accumulation is an accumulation-precision knob, not an overflow one:
across 276 linears, including Z-Image's ~1.0e6 activation peaks (which overflow FP16),
0 non-finite elements and identical max-abs for both modes.

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

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

* Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override

Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want
fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and
prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match
on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40;
GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it.

Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s),
no overflow, quality below the quant noise floor. So the default leans to accuracy on
data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force)
lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off).

187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override
threading).

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

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

* Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9

scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer.
Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with
use_triton_kernel=False (the default triton path needs the missing MSLK library), but
only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than
using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's
CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto
on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast.

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

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

* Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder

Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4
tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS,
2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the
diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both
slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166
vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is
correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add
scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end).

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

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

* Studio diffusion (Phase 9): pre-quantized transformer loading

The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on
the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs
13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline
(scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton
on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True)
the quantized weights, so the dense bf16 never touches the GPU.

Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching
GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same
torchao config + min_features filter the runtime path uses, applied ahead of time.

New core/inference/diffusion_prequant.py (resolve_prequant_source +
load_prequantized_transformer, best-effort, lazy imports). diffusion.py
_load_dense_quant_pipeline tries the pre-quant source first and falls back to the
dense materialise+quantise path, then to GGUF, so the default is unchanged.
DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty
prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for
the resolver, the meta-init+assign loader, and the backend branch selection +
fallbacks; GPU verification via scripts/verify_prequant_backend.py.

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

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

* Studio diffusion (Phase 10): attention-backend selection

Add a selectable attention kernel via the diffusers set_attention_backend
dispatcher. Attention is memory-bandwidth bound, so a better kernel is an
end-to-end win orthogonal to the linear-weight quantisation (it speeds the QK/PV
matmuls torchao never touches) and composes with torch.compile.

auto picks the best exact backend for the device: cuDNN fused attention
(_native_cudnn) on NVIDIA when a speed profile is active, measured ~1.18x
end-to-end on a B200 (Z-Image 1024px/8 steps) with LPIPS ~0.004 vs the default
(below the compile/quant noise floor); native SDPA elsewhere and when speed=off
(so off stays bit-identical). Explicit native/cudnn/flash/flash3/flash4/sage/
xformers/aiter are honored, and an unavailable kernel falls back to the default
rather than failing the load.

New core/inference/diffusion_attention.py (normalize + per-device select + apply,
best-effort, lazy imports). Set on pipe.transformer BEFORE compile in load_pipeline;
attention_backend threads through begin_load / load_pipeline / status like the other
load knobs. New request field attention_backend + status field. Hermetic CPU tests
for normalize / select policy / apply fallback, plus route threading + 422. Measured
via scripts/perf_levers_probe.py.

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

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

* Studio diffusion (Phase 11): prefer int8 on consumer GPUs in the auto ladder

Consumer / workstation GPUs halve fp8 (and fp16/bf16) FP32-accumulate tensor-core
throughput, while int8 runs at full rate (int32 accumulate is not nerfed). Public
benchmarks (SDNQ across RTX 3090/4090/5090, AMD, Intel) confirm int8 via torch._int_mm
is as fast or faster than fp8 on every consumer part, and the only path on pre-Ada
consumer cards without fp8 tensor cores. So when transformer_quant=auto, reorder the
arch tier to put int8 first on a consumer/workstation GPU (detected by the existing
_is_consumer_gpu name heuristic), while data-center HBM parts keep fp8 first.

Pure ladder reorder via _prefer_consumer_scheme; no new flags. Verified non-regression
on a B200 (still picks fp8). Hermetic tests for consumer Blackwell/Ada/workstation
(-> int8) and data-center Ada/Hopper/Blackwell (-> fp8).

* Studio diffusion (Phase 12): First-Block-Cache step caching for many-step DiT

Add opt-in step caching (First-Block-Cache) for the diffusion transformer. Across
denoise steps a DiT's output settles, so once the first block's residual barely
changes the remaining blocks are skipped and their cached output reused. diffusers
ships it natively (FirstBlockCacheConfig + transformer.enable_cache, with the
standalone apply_first_block_cache hook as a fallback).

Measured on Flux.1-dev (28 steps, 1024px): ~1.4x on top of torch.compile (2.83 ->
2.03s) at LPIPS ~0.08 vs the no-cache output, well inside the quality bar.

OFF by default and a per-load opt-in: the win scales with step count, so it is for
many-step models (Flux / Qwen-Image) and pointless for few-step distilled models
(e.g. Z-Image-Turbo at ~8 steps), where a single skipped step is a large fraction
of the trajectory. It composes with regional compile only with fullgraph=False (the
cache's per-step decision is a torch.compiler.disable graph break), which the speed
layer now switches to automatically when a cache is engaged. Best-effort: a model
whose block signature the hook does not recognise is caught and the load proceeds
uncached.

- new core/inference/diffusion_cache.py: normalize_transformer_cache + apply_step_cache
  (enable_cache / apply_first_block_cache fallback; threshold auto-raised for a
  quantised transformer per ParaAttention's fp8 guidance; lazy diffusers import).
- diffusion_speed.py: apply_speed_optims takes cache_active; compile drops fullgraph
  when a cache is engaged.
- diffusion.py: apply_step_cache before compile; thread transformer_cache /
  transformer_cache_threshold through begin_load -> load_pipeline and report the
  engaged mode in status().
- models/inference.py + routes/inference.py: transformer_cache (off | fbcache) and
  transformer_cache_threshold request fields, engaged mode in the status response.
- hermetic tests for normalisation, the enable_cache / hook-fallback paths, threshold
  selection, and best-effort failure handling, plus route threading + validation.
- scripts/fbcache_flux_probe.py: the Flux validation probe (latency / speedup / VRAM /
  LPIPS vs the compiled no-cache baseline).

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

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

* Studio diffusion (Phase 14): fix int8 dense quant on Flux / Qwen (skip M=1 modulation linears)

The opt-in dense int8 transformer path crashed on Flux.1 and Qwen-Image with
'torch._int_mm: self.size(0) needs to be greater than 16, but got 1'. int8 dynamic quant
goes through torch._int_mm, which requires the activation row count M > 16. A DiT's AdaLN
modulation projections (Flux norm1.linear 3072->18432, Qwen img_mod.1 / txt_mod.1, Flux.2
*_modulation.linear) and its timestep / guidance / pooled-text conditioning embedders are
computed once from the [batch, dim] conditioning vector (M = batch = 1), not per token, so
they hit _int_mm at M=1 and crash. Their feature dims are large, so the existing
min_features filter did not exclude them.

Fix: the int8 filter now also skips any Linear whose fully-qualified name matches a
modulation / conditioning-embedder token (norm, _mod, modulation, timestep_embed,
guidance_embed, time_text_embed, pooled). These layers run at M=1 once per block and are a
negligible share of the FLOPs, so int8 keeps the full speedup on the attention / FFN layers
(M = sequence length). fp8 / nvfp4 / mxfp8 use scaled_mm, which has no M>16 limit and
quantises these layers fine, so the exclusion is int8-only. Sequence embedders
(context_embedder / x_embedder / txt_in, M = seq) are deliberately not excluded -- note
'context_embedder' contains the substring 'text_embed', which is why the token is the
specific 'time_text_embed', not 'text_embed'.

Measured on a B200 (1024px, transformer_quant=int8 + speed=default), int8 now runs on every
supported model and is the fastest dense path on Flux/Qwen (int8 runs full-rate vs fp8's
FP32-accumulate): FLUX.1-dev 9.62s eager -> 1.98s (4.86x, vs fp8 2.15s), Qwen-Image -> 1.87s
(5.57x, vs fp8 2.09s), FLUX.1-schnell -> 0.41s (3.59x). Z-Image and Flux.2-klein (already
working) are unchanged.

- diffusion_transformer_quant.py: add _INT8_EXCLUDE_NAME_TOKENS; make_filter_fn takes
  exclude_name_tokens; quantize_transformer passes it for int8 only.
- hermetic test that the int8 filter excludes the modulation / embedder linears (and keeps
  attention / FFN / sequence-embedder linears), while fp8 keeps them.
- scripts/int8_linear_probe.py: the meta-device probe used to enumerate each transformer's
  Linear layers and derive the exclusion list.

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

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

* Studio diffusion (Phase 15): build int8 pre-quantized checkpoints (skip M=1 modulation linears)

The prequant-checkpoint builder applied the dense quant filter without the int8-only
M=1 modulation / conditioning-embedder exclusion the runtime path uses, so a built int8
checkpoint baked those projections as int8 and crashed (torch._int_mm needs M>16) at the
first denoise step on Flux / Qwen. Factor the scheme->exclusion decision into a shared
exclude_tokens_for_scheme() used by both the runtime quantise path and the offline builder
so they can never drift, and apply it in build_prequant_checkpoint.py. int8 prequant now
produces a working checkpoint on every supported model, giving int8 (the consumer-preferred
scheme) the same ~2x load-VRAM and download reduction fp8 already had.

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

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

* Studio diffusion (Phase 16): route no-GPU loads to the native sd.cpp engine

When no CUDA/ROCm/XPU GPU is available, route diffusion load/generate to the
native stable-diffusion.cpp engine instead of diffusers, with diffusers as the
guaranteed fallback. On CPU sd.cpp is 1.4-2.8x faster and uses 1.5-2.2x less RAM.

- diffusion_engine_router: centralised engine selection (built on the existing
  select_diffusion_engine), env opt-outs, MPS gating, recorded fallback reason.
- sd_cpp_backend (SdCppDiffusionBackend): the diffusers backend method surface
  backed by sd-cli, with lazy binary install, registry-driven asset fetch,
  step-progress parsing, and cancellation.
- diffusion_families: per-family single-file VAE + text-encoder asset mapping.
- sd_cpp_engine: cancellation support (process-group kill + SdCppCancelled).
- routes/inference + gpu_arbiter: drive the active engine via the router; the
  API now reports the active engine and any fallback reason.
- tests for the backend, router, route selection, and cancellation.

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

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

* Phase 16 review fixes: engine-switch unload, sd.cpp error mapping, per-image seeds, Qwen sampler

Address review feedback on #6724:
- engine router: unload the engine being deactivated on a switch, so the old
  model is not left resident-but-unreachable (the evictor only targets the active
  engine).
- generate route: sd.cpp execution errors (nonzero exit / timeout / missing
  output) now map to 500, not 409 (which only means not-loaded / cancelled).
- native batch: return per-image seeds and persist the actual seed for each image
  so every batch image is reproducible.
- Qwen-Image native path: apply --sampling-method euler --flow-shift 3 per the
  stable-diffusion.cpp docs; other families keep sd-cli defaults.
- honor speed_mode (native --diffusion-fa) and, off-CPU, memory_mode/cpu_offload
  offload flags on the native load instead of hardcoding them off.
- fail the load when the sd-cli binary is present but not runnable (version()
  now returns None on exec error / nonzero exit).
- size estimate: only treat the transformer asset as a possible local path.

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

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

* Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in

load_prequantized_transformer ends in torch.load(weights_only=False), which executes
arbitrary code from the pickle. The transformer_prequant_path load-request field reached
that unpickle for any local file an authenticated caller named, so a request could trigger
remote code execution. Refuse the source.kind=='path' branch unless the operator sets
UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted
and unaffected. Document the requirement on the API field and add gate tests.

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

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

* Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa

- apply_attention_backend now restores the native default when no backend is requested or a
  kernel fails. diffusers keeps a process-wide active attention backend that
  set_attention_backend updates, and a fresh transformer's processors follow it, so a load
  that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile
  load pinned, breaking the bit-identical/off guarantee.
- select_attention_backend drops flash3/flash4 up front when the CUDA capability is below
  Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit
  request on the wrong card set fine then crashed mid-generation; it now falls back to native.
- Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a
  valid alias of native) is accepted instead of 422-rejected by Pydantic.
- Drop the dead replace('-','_') normalization (no alias uses dashes/underscores).
- perf_levers_probe.py output dir is now relative to the script, not a hardcoded path.

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

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

* Studio diffusion (Phase 12): only engage FBCache on context-aware transformers; quantized threshold for GGUF

- apply_step_cache now engages only via the transformer's native enable_cache (the diffusers
  CacheMixin path), which exists exactly when the pipeline wraps the transformer call in a
  cache_context. The standalone apply_first_block_cache fallback installed on non-CacheMixin
  transformers too (e.g. Z-Image), whose pipeline opens no cache_context, so the load reported
  transformer_cache=fbcache and then the first generation crashed inside the hook. Such a model
  now runs uncached per the best-effort contract.
- GGUF transformers are quantized (the default Studio load path), so they now use the higher
  quantized FBCache threshold when the caller leaves it unset, instead of the dense default
  that could keep the cache from triggering.
- fbcache_flux_probe.py: compile cached runs with fullgraph=False (FBCache is a graph break, so
  fullgraph=True failed warmup and silently measured an eager cached run); output dir is now
  relative to the script, not a hardcoded path.

* Studio diffusion (Phase 11): keep professional RTX cards on the fp8 ladder

_is_consumer_gpu treated professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) as
consumer because their names carry no datacenter token, so the auto ladder moved int8 ahead
of fp8 and the fp8 path chose fast accumulate for them. The rest of the backend already
classifies these as datacenter/professional (llama_cpp.py _DATACENTER_GPU_RE), so detect the
same RTX PRO 6000 / RTX 6000 Ada markers here and keep fp8 first with precise accumulate.

Also fix the consumer-Blackwell test to use compute capability (10, 0) instead of (12, 0).

* Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config

Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just
TypeError on older torchao), which would break the mxfp8 config helper instead of falling
back to the default. Catch both so the fallback is robust.

quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds
CUDA memory during the per-row VRAM probe; output dir relative to the script.

* Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load

- snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform
  missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the
  whole snapshot. restore_backend_flags restores each flag independently so one failure can't
  leave the others leaked process-wide.
- load_pipeline restores the flags (and clears the GPU cache) when the build fails after
  apply_speed_optims mutated the process-wide flags but before _state captured them for unload
  to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and
  contaminated later off generations.

* [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 14): guard the int8 exclusion filter against a None fqn

The filter callback can be invoked without a module name, so fqn.lower() would raise
AttributeError on None. Fall back to an empty name (nothing matches the exclusion tokens,
so the linear is kept) instead of crashing the quantise pass.

* Studio diffusion (Phase 16) review fixes: native engine robustness

- sd_cpp_backend: stop truncating explicit seeds to 53 bits (mask to int64);
  a large requested seed was silently collapsed (2**53 -> 0) and distinct seeds
  aliased to the same image. Random seeds stay 53-bit (JS-safe).
- sd_cpp_backend: sanitize empty/whitespace hf_token to None so HfApi/hf_hub
  fall back to anonymous instead of failing auth on a blank token.
- sd_cpp_backend: a superseding load now cancels the in-flight generation, so the
  old sd-cli can no longer return/persist an image from the previous model.
- diffusion_engine_router: run the previous engine's unload() OUTSIDE the lock so a
  slow 10+ GB free / CUDA sync does not block engine selection.
- diffusion_engine_router: probe sd-cli runnability (version()) before committing to
  native, so a present-but-unrunnable binary falls back to diffusers at selection.
- diffusion_device: resolve a torch-free CPU target when torch is unavailable, so a
  CPU-only install can still reach the native sd.cpp engine instead of failing load.
- tests updated for the runnability probe + a not-runnable fallback case.

* Studio diffusion (Phase 9) review fixes: prequant safety + validation

- SECURITY: a request-supplied local pre-quant path is now unpickled only when it
  resolves inside an operator-configured ALLOWLIST of directories
  (UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in,
  once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on
  any path a load request named (arbitrary code execution). realpath() blocks symlink
  escapes; a bare on/off toggle is no longer a wildcard.
- Validate the checkpoint's min_features against the runtime Linear filter, so a
  checkpoint that quantised a different layer set is rejected instead of silently
  loading a model that mismatches the dense path while reporting the same scheme.
- Tolerant base_model_id compare (exact or same final path/repo segment), so a local
  path or fork of the canonical base is accepted instead of falling back to dense.
- _has_meta_tensors uses any(chain(...)) (no intermediate lists).
- prequant verify/probe scripts use repo-relative paths (+ env overrides), not the
  author's absolute /mnt paths.
- tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail.

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

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

* Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts

- diffusion_memory: when group offload is unavailable and the plan falls back to
  whole-module offload, enable VAE tiling (the group plan left it off, but the fallback
  is the low-VRAM path where the decode spike can OOM). Covers both the group and
  sequential fallback branches.
- perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a
  balanced bit-identity regression actually fails the check instead of exiting 0.
- compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and
  the load-progress poll has a 30 min deadline instead of looping forever on a hang.
- test for the group->model fallback enabling VAE tiling.

* Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path

- diffusion: a torchao-quantized transformer is committed only compiled. A dense model
  resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF
  it replaced), so when transformer_quant engaged and speed resolved to off, promote to
  default (regional compile); warn loudly if compile still does not engage.
- diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the
  CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK);
  otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF.
- nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path).
- test asserts the eager-quant -> default-compile promotion.

* Studio diffusion (Phase 10) review fixes: attention gating + probe isolation

- diffusion_attention: gate the auto cuDNN-attention upgrade on SM80+; on pre-Ampere
  NVIDIA (T4/V100) cuDNN fused SDPA is accepted at set time but fails at first generation,
  so auto now stays on native SDPA there.
- diffusion_attention: _active_attention_backend handles get_active_backend() returning an
  enum/None (not a tuple); the old  unpack always raised and was swallowed, so
  the native-restore short-circuit never fired.
- perf_levers_probe: free the resident pipe on a skipped (attn/fbcache) variant; run LPIPS
  on CPU so it isn't charged to every variant's peak VRAM; reset force_fuse_int_mm_with_mul
  so the inductor_flags variant doesn't leak into later compiled rows.
- tests for the SM80 cuDNN gate.

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

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

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

* Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc

Codex review: the transformer_prequant_path field description still told operators
to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the
prior security fix made that variable a directory allowlist -- _allowed_prequant_roots
deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator
following the documented =1 would have every transformer_prequant_path request
silently refused. The description now states it must name one or more allowlisted
directories and that a bare on/off value is not accepted.

Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does
not say =1, and describes an allowlist/directory (guards against doc drift).

* Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset

Codex review on attention-backend selection:

- Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on
  pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first
  generation with no fallback. select_attention_backend now applies
  _cudnn_attention_supported() to an explicit cuDNN request too.

- flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a
  Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3
  is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a
  (min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+.

- apply_attention_backend's success path left diffusers' process-wide active
  backend pinned to the kernel it set; a later component whose processors are
  unconfigured (backend None) would inherit it. It now resets the global registry
  to native after a successful per-transformer set (the transformer keeps its own
  backend), best-effort. Also fixed _active_attention_backend: get_active_backend()
  returns a (name, fn) tuple, so the prior code stringified the tuple and never
  matched a name, defeating the native-restore short-circuit.

Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on
SM90; global registry reset after a successful set; _active_attention_backend
reads the tuple return.

* Studio diffusion (Phase 11) review round 2: keep GH200/B300 on the fp8 ladder

Codex review: _DATACENTER_GPU_TOKENS omitted GH200 (Grace-Hopper) and B300
(Blackwell Ultra), though it has the distinct GB200/GB300 superchip tokens. So
_is_consumer_gpu returned True for 'NVIDIA GH200 480GB' / 'NVIDIA B300', and the
auto ladder moved int8 ahead of fp8 on those data-center parts -- contradicting
llama_cpp.py's datacenter regex, which lists both. Added GH200 and B300 so they
are treated as data-center class and keep the intended fp8-first behavior.

Test: extends the datacenter parametrize with 'NVIDIA B300' and
'NVIDIA GH200 480GB' (now _is_consumer_gpu False).

* Studio diffusion (Phase 14) review round 2: apply int8 M=1 exclusion in the builder

Codex review: the M=1 modulation/embedder exclusion was wired only into the dense
runtime quantiser; the offline builder scripts/build_prequant_checkpoint.py called
make_filter_fn(min_features) with no exclusion. So an int8 prequant checkpoint
quantised the AdaLN modulation and conditioning-embedder linears, and loading it
via transformer_prequant_path (the load path only loads already-quantised tensors,
it can't re-skip them) reintroduced the torch._int_mm M=1 crash this phase fixes
for the runtime path.

Extracted int8_exclude_name_tokens(scheme) as the single source of truth (int8 ->
the M=1 exclusion, every other scheme -> none) and use it in both the runtime
quantiser and the builder, so a prequant artifact's quantised-layer set always
matches the runtime. fp8/fp4/mx artifacts are byte-identical (empty exclusion).

Test: int8_exclude_name_tokens returns the exclusion for int8 and () for
fp8/nvfp4/mxfp8.

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

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

* Studio diffusion (Phase 16) review round 2: native CPU arbiter, status offload, load race

Codex review on the native-engine routing:

- The /images/load route took the GPU arbiter (acquire_for(DIFFUSION) -> evict chat)
  unconditionally after engine selection. A native sd.cpp load on a pure-CPU host
  never touches the GPU, so that needlessly tore down the resident chat model. The
  handoff is now gated: diffusers always takes it, a force-native sd.cpp load on a
  CUDA/XPU/MPS box still takes it, but a native sd.cpp load on a CPU host skips it.

- sd_cpp status() hardcoded offload_policy 'none' / cpu_offload False even when
  _run_load computed real offload flags (balanced/low_vram/cpu_offload off-CPU), so
  the setting was unverifiable. status now derives them from state.offload_flags
  (still 'none' on CPU, where the flags are empty).

- _run_load committed the new state without cancelling/waiting on a generation that
  started during the (slow) asset download, so a stale sd-cli run against the OLD
  model could finish afterward and persist an image from the previous model once the
  new load reported ready. The commit now signals the in-flight cancel and waits on
  _generate_lock before swapping _state (taken only at commit, so the download never
  serialises against generation), mirroring the diffusers load path.

Tests: CPU native load skips the arbiter while a GPU native load takes it; status
reports offload active when flags are set; _run_load cancels and waits for an
in-flight generation before committing.

* Studio diffusion (Phase 14) review round 2: align helper name with the stack

Rename the int8 exclusion helper to exclude_tokens_for_scheme, matching the
identical helper already present higher in the diffusion stack (Phase 16). The
helper definition, the runtime quantiser call, and the offline builder are now
byte-identical to that version, so the two branches no longer introduce a
divergent name for the same single-source-of-truth and the stack merges without
a conflict on this fix. No behavior change.

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

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

* Studio diffusion: persistent sd-server for the native engine (load once, serve many)

The native sd.cpp tier ran sd-cli one-shot per image, so begin_load only resolved
asset paths and every generation re-spawned sd-cli and reloaded the multi-GB GGUF
from disk (a batch of N = N full reloads). This makes it a resident backend backed by
stable-diffusion.cpp's persistent sd-server, mirroring the chat backend's llama-server
lifecycle:

- begin_load spawns sd-server once (the model loads there) and polls /v1/models until
  ready; unload kills it.
- generate submits ONE async /sdcpp/v1/img_gen job for the whole batch (no reload),
  polls it to completion, and decodes the returned images. Step progress and ETA come
  from the server's stdout (the job JSON has no per-step field).
- The one-shot sd-cli path is kept as an automatic fallback: it is used when sd-server
  is absent, and also when a present sd-server fails to start, so behavior is never
  worse than before. The public backend surface is unchanged, so routes/router need no
  change.

New: sd_cpp_server.py (SdCppServer manager: spawn/readiness/job-submit-poll/cancel/stop,
process spawned inside the drain thread so PR_SET_PDEATHSIG binds to the interpreter, not
a transient thread; empty scratch dir for the server's per-request LoRA/upscaler/embd
scans). Extended: sd_cpp_engine.py (find_sd_server_binary), sd_cpp_args.py
(build_sd_cpp_server_command + build_img_gen_request), sd_cpp_backend.py (server/one-shot
modes, ensure_sd_server_binary upgrades existing sd-cli-only installs), and the prebuilt
installer (locate + chmod sd-server, which ships in the same archive as sd-cli).

Verified on a B200 (Z-Image-Turbo-GGUF, CUDA sd-server): one model load across multiple
generations (server pid stable, a single 'listening on:'), a batch served from one job
with distinct per-image seeds, the second generation faster than the first, and
unload/reload spawning a fresh process. 105 sd.cpp + 81 diffusion tests pass.

Addresses the review of the Phase 16 native-engine PR.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

* Studio native diffusion: harden the persistent sd-server path

Addresses review findings on the sd-server backend:

- Router: treat a runnable sd-server as native availability, so an
  sd-server-only install (no sd-cli) still routes to the native engine
  instead of silently falling back to diffusers.
- Backend: probe the sd-server binary before the multi-GB asset download,
  falling back to one-shot sd-cli up front when it cannot run.
- Backend: a lazily cached one-shot fallback engine no longer pins the
  backend to one-shot; only an explicitly injected engine does, so a
  now-available server can be used on the next load.
- Backend: mask explicit seeds to sd.cpp's signed int64 range before
  submitting a server job (large seeds were rejected/wrapped in server
  mode only), and split batches above the server's per-job limit into
  chunks, each with a timeout proportional to its image count.
- Backend/server: make server startup cancellable. stop() signals an abort
  event before taking the lifecycle lock so a blocking readiness wait bails
  promptly; unload() stops a not-yet-committed pending server.
- Backend: status() clears stale loaded state when the resident server has
  exited, so clients reload instead of hammering a dead process with 500s.
- Server: abandon a poll whose best-effort cancel is not honored within a
  grace window (releasing the generate lock), report a pre-submit
  stop/cancel as cancellation (409, not 500), and verify JSON responses are
  the expected type before indexing.
- Server: use a bounded deque for the stdout tail buffer.
- Add native_mode to DiffusionStatusResponse so the field is not dropped by
  the response model.

Adds regression tests for each behavioral change.

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

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

* sd-server: harden native lifecycle and GPU install path

- Treat a crashed sd-server probe (signal death / non-127 nonzero) as
  unavailable so a broken prebuilt falls back to diffusers instead of
  routing to a server that dies on startup.
- Drop stale loaded state when a resident server has exited before a
  generate, returning the recoverable not-loaded path rather than a 500.
- Reject incomplete server batches (fewer blobs than requested) like the
  one-shot path instead of silently dropping images.
- Bound the server log tail in place (keep the deque(maxlen)) and bypass
  HTTP(S) proxies for the loopback client (trust_env=False).
- Honor a stop() that arrives after the server is published but before
  start() takes the lock, so a cancelled load cannot leak a spawned model
  process.
- Map a closed-client RuntimeError during poll to a cancellation when the
  generation is being cancelled, so unload races surface as 409 not 500.
- Stop a timed-out server job (best-effort cancel then teardown) so an
  abandoned generation cannot keep denoising and block later loads.
- Install the accelerator-matched sd-server build (ROCm/Vulkan/CUDA) and
  probe the resident server before auto-installing sd-cli, so a server-only
  or GPU host does not fetch the wrong or an unused binary.

---------

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>
This commit is contained in:
Daniel Han 2026-07-03 08:37:32 -07:00 • committed by GitHub
commit 7d8b2db236
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 2108 additions and 126 deletions
Download patch file Download diff file Expand all files Collapse all files

42
studio/backend/core/inference/diffusion_engine_router.py
View file

@ -29,7 +29,12 @@ from typing import Any, Optional
from core.inference.diffusion_device import resolve_diffusion_device_target
from core.inference.diffusion_families import DiffusionFamily, family_sd_cpp_supported
from core.inference.sd_cpp_backend import _install_allowed, ensure_sd_cpp_binary
from core.inference.sd_cpp_backend import (
_install_allowed,
_server_binary_runnable,
ensure_sd_cpp_binary,
ensure_sd_server_binary,
)
from core.inference.sd_cpp_engine import (
ENGINE_DIFFUSERS,
ENGINE_SD_CPP,
@ -137,20 +142,37 @@ def select_and_activate_engine(fam: DiffusionFamily, *, hf_token: Optional[str]
fam_ok = family_sd_cpp_supported(fam)
binary = None
server_binary = None
if policy_eligible and fam_ok:
binary = ensure_sd_cpp_binary(
# Probe the resident sd-server FIRST: the backend PREFERS it, and an sd-server-only
# install (no sd-cli) must still route to native rather than silently falling back to
# diffusers. Checking it before the sd-cli install also means a server-only host does
# not pay an avoidable sd-cli download. Install the accelerator-matched build (ROCm /
# Vulkan / CUDA) so a forced-native GPU load gets the GPU server, not the CPU one.
server_binary = ensure_sd_server_binary(
allow_install = _install_allowed(),
accelerator = _install_accelerator_for(backend),
)
# Probe runnability here, before committing the route to native: a present but
# non-runnable binary (wrong arch, missing shared libs, no execute bit) would
# otherwise pass as available and only fail inside the background load, instead
# of falling back to diffusers now.
if server_binary and not _server_binary_runnable(server_binary):
logger.warning(
"sd-server at %s is present but not runnable; not using it", server_binary
)
server_binary = None
# sd-cli is the one-shot fallback. Always LOCATE an existing binary, but only
# auto-INSTALL it when there is no usable server, so a server-only install is not
# forced to also download a CLI it will never use. Probe runnability before
# committing native: a present but non-runnable binary (wrong arch, missing shared
# libs, no execute bit) would otherwise pass as available and only fail inside the
# background load, instead of falling back to diffusers now.
binary = ensure_sd_cpp_binary(
allow_install = _install_allowed() and server_binary is None,
accelerator = _install_accelerator_for(backend),
)
if binary and SdCppEngine(binary = binary).version() is None:
logger.warning("sd-cli at %s is present but not runnable; using diffusers", binary)
logger.warning("sd-cli at %s is present but not runnable; not using it", binary)
binary = None
native_available = bool(binary) and policy_eligible and fam_ok
native_available = bool(binary or server_binary) and policy_eligible and fam_ok
choice = select_diffusion_engine(
backend, native_available = native_available, prefer_native = prefer_native
)
@ -162,8 +184,8 @@ def select_and_activate_engine(fam: DiffusionFamily, *, hf_token: Optional[str]
reason = f"GPU backend '{backend}' uses diffusers"
elif not fam_ok:
reason = f"family '{fam.name}' has no native sd.cpp asset mapping"
elif not binary:
reason = "sd-cli binary unavailable"
elif not (binary or server_binary):
reason = "native sd.cpp binary unavailable"
else:
reason = "diffusers selected"
return _activate(ENGINE_DIFFUSERS, reason)

130
studio/backend/core/inference/sd_cpp_args.py
View file

@ -325,6 +325,136 @@ def build_sd_cpp_upscale_command(
return cmd
def build_sd_cpp_server_command(
binary: str,
files: SdCppModelFiles,
*,
host: str,
port: int,
vae_format: Optional[str] = None,
offload: Optional[list[str]] = None,
native_speed: Optional[str] = None,
threads: Optional[int] = None,
scratch_dir: Optional[str] = None,
verbose: bool = False,
extra_args: Optional[list[str]] = None,
) -> list[str]:
"""Build the ``sd-server`` argv: model + hardware/server flags only.
``sd-server`` (stable-diffusion.cpp ``examples/server``) loads the model once at
spawn from the SAME flags ``sd-cli`` takes (``--diffusion-model`` / ``--vae`` /
the text encoders / ``--vae-format`` / offload + speed), and adds ``--listen-ip``
/ ``--listen-port``. Per-generation parameters (prompt, size, steps, seed, cfg,
sampler, batch) are NOT here -- they go in each ``/sdcpp/v1/img_gen`` request, so
one resident process serves many generations without reloading the weights.
``offload`` / ``native_speed`` map to the exact same sd.cpp flags as the one-shot
engine (``--offload-to-cpu`` / ``--diffusion-fa`` / ...), verified to be accepted
by ``sd-server --help``. ``scratch_dir`` (if given) is pointed at by the LoRA /
hires-upscaler / embeddings directory flags: sd-server's img_gen handler recursively
iterates those dirs, and an unset / missing dir makes it fail the request, so we give
it a real (empty) directory. ``extra_args`` is appended last so a power user can
override anything (sd.cpp's parser is last-wins).
"""
if not files.diffusion_model:
raise ValueError("diffusion_model path is required")
cmd: list[str] = [binary, "--diffusion-model", files.diffusion_model]
for flag, value in (
("--vae", files.vae),
("--clip_l", files.clip_l),
("--clip_g", files.clip_g),
("--t5xxl", files.t5xxl),
("--llm", files.llm),
("--qwen2vl", files.qwen2vl),
):
if value:
cmd += [flag, value]
if vae_format:
cmd += ["--vae-format", vae_format]
cmd += ["--listen-ip", str(host), "--listen-port", str(int(port))]
if scratch_dir:
cmd += [
"--lora-model-dir",
scratch_dir,
"--hires-upscalers-dir",
scratch_dir,
"--embd-dir",
scratch_dir,
]
if threads is not None:
cmd += ["--threads", str(int(threads))]
offload = list(offload or [])
if offload:
cmd += offload
# De-dup speed flags against offload (offload may already include --diffusion-fa).
cmd += [f for f in native_speed_flags(native_speed) if f not in offload]
if verbose:
cmd += ["-v"]
if extra_args:
cmd += list(extra_args)
return cmd
def build_img_gen_request(
*,
prompt: str,
negative_prompt: Optional[str] = None,
width: int = 1024,
height: int = 1024,
steps: Optional[int] = None,
seed: Optional[int] = None,
batch_count: int = 1,
sample_method: Optional[str] = None,
flow_shift: Optional[float] = None,
cfg_scale: Optional[float] = None,
distilled_guidance: Optional[float] = None,
output_format: str = "png",
) -> dict:
"""Build the ``POST /sdcpp/v1/img_gen`` JSON body for one text-to-image request.
The native ``sdcpp`` API takes the whole batch in one request (``batch_count``),
so a batch reuses the resident model with no reload. Sampling lives under
``sample_params``; guidance is split exactly like the one-shot engine's
``_map_guidance``: a FLUX distilled value goes to ``guidance.distilled_guidance``,
a real classifier-free scale goes to ``guidance.txt_cfg``. Only set keys are
emitted so the server applies its own defaults for the rest.
"""
if not str(prompt).strip():
raise ValueError("prompt is required")
guidance: dict = {}
if cfg_scale is not None:
guidance["txt_cfg"] = float(cfg_scale)
if distilled_guidance is not None:
guidance["distilled_guidance"] = float(distilled_guidance)
sample_params: dict = {}
if steps is not None:
sample_params["sample_steps"] = int(steps)
if sample_method:
sample_params["sample_method"] = str(sample_method)
if flow_shift is not None:
sample_params["flow_shift"] = float(flow_shift)
if guidance:
sample_params["guidance"] = guidance
req: dict = {
"prompt": prompt,
"negative_prompt": negative_prompt or "",
"width": int(width),
"height": int(height),
"batch_count": max(1, int(batch_count)),
"output_format": output_format,
}
if seed is not None:
req["seed"] = int(seed)
if sample_params:
req["sample_params"] = sample_params
return req
def _fmt_float(value: float) -> str:
"""Compact float -> str: drop a trailing ``.0`` so ``1.0`` -> ``1`` (sd-cli
accepts both, but the tidy form keeps logged commands readable)."""

543
studio/backend/core/inference/sd_cpp_backend.py
View file

@ -49,13 +49,22 @@ from core.inference.diffusion_memory import (
OFFLOAD_NONE,
OFFLOAD_SEQUENTIAL,
)
from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles, offload_flags
from core.inference.sd_cpp_args import (
SdCppGenParams,
SdCppModelFiles,
build_img_gen_request,
offload_flags,
)
from core.inference.sd_cpp_engine import (
SdCppCancelled,
SdCppEngine,
find_sd_cpp_binary,
find_sd_server_binary,
runtime_env,
)
from core.inference.sd_cpp_server import SdCppServer
from loggers import get_logger
from utils.subprocess_compat import windows_hidden_subprocess_kwargs
logger = get_logger(__name__)
@ -68,6 +77,45 @@ _STEP_RE = re.compile(r"(\d+)\s*/\s*(\d+)")
# download / extract / chmod.
_install_lock = threading.Lock()
# sd-server accepts at most this many images per img_gen job; larger Studio batches
# (the request model allows up to 32) are split into chunks of this size, the way the
# one-shot path did them one image at a time.
_MAX_SERVER_BATCH = 8
# Per-image wall-clock budget for a server job, so a batch gets a timeout proportional to
# its image count (matching the one-shot path, where each image had its own budget) rather
# than one fixed deadline the whole batch has to finish within.
_SERVER_PER_IMAGE_TIMEOUT_S = 1800.0
def _server_binary_runnable(binary: str) -> bool:
"""Best-effort probe that ``binary`` can actually execute (not just exist).
Runs ``<binary> --help`` with the same runtime env the server will use, so a present
but unrunnable build (wrong arch, missing shared libs, no execute bit) is caught before
a multi-GB asset download. Conservative: only a clear "cannot launch" signal (OSError,
or the dynamic-loader exit codes 126/127) returns False; anything else is treated as
runnable so a quirky ``--help`` exit code never blocks a working binary."""
import subprocess
try:
proc = subprocess.run(
[binary, "--help"],
capture_output = True,
timeout = 20,
env = runtime_env(binary),
**windows_hidden_subprocess_kwargs(),
)
except OSError:
return False # cannot exec at all (wrong arch / no execute bit / missing loader)
except Exception: # noqa: BLE001 -- timeout or anything odd: don't block on a flaky probe
return True
# A negative return code is a signal death (e.g. -4 SIGILL from an incompatible
# prebuilt on an older CPU): the binary launches but immediately crashes, so treat it
# as unavailable and let the load fall back to diffusers instead of routing to a
# server that will die on startup.
return proc.returncode >= 0 and proc.returncode not in (126, 127)
def ensure_sd_cpp_binary(*, allow_install: bool = True, accelerator: str = "cpu") -> Optional[str]:
"""Path to a usable ``sd-cli`` binary, installing the prebuilt once if needed.
@ -105,9 +153,51 @@ def ensure_sd_cpp_binary(*, allow_install: bool = True, accelerator: str = "cpu"
return None
def ensure_sd_server_binary(
*, allow_install: bool = True, accelerator: str = "cpu"
) -> Optional[str]:
"""Path to a usable ``sd-server`` binary, installing the prebuilt once if needed.
Unlike ``ensure_sd_cpp_binary``, this installs when *sd-server specifically* is
missing -- even if an ``sd-cli`` from an older install is already present -- so an
existing one-shot install is upgraded to the persistent server (the prebuilt archive
ships both). Returns None when it is absent and cannot be installed; the backend then
uses the one-shot fallback. Never raises.
"""
found = find_sd_server_binary()
if found:
return found
if not allow_install:
return None
with _install_lock:
found = find_sd_server_binary()
if found:
return found
try:
import sys
studio_dir = Path(__file__).resolve().parents[3] # .../studio
if str(studio_dir) not in sys.path:
sys.path.insert(0, str(studio_dir))
from install_sd_cpp_prebuilt import install as _install
except Exception as exc: # noqa: BLE001 -- import path / module issues are non-fatal
logger.warning("sd-server installer import failed: %s", exc)
return None
try:
_install(accelerator = accelerator) # extracts sd-cli AND sd-server
except Exception as exc: # noqa: BLE001 -- download/extract failure -> fall back
logger.warning("sd-server auto-install failed: %s", exc)
return None
return find_sd_server_binary()
@dataclass(frozen = True)
class _SdState:
"""The loaded native checkpoint: resolved asset paths + run settings."""
"""The loaded native checkpoint: resolved asset paths + run settings.
``server`` is the resident ``sd-server`` process (the model is loaded once, inside
it) when ``mode == "server"``; in the ``"oneshot"`` fallback it is ``None`` and each
generation re-runs ``sd-cli``."""
repo_id: str
base_repo: str
@ -120,6 +210,8 @@ class _SdState:
threads: Optional[int] = None
sampling_method: Optional[str] = None
flow_shift: Optional[float] = None
server: Optional[SdCppServer] = None
mode: str = "server"
def _memory_policy(memory_mode: Optional[str], cpu_offload: bool) -> str:
@ -191,11 +283,19 @@ class SdCppDiffusionBackend:
self._lock = threading.Lock()
self._generate_lock = threading.Lock()
self._engine = engine # resolved lazily on first load so import stays cheap
# An engine passed in is an EXPLICIT injection (the test seam / escape hatch) and
# pins one-shot mode; an engine cached later by a runtime fallback must NOT, so a
# now-available server can still be used on the next load.
self._engine_injected = engine is not None
self._state: Optional[_SdState] = None
self._loading: Optional[_SdLoading] = None
self._load_token = 0
self._cancel_event = threading.Event()
self._active_generate_cancel: Optional[threading.Event] = None
# The sd-server being started for an in-flight load, before it is committed to
# _state. Tracked so an unload / superseding load can stop it mid-startup instead
# of leaving it loading (and holding the generate lock) for the whole timeout.
self._pending_server: Optional[SdCppServer] = None
self._gen: Optional[_SdGen] = None
@property
@ -212,6 +312,38 @@ class SdCppDiffusionBackend:
self._engine = SdCppEngine(binary = binary)
return self._engine
def _resolve_backend(self) -> tuple[str, Optional[str], Optional[SdCppEngine]]:
"""Pick the native execution mode: ("server", binary, None) or ("oneshot", None, engine).
The persistent ``sd-server`` is preferred (load once, serve many). The one-shot
``sd-cli`` is the fallback for older / custom builds that lack the server target.
An explicitly injected engine forces one-shot (the unit-test seam and an escape
hatch), so a test never spawns a real server or triggers an install. A lazily
cached fallback engine does NOT force one-shot: once a resident server becomes
available (installed, or a per-model start that previously failed now works), the
next load can use it, instead of being pinned to one-shot for the whole session.
"""
if self._engine_injected and self._engine is not None:
return "oneshot", None, self._resolve_engine()
# Install the sd-server build matching the resolved device backend (ROCm / Vulkan /
# CUDA), not the default CPU build: a forced/enabled native load on a GPU host must
# not silently fetch the plain-CPU server. Lazy import avoids an import cycle with
# the router, which imports this backend during engine selection.
from core.inference.diffusion_engine_router import _install_accelerator_for
accelerator = _install_accelerator_for(
getattr(resolve_diffusion_device_target(), "backend", "cpu")
)
server_binary = ensure_sd_server_binary(
allow_install = _install_allowed(), accelerator = accelerator
)
if server_binary is not None:
return "server", server_binary, None
logger.warning(
"sd-server not found; falling back to one-shot sd-cli (reloads the model per image)."
)
return "oneshot", None, self._resolve_engine()
# ── Background load + progress ─────────────────────────────────────────
def begin_load(
@ -298,10 +430,35 @@ class SdCppDiffusionBackend:
_load_token: int,
) -> None:
try:
# Ensure the binary up front so an install failure surfaces before the
# multi-GB asset pull (the router also pre-checks, but a forced reload here
# must not silently download then fail at generate).
engine = self._resolve_engine()
# Resolve the backend mode (persistent sd-server preferred, one-shot sd-cli
# fallback) and binary up front so an install / missing-binary failure
# surfaces before the multi-GB asset pull.
mode, server_binary, engine = self._resolve_backend()
if mode == "server":
# Probe the server binary before the multi-GB asset pull: a present but
# unrunnable build (wrong arch / missing libs) would otherwise download
# everything and only then fail to start. If it cannot run, fall back to
# the one-shot engine now (when it is usable), else surface the failure.
assert server_binary is not None
if not _server_binary_runnable(server_binary):
logger.warning(
"sd-server at %s is present but not runnable; trying one-shot sd-cli.",
server_binary,
)
try:
usable = self._resolve_engine().version() is not None
except Exception: # noqa: BLE001
usable = False
if not usable:
raise RuntimeError("sd-server binary is present but not runnable.")
mode, server_binary, engine = "oneshot", None, self._resolve_engine()
if mode == "oneshot":
# Probe the binary: version() returns None when the present binary cannot
# run (bad perms / missing libs), so fail now rather than commit a "ready"
# state that crashes on the first generation.
assert engine is not None
if engine.version() is None:
raise RuntimeError("sd-cli binary is present but not runnable.")
assets = self._asset_specs(repo_id, gguf_filename, fam)
self._set_expected_bytes(assets, hf_token)
@ -318,37 +475,21 @@ class SdCppDiffusionBackend:
)
device = resolve_diffusion_device_target().device
# Honor the requested speed everywhere; offload only off-CPU (forced
# sd_cpp / MPS), since on CPU sd-cli is resident in RAM and the offload
# flags are no-ops.
# sd_cpp / MPS), since on CPU the weights are resident in RAM and the
# offload flags are no-ops.
offload: tuple[str, ...] = ()
if device != "cpu":
offload = tuple(offload_flags(_memory_policy(memory_mode, cpu_offload)))
state = _SdState(
repo_id = repo_id,
base_repo = base,
family = fam,
device = device,
files = files,
vae_format = fam.sd_cpp_vae_format,
native_speed = _native_speed_for(speed_mode),
offload_flags = offload,
threads = None,
sampling_method = fam.sd_cpp_sampling_method,
flow_shift = fam.sd_cpp_flow_shift,
)
# Probe the binary: version() returns None when the present binary cannot
# run (bad permissions / missing shared libs), so fail the load now rather
# than commit a "ready" state that crashes on the first generation.
if engine.version() is None:
raise RuntimeError("sd-cli binary is present but not runnable.")
# A generation that started during the (slow) asset download is still running
# against the OLD model. Abort it, then WAIT on _generate_lock for it to exit
# before publishing the new state -- otherwise that stale sd-cli run can finish
# afterward and persist an image from the previous model once this load reports
# ready (mirrors the diffusers load_pipeline commit). _generate_lock is taken
# only here, not during the download, so the long fetch never serialises against
# generation; the inner token re-check guards an unload/newer load arriving while
# we waited.
native_speed = _native_speed_for(speed_mode)
# Tear down any previously-loaded model, then commit the new one. A generation
# that started during the (slow) asset download is still running against the OLD
# model: abort it and WAIT on _generate_lock for it to exit before swapping, or
# a stale run could finish afterward and persist an image from the previous
# model. For server mode we stop the old server and start (load) the new one
# HERE, under _generate_lock, so generation never races a half-loaded server and
# two resident models never coexist. _generate_lock is taken only now, not during
# the download, so the long fetch never serialises against generation.
with self._lock:
if self._load_token != _load_token:
return # superseded / cancelled
@ -358,6 +499,78 @@ class SdCppDiffusionBackend:
with self._lock:
if self._load_token != _load_token:
return # superseded / cancelled while waiting
old_state = self._state
self._state = None # the old model is being torn down
if old_state is not None and old_state.server is not None:
old_state.server.stop()
server: Optional[SdCppServer] = None
if mode == "server":
assert server_binary is not None
server = SdCppServer(server_binary)
# Publish the not-yet-committed server so unload() / a superseding load
# can stop it mid-startup (SdCppServer.stop aborts the readiness wait
# without waiting on the lifecycle lock), instead of it loading for the
# full startup timeout while holding the generate lock.
with self._lock:
self._pending_server = server
try:
# Blocks until the server has loaded the model and is answering
# (its readiness check); raises with the log tail on a failed load.
server.start(
files,
vae_format = fam.sd_cpp_vae_format,
offload = list(offload),
native_speed = native_speed,
threads = None,
)
except SdCppCancelled:
# Startup was aborted by an unload / superseding load: stop the
# half-started server and bail (the outer handler returns cleanly).
server.stop()
raise
except Exception as start_exc: # noqa: BLE001
# A present-but-unusable sd-server must be no worse than the
# one-shot engine: fall back to sd-cli when it is usable, else
# surface the server error.
logger.warning(
"sd-server failed to start (%s); falling back to one-shot sd-cli.",
start_exc,
)
server.stop()
server = None
try:
usable = self._resolve_engine().version() is not None
except Exception: # noqa: BLE001
usable = False
if not usable:
raise start_exc
mode = "oneshot"
finally:
with self._lock:
if self._pending_server is server:
self._pending_server = None
state = _SdState(
repo_id = repo_id,
base_repo = base,
family = fam,
device = device,
files = files,
vae_format = fam.sd_cpp_vae_format,
native_speed = native_speed,
offload_flags = offload,
threads = None,
sampling_method = fam.sd_cpp_sampling_method,
flow_shift = fam.sd_cpp_flow_shift,
server = server,
mode = mode,
)
with self._lock:
if self._load_token != _load_token:
# Superseded / unloaded while we were loading: discard the server
# we just started so it doesn't leak (and keep _state unloaded).
if server is not None:
server.stop()
return
self._state = state
self._loading = None
except SdCppCancelled:
@ -475,76 +688,62 @@ class SdCppDiffusionBackend:
seed: Optional[int] = None,
batch_size: int = 1,
) -> dict[str, Any]:
import tempfile
from PIL import Image
cancel = threading.Event()
with self._generate_lock:
with self._lock:
state = self._state
if state is None:
raise RuntimeError(DIFFUSION_NOT_LOADED_MSG)
# A resident server can exit while idle; if a client generates without first
# polling status, drop the stale loaded state and report not-loaded so it gets
# the recoverable reload path instead of a 500 from img_gen (not running).
if (
state.mode == "server"
and state.server is not None
and not state.server.is_alive()
):
self._state = None
raise RuntimeError(DIFFUSION_NOT_LOADED_MSG)
self._active_generate_cancel = cancel
engine = self._resolve_engine()
try:
if seed is None:
seed = int.from_bytes(os.urandom(6), "big") & ((1 << 53) - 1)
else:
seed = int(seed)
cfg_scale, flux_guidance = _map_guidance(state.family, guidance)
extra_args: list[str] = []
if state.vae_format:
extra_args += ["--vae-format", state.vae_format]
if state.flow_shift is not None:
extra_args += ["--flow-shift", repr(float(state.flow_shift))]
self._gen = _SdGen(total_steps = int(steps))
images = []
seeds: list[int] = []
with tempfile.TemporaryDirectory(prefix = "sdcpp_gen_") as tmpdir:
for index in range(max(1, int(batch_size))):
if cancel.is_set():
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
# Distinct seed per batch image (sd-cli is one image/run here),
# so a batch is reproducible image-by-image from the base seed.
# Mask to sd-cli's int64 range, NOT 53 bits: the request model and
# the diffusers backend both accept large explicit seeds, so a tight
# 2**53 mask would silently truncate them (2**53 -> 0) and collide
# distinct requested seeds onto the same image. Randomly-drawn seeds
# above are already 53-bit (JS-safe); explicit seeds pass through.
seed_i = (seed + index) & ((1 << 63) - 1)
out_path = str(Path(tmpdir) / f"img_{index}.png")
params = SdCppGenParams(
prompt = prompt,
negative_prompt = negative_prompt or None,
width = int(width),
height = int(height),
steps = int(steps),
cfg_scale = cfg_scale,
guidance = flux_guidance,
seed = seed_i,
sampling_method = state.sampling_method,
batch_count = 1,
)
engine.generate(
state.files,
params,
output_path = out_path,
offload = list(state.offload_flags) or None,
native_speed = state.native_speed,
threads = state.threads,
extra_args = extra_args or None,
on_log = self._on_log,
cancel_event = cancel,
)
with Image.open(out_path) as im:
images.append(im.copy())
seeds.append(seed_i)
if state.mode == "server" and state.server is not None:
images, seeds = self._generate_server(
state,
prompt = prompt,
negative_prompt = negative_prompt,
width = width,
height = height,
steps = steps,
seed = seed,
batch_size = batch_size,
cfg_scale = cfg_scale,
flux_guidance = flux_guidance,
cancel = cancel,
)
else:
images, seeds = self._generate_oneshot(
state,
prompt = prompt,
negative_prompt = negative_prompt,
width = width,
height = height,
steps = steps,
seed = seed,
batch_size = batch_size,
cfg_scale = cfg_scale,
flux_guidance = flux_guidance,
cancel = cancel,
)
if cancel.is_set():
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
# ``seeds`` is the per-image seed (each sd-cli run used seed+index), so
# the route can persist the real seed for every image in the batch.
# ``seeds`` is the per-image seed (image i used seed+i), so the route can
# persist the real seed for every image in the batch.
return {
"images": images,
"seed": int(seed),
@ -559,6 +758,144 @@ class SdCppDiffusionBackend:
if self._active_generate_cancel is cancel:
self._active_generate_cancel = None
def _generate_server(
self,
state: _SdState,
*,
prompt: str,
negative_prompt: Optional[str],
width: int,
height: int,
steps: int,
seed: int,
batch_size: int,
cfg_scale: Optional[float],
flux_guidance: Optional[float],
cancel: threading.Event,
) -> tuple[list, list[int]]:
"""Generate via the resident sd-server (no model reload).
A batch larger than the server's per-job limit is split into chunks: the server
rejects a batch_count above _MAX_SERVER_BATCH, and the one-shot path served large
batches image-by-image, so preserve that. The base seed is masked to sd.cpp's
signed-int64 range (the request model / diffusers accept larger seeds), and each
chunk is submitted at base+offset so the per-image seeds stay reproducible. Each
chunk gets a timeout proportional to its image count so a slow CPU batch is not
cancelled partway through on one fixed deadline."""
import io
from PIL import Image
assert state.server is not None
total = max(1, int(batch_size))
# sd.cpp's image seed is signed int64; mask the base (and every derived seed) so a
# large explicit seed is not rejected / wrapped inconsistently by the server.
base_seed = int(seed) & ((1 << 63) - 1)
images: list = []
seeds: list[int] = []
for offset in range(0, total, _MAX_SERVER_BATCH):
if cancel.is_set():
raise SdCppCancelled("sd-server generation was cancelled.")
count = min(_MAX_SERVER_BATCH, total - offset)
chunk_seed = (base_seed + offset) & ((1 << 63) - 1)
payload = build_img_gen_request(
prompt = prompt,
negative_prompt = negative_prompt or None,
width = int(width),
height = int(height),
steps = int(steps),
seed = chunk_seed,
batch_count = count,
sample_method = state.sampling_method,
flow_shift = state.flow_shift,
cfg_scale = cfg_scale,
distilled_guidance = flux_guidance,
)
blobs = state.server.img_gen(
payload,
on_step = self._on_log,
cancel_event = cancel,
total_timeout = _SERVER_PER_IMAGE_TIMEOUT_S * count,
)
# All-or-nothing per chunk, like the one-shot path: if the server returns fewer
# blobs than requested (e.g. one image in the batch failed to encode), fail
# rather than silently dropping images from the user's requested batch.
if not cancel.is_set() and len(blobs) != count:
raise RuntimeError(
f"sd-server returned {len(blobs)} of {count} requested images in the batch."
)
images.extend(Image.open(io.BytesIO(b)).convert("RGB") for b in blobs)
# sd.cpp advances the seed per image within a job, so report chunk_seed+i.
seeds.extend((chunk_seed + i) & ((1 << 63) - 1) for i in range(len(blobs)))
return images, seeds
def _generate_oneshot(
self,
state: _SdState,
*,
prompt: str,
negative_prompt: Optional[str],
width: int,
height: int,
steps: int,
seed: int,
batch_size: int,
cfg_scale: Optional[float],
flux_guidance: Optional[float],
cancel: threading.Event,
) -> tuple[list, list[int]]:
"""Fallback path: re-run one-shot sd-cli per image (reloads the model each time)."""
import tempfile
from PIL import Image
engine = self._resolve_engine()
extra_args: list[str] = []
if state.vae_format:
extra_args += ["--vae-format", state.vae_format]
if state.flow_shift is not None:
extra_args += ["--flow-shift", repr(float(state.flow_shift))]
images = []
seeds: list[int] = []
with tempfile.TemporaryDirectory(prefix = "sdcpp_gen_") as tmpdir:
for index in range(max(1, int(batch_size))):
if cancel.is_set():
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
# Distinct seed per batch image, reproducible image-by-image from the base
# seed. Mask to int64, NOT 53 bits: the request model and the diffusers
# backend both accept large explicit seeds, so a tight 2**53 mask would
# truncate them and collide distinct requested seeds onto the same image.
seed_i = (seed + index) & ((1 << 63) - 1)
out_path = str(Path(tmpdir) / f"img_{index}.png")
params = SdCppGenParams(
prompt = prompt,
negative_prompt = negative_prompt or None,
width = int(width),
height = int(height),
steps = int(steps),
cfg_scale = cfg_scale,
guidance = flux_guidance,
seed = seed_i,
sampling_method = state.sampling_method,
batch_count = 1,
)
engine.generate(
state.files,
params,
output_path = out_path,
offload = list(state.offload_flags) or None,
native_speed = state.native_speed,
threads = state.threads,
extra_args = extra_args or None,
on_log = self._on_log,
cancel_event = cancel,
)
with Image.open(out_path) as im:
images.append(im.copy())
seeds.append(seed_i)
return images, seeds
def _on_log(self, line: str) -> None:
gen = self._gen
if gen is None or gen.total_steps <= 0:
@ -596,13 +933,40 @@ class SdCppDiffusionBackend:
with self._lock:
if self._active_generate_cancel is not None:
self._active_generate_cancel.set()
state = self._state
self._state = None
self._load_token += 1
self._loading = None
# A load may be mid server.start() with the server not yet committed to _state;
# grab it too so we can stop it (its startup is abortable) instead of leaving it
# loading for the full startup timeout.
pending = self._pending_server
self._pending_server = None
# Stop the resident server outside the lock (terminate can take a few seconds). A
# mid-flight generation had its cancel event set above, so its poll loop unwinds
# as the process goes away.
if state is not None and state.server is not None:
state.server.stop()
if pending is not None and pending is not (state.server if state else None):
pending.stop()
return self.status()
def status(self) -> dict[str, Any]:
state = self._state
# A resident sd-server can exit after load (OOM-killed / crashed while idle). If so,
# drop the stale loaded state so status reports not-loaded and clients reload,
# instead of every generation failing with a 500 against a dead process.
if (
state is not None
and state.mode == "server"
and state.server is not None
and not state.server.is_alive()
):
logger.warning("sd-server exited after load; clearing loaded state")
with self._lock:
if self._state is state:
self._state = None
state = None
if state is None:
return {
"loaded": False,
@ -622,6 +986,7 @@ class SdCppDiffusionBackend:
"attention_backend": None,
"transformer_cache": None,
"engine": "sd_cpp",
"native_mode": None,
}
return {
"loaded": True,
@ -645,6 +1010,8 @@ class SdCppDiffusionBackend:
"attention_backend": None,
"transformer_cache": None,
"engine": "sd_cpp",
# "server" = resident sd-server (load once); "oneshot" = legacy per-image sd-cli.
"native_mode": state.mode,
}

94
studio/backend/core/inference/sd_cpp_engine.py
View file

@ -51,6 +51,9 @@ logger = logging.getLogger(__name__)
# target is ``sd-cli``; older builds shipped ``sd`` -- both are probed on PATH.
_BINARY_STEM = "sd-cli"
_LEGACY_STEM = "sd"
# The persistent HTTP server target (stable-diffusion.cpp ``examples/server``). It
# ships next to ``sd-cli`` in both the prebuilt archives and the cmake build tree.
_SERVER_STEM = "sd-server"
class SdCppCancelled(RuntimeError):
@ -119,11 +122,11 @@ def runtime_env(binary: str, base_env: Optional[dict[str, str]] = None) -> dict[
return env
def _layout_candidates(root: Path) -> list[Path]:
"""sd-cli locations under a stable-diffusion.cpp checkout/install ``root``,
def _layout_candidates(root: Path, stem: str = _BINARY_STEM) -> list[Path]:
"""``stem`` locations under a stable-diffusion.cpp checkout/install ``root``,
highest priority first: the cmake ``build/bin`` tree, then a Windows Release
subdir, then the root itself."""
name = _binary_name(_BINARY_STEM)
name = _binary_name(stem)
cands = [
root / "build" / "bin" / name,
root / "build" / "bin" / "Release" / name,
@ -133,66 +136,101 @@ def _layout_candidates(root: Path) -> list[Path]:
return cands
def find_sd_cpp_binary() -> Optional[str]:
"""Locate the ``sd-cli`` binary, or None.
def _first_file(paths: list[Path]) -> Optional[str]:
for p in paths:
try:
if p.is_file():
return str(p)
except OSError:
continue
return None
Search order (mirrors the llama.cpp finder so a Studio install lands where
both engines look):
1. ``SD_CLI_PATH`` env -- a direct path to the binary.
def _find_binary(
*, direct_env: str, path_stems: tuple[str, ...], layout_stem: str
) -> Optional[str]:
"""Shared finder for the stable-diffusion.cpp binaries.
Search order (mirrors the llama.cpp finder so a Studio install lands where every
binary is looked for):
1. ``direct_env`` -- a direct path to the binary.
2. ``UNSLOTH_SD_CPP_PATH`` env -- a stable-diffusion.cpp install dir.
3. the installer target: ``<UNSLOTH_STUDIO_HOME>/../stable-diffusion.cpp`` when
that env (or ``STUDIO_HOME``) is set, else ``~/.unsloth/stable-diffusion.cpp``.
4. ``./stable-diffusion.cpp`` in-tree build (developer checkout).
5. ``sd-cli`` (then legacy ``sd``) on PATH.
5. ``path_stems`` on PATH (in order).
"""
def _first_file(paths: list[Path]) -> Optional[str]:
for p in paths:
try:
if p.is_file():
return str(p)
except OSError:
continue
return None
# 1. Direct binary path.
env_bin = os.environ.get("SD_CLI_PATH")
env_bin = os.environ.get(direct_env)
if env_bin and Path(env_bin).is_file():
return env_bin
# 2. Custom install dir.
custom = os.environ.get("UNSLOTH_SD_CPP_PATH")
if custom:
hit = _first_file(_layout_candidates(Path(custom)))
hit = _first_file(_layout_candidates(Path(custom), layout_stem))
if hit:
return hit
# 3. Default install root: the installer's default_install_dir() -- a sibling of
# the llama.cpp install under UNSLOTH_STUDIO_HOME / STUDIO_HOME when set, else
# ~/.unsloth. Mirror that env resolution or a custom Studio home never resolves.
# 3. Default install root. Honors UNSLOTH_STUDIO_HOME / STUDIO_HOME the same way
# the installer's default_install_dir does (base = the Studio home's parent), so
# a binary installed under a custom Studio root is discovered and side-by-side
# Studios stay isolated; falls back to the sibling of ~/.unsloth/llama.cpp.
studio_home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME")
default_base = Path(studio_home).parent if studio_home else Path.home() / ".unsloth"
hit = _first_file(_layout_candidates(default_base / "stable-diffusion.cpp"))
default_root = (
Path(studio_home).parent / "stable-diffusion.cpp"
if studio_home
else Path.home() / ".unsloth" / "stable-diffusion.cpp"
)
hit = _first_file(_layout_candidates(default_root, layout_stem))
if hit:
return hit
# 4. In-tree developer build: <repo_root>/stable-diffusion.cpp.
try:
project_root = Path(__file__).resolve().parents[4]
hit = _first_file(_layout_candidates(project_root / "stable-diffusion.cpp"))
hit = _first_file(_layout_candidates(project_root / "stable-diffusion.cpp", layout_stem))
if hit:
return hit
except (OSError, IndexError):
pass
# 5. PATH.
for stem in (_BINARY_STEM, _LEGACY_STEM):
for stem in path_stems:
on_path = shutil.which(stem)
if on_path:
return on_path
return None
def find_sd_cpp_binary() -> Optional[str]:
"""Locate the one-shot ``sd-cli`` binary (env ``SD_CLI_PATH``), or None.
Probes ``sd-cli`` then legacy ``sd`` on PATH. This is the fallback engine once the
persistent ``sd-server`` exists; it also still backs the ESRGAN upscale mode.
"""
return _find_binary(
direct_env = "SD_CLI_PATH",
path_stems = (_BINARY_STEM, _LEGACY_STEM),
layout_stem = _BINARY_STEM,
)
def find_sd_server_binary() -> Optional[str]:
"""Locate the persistent ``sd-server`` binary (env ``SD_SERVER_PATH``), or None.
Same precedence as ``find_sd_cpp_binary`` but keyed to the ``sd-server`` stem, so
a Studio install (prebuilt archive or cmake build, both of which ship ``sd-server``
next to ``sd-cli``) is found in the same places. Preferred over the one-shot CLI:
it loads the model once and serves many generations without reloading from disk.
"""
return _find_binary(
direct_env = "SD_SERVER_PATH",
path_stems = (_SERVER_STEM,),
layout_stem = _SERVER_STEM,
)
class SdCppEngine:
"""A thin handle over a located ``sd-cli`` binary.

502
studio/backend/core/inference/sd_cpp_server.py Normal file
View file

@ -0,0 +1,502 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Persistent ``sd-server`` (stable-diffusion.cpp) process manager.
The native diffusion tier used to shell out to one-shot ``sd-cli`` per image, which
reloaded the multi-GB GGUF from disk every generation. ``sd-server`` (the upstream
``examples/server`` target) loads the model once at spawn and serves many generations
over HTTP, exactly like the chat backend's persistent ``llama-server``. This manager
owns ONLY the process + HTTP lifecycle; the backend (``sd_cpp_backend.py``) still owns
asset resolution, request validation, and the public Studio surface.
Shape mirrors ``core/rag/embed_llama_server.py``:
* ``start`` -- pick a free loopback port, spawn the server (model loads here),
drain stdout on a daemon thread, poll until ready.
* ``img_gen`` -- POST ``/sdcpp/v1/img_gen`` (the whole batch in one request),
poll the async job to a terminal state, return image bytes.
* ``stop`` -- SIGTERM -> wait -> SIGKILL, join the drain thread (idempotent).
Readiness is real: upstream ``main.cpp`` loads the model BEFORE it binds the port and
prints ``listening on:``, so a 200 from ``GET /v1/models`` (a trivial handler) means the
model is loaded; a load failure exits the process before listening and is surfaced with
the captured log tail. (The richer ``/sdcpp/v1/capabilities`` handler can block in some
builds, so it is not used for readiness.) The job JSON has no per-step field, so step progress is recovered by
parsing the server's stdout (the same ``N/M`` lines ``sd-cli`` emits), routed to the
active generation's callback.
Import-light on purpose (no torch / diffusers / PIL), so selecting the native tier on
a CPU box never drags the GPU stack into the process.
"""
from __future__ import annotations
import atexit
import base64
import logging
import shutil
import socket
import subprocess
import tempfile
import threading
import time
from collections import deque
from typing import Any, Callable, Optional
import httpx
from core.inference.sd_cpp_args import SdCppModelFiles, build_sd_cpp_server_command
from core.inference.sd_cpp_engine import SdCppCancelled, runtime_env
from utils.native_path_leases import child_env_without_native_path_secret
from utils.process_lifetime import adopt_pid, child_popen_kwargs, forget_pid
from utils.subprocess_compat import windows_hidden_subprocess_kwargs
logger = logging.getLogger(__name__)
# httpx transport errors meaning "the server is gone / connection refused" -- treated
# as "not ready yet" while polling readiness, and as a fatal "server died" mid-request.
_TRANSPORT_ERRORS = (
httpx.ConnectError,
httpx.ReadError,
httpx.RemoteProtocolError,
httpx.WriteError,
)
# Readiness probe. Upstream binds the port only AFTER the model is loaded, so any 200
# means ready. We use /v1/models (a trivial, always-fast handler) rather than
# /sdcpp/v1/capabilities: the capabilities handler can block in some builds (it enumerates
# model metadata), which would stall readiness even though the server is up.
_READY_PATH = "/v1/models"
# Native async sdcpp API.
_IMG_GEN_PATH = "/sdcpp/v1/img_gen"
_JOBS_PATH = "/sdcpp/v1/jobs"
_TERMINAL_OK = "completed"
_TERMINAL_FAIL = "failed"
_TERMINAL_CANCELLED = "cancelled"
# After a cancel is requested, how long to let the server reflect it in job status before
# abandoning the poll. The native cancel is best-effort, so without this cap a server that
# ignores/loses the cancel would keep this call (and the backend's generate lock) alive
# until the job finishes naturally, blocking a superseding load from swapping the model.
_CANCEL_GRACE_S = 5.0
class SdCppServer:
"""A resident ``sd-server`` subprocess plus the HTTP client that drives it."""
def __init__(
self,
binary: str,
*,
host: str = "127.0.0.1",
) -> None:
self.binary = binary
self.host = host
self.port: Optional[int] = None
self._process: Optional[subprocess.Popen] = None
# Fixed-size, thread-safe tail buffer: the drain thread appends while lifecycle /
# request threads read it for diagnostics, so a deque(maxlen) is safer and cheaper
# than a list with manual slicing.
self._tail: deque[str] = deque(maxlen = 200)
self._stdout_thread: Optional[threading.Thread] = None
self._lifecycle_lock = threading.Lock()
# Set (lock-free) by stop() so a blocking start()/readiness wait can be aborted
# promptly without waiting on the lifecycle lock start() holds.
self._abort = threading.Event()
# Set for the duration of a generation so the continuous stdout drain can feed
# the active request's step-progress callback; cleared in img_gen's finally.
self._step_listener: Optional[Callable[[str], None]] = None
# trust_env=False: this client only ever talks to the loopback sd-server, so it must
# not route through HTTP_PROXY/HTTPS_PROXY (a proxy without 127.0.0.1 in NO_PROXY
# would break readiness/generation). Matches the local llama-server clients.
self._client = httpx.Client(timeout = 30.0, trust_env = False)
self._scratch_dir: Optional[str] = None
self._stopped = False
atexit.register(self.stop)
# ── lifecycle ────────────────────────────────────────────────────────────
@property
def base_url(self) -> str:
return f"http://{self.host}:{self.port}"
def is_alive(self) -> bool:
return self._process is not None and self._process.poll() is None
@staticmethod
def _find_free_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return int(s.getsockname()[1])
def start(
self,
files: SdCppModelFiles,
*,
vae_format: Optional[str] = None,
offload: Optional[list[str]] = None,
native_speed: Optional[str] = None,
threads: Optional[int] = None,
env: Optional[dict[str, str]] = None,
startup_timeout: float = 600.0,
) -> None:
"""Spawn the server (which loads the model) and block until it is ready.
Raises ``RuntimeError`` (with the captured log tail) if the process exits during
startup or never answers within ``startup_timeout``. Holds the lifecycle lock so
a concurrent start/stop can't interleave.
"""
with self._lifecycle_lock:
# A stop()/unload that raced in AFTER the backend published this server as
# _pending_server but BEFORE start() took the lock has already set _abort and
# closed the httpx client. Honor that delivered stop instead of clearing the
# abort and spawning a model process the cancelled load would then leak.
if self._stopped or self._abort.is_set():
raise SdCppCancelled("sd-server start was cancelled before launch.")
self._abort.clear()
port = self._find_free_port()
# An empty scratch dir for sd-server's LoRA / upscaler / embeddings scans
# (it recursively iterates them per request and errors on a missing dir).
self._scratch_dir = tempfile.mkdtemp(prefix = "sdcpp_dirs_")
cmd = build_sd_cpp_server_command(
self.binary,
files,
host = self.host,
port = port,
vae_format = vae_format,
offload = list(offload or []),
native_speed = native_speed,
threads = threads,
scratch_dir = self._scratch_dir,
verbose = True, # sd-server prints the per-step sampling lines we parse
)
run_env = runtime_env(self.binary, child_env_without_native_path_secret())
if env:
run_env.update(env)
logger.info("starting sd-server: %s", " ".join(cmd))
# Clear in place: reassigning to [] drops the deque(maxlen=200) bound, so the
# continuous stdout drain would then grow the tail without limit for the whole
# resident-server lifetime.
self._tail.clear()
self._spawn_error: Optional[Exception] = None
spawned = threading.Event()
# Spawn INSIDE the drain thread, which then reads stdout for the process's whole
# lifetime. child_popen_kwargs() sets PR_SET_PDEATHSIG, which on Linux is bound to
# the CREATING THREAD -- so the child must be created by a thread that outlives it,
# or a transient spawner thread ending would kill the server. The drain thread is
# exactly that long-lived owner; it dies only when the process exits or the
# interpreter goes away (the case we DO want to reap the GPU-resident server).
def _own_process() -> None:
try:
proc = subprocess.Popen(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
errors = "replace",
env = run_env,
**windows_hidden_subprocess_kwargs(),
**child_popen_kwargs(),
)
except Exception as exc: # noqa: BLE001 -- surface the spawn failure to start()
self._spawn_error = exc
spawned.set()
return
self._process = proc
self.port = port
adopt_pid(proc.pid) # so a global shutdown sweep also reaps it
spawned.set()
self._drain_stdout(proc)
# stdout closed == the process exited; reap it so it is not left a zombie
# until the next stop()/reload.
try:
proc.wait(timeout = 5)
except Exception: # noqa: BLE001
pass
self._stdout_thread = threading.Thread(
target = _own_process, daemon = True, name = "sd-server-owner"
)
self._stdout_thread.start()
spawned.wait()
if self._spawn_error is not None:
self._dispose()
raise RuntimeError(f"failed to spawn sd-server: {self._spawn_error}")
if not self._wait_ready(startup_timeout):
tail = "\n".join(list(self._tail)[-30:])
aborted = self._abort.is_set()
self._kill_locked()
self._dispose()
if aborted:
raise SdCppCancelled("sd-server startup was cancelled.")
raise RuntimeError("sd-server failed to become ready. Last output:\n" + tail[:2000])
def _wait_ready(
self,
timeout: float,
interval: float = 0.5,
) -> bool:
"""Poll ``/v1/models`` until 200; bail early if the process exits.
Upstream binds the port only AFTER the model is loaded, so a 200 here is a true
ready signal (no half-loaded race)."""
deadline = time.monotonic() + timeout
url = f"{self.base_url}{_READY_PATH}"
while time.monotonic() < deadline:
# A concurrent stop() (unload / superseding load) sets _abort so this wait can
# bail without holding the model-load hostage for the full startup_timeout.
if self._abort.is_set():
logger.info("sd-server startup aborted before ready")
return False
if not self.is_alive():
code = None if self._process is None else self._process.returncode
logger.error("sd-server exited early during load (code %s)", code)
return False
try:
if self._client.get(url, timeout = 2.0).status_code == 200:
return True
except (*_TRANSPORT_ERRORS, httpx.TimeoutException):
pass
time.sleep(interval)
logger.error("sd-server readiness timed out after %ss", timeout)
return False
def _drain_stdout(self, proc: subprocess.Popen) -> None:
"""Drain stdout so the pipe never deadlocks; keep a tail for diagnostics and
feed each line to the active generation's step callback."""
try:
assert proc.stdout is not None
for raw in proc.stdout:
line = raw.rstrip()
if not line:
continue
self._tail.append(line) # deque(maxlen) discards the oldest automatically
logger.debug("[sd-server] %s", line)
cb = self._step_listener
if cb is not None:
try:
cb(line)
except Exception: # noqa: BLE001 -- a progress callback must never break drain
pass
except Exception: # noqa: BLE001 -- drain thread must never raise (pipe closed at teardown)
pass
def stop(self) -> None:
"""Terminate the server (SIGTERM -> SIGKILL), join the drain, and release the HTTP
client + atexit handler. Idempotent."""
# Signal abort BEFORE contending for the lifecycle lock: a concurrent start() holds
# that lock for the whole (up to startup_timeout) readiness wait, so setting the
# event lets that wait bail immediately instead of stop() blocking behind it.
self._abort.set()
self._stopped = True
with self._lifecycle_lock:
self._kill_locked()
self._dispose()
def _dispose(self) -> None:
"""Release per-instance resources (on stop / failed start). The backend never
reuses a disposed server, so this closes the pooled httpx client and drops the
atexit handler that would otherwise pin every reloaded instance for the session."""
try:
atexit.unregister(self.stop)
except Exception: # noqa: BLE001
pass
try:
self._client.close()
except Exception: # noqa: BLE001
pass
if self._scratch_dir:
shutil.rmtree(self._scratch_dir, ignore_errors = True)
self._scratch_dir = None
def _kill_locked(self) -> None:
proc = self._process
if proc is None:
return
pid = proc.pid
try:
proc.terminate()
proc.wait(timeout = 5)
except subprocess.TimeoutExpired:
logger.warning("sd-server did not exit on SIGTERM; killing")
try:
proc.kill()
proc.wait(timeout = 5)
except Exception: # noqa: BLE001 -- best-effort teardown
pass
except Exception as exc: # noqa: BLE001
logger.warning("error terminating sd-server: %s", exc)
finally:
forget_pid(pid)
self._process = None
self.port = None
if self._stdout_thread is not None:
self._stdout_thread.join(timeout = 2)
self._stdout_thread = None
# ── generation ───────────────────────────────────────────────────────────
def img_gen(
self,
payload: dict[str, Any],
*,
on_step: Optional[Callable[[str], None]] = None,
cancel_event: Optional[threading.Event] = None,
poll_interval: float = 0.4,
submit_timeout: float = 60.0,
total_timeout: float = 1800.0,
) -> list[bytes]:
"""Submit one async ``img_gen`` job, poll it to completion, return image bytes.
``on_step`` receives each server stdout line (for the step bar). ``cancel_event``,
when set, cancels the job via the native endpoint and raises ``SdCppCancelled``.
Raises ``RuntimeError`` on submit/poll failures (including the server dying), with
the log tail attached.
"""
# If the server was already stopped for a cancel/unload/superseding load that set
# the cancel event before this submit began, report it as a cancellation (which the
# route maps to a client-state 409) rather than a generic "server died" 500.
if self._stopped or not self.is_alive():
if cancel_event is not None and cancel_event.is_set():
raise SdCppCancelled("sd-server generation was cancelled.")
raise RuntimeError("sd-server is not running.")
self._step_listener = on_step
job_id: Optional[str] = None
try:
# Submit -> 202 Accepted + job id.
try:
resp = self._client.post(
f"{self.base_url}{_IMG_GEN_PATH}", json = payload, timeout = submit_timeout
)
except (*_TRANSPORT_ERRORS, httpx.TimeoutException) as exc:
raise RuntimeError(self._died_message("img_gen submit", exc)) from exc
if resp.status_code == 429:
raise RuntimeError("sd-server job queue is full (HTTP 429).")
if resp.status_code not in (200, 202):
raise RuntimeError(
f"sd-server img_gen submit -> {resp.status_code}: {resp.text[:500]}"
)
try:
job = resp.json()
except ValueError as exc:
raise RuntimeError(
f"sd-server img_gen returned a non-JSON submit response: {exc}"
) from exc
if not isinstance(job, dict):
raise RuntimeError(
f"sd-server img_gen returned an unexpected submit response type: {type(job)}"
)
job_id = job.get("id")
if not job_id:
raise RuntimeError(f"sd-server img_gen returned no job id: {job}")
# Poll the job to a terminal state.
deadline = time.monotonic() + total_timeout
cancel_sent_at: Optional[float] = None
while True:
if cancel_event is not None and cancel_event.is_set():
if cancel_sent_at is None:
self.cancel(job_id)
cancel_sent_at = time.monotonic()
elif time.monotonic() - cancel_sent_at > _CANCEL_GRACE_S:
# The best-effort cancel was not reflected in job status within the
# grace window; abandon the poll so the caller can stop the server
# instead of holding the generate lock until the job finishes.
raise SdCppCancelled("sd-server generation was cancelled.")
if not self.is_alive():
# If we're unwinding a cancel (e.g. unload killed the server), surface a
# clean cancellation rather than a generic "server died" error.
if cancel_event is not None and cancel_event.is_set():
raise SdCppCancelled("sd-server generation was cancelled.")
raise RuntimeError(self._died_message("img_gen poll", None))
if time.monotonic() > deadline:
# Best-effort cancel, then tear the server down: current sd-server does
# not interrupt an already-generating job (cancel_generating=false / 409),
# so leaving it up would keep denoising the abandoned job and block later
# generations/reloads behind it. Stopping frees the slot; the backend sees
# the dead server on the next generate and takes the recoverable reload path.
self.cancel(job_id)
self.stop()
raise RuntimeError(f"sd-server generation timed out after {total_timeout}s")
try:
jr = self._client.get(f"{self.base_url}{_JOBS_PATH}/{job_id}", timeout = 10.0)
except (*_TRANSPORT_ERRORS, httpx.TimeoutException):
time.sleep(poll_interval)
continue
except RuntimeError as exc:
# A concurrent stop()/unload closes the shared httpx client; httpx then
# raises a plain RuntimeError ("client has been closed") that is NOT a
# transport error. When we are being cancelled, report it as a clean
# cancellation (route -> 409) instead of a generic 500 generation failure.
if cancel_event is not None and cancel_event.is_set():
raise SdCppCancelled("sd-server generation was cancelled.") from exc
raise
if jr.status_code in (404, 410):
raise RuntimeError(f"sd-server job {job_id} is gone (HTTP {jr.status_code}).")
if jr.status_code != 200:
time.sleep(poll_interval)
continue
try:
jd = jr.json()
except ValueError as exc:
raise RuntimeError(f"sd-server job status was not JSON: {exc}") from exc
if not isinstance(jd, dict):
raise RuntimeError(
f"sd-server job status returned an unexpected response type: {type(jd)}"
)
status = jd.get("status")
if status == _TERMINAL_OK:
return self._decode_images(jd)
if status == _TERMINAL_FAIL:
err = jd.get("error") or {}
raise RuntimeError(
"sd-server generation failed: "
f"{err.get('code', 'error')}: {err.get('message', '')}".strip()
)
if status == _TERMINAL_CANCELLED:
raise SdCppCancelled("sd-server generation was cancelled.")
time.sleep(poll_interval)
finally:
self._step_listener = None
def cancel(self, job_id: str) -> None:
"""Best-effort native cancel of an in-flight job."""
try:
self._client.post(f"{self.base_url}{_JOBS_PATH}/{job_id}/cancel", timeout = 5.0)
except Exception: # noqa: BLE001 -- cancel is best-effort
pass
@staticmethod
def _decode_images(job: dict[str, Any]) -> list[bytes]:
# Defensive against an unexpected response shape (a misbehaving/older server):
# verify each level is the type we index before calling dict/list methods.
result = job.get("result") if isinstance(job, dict) else None
images = result.get("images") if isinstance(result, dict) else None
items = [it for it in images if isinstance(it, dict)] if isinstance(images, list) else []
out: list[bytes] = []
for item in sorted(items, key = lambda d: d.get("index", 0)):
b64 = item.get("b64_json")
if not b64:
continue
try:
out.append(base64.b64decode(b64))
except Exception as exc: # noqa: BLE001
raise RuntimeError(f"sd-server returned an undecodable image: {exc}") from exc
if not out:
raise RuntimeError("sd-server completed the job but returned no images.")
return out
def _died_message(self, where: str, exc: Optional[Exception]) -> str:
tail = "\n".join(list(self._tail)[-20:])
base = f"sd-server connection lost during {where}"
if not self.is_alive():
code = None if self._process is None else self._process.returncode
base += f" (process exited, code {code})"
if exc is not None:
base += f": {exc}"
if tail:
base += "\nLast output:\n" + tail[:1500]
return base

5
studio/backend/models/inference.py
View file

@ -1924,6 +1924,11 @@ class DiffusionStatusResponse(BaseModel):
)
transformer_cache: Optional[str] = Field(None, description = "Step cache engaged: fbcache | null")
engine: Optional[str] = Field(None, description = "Active diffusion engine: diffusers | sd_cpp")
native_mode: Optional[str] = Field(
None,
description = "Native sd.cpp execution mode: server (resident sd-server) | oneshot "
"(per-image sd-cli) | null (diffusers engine)",
)
fallback_reason: Optional[str] = Field(
None,
description = "Why diffusers was chosen over the native sd.cpp engine (null when none)",

14
studio/backend/tests/test_diffusion_engine_router.py
View file

@ -32,6 +32,10 @@ def _clean_env_and_state(monkeypatch):
"get_active_diffusion_engine",
lambda: SimpleNamespace(status = lambda: {"loaded": False, "repo_id": None}),
)
# Default: no resident sd-server (so existing tests exercise the sd-cli path only) and
# a stubbed runnability probe, so neither reaches the real install/exec path.
monkeypatch.setattr(r, "ensure_sd_server_binary", lambda **_: None)
monkeypatch.setattr(r, "_server_binary_runnable", lambda *_a, **_k: True)
yield
@ -70,6 +74,16 @@ def test_cpu_with_binary_and_supported_family_picks_sd_cpp(monkeypatch):
assert r.active_engine_name() == ENGINE_SD_CPP
def test_cpu_with_only_sd_server_picks_sd_cpp(monkeypatch):
# An sd-server-only install (no runnable sd-cli) must still route to native: the
# backend prefers the resident server, so a runnable sd-server is native availability.
_set_device(monkeypatch, "cpu")
_set_binary(monkeypatch, None) # no sd-cli
monkeypatch.setattr(r, "SdCppEngine", lambda **_: SimpleNamespace(version = lambda: None))
monkeypatch.setattr(r, "ensure_sd_server_binary", lambda **_: "/usr/bin/sd-server")
assert _select() == ENGINE_SD_CPP
def test_present_but_not_runnable_binary_falls_back(monkeypatch):
# A binary that exists but cannot run (version() -> None) must fall back to
# diffusers at selection, not commit native and fail inside the load.

103
studio/backend/tests/test_sd_cpp_args.py
View file

@ -21,7 +21,9 @@ from core.inference.sd_cpp_args import (
SdCppGenParams,
SdCppModelFiles,
SdCppUpscaleParams,
build_img_gen_request,
build_sd_cpp_command,
build_sd_cpp_server_command,
build_sd_cpp_upscale_command,
native_speed_flags,
offload_flags,
@ -364,3 +366,104 @@ def test_build_upscale_requires_input_and_model():
SdCppUpscaleParams(input_image = "/i.png", upscale_model = ""),
output_path = "/o.png",
)
# ── sd-server spawn command ──────────────────────────────────────────────────
def test_server_command_has_model_and_listen_but_no_request_params():
files = SdCppModelFiles(diffusion_model = "/m/z.gguf", vae = "/m/ae.sft", llm = "/m/q.gguf")
cmd = build_sd_cpp_server_command(
"/bin/sd-server", files, host = "127.0.0.1", port = 5678, vae_format = "flux2"
)
assert _pair(cmd, "--diffusion-model") == "/m/z.gguf"
assert _pair(cmd, "--vae") == "/m/ae.sft"
assert _pair(cmd, "--llm") == "/m/q.gguf"
assert _pair(cmd, "--vae-format") == "flux2"
assert _pair(cmd, "--listen-ip") == "127.0.0.1"
assert _pair(cmd, "--listen-port") == "5678"
# Per-request parameters must NOT be baked into the spawn command.
for flag in (
"--prompt",
"--seed",
"--steps",
"--cfg-scale",
"--guidance",
"--width",
"--height",
"--batch-count",
):
assert flag not in cmd
def test_server_command_maps_offload_and_speed_and_dedupes():
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
cmd = build_sd_cpp_server_command(
"/bin/sd-server",
files,
host = "127.0.0.1",
port = 1,
offload = ["--offload-to-cpu", "--diffusion-fa"],
native_speed = "default", # would add --diffusion-fa again
threads = 8,
)
assert _pair(cmd, "--threads") == "8"
assert cmd.count("--diffusion-fa") == 1 # de-duped against offload
assert "--offload-to-cpu" in cmd
def test_server_command_scratch_dir_expands_to_lora_upscaler_embd():
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
cmd = build_sd_cpp_server_command(
"/bin/sd-server", files, host = "127.0.0.1", port = 1, scratch_dir = "/tmp/scratch"
)
assert _pair(cmd, "--lora-model-dir") == "/tmp/scratch"
assert _pair(cmd, "--hires-upscalers-dir") == "/tmp/scratch"
assert _pair(cmd, "--embd-dir") == "/tmp/scratch"
# Absent when not requested.
bare = build_sd_cpp_server_command("/bin/sd-server", files, host = "127.0.0.1", port = 1)
assert "--lora-model-dir" not in bare and "--hires-upscalers-dir" not in bare
def test_server_command_requires_diffusion_model():
with pytest.raises(ValueError):
build_sd_cpp_server_command(
"/bin/sd-server", SdCppModelFiles(diffusion_model = ""), host = "127.0.0.1", port = 1
)
# ── img_gen request body ─────────────────────────────────────────────────────
def test_img_gen_request_maps_core_fields():
req = build_img_gen_request(
prompt = "a fox",
negative_prompt = "blurry",
width = 512,
height = 768,
steps = 8,
seed = 42,
batch_count = 3,
sample_method = "euler",
cfg_scale = 4.0,
)
assert req["prompt"] == "a fox" and req["negative_prompt"] == "blurry"
assert req["width"] == 512 and req["height"] == 768
assert req["seed"] == 42 and req["batch_count"] == 3
assert req["sample_params"]["sample_steps"] == 8
assert req["sample_params"]["sample_method"] == "euler"
assert req["sample_params"]["guidance"]["txt_cfg"] == 4.0
assert req["output_format"] == "png"
def test_img_gen_request_flux_uses_distilled_guidance():
req = build_img_gen_request(prompt = "x", steps = 4, distilled_guidance = 3.5, flow_shift = 3.0)
g = req["sample_params"]["guidance"]
assert g["distilled_guidance"] == 3.5
assert "txt_cfg" not in g
assert req["sample_params"]["flow_shift"] == 3.0
def test_img_gen_request_requires_prompt():
with pytest.raises(ValueError):
build_img_gen_request(prompt = " ", steps = 4)

313
studio/backend/tests/test_sd_cpp_backend.py
View file

@ -77,10 +77,69 @@ def _loaded_backend(fam_name = "z-image", engine = None):
vae_format = fam.sd_cpp_vae_format,
sampling_method = fam.sd_cpp_sampling_method,
flow_shift = fam.sd_cpp_flow_shift,
mode = "oneshot", # this fixture injects an engine, so it exercises the one-shot path
)
return b
class _FakeServer:
"""Stands in for SdCppServer: records the spawn + one img_gen per whole batch."""
def __init__(self, binary):
self.binary = binary
self.started = None
self.stopped = False
self.payloads = []
self.timeouts = []
self.alive = True
def is_alive(self):
return self.alive and not self.stopped
def start(
self,
files,
*,
vae_format = None,
offload = None,
native_speed = None,
threads = None,
):
self.started = dict(
files = files,
vae_format = vae_format,
offload = offload,
native_speed = native_speed,
threads = threads,
)
def img_gen(
self,
payload,
*,
on_step = None,
cancel_event = None,
total_timeout = None,
):
import io as _io
self.payloads.append(payload)
self.timeouts.append(total_timeout)
if on_step is not None:
steps = payload.get("sample_params", {}).get("sample_steps", 0)
on_step(f" {steps}/{steps}")
n = int(payload.get("batch_count", 1))
blobs = []
for i in range(n):
buf = _io.BytesIO()
Image.new("RGB", (1, 1), (i, i, i)).save(buf, format = "PNG")
blobs.append(buf.getvalue())
return blobs
def stop(self):
self.stopped = True
# ── asset resolution ──────────────────────────────────────────────────────────
@ -321,6 +380,260 @@ def test_run_load_cancels_and_waits_for_inflight_generation(monkeypatch):
assert b._state is not None and b._state.repo_id == "unsloth/Z-Image-Turbo-GGUF"
# ── persistent sd-server mode ──────────────────────────────────────────────────
def test_resolve_backend_prefers_server(monkeypatch):
b = SdCppDiffusionBackend() # no injected engine
monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server")
mode, binary, engine = b._resolve_backend()
assert mode == "server" and binary == "/x/sd-server" and engine is None
def test_resolve_backend_injected_engine_forces_oneshot():
b = SdCppDiffusionBackend(engine = _FakeEngine())
mode, binary, engine = b._resolve_backend()
assert mode == "oneshot" and binary is None and engine is not None
def test_resolve_backend_falls_back_to_oneshot_without_server(monkeypatch):
b = SdCppDiffusionBackend()
monkeypatch.setattr(bk, "find_sd_server_binary", lambda: None)
monkeypatch.setattr(bk, "_install_allowed", lambda: False) # don't attempt a real install
monkeypatch.setattr(bk, "find_sd_cpp_binary", lambda: "/usr/bin/sd-cli")
mode, binary, engine = b._resolve_backend()
assert mode == "oneshot" and engine is not None
def test_resolve_backend_cached_fallback_engine_does_not_pin_oneshot(monkeypatch):
# A lazily cached fallback engine (NOT an explicit injection) must not force one-shot:
# once a server is available again, the next load can use it.
b = SdCppDiffusionBackend() # no injected engine
b._engine = _FakeEngine() # simulate a prior lazy one-shot fallback caching the engine
monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server")
mode, binary, engine = b._resolve_backend()
assert mode == "server" and binary == "/x/sd-server" and engine is None
def _run_server_load(
monkeypatch,
b,
servers,
fam_name = "z-image",
):
fam = detect_family(fam_name)
monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server")
# The fake binary path is not a real executable; skip the up-front runnability probe.
monkeypatch.setattr(bk, "_server_binary_runnable", lambda *_a, **_k: True)
def _factory(binary):
s = _FakeServer(binary)
servers.append(s)
return s
monkeypatch.setattr(bk, "SdCppServer", _factory)
monkeypatch.setattr(b, "_asset_specs", lambda *a, **k: [])
monkeypatch.setattr(b, "_set_expected_bytes", lambda *a, **k: None)
monkeypatch.setattr(
b,
"_fetch_assets",
lambda *a, **k: {"diffusion_model": "/m/z.gguf", "vae": "/m/vae.sft", "llm": "/m/llm.sft"},
)
monkeypatch.setattr(
bk, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cpu")
)
b._load_token = 1
b._run_load(
repo_id = "unsloth/Z-Image-Turbo-GGUF",
gguf_filename = "z.gguf",
base = fam.base_repo,
fam = fam,
hf_token = None,
_load_token = 1,
)
def test_server_load_spawns_once_and_status_reports_mode(monkeypatch):
b = SdCppDiffusionBackend()
servers: list = []
_run_server_load(monkeypatch, b, servers)
assert len(servers) == 1
assert servers[0].started is not None # the model is loaded once, at spawn
assert b._state is not None and b._state.mode == "server" and b._state.server is servers[0]
assert b.status()["native_mode"] == "server"
def test_server_generate_uses_one_request_for_whole_batch(monkeypatch):
b = SdCppDiffusionBackend()
servers: list = []
_run_server_load(monkeypatch, b, servers)
out = b.generate(prompt = "a fox", width = 64, height = 64, steps = 8, seed = 7, batch_size = 3)
assert len(out["images"]) == 3
assert all(isinstance(im, Image.Image) for im in out["images"])
# ONE job for the whole batch (no per-image model reload), unlike the one-shot path.
assert len(servers[0].payloads) == 1
assert servers[0].payloads[0]["batch_count"] == 3
assert out["seed"] == 7 and out["seeds"] == [7, 8, 9]
# step progress was driven from the server's stdout line.
assert b._gen is None # cleared after generate
def test_server_generate_splits_batches_above_server_limit(monkeypatch):
# A batch above the server's per-job limit is chunked (the one-shot path did these
# image-by-image); each chunk gets a timeout proportional to its image count.
b = SdCppDiffusionBackend()
servers: list = []
_run_server_load(monkeypatch, b, servers)
out = b.generate(prompt = "x", width = 64, height = 64, steps = 4, seed = 100, batch_size = 10)
assert len(out["images"]) == 10
counts = [p["batch_count"] for p in servers[0].payloads]
assert counts == [bk._MAX_SERVER_BATCH, 10 - bk._MAX_SERVER_BATCH] # [8, 2]
# Each chunk's timeout scales with its image count, not one fixed batch deadline.
assert servers[0].timeouts == [
bk._SERVER_PER_IMAGE_TIMEOUT_S * 8,
bk._SERVER_PER_IMAGE_TIMEOUT_S * 2,
]
# Seeds run contiguously across chunks (chunk 2 submitted at base + 8).
assert out["seeds"] == list(range(100, 110))
assert servers[0].payloads[1]["seed"] == 108
def test_server_generate_masks_large_seed(monkeypatch):
# sd.cpp's image seed is signed int64; a larger explicit seed must be masked before it
# reaches the server (the request model / diffusers accept up to 2**64 - 1).
b = SdCppDiffusionBackend()
servers: list = []
_run_server_load(monkeypatch, b, servers)
out = b.generate(prompt = "x", width = 64, height = 64, steps = 4, seed = 2**64 - 1, batch_size = 1)
assert servers[0].payloads[0]["seed"] <= (1 << 63) - 1
assert all(s <= (1 << 63) - 1 for s in out["seeds"])
def test_status_clears_when_server_died(monkeypatch):
b = SdCppDiffusionBackend()
servers: list = []
_run_server_load(monkeypatch, b, servers)
assert b.status()["loaded"] is True
servers[0].alive = False # the resident server crashed / was OOM-killed
st = b.status()
assert st["loaded"] is False
assert b._state is None # stale state was dropped so clients reload
def test_server_generate_progress_from_stdout(monkeypatch):
b = SdCppDiffusionBackend()
servers: list = []
_run_server_load(monkeypatch, b, servers)
seen = {}
class _WatchServer(_FakeServer):
def img_gen(
self,
payload,
*,
on_step = None,
cancel_event = None,
total_timeout = None,
):
on_step(" 4/8")
seen["mid"] = b.generate_progress()
return super().img_gen(
payload, on_step = on_step, cancel_event = cancel_event, total_timeout = total_timeout
)
b._state = bk._SdState(
repo_id = b._state.repo_id,
base_repo = b._state.base_repo,
family = b._state.family,
device = b._state.device,
files = b._state.files,
vae_format = b._state.vae_format,
sampling_method = b._state.sampling_method,
flow_shift = b._state.flow_shift,
server = _WatchServer("/x/sd-server"),
mode = "server",
)
b.generate(prompt = "x", steps = 8, seed = 1)
assert seen["mid"]["step"] == 4 and seen["mid"]["total_steps"] == 8
def test_server_unload_stops_server(monkeypatch):
b = SdCppDiffusionBackend()
servers: list = []
_run_server_load(monkeypatch, b, servers)
st = b.unload()
assert st["loaded"] is False
assert servers[0].stopped is True
assert b._state is None
def test_server_reload_stops_old_server_before_new(monkeypatch):
b = SdCppDiffusionBackend()
servers: list = []
_run_server_load(monkeypatch, b, servers)
# A second load must tear down the first server and start a fresh one.
b._load_token = 2
fam = detect_family("z-image")
b._run_load(
repo_id = "unsloth/Z-Image-Turbo-GGUF",
gguf_filename = "z.gguf",
base = fam.base_repo,
fam = fam,
hf_token = None,
_load_token = 2,
)
assert len(servers) == 2
assert servers[0].stopped is True # old server stopped
assert b._state.server is servers[1] and servers[1].stopped is False
def test_server_start_failure_falls_back_to_oneshot(monkeypatch):
# A present-but-broken sd-server must not fail the load when sd-cli works.
b = SdCppDiffusionBackend()
monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server")
# Probe passes; the failure we exercise here is in start(), not the up-front probe.
monkeypatch.setattr(bk, "_server_binary_runnable", lambda *_a, **_k: True)
class _BadServer:
def __init__(self, binary):
self.stopped = False
def start(self, *a, **k):
raise RuntimeError("sd-server broken")
def stop(self):
self.stopped = True
monkeypatch.setattr(bk, "SdCppServer", _BadServer)
fake = _FakeEngine()
monkeypatch.setattr(b, "_resolve_engine", lambda: fake)
monkeypatch.setattr(b, "_asset_specs", lambda *a, **k: [])
monkeypatch.setattr(b, "_set_expected_bytes", lambda *a, **k: None)
monkeypatch.setattr(
b,
"_fetch_assets",
lambda *a, **k: {"diffusion_model": "/m/z.gguf", "vae": "/m/vae.sft", "llm": "/m/llm.sft"},
)
monkeypatch.setattr(
bk, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cpu")
)
fam = detect_family("z-image")
b._load_token = 1
b._run_load(
repo_id = "unsloth/Z-Image-Turbo-GGUF",
gguf_filename = "z.gguf",
base = fam.base_repo,
fam = fam,
hf_token = None,
_load_token = 1,
)
assert b._state is not None and b._state.mode == "oneshot" and b._state.server is None
# and it can still generate via the one-shot engine
out = b.generate(prompt = "x", steps = 4, seed = 1)
assert len(out["images"]) == 1 and len(fake.calls) == 1
def test_run_load_redacts_paths_in_progress_error(monkeypatch):
# A load failure surfaced via load_progress() must run through redact_native_paths, the
# same scrub the diffusers load path applies, so a registered native path can't leak.

53
studio/backend/tests/test_sd_cpp_engine.py
View file

@ -24,6 +24,7 @@ from core.inference.sd_cpp_engine import (
ENGINE_SD_CPP,
SdCppEngine,
find_sd_cpp_binary,
find_sd_server_binary,
runtime_env,
select_diffusion_engine,
)
@ -75,6 +76,58 @@ def test_find_returns_none_when_absent(tmp_path, monkeypatch):
assert find_sd_cpp_binary() is None
# ── sd-server discovery ──────────────────────────────────────────────────────
def _clear_server_env(monkeypatch):
monkeypatch.delenv("SD_SERVER_PATH", raising = False)
monkeypatch.delenv("SD_CLI_PATH", raising = False)
monkeypatch.delenv("UNSLOTH_SD_CPP_PATH", raising = False)
def test_find_server_prefers_sd_server_path_env(tmp_path, monkeypatch):
_clear_server_env(monkeypatch)
binary = tmp_path / "sd-server"
binary.write_text("x")
monkeypatch.setenv("SD_SERVER_PATH", str(binary))
monkeypatch.setattr(eng.shutil, "which", lambda *_a: None)
assert find_sd_server_binary() == str(binary)
def test_find_server_build_layout(tmp_path, monkeypatch):
_clear_server_env(monkeypatch)
root = tmp_path / "sdcpp"
built = root / "build" / "bin" / "sd-server"
built.parent.mkdir(parents = True)
built.write_text("x")
monkeypatch.setenv("UNSLOTH_SD_CPP_PATH", str(root))
monkeypatch.setattr(eng.shutil, "which", lambda *_a: None)
assert find_sd_server_binary() == str(built)
def test_find_server_path_fallback(tmp_path, monkeypatch):
_clear_server_env(monkeypatch)
monkeypatch.setattr(eng.Path, "home", staticmethod(lambda: tmp_path / "nohome"))
monkeypatch.setattr(
eng.shutil, "which", lambda stem: "/usr/bin/sd-server" if stem == "sd-server" else None
)
assert find_sd_server_binary() == "/usr/bin/sd-server"
def test_find_server_not_confused_with_sd_cli(tmp_path, monkeypatch):
# A tree that has only sd-cli must NOT be reported as an sd-server (and vice versa),
# so the backend correctly falls back to one-shot when only the CLI is present.
_clear_server_env(monkeypatch)
root = tmp_path / "sdcpp"
(root / "build" / "bin").mkdir(parents = True)
(root / "build" / "bin" / "sd-cli").write_text("x")
monkeypatch.setenv("UNSLOTH_SD_CPP_PATH", str(root))
monkeypatch.setattr(eng.Path, "home", staticmethod(lambda: tmp_path / "nohome"))
monkeypatch.setattr(eng.shutil, "which", lambda *_a: None)
assert find_sd_server_binary() is None
assert find_sd_cpp_binary() == str(root / "build" / "bin" / "sd-cli")
# ── availability / version ──────────────────────────────────────────────────

417
studio/backend/tests/test_sd_cpp_server.py Normal file
View file

@ -0,0 +1,417 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for the persistent sd-server process manager (SdCppServer).
Hermetic: subprocess.Popen and the httpx client are faked, so nothing spawns a real
binary or opens a socket beyond the free-port probe."""
from __future__ import annotations
import base64
import io
import threading
import pytest
from PIL import Image
from core.inference import sd_cpp_server as srv
from core.inference.sd_cpp_args import SdCppModelFiles
from core.inference.sd_cpp_engine import SdCppCancelled
from core.inference.sd_cpp_server import SdCppServer
_FILES = SdCppModelFiles(diffusion_model = "/m/z.gguf", vae = "/m/vae.sft", llm = "/m/llm.sft")
def _png_b64(shade: int) -> str:
buf = io.BytesIO()
Image.new("RGB", (1, 1), (shade, shade, shade)).save(buf, format = "PNG")
return base64.b64encode(buf.getvalue()).decode()
class _FakePopen:
"""Minimal Popen stand-in. stdout yields the scripted lines then BLOCKS until the
process is terminated/killed/exited -- mirroring a real child that holds its pipe
open for its lifetime (so the owner/drain thread stays alive, as in production)."""
def __init__(
self,
lines = (),
exit_code = None,
):
self.pid = 4242
self._lines = list(lines)
self._exit = exit_code # None == alive
self.returncode = exit_code
self.terminated = False
self.killed = False
self._done = threading.Event()
if exit_code is not None:
self._done.set()
@property
def stdout(self):
def _gen():
for ln in self._lines:
yield ln
self._done.wait() # hold the pipe open until the process ends
return _gen()
def poll(self):
return self._exit
def terminate(self):
self.terminated = True
self._exit = 0
self.returncode = 0
self._done.set()
def wait(self, timeout = None):
self._done.wait(timeout)
if self._exit is None:
self._exit = 0
self.returncode = 0
return self.returncode
def kill(self):
self.killed = True
self._exit = -9
self.returncode = -9
self._done.set()
class _Resp:
def __init__(
self,
status_code,
payload = None,
text = "",
bad_json = False,
):
self.status_code = status_code
self._payload = payload if payload is not None else {}
self.text = text
self._bad_json = bad_json
def json(self):
if self._bad_json:
raise ValueError("not json")
return self._payload
class _FakeClient:
def __init__(
self,
*,
get = None,
post = None,
):
self._get = get or (lambda url: _Resp(200, {}))
self._post = post or (lambda url, json: _Resp(202, {"id": "job1"}))
self.get_urls = []
self.post_calls = []
self.closed = False
def get(
self,
url,
timeout = None,
):
self.get_urls.append(url)
return self._get(url)
def post(
self,
url,
json = None,
timeout = None,
):
self.post_calls.append((url, json))
return self._post(url, json)
def close(self):
self.closed = True
@pytest.fixture
def patched(monkeypatch):
"""Neutralise process-lifetime side effects for the manager under test."""
monkeypatch.setattr(srv, "adopt_pid", lambda pid: None)
monkeypatch.setattr(srv, "forget_pid", lambda pid: None)
monkeypatch.setattr(srv, "child_popen_kwargs", lambda: {})
monkeypatch.setattr(srv, "windows_hidden_subprocess_kwargs", lambda: {})
return monkeypatch
def _server_with(popen, client):
s = SdCppServer("/x/sd-server")
s._client = client
# Attach the fake process + port so generation tests can run without start().
s._process = popen
s.port = 1234
return s
# ── start / readiness ──────────────────────────────────────────────────────────
def test_start_becomes_ready_when_capabilities_200(patched):
popen = _FakePopen(lines = ["loading model", "listening on: http://127.0.0.1:1"])
patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen)
s = _server_with(
popen, _FakeClient(get = lambda url: _Resp(200, {"model": {"path": "/m/z.gguf"}}))
)
s.start(_FILES, startup_timeout = 5.0)
assert s.is_alive() is True
assert s.port is not None
def test_start_fails_fast_when_process_exits(patched):
# Model load failed -> process exits before listening; start must raise with the tail.
popen = _FakePopen(lines = ["error: bad model"], exit_code = 1)
patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen)
# Capabilities never answers (connection refused) -> readiness relies on exit detection.
s = _server_with(
popen, _FakeClient(get = lambda url: (_ for _ in ()).throw(srv.httpx.ConnectError("refused")))
)
with pytest.raises(RuntimeError, match = "failed to become ready"):
s.start(_FILES, startup_timeout = 2.0)
# ── generation ───────────────────────────────────────────────────────────────
def _completed_job(images_b64):
return _Resp(
200,
{
"status": "completed",
"result": {"images": [{"index": i, "b64_json": b} for i, b in enumerate(images_b64)]},
},
)
def test_img_gen_returns_image_bytes_in_index_order(patched):
popen = _FakePopen()
s = _server_with(
popen,
_FakeClient(
post = lambda url, json: _Resp(202, {"id": "jobA"}),
# result images deliberately out of order -> manager must sort by index.
get = lambda url: _Resp(
200,
{
"status": "completed",
"result": {
"images": [
{"index": 1, "b64_json": _png_b64(200)},
{"index": 0, "b64_json": _png_b64(50)},
]
},
},
),
),
)
blobs = s.img_gen({"prompt": "x", "batch_count": 2, "sample_params": {"sample_steps": 4}})
assert len(blobs) == 2
first = Image.open(io.BytesIO(blobs[0])).convert("RGB").getpixel((0, 0))
assert first == (50, 50, 50) # index 0 first
def test_img_gen_failed_job_raises(patched):
popen = _FakePopen()
s = _server_with(
popen,
_FakeClient(
post = lambda url, json: _Resp(202, {"id": "jobF"}),
get = lambda url: _Resp(
200, {"status": "failed", "error": {"code": "x", "message": "boom"}}
),
),
)
with pytest.raises(RuntimeError, match = "generation failed.*boom"):
s.img_gen({"prompt": "x"})
def test_img_gen_queue_full_raises(patched):
popen = _FakePopen()
s = _server_with(popen, _FakeClient(post = lambda url, json: _Resp(429, text = "busy")))
with pytest.raises(RuntimeError, match = "queue is full"):
s.img_gen({"prompt": "x"})
def test_img_gen_cancel_posts_cancel_and_raises(patched):
popen = _FakePopen()
cancel = threading.Event()
cancel.set() # already cancelled before the first poll
client = _FakeClient(
post = lambda url, json: _Resp(202, {"id": "jobC"}),
get = lambda url: _Resp(
200, {"status": "cancelled", "error": {"code": "cancelled", "message": "c"}}
),
)
s = _server_with(popen, client)
with pytest.raises(SdCppCancelled):
s.img_gen({"prompt": "x"}, cancel_event = cancel)
assert any(url.endswith("/cancel") for url, _ in client.post_calls)
def test_img_gen_detects_server_death(patched):
popen = _FakePopen()
def _die_get(url):
popen._exit = 137 # the process died between submit and poll
return _Resp(200, {"status": "generating"})
s = _server_with(
popen, _FakeClient(post = lambda url, json: _Resp(202, {"id": "jobD"}), get = _die_get)
)
with pytest.raises(RuntimeError, match = "connection lost|process exited"):
s.img_gen({"prompt": "x"})
# ── stdout routing + stop ──────────────────────────────────────────────────────
def test_drain_routes_lines_to_step_listener_and_tail(patched):
s = SdCppServer("/x/sd-server")
seen = []
s._step_listener = seen.append
# exit_code set so stdout ends after the scripted lines (a live fake would block).
s._drain_stdout(_FakePopen(lines = ["sampling 1/8", "", "sampling 8/8", "done"], exit_code = 0))
assert "sampling 1/8" in seen and "sampling 8/8" in seen
assert "" not in seen # blank lines skipped
assert s._tail[-1] == "done"
def test_stop_is_idempotent_and_terminates(patched):
popen = _FakePopen()
patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen)
client = _FakeClient(get = lambda url: _Resp(200, {}))
s = _server_with(popen, client)
s.start(_FILES, startup_timeout = 5.0)
s.stop()
assert popen.terminated is True
assert s.is_alive() is False
assert client.closed is True # stop() releases the pooled HTTP client
s.stop() # second call must not raise
def test_img_gen_submit_error_raises(patched):
popen = _FakePopen()
s = _server_with(popen, _FakeClient(post = lambda url, json: _Resp(400, text = "bad params")))
with pytest.raises(RuntimeError, match = "submit -> 400"):
s.img_gen({"prompt": "x"})
def test_img_gen_malformed_submit_json_raises(patched):
popen = _FakePopen()
s = _server_with(popen, _FakeClient(post = lambda url, json: _Resp(202, bad_json = True)))
with pytest.raises(RuntimeError, match = "non-JSON submit"):
s.img_gen({"prompt": "x"})
def test_img_gen_empty_result_raises(patched):
popen = _FakePopen()
s = _server_with(
popen,
_FakeClient(
post = lambda url, json: _Resp(202, {"id": "jobE"}),
get = lambda url: _Resp(200, {"status": "completed", "result": {"images": []}}),
),
)
with pytest.raises(RuntimeError, match = "no images"):
s.img_gen({"prompt": "x"})
def test_img_gen_rejected_after_stop(patched):
popen = _FakePopen()
patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen)
s = _server_with(popen, _FakeClient(get = lambda url: _Resp(200, {})))
s.start(_FILES, startup_timeout = 5.0)
s.stop()
with pytest.raises(RuntimeError, match = "not running"):
s.img_gen({"prompt": "x"})
# ── cancellation + defensive parsing (review follow-ups) ───────────────────────
def test_img_gen_cancelled_before_submit_reports_cancellation(patched):
# The server was stopped for a cancel/unload before submit; with the cancel event set
# this must surface as a cancellation (route -> 409), not a generic "not running" 500.
popen = _FakePopen()
patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen)
s = _server_with(popen, _FakeClient(get = lambda url: _Resp(200, {})))
s.start(_FILES, startup_timeout = 5.0)
s.stop()
cancel = threading.Event()
cancel.set()
with pytest.raises(SdCppCancelled):
s.img_gen({"prompt": "x"}, cancel_event = cancel)
def test_img_gen_abandons_when_cancel_not_honored(patched):
# A best-effort cancel the server ignores must not pin this call (and the generate
# lock) until natural completion: after the grace window it raises cancellation.
patched.setattr(srv, "_CANCEL_GRACE_S", 0.0)
popen = _FakePopen()
cancel = threading.Event()
cancel.set()
client = _FakeClient(
post = lambda url, json: _Resp(202, {"id": "jobG"}),
get = lambda url: _Resp(200, {"status": "generating"}), # never terminal
)
s = _server_with(popen, client)
with pytest.raises(SdCppCancelled):
s.img_gen({"prompt": "x"}, cancel_event = cancel, poll_interval = 0.01)
def test_img_gen_non_dict_submit_json_raises(patched):
popen = _FakePopen()
s = _server_with(popen, _FakeClient(post = lambda url, json: _Resp(202, ["not", "a", "dict"])))
with pytest.raises(RuntimeError, match = "unexpected submit response"):
s.img_gen({"prompt": "x"})
def test_img_gen_non_dict_status_json_raises(patched):
popen = _FakePopen()
s = _server_with(
popen,
_FakeClient(
post = lambda url, json: _Resp(202, {"id": "jobH"}),
get = lambda url: _Resp(200, ["unexpected"]),
),
)
with pytest.raises(RuntimeError, match = "unexpected response type"):
s.img_gen({"prompt": "x"}, poll_interval = 0.01)
def test_decode_images_tolerates_unexpected_shapes():
# A misbehaving/older server can return non-dict result/images/items; _decode_images
# must raise a clean "no images" rather than an AttributeError on .get().
for job in ({"result": ["x"]}, {"result": {"images": "nope"}}, {"result": {"images": [1, 2]}}):
with pytest.raises(RuntimeError, match = "no images"):
SdCppServer._decode_images(job)
def test_start_aborted_by_concurrent_stop(patched):
# A stop() during the readiness wait must abort start() promptly (without waiting out
# the startup timeout) and surface as a cancellation.
popen = _FakePopen(lines = ["loading model"])
patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen)
def _never_ready(url):
raise srv.httpx.ConnectError("refused")
s = _server_with(popen, _FakeClient(get = _never_ready))
def _stop_soon():
import time as _t
_t.sleep(0.2)
s.stop()
threading.Thread(target = _stop_soon, daemon = True).start()
with pytest.raises(SdCppCancelled):
s.start(_FILES, startup_timeout = 30.0)

18
studio/install_sd_cpp_prebuilt.py
View file

@ -140,6 +140,17 @@ def _locate_sd_cli(root: Path) -> Optional[Path]:
return None
def _locate_sd_server(root: Path) -> Optional[Path]:
"""The persistent ``sd-server`` binary in the extracted tree, if the archive ships
one (modern stable-diffusion.cpp releases do). Best-effort: the native backend
falls back to one-shot ``sd-cli`` when it is absent."""
name = "sd-server.exe" if sys.platform == "win32" else "sd-server"
for p in root.rglob(name):
if p.is_file():
return p
return None
def _download(
url: str,
dest: Path,
@ -237,6 +248,13 @@ def install(
if sys.platform != "win32":
_make_executable(sd_cli)
print(f"installed sd-cli -> {sd_cli}", flush = True)
# The same archive ships the persistent sd-server; make it runnable too so the
# native backend can prefer it (load once, serve many) over one-shot sd-cli.
sd_server = _locate_sd_server(target)
if sd_server is not None and sys.platform != "win32":
_make_executable(sd_server)
if sd_server is not None:
print(f"installed sd-server -> {sd_server}", flush = True)
return sd_cli