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.
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).
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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
* 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>
ensure_diffusion_visual_server() downloaded the visual-server release
asset with the unverified download_file() and marked it executable,
bypassing the approved-checksum manifest that gates every other prebuilt
llama.cpp artifact. The backend later auto-discovers that binary and
launches it through DG_VISUAL_BIN, so a compromised or substituted
release asset could place attacker-controlled native code in the install
tree and have it executed under the Studio user.
Require the matched asset to be present in the approved checksum manifest
and download it through download_file_verified() with the published
sha256. A name-matching asset that is absent from the manifest is refused
rather than executed.
Add regression tests covering the verified-download path and the refusal
of an unapproved asset.
