Commit graph

325 commits

Author SHA1 Message Date
Daniel Han
1d289aab61 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.
2026-06-28 06:02:12 +00:00
pre-commit-ci[bot]
ca52807680 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-26 12:36:49 +00:00
Daniel Han
efc7a44747 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).
2026-06-26 12:35:50 +00:00
Daniel Han
764a3e1dc9 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).
2026-06-26 12:09:16 +00:00
pre-commit-ci[bot]
c7e30e9ea8 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-26 12:03:36 +00:00
Daniel Han
b923675549 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.
2026-06-26 12:02:27 +00:00
pre-commit-ci[bot]
520b80a082 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-26 11:24:35 +00:00
Daniel Han
b90f833469 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.
2026-06-26 11:23:20 +00:00
Daniel Han
3a21f12500 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).
2026-06-26 08:56:23 +00:00
pre-commit-ci[bot]
8c3554e38c [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-26 08:01:47 +00:00
Daniel Han
10db6a4777 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).
2026-06-26 08:01:15 +00:00
pre-commit-ci[bot]
3cfb29dd6b [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-26 06:17:18 +00:00
Daniel Han
983f2c14a0 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.
2026-06-26 06:15:18 +00:00
Daniel Han
ede94176f6 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).
2026-06-26 05:04:10 +00:00
pre-commit-ci[bot]
141cb3ae10 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-26 03:20:02 +00:00
Daniel Han
395816cf7e 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.
2026-06-26 03:18:44 +00:00
pre-commit-ci[bot]
f748d97702 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-25 16:08:55 +00:00
Daniel Han
703d2df687 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.
2026-06-25 16:08:08 +00:00
pre-commit-ci[bot]
33977e0991 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-25 15:51:26 +00:00
Daniel Han
a7b8f825da Studio diffusion (Phase 4): native stable-diffusion.cpp engine for CPU/Mac
Adds the CPU / Apple-Silicon tier of the two-engine strategy, mirroring the
chat backend's llama.cpp shell-out. Diffusers stays the default on CUDA / ROCm
/ XPU; this covers the hardware diffusers serves poorly, consuming the same
split GGUF assets Studio already curates.

- sd_cpp_args.py: pure sd-cli command builder. Maps the family to its
  text-encoder flag (Z-Image Qwen3 to --llm, Qwen-Image to --qwen2vl, FLUX.1
  CLIP-L + T5), and the diffusers memory policy (none/group/model/sequential)
  to sd.cpp's offload flags (--offload-to-cpu / --clip-on-cpu / --vae-on-cpu /
  --vae-tiling / --diffusion-fa), so one user knob drives both engines.
- sd_cpp_engine.py: SdCppEngine over a located sd-cli. find_sd_cpp_binary()
  with the same precedence as the llama finder (env override, then the Studio
  install root, then in-tree, then PATH), an is_available/version probe, and a
  one-shot subprocess generate that streams progress and returns the PNG.
  runtime_env() prepends the binary's directory to the platform library path
  so a prebuilt's bundled libstable-diffusion.so resolves.
  select_diffusion_engine() is the pure routing decision (GPU backends to
  diffusers, CPU/MPS to native when present).
- install_sd_cpp_prebuilt.py: resolve + download the per-host prebuilt
  (macOS-arm64/Metal, Linux x86_64 CPU, Vulkan/ROCm/Windows variants) into the
  Studio install root. resolve_release_asset() is a pure, unit-tested
  host-to-asset matrix.
- scripts/sd_cpp_smoke.py: end-to-end native generation harness.

Tests (CPU-only, subprocess/filesystem stubbed): 49 new across args, engine,
routing, runtime env, and the installer resolver. Full diffusion suite 166
passing.

Verified on a B200 box: built sd-cli (CUDA) and the prebuilt (CPU) both
generate Z-Image-Turbo Q4_K end to end through SdCppEngine: balanced (group
offload, 5.0s gen), low_vram (full CPU offload + VAE tiling, 13.4s), and the
dynamically-linked CPU prebuilt (50.4s on CPU), all producing coherent images.
2026-06-25 15:50:10 +00:00
pre-commit-ci[bot]
7ca5573498 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-25 14:43:45 +00:00
Daniel Han
dbb0292561 Studio diffusion (Phase 2C): NVFP4 text-encoder quant (+ generalise fp8 knob)
Generalise the text-encoder precision knob from a fp8 bool to text_encoder_quant
(fp8 | nvfp4). nvfp4 quantises the companion text encoder to 4-bit via torchao
NVFP4 weight-only (two-level microscaling) on Blackwell's FP4 tensor cores; fp8
stays the broader-hardware path (cc>=8.9). Both are gated, best-effort, and run
before placement; status reports the mode actually engaged. This is the lean
realisation of GGUF-native text-encoder quant: 4-bit on the encoder without the
3045-line port.

Verified on Z-Image (B200, balanced/group where the encoder stays resident), vs the
bf16 encoder: nvfp4 cut generation peak VRAM 48% (10840 -> 5593 MB, the lowest TE
option, below whole-model offload) at near-fp8 quality (16.4 vs 17.1 dB PSNR), and
both quants ran faster than bf16. A memory-vs-quality tradeoff (off by default);
size it per model with the Phase 5 quality harness. diffusion_bench gains
--text-encoder-quant.

129 CPU tests pass.
2026-06-25 14:42:54 +00:00
pre-commit-ci[bot]
54eea13036 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-25 13:56:37 +00:00
Daniel Han
817e14ad19 Studio diffusion (Phase 2B): opt-in fp8 text-encoder layerwise casting
Add a text_encoder_fp8 knob that casts the companion text encoder(s) to fp8 (e4m3)
storage via diffusers apply_layerwise_casting, upcasting per layer to the bf16
compute dtype while normalisations and embeddings stay full precision. Applied
before placement, gated to CUDA + bf16, best-effort (a failure leaves the encoder
dense). status reports which encoders were cast.

Verified on Z-Image (B200, balanced/group mode where the encoder stays resident):
generation peak VRAM dropped 37% (10840 -> 6791 MB, below the lowest-VRAM offload)
at near-resident speed. It is a memory-vs-quality tradeoff, not free -- ~20 dB PSNR
vs the bf16 encoder, a larger shift than one transformer quant step -- so it is off
by default and documented as such, with the Phase 5 harness to size the cost.

127 CPU tests pass.
2026-06-25 13:55:08 +00:00
Daniel Han
9de8684a98 Studio diffusion (Phase 3): opt-in speed layer (channels_last / compile / TF32)
Add a speed_mode knob (off by default, so the render path stays bit-identical):
default applies channels_last VAE + regional torch.compile of the denoiser's
repeated block where eligible; max also enables TF32 matmul and fused QKV. Regional
compile is gated off for the GGUF transformer (dequantises per-op) and for families
flagged not compile-friendly (a new supports_torch_compile flag, False for Z-Image),
so it activates automatically only once a non-GGUF bf16 transformer is loaded. Speed
optims run before placement/offload, per the diffusers composition order. status now
reports speed_mode + the optims actually engaged.

Verified on Z-Image (B200): default -> ['channels_last'], max -> ['channels_last',
'tf32'], compile correctly skipped for GGUF; generation works in every mode.

121 CPU tests pass.
2026-06-25 13:55:08 +00:00
Daniel Han
2f00c77e75 Studio diffusion (Phase 2D): streamed block-level offload + functional VAE tiling
Add a streamed 'group' offload tier (diffusers apply_group_offloading, block_level,
use_stream) that keeps the transformer flowing through the GPU a few blocks at a
time while the text encoder / VAE stay resident, and fix VAE tiling to drive the
VAE submodule (pipelines like Z-Image expose enable_tiling on pipe.vae, not the
pipeline). apply_memory_plan now returns the (policy, tiling) actually engaged so
status never overstates either, and group falls back to whole-module offload when
the transformer can't be streamed.

Measured on Z-Image (B200), all lossless (PSNR inf vs resident): balanced/group
cuts generation peak VRAM 32% (15951 -> 10840 MB) at near-resident speed (2.07 ->
2.99s); low_vram/model cuts it 48% (-> 8318 MB) but is slower (7.99s). Mode names
now match that tradeoff: balanced = stream the transformer, low_vram = offload
every component. auto picks group when the companions fit resident, else model.

112 CPU tests pass.
2026-06-25 13:55:08 +00:00
Daniel Han
3a0bd55bdf Studio diffusion (Phase 2A): measured-budget memory planner + offload/VAE policy
Add a lean, backend-agnostic memory policy that picks a CPU-offload policy and
VAE tiling/slicing from measured free device memory vs the model's estimated
resident footprint, then applies it to the built pipeline. auto stays resident
when the model fits (byte-identical to the prior resident path), and falls to
whole-module offload when tight; fast/balanced/low_vram are explicit overrides.
Sequential submodule offload is unreliable for GGUF transformers on diffusers
0.38, so it falls back to whole-module offload and status reports the policy
actually engaged.

Verified on Z-Image-Turbo Q4_K_M (B200): auto reproduces the resident image with
no VRAM/latency regression (PSNR inf); balanced/low_vram cut generation peak VRAM
47.9% (15951 -> 8318 MB) with byte-identical output, at the expected latency cost.

73 prior + 35 new CPU tests pass.
2026-06-25 13:55:08 +00:00
pre-commit-ci[bot]
7f6e69dd5c [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-25 11:30:41 +00:00
Daniel Han
8ef51d744b Studio diffusion: cross-platform device policy, fp16 guard, lock split, validate-before-evict
Phase 1 of porting the richer diffusion stack onto the image-generation backend.

- Add a compartmentalized device/dtype policy module (diffusion_device.py)
  resolving CUDA/ROCm/XPU/MPS/CPU with capability flags. Keeps the NVIDIA
  capability-based bf16 choice; ROCm and XPU are isolated; MPS uses bf16 or
  fp32, never a silent fp16 that renders a black image.
- Add a per-family fp16_incompatible flag (Z-Image) and promote a resolved
  float16 to float32 for those families so they do not produce black images.
- Split the backend locks: a generation holds only _generate_lock, so status,
  unload, and a new load are never blocked by a long denoise. Add per-generation
  cancellation via callback_on_step_end so an eviction or a superseding load
  preempts a running generation; a replacement load waits for it to stop before
  allocating, so two pipelines never sit in VRAM at once.
- Validate a load request before the GPU handoff so an unloadable pick never
  evicts a working chat model, and reject missing local paths up front.
- Add CPU-only tests for the device policy, dtype guard, lock split and
  cancellation, and validate-before-evict, plus a GPU benchmark/regression
  script (scripts/diffusion_bench.py) measuring latency, peak VRAM, and PSNR
  against a saved reference.
2026-06-25 11:14:39 +00:00
pre-commit-ci[bot]
39eb022b52 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-25 04:11:49 +00:00
oobabooga
002b1216ee Address Codex review: FP16 fallback on pre-Ampere, forward HF token, hide edit models, download outside the load lock 2026-06-25 01:08:18 -03:00
pre-commit-ci[bot]
0a2a423b08 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-25 03:21:55 +00:00
oobabooga
2a5de505df Cancel diffusion loads on unload, mask seeds to JS-safe range, skip malformed gallery records 2026-06-24 23:30:29 -03:00
oobabooga
818ddb828f Add FLUX.2-klein-9B to the curated image models 2026-06-24 22:31:57 -03:00
oobabooga
6411f9464a Add FLUX.2-klein-4B via Flux2KleinPipeline 2026-06-24 22:09:59 -03:00
oobabooga
5d80d30168 Support Qwen-Image and FLUX.1 image models alongside Z-Image 2026-06-24 21:42:57 -03:00
oobabooga
5c4d09a1c9 Paginate the image gallery with infinite scroll 2026-06-24 20:48:07 -03:00
oobabooga
2bf9d73b4a Reuse formatEta, compute generation ETA once per step, guard re-renders 2026-06-24 19:50:10 -03:00
oobabooga
54d01a505f Add a live per-step generation progress bar with ETA 2026-06-24 19:41:15 -03:00
oobabooga
a87e092f59 Drop dead PNG-base64 helper and dedupe the image-gen task list 2026-06-24 19:17:49 -03:00
oobabooga
eced3620fa Add batch size and sequential count to image generation with reproducible per-run seeds 2026-06-24 18:02:38 -03:00
oobabooga
caa7efecc0 Persist the image gallery to disk with recipes embedded in each PNG, and match official Z-Image defaults 2026-06-24 17:50:11 -03:00
oobabooga
adebdfc9f4 Fix image load progress, hide diffusion models from chat picker, and gallery polish 2026-06-24 17:17:28 -03:00
oobabooga
453efa7ae7 drop unused diffusion status 'loading' field; cap session gallery 2026-06-24 16:18:14 -03:00
oobabooga
425d61e47a Z-Image-Turbo image generation with chat-style model picker 2026-06-24 15:32:56 -03:00
oobabooga
e8d3cfe2b5 Merge branch 'main' into image-generation 2026-06-24 14:07:59 -03:00
oobabooga
ab6c9ecfee
Studio: honor stream=false on the GGUF agentic tool path (#6570) (#6618)
* Studio: honor stream=false on the GGUF agentic tool path (#6570)

* Studio: dedup the #6570 non-streaming tool tests and cover cached_tokens

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

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

* Studio: cover the cached_tokens metadata fix and clarify the drain comment (#6570)

* Studio: align the GGUF tool drain naming and tighten its comment (#6570)

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
2026-06-24 15:37:08 +01:00
Nilay
e5cf956601
Studio: shareable per-checkpoint preview links (#6486)
* checkpoint preview endpoint

* harden new preview endpoints

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

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

* address review

* Studio preview: pin adapter, guard streaming submit, robust copy-link

Harden the public per-checkpoint preview surface:

- Pin use_adapter=True in the preview payload sanitizer. Otherwise an
  unauthenticated /p caller can POST use_adapter=false, which calls
  disable_adapter_layers() on the shared in-memory model without restoring
  it; since load_model skips reloads for the same checkpoint, every later
  visitor (the page never sends the field) keeps getting base-model output
  instead of the fine-tuned checkpoint. Forcing it on also re-enables a
  previously disabled adapter and no-ops on merged checkpoints.
- Ignore preview-page submits while a response is streaming. The send
  button was disabled but the Enter handler still called requestSubmit(),
  so a second request could start before the first reply landed in msgs and
  reorder the chat history. Both the keydown and submit handlers now honor
  the disabled button.
- Keep the cloudflare-URL polling loop alive across transient startup fetch
  errors instead of letting one rejection halt it.
- Build the copy-link from a backend preview_ref (output dir relative to
  outputs_root, gated on previewability and the two-segment /p route limit)
  so a nested output dir no longer copies a basename-only link that 404s.
  Expose preview_ref on training run summaries.

Add route-level security tests (path traversal, payload sanitization,
asset containment, CSP header, HTML title escaping, streaming lock held
until drained) and preview_ref unit tests.

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

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

* Studio preview: Safari-safe submit and adapter pin only for LoRA

Follow-ups from cross-browser and route simulations:

- Preview page: send the message from a shared send() helper called by both
  the form submit and the Enter key, instead of form.requestSubmit(). The
  latter throws on Safari < 16 and older iOS, which broke Enter-to-send there.
  Verified across Chromium, Firefox and WebKit with Playwright.
- Only pin use_adapter=True when the resolved checkpoint is a LoRA adapter
  (adapter_config.json present); for a merged checkpoint strip it to None.
  A merged model has no adapter to toggle, so forcing it on only produced a
  per-request "not a PeftModel" warning. The cross-request base-model
  contamination fix still holds for LoRA previews.

Add a merged-checkpoint test asserting use_adapter is stripped to None.

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

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

* Studio preview: trim verbose comments

Tighten comments across the preview routes, page, checkpoint helpers, and tests
to short single-line notes; drop ones that just restate the code. No behavior
change (verified comment/docstring-only with comment_tools.py check).

* Harden preview routes for PR #6486

- Return a generic 400 detail on a rejected preview path so the public /p
  route never echoes the absolute install path (the real reason is logged
  server-side instead).
- Strip confirm_tool_calls, session_id and rag_scope in the preview payload
  sanitizer so the public surface stays inert regardless of the tool gate.
- Use Path.is_relative_to for the asset containment check, matching the rest
  of the codebase.
- Add img-src 'self' and font-src 'self' to the preview page CSP.
- Preview page: on a mid-stream error keep the streamed text, flag the break,
  and restore the prompt so the user can retry; drop the unused --font-sans var.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-06-24 06:31:53 -07:00
oobabooga
346d96d7f2
Studio: cap GGUF context to unified memory on Apple Silicon (#6622)
* Studio: cap GGUF context to unified memory on Apple Silicon

* Studio: tighten Apple ctx-cap comments and drop the overstated MLX-sync claim

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

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

* Studio: reserve flat MTP fraction and floor sparse-KV ctx in the Apple unified-memory cap

The Apple Silicon GGUF context cap mirrored the discrete-GPU auto-fit branch but
missed two protections the discrete path already applies:

- It passed the full unified-memory budget with budget_frac=1.0 without first
  reserving the flat MTP fraction the discrete path takes off via _pin_fraction.
  With an MTP draft whose KV cannot be byte-sized (e.g. Qwen3.6-MTP, #6529), the
  cap filled the whole budget and left nothing for the draft, so unified memory
  could still over-commit. Reserve _flat_mtp_reserve up front; this is a no-op
  when MTP is not engaged.

- It required _can_estimate_kv(), so a GGUF with sparse KV metadata skipped the
  cap entirely and launched at full native context. Mirror the discrete
  file-size-only fallback and floor the auto context to 4096 when the cache
  cannot be sized.

Adds regression tests for both paths.

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

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

* Studio: tighten comments in the Apple unified-memory context cap

Condense the verbose comment blocks in the Apple budget helper, the no-GPU
Metal branch, and the context-fit tests. Comments only, no code change
(verified with ast-based comment_tools check); suite still green.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <danielhanchen@gmail.com>
2026-06-24 04:39:33 -07:00
Saicharan Ramineni
9d53656614
Make _uv_safe_path space-safe on macOS/Linux (#6503) (#6534)
* Copy uv `-c`/`-r` paths to a space-free temp dir on macOS/Linux

uv 0.11.x truncates a constraints/requirements path passed via `-c`/`-r`
at the first space, so `unsloth studio` setup from a repo cloned under a
path containing a space (e.g. `/Users/me/Open Source/unsloth`) fails with:

    error: File not found: `/Users/me/Open`

_uv_safe_path() already worked around this on Windows via the 8.3 short
path but returned the space-containing path unchanged on macOS/Linux,
which have no 8.3 equivalent. Extend it to copy the (small, flat)
constraints/requirements file into a space-free temp dir and hand uv the
copy; the temp dirs are removed at process exit. Falls back to the
original path on any error, so it is never worse than before.

Refs unslothai/unsloth#6503

* Route UV_OVERRIDE through _uv_safe_path and fix temp-dir leak (#6503)

The -c/-r fix did not cover UV_OVERRIDE, which uv also truncates at the first
space. On Apple Silicon the overrides file is handed to uv via UV_OVERRIDE at
install time (install_python_stack.py) and during the MLX self-heal
(utils.mlx_repair), so a repo under a path containing a space still broke every
uv call there. Move _uv_safe_path into backend.utils.uv_path_safety so both
sites share it, and route UV_OVERRIDE through it.

Also stop leaking the temp dir when shutil.copyfile fails after mkdtemp, and add
tests for the UV_OVERRIDE channel, the TMPDIR-with-space fallback, the atexit
cleanup, and the no-leak path.

---------

Co-authored-by: danielhanchen <danielhanchen@gmail.com>
2026-06-24 04:02:24 -07:00