Commit graph

139 commits

Author SHA1 Message Date
Daniel Han
cdfc6f0f56 ControlNet: address review findings on the diffusers path
- resolve_controlnet enforces catalog family compatibility so a direct API call
  cannot load a ControlNet built for another family through the wrong pipeline.
- Unknown ControlNet ids now surface as a 400 (call site maps FileNotFoundError
  to ValueError) instead of a generic 500.
- strength 0 disables ControlNet entirely, so a no-op selection never pays the
  download / VRAM cost; the control image is decoded and validated BEFORE the
  ControlNet is resolved or built, so a malformed image fails fast for the same reason.
- ControlNet loads use the base compute dtype (state.dtype is a display string,
  not a torch.dtype, so it silently fell back to float32) and honor the base
  offload policy via group offloading instead of forcing the module resident.
- Empty/malformed HF token coerced to anonymous access.
- Flux Union ControlNet control_mode mapped from the selected control type.
- resolve_controlnet drops the unused hf_token/cancel_event params.
- ControlNetSpec validates guidance_start <= guidance_end (clean 422).
- Images UI ControlNet Select shows its placeholder when nothing is selected.

Adds regression tests for family enforcement and the union control-mode map.
2026-07-02 01:08:02 +00:00
Daniel Han
355f54f6e4 Merge remote-tracking branch 'origin/diffusion-lora' into diffusion-controlnet 2026-07-01 11:57:56 +00:00
Daniel Han
9637e85311 Merge remote-tracking branch 'origin/diffusion-image-workflows' into diffusion-lora 2026-07-01 11:56:47 +00:00
Daniel Han
38ed3ce5b5 Merge remote-tracking branch 'origin/diffusion-phase16-native-engine-routing' into diffusion-image-workflows
# Conflicts:
#	studio/backend/core/inference/diffusion.py
#	studio/backend/core/inference/diffusion_families.py
#	studio/backend/tests/test_sd_cpp_install.py
#	studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
#	studio/frontend/src/features/images/api.ts
#	studio/frontend/src/features/images/images-page.tsx
#	studio/install_sd_cpp_prebuilt.py
2026-07-01 11:48:33 +00:00
Daniel Han
91edf06748 Merge branch 'diffusion-phase15-int8-prequant' into diffusion-phase16-native-engine-routing
# Conflicts:
#	studio/backend/routes/inference.py
2026-07-01 11:31:54 +00:00
Daniel Han
2a5713aff5 Merge branch 'diffusion-phase11-consumer-int8' into diffusion-phase12-fbcache 2026-07-01 11:30:26 +00:00
Daniel Han
b84d7ea7f3 Merge branch 'diffusion-phase9-prequant' into diffusion-phase10-attention
# Conflicts:
#	studio/backend/tests/test_diffusion_routes.py
2026-07-01 11:30:23 +00:00
Daniel Han
4fec0bb79c Merge branch 'diffusion-phase8-quant' into diffusion-phase9-prequant
# Conflicts:
#	scripts/compare_engines.py
2026-07-01 11:29:36 +00:00
Daniel Han
5d86a7685a Merge branch 'diffusion-phase7-perf' into diffusion-phase8-quant
# Conflicts:
#	studio/backend/core/inference/diffusion.py
2026-07-01 11:29:00 +00:00
Daniel Han
48628252bd Merge remote-tracking branch 'origin/image-generation' into diffusion-phase4-native
# Conflicts:
#	scripts/diffusion_bench.py
#	scripts/diffusion_quality.py
#	studio/backend/core/inference/diffusion.py
#	studio/backend/core/inference/diffusion_device.py
#	studio/backend/core/inference/diffusion_families.py
#	studio/backend/core/inference/diffusion_memory.py
#	studio/backend/core/inference/diffusion_precision.py
#	studio/backend/core/inference/diffusion_speed.py
#	studio/backend/models/inference.py
#	studio/backend/routes/inference.py
#	studio/backend/tests/test_diffusion_backend.py
#	studio/backend/tests/test_diffusion_device.py
#	studio/backend/tests/test_diffusion_memory.py
#	studio/backend/tests/test_diffusion_precision.py
#	studio/backend/tests/test_diffusion_speed.py
2026-07-01 11:19:15 +00:00
Daniel Han
bce892a7f7 Studio diffusion: ControlNet for the Images workflow (diffusers)
Add ControlNet conditioning, the #2 most-used diffusion workflow after
LoRA, on the diffusers backend for the families with ControlNet pipelines
(FLUX.1 and Qwen-Image), with Union models as the default picks.

Backend
- New core/inference/diffusion_controlnet.py: family-gated discovery
  (curated Union models + local dirs + bare owner/name repos), resolution
  to a loadable repo/dir, control-image preprocessing (passthrough +
  a dependency-free canny edge map), and a supports_controlnet gate.
- diffusion.py: a ControlNet manager parallel to the LoRA one. Loads the
  (small) ControlNet model once via from_pretrained (cached by id) and
  builds the family's ControlNet pipeline via Pipeline.from_pipe(base,
  controlnet=model), reusing the resident base modules at their loaded
  dtype (no reload, no recast). Passes the control image + conditioning
  scale + guidance start/end at generate time; cleared on unload.
- Families: FLUX.1 -> FluxControlNetPipeline/Model, Qwen-Image ->
  QwenImageControlNetPipeline/Model. Others declare none (gated off).
- Gated off for the native engine, GGUF-via-diffusers, and torchao
  fp8/int8 dense (same rule as LoRA). v1 conditions txt2img only.
- Request contract: optional controlnet on DiffusionGenerateRequest;
  supports_controlnet in status; the choice persisted in gallery meta.
- New GET /api/models/diffusion-controlnets for the picker.

Frontend
- A ControlNet control in the Images rail (model select + control-image
  upload + control-type select + strength slider), gated by the loaded
  model's supports_controlnet + family, shown for text-to-image.

Tests
- New test_diffusion_controlnet.py (10): discovery/resolve/preprocess/gate
  helpers, request validation, family wiring, and the diffusers pipe
  manager (loads once, caches, from_pipe with controlnet, rejects
  unsupported families).
2026-07-01 05:06:02 +00:00
Daniel Han
9ac847b9c0 Studio diffusion: LoRA adapters for the Images workflow
Add community LoRA support across both diffusion backends, the single
biggest step toward broad image-workflow coverage.

Backend
- New shared module core/inference/diffusion_lora.py: adapter discovery
  (local scan + curated catalog + owner/name[:file] Hub refs), download
  via hf_hub_download_with_xet_fallback, alias sanitization, native
  managed-dir materialization with collision-broken aliases, prompt-tag
  injection (deduped against user-typed tags), and a supports_lora gate.
- Native sd-cli: resolve + materialize selected LoRAs into a per-run
  managed dir, inject <lora:ALIAS:w> tags, pass --lora-model-dir with
  --lora-apply-mode auto. The arg builder already emitted these flags.
- Diffusers: non-fused load_lora_weights + set_adapters manager, tracked
  on the pipe so an unchanged selection is a no-op and a model swap
  resets; cleared on unload. Never fuses (breaks quantized transformers
  and blocks live weight tweaks).
- Gated off where unsupported: torchao fp8/int8 dense, GGUF-via-diffusers,
  and native Qwen-Image (no LoRA name-conversion branch upstream).
- Request contract: optional loras on DiffusionGenerateRequest; empty or
  omitted is identical to today. supports_lora surfaced in status; chosen
  LoRAs persisted in gallery recipe metadata.
- New GET /api/models/diffusion-loras for the picker (family-filtered).

Frontend
- Repeatable multi-LoRA picker (adapter select + weight slider 0..2 +
  remove), gated by the loaded model's supports_lora and family, max 8.

Tests
- New test_diffusion_lora.py (14): helpers, request validation, native
  tag/dir wiring, diffusers set_adapters manager, supports_lora matrix.
2026-07-01 04:00:34 +00:00
Daniel Han
b14e2f9be7 Studio diffusion: image workflows (safetensors, image-conditioned, editing) + Images UI
Backend:
- Load non-GGUF safetensors models: full bnb-4bit pipelines and single-file
  fp8 transformers, gated to the unsloth org plus a curated allowlist.
- Image-conditioned workflows built with Pipeline.from_pipe so they reuse the
  loaded transformer/VAE/text-encoder with no extra VRAM: img2img, inpaint,
  outpaint, and a hires-fix upscale pass.
- Instruction editing as its own family kind (Qwen-Image-Edit-2511,
  FLUX.1-Kontext-dev) and FLUX.2-klein reference conditioning (single and
  multi-reference) plus klein inpaint.
- Auto-resize odd-sized inputs to a multiple of 16 (and resize the matched
  mask) so img2img/inpaint/edit no longer reject non-/16 uploads. Bound the
  decoded image size and cap upscale output to avoid OOM on large inputs.
- Fixes: from_pipe defaulting to a float32 recast that crashed torchao
  quantized transformers; image-conditioned calls forcing the slider size
  onto the input image. Native sd.cpp engine rejects image-conditioned and
  reference requests it cannot serve.

Frontend:
- Redesigned Images page with capability-gated workflow tabs (Create,
  Transform, Inpaint, Extend, Upscale, Reference, Edit), a brush mask editor,
  client-side outpaint, and a multi-reference picker.
- Advanced options moved to a right-docked panel mirroring Chat: closed by
  default, toggled by a single fixed top-bar button that stays in place.

sd.cpp installer: pin the release, verify each download's sha256, add a
download timeout, and make the source repo configurable for a future mirror.
2026-07-01 01:23:52 +00:00
Daniel Han
6ae6fb1c45
Studio diffusion (Phase 2): memory planner, streamed offload, fp8 TE, speed layer, quality harness (#6675)
---------

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

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-06-30 16:33:47 -03:00
Daniel Han
f0ba9fa2e1 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).
2026-06-29 10:38:17 +00:00
Daniel Han
cf5f23996f 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.
2026-06-28 05:58:18 +00:00
Daniel Han
4c8137a7af 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.
2026-06-28 05:51:34 +00:00
pre-commit-ci[bot]
86b3a5d2d1 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-28 04:31:39 +00:00
Daniel Han
7f3c206fa1 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.
2026-06-28 04:30:21 +00:00
oobabooga
0607351a81 Merge remote-tracking branch 'origin/main' into image-generation
# Conflicts:
#	studio/backend/routes/models.py
2026-06-26 23:23:02 -03:00
Avaya Aggarwal
cb274484a6
Add GGUF --tensor-parallel CLI option (#6561)
---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-06-26 15:30:10 -03: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
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
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
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
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
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]
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
5c4d09a1c9 Paginate the image gallery with infinite scroll 2026-06-24 20:48:07 -03:00
oobabooga
54d01a505f Add a live per-step generation progress bar with ETA 2026-06-24 19:41:15 -03:00
oobabooga
571d506c20 Refine image generation UI: aspect ratio sliders, info hints, action toolbar, and timestamped export filenames 2026-06-24 18:47:33 -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
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
Wasim Yousef Said
37166efcfc
Fix Gemma 4 GGUF OpenAI API streams (#6476)
* Fix Gemma 4 GGUF OpenAI API streams

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

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

* Avoid duplicate Responses stream disconnect watcher

* Keep reasoning-only Responses output hidden

* Address Gemma stream review comments

* Avoid Responses stream task-group cleanup

* Harden OpenAI chat completion streams

* Address OpenAI stream review issues

* Clean up Studio OpenAI stream helpers

* Fix Studio passthrough cold stream timeout

* Fix tool parser compatibility exports lint

* Preserve audio stream disconnect cancellation

* Avoid synthetic finish after passthrough errors

* Address stream cleanup and Gemma parser reviews

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

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

* Gemma 4: parse bare-string tool args and keep safetensors tools for native <|tool_call>

- Quote bare unquoted string values in Gemma native tool-call args (e.g.
  {location:Tokyo,unit:celsius}) so they parse; JSON scalars stay typed.
- Stop _detect_safetensors_features from suppressing supports_tools for
  templates that emit Gemma native <|tool_call>, which the shared parser
  now reads.
- Add tests for both.

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

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

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

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

* Harden Gemma tool-call parsing and stream-error detection

Address three issues in the Gemma-native tool-call path:

- _quote_gemma_object_keys stopped a bare (unquoted) string value at the
  first comma, so an argument like `location:New York, NY` was split
  mid-value and the synthesized JSON failed to parse, dropping the whole
  tool call. A bare value now ends only at `}` or a comma that begins the
  next `key:` pair.

- parse_tool_calls_from_text scanned the entire response for Gemma markers
  even inside a tool call already parsed from a `<tool_call>{...}` JSON
  block, so a marker-like string inside an argument (data) was promoted to
  a second, unintended tool call. Matches inside an already-consumed call
  span are now skipped.

- _openai_passthrough_stream relied on _monitor_openai_sse_line to flag a
  stream error, which returns early when monitor_id is None
  (skip_api_monitor), so an upstream error chunk left saw_stream_error
  unset and the synthetic-finish guard emitted a successful finish_reason
  after a failed stream. Error chunks are now detected independently of API
  monitoring.

Adds tests/test_gemma_tool_parse_edge_cases.py covering the comma and
marker-injection cases.

* Emit the terminal finish_reason chunk in GGUF streams

The OpenAI chat-completions GGUF tool stream and plain stream both built a
final ChatCompletionChunk carrying finish_reason but never yielded it, so
clients received the optional usage chunk and [DONE] with no chunk carrying
finish_reason. OpenAI-compatible consumers rely on that terminal choice to
distinguish stop/length/tool_calls. Yield it before the usage chunk and
[DONE], matching the other streaming paths.

* Parse tool calls in document order and skip nested markers both ways

Unify the JSON- and Gemma-format tool-call passes into a single
position-ordered scan:

- Calls are now emitted in byte order across both formats, so a mixed
  output like `<|tool_call>call:create{...}<tool_call|> ... <tool_call>
  {"name":"read",...}</tool_call>` executes create before read, matching
  the order they appear in (tools run in returned order).

- A candidate that starts inside an already-accepted call's span is
  skipped, in both directions: a JSON marker inside a Gemma argument and a
  Gemma marker inside a JSON argument are treated as data, not promoted to
  a second executable tool call.

Extends tests/test_gemma_tool_parse_edge_cases.py with the ordering and
JSON-in-Gemma nesting cases.

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

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

* Quote bare Gemma array elements; order finish before trailing usage

- _quote_gemma_object_keys skipped array values, so a Gemma call with a
  bare-string array argument like labels:[bug,ui] produced invalid JSON and
  the whole tool call was dropped. Array values are now scanned and bare
  string elements quoted, while numbers, quoted strings, and JSON literals
  are preserved.

- In the OpenAI passthrough stream, a trailing usage-only chunk
  (stream_options.include_usage) that arrived before any finish chunk was
  relayed before the synthetic finish, producing usage -> finish -> [DONE].
  Emit the synthetic finish before that usage chunk so the order matches the
  other streams (finish -> usage -> [DONE]).

Extends tests/test_gemma_tool_parse_edge_cases.py with the bare-array cases.

* Harden Gemma array parsing, XML-parameter guard, and stream teardown

Address five review findings on the Gemma tool-call and OpenAI passthrough
streaming paths:

- parse_tool_calls_from_text collected JSON and Gemma markers without the
  _inside_open_parameter guard, so a marker embedded in an existing
  <function=...><parameter=...> value was promoted to a separate tool call.
  Candidates that start inside an open XML parameter are now skipped, matching
  the guard the XML-style parser already applies.

- _quote_gemma_array_elements preserved array elements starting with { or [
  verbatim, so an array of objects (items:[{path:a}]) or a nested array failed
  json.loads and the whole call was dropped. Object and nested-array elements
  are now normalised recursively.

- _openai_passthrough_stream synthesized a finish chunk before a trailing
  usage-only chunk and set saw_finish_reason, which made the EOF guard skip the
  [DONE] sentinel. The EOF path now emits [DONE] whenever the upstream omitted
  it, even after a finish chunk was already synthesized.

- /generate/stream drove generation through asyncio.to_thread with no
  disconnect watcher, so a client disconnect during a long generation went
  unnoticed until the next send. It now runs _await_disconnect_then_cancel
  against the request, matching the other local streaming endpoints.

- _SameTaskStreamingResponse closed the body iterator with aclose() on a
  send-side disconnect, raising GeneratorExit so the generators' cancellation
  handlers (which finish the api_monitor entry) never ran. It now throws
  CancelledError, falling back to aclose() when athrow is unavailable.

Extends tests/test_gemma_tool_parse_edge_cases.py with array-of-objects,
nested-array, and marker-inside-XML-parameter cases.

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

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

* Watch disconnects on Anthropic streams; keep timestamps in Gemma values

Two follow-ups on the streaming and tool-parse paths:

- _anthropic_tool_stream and _anthropic_plain_stream drove generation through
  asyncio.to_thread(next, gen, ...) and only polled is_disconnected() between
  events, so a client disconnect during prefill or a long generation/tool step
  held the decode slot until the next event or a failed send. Both now run the
  _await_disconnect_then_cancel watcher used by the other local streams, stop it
  in finally, and break promptly when cancel_event is set.

- _GEMMA_NEXT_KEY_RE treated any comma followed by word-chars-then-colon as the
  next key, so a bare value such as "meet at 10:00, 11:00 tomorrow" was split
  into bogus keys. The next-key token must now be identifier-shaped (start with
  a letter or underscore), so a comma before a timestamp, ratio, or other
  numeric-then-colon text stays part of the value.

Adds a timestamp-in-bare-value regression test.

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

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

* Guard nested markers, reset on disconnect, clean unstarted streams

Three follow-ups on the tool-parse and streaming paths:

- parse_tool_calls_from_text only skipped markers that fell inside a span it
  had already parsed successfully, so when an unquoted Gemma argument contained
  a literal marker (code:<|tool_call>call:terminal{...}<tool_call|>) the outer
  object failed to normalize, its span was never recorded, and the inner marker
  was promoted to a standalone terminal call. Candidates nested inside any other
  candidate's brace span are now skipped regardless of whether the enclosing
  candidate parsed, so a marker in malformed outer data is never executed.

- /generate/stream skipped backend.reset_generation_state() when the disconnect
  watcher set cancel_event between chunks: the loop broke and the finally's reset
  is guarded on cancel_event being unset. A subprocess backend kept decoding
  after the client left. The cancel-break path now resets the backend.

- _SameTaskStreamingResponse threw CancelledError / called aclose() on the body
  iterator on a send-side disconnect, but neither runs the try/finally of a
  generator that never started (early disconnect on http.response.start), so the
  passthrough's eagerly-opened upstream httpx stream and cancel-registry entry
  leaked. It now tracks whether the body started and, when it did not, runs an
  optional unstarted_cleanup hook; the OpenAI passthrough wires it to close the
  upstream resp/client and exit the cancel tracker.

Adds a nested-unquoted-marker regression test.

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

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

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-06-23 06:13:56 -07:00
oobabooga
a1423e05d2 Studio: local diffusion image generation page 2026-06-23 00:39:36 -03:00
UmranPros
9e83399f9e
Studio: fix Gemma 4 separate-drafter MTP detection and fallback (#6459)
Recognise the Gemma 4 separate-drafter MTP family, auto-download the drafter with retry, fall back to n-gram with a clear reason when it cannot be resolved, and retry the download on reload. Gemma 3n (ships no drafter) and embedded-MTP models (Qwen) are unaffected.

Fixes #6406
2026-06-20 04:43:37 -07:00
Daniel Han
76a2b9edf1
Studio: Auto disables MTP for MLA models (GLM-5.2 et al.); UNSLOTH_MLA_MTP_ENABLED to re-enable (#6468)
* Studio: Auto disables MTP for MLA models (GLM-5.2 et al.); UNSLOTH_MLA_MTP_ENABLED to re-enable

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

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

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

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

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

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

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

---------

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

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

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

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

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

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

GLM-5.2 reasoning levels:

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

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

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

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

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

* Studio: address review feedback on reasoning effort and formatBytes

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

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

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

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

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

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

Review follow-ups on the enable_thinking_effort work:

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

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-18 10:10:01 -07:00