Commit graph

244 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
2cae1c1e7c Merge remote-tracking branch 'origin/main' into image-generation
# Conflicts:
#	studio/backend/routes/models.py
#	studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
2026-07-01 10:47:21 +00:00
Daniel Han
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
Anish Umale
d0f8d40c36
studio: allow updating HF models through UI (#5388)
* add models for /update endpoint

* add logic for identifying out of date hf models

* add endpoint for updating hf models

* add relevant field to GgufVariantDetail

* make exception handling better

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

* hook up /update endpoint on the frontend

* implement update scenarios for the model picker

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

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

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

* fix import and make hf calls async

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

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

* remove has_vision from UpdateRequest

* fix ci

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

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

* clear cancel event before updating gguf variant

* set _cancel_event back if it was set initially

* add hf_token to get_paths_info

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

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

* studio: harden model update endpoint and update checks

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

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

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

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

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

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

Adds regression tests for the multi-revision update check.

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

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

* Studio: accept force_download kwarg in hf_xet_fallback test double

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

* Fix Studio model update regressions

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

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

* Address Studio update review feedback

* Address Studio update edge cases

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

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

* Share GGUF update status helper

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

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

* Fix GGUF update detection and cache cleanup

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

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

* Fix cached GGUF update badges

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: shimmyshimmer <107991372+shimmyshimmer@users.noreply.github.com>
Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
2026-07-01 01:54:57 +03:00
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
43d3caf38b
Studio: imatrix GGUF option and FP8/NVFP4 compressed export in the export UI (#6729)
* Studio: wire imatrix GGUF option and FP8/NVFP4 compressed export into the export UI

GGUF export gains an importance-matrix toggle. When enabled it auto-downloads the
upstream Unsloth imatrix for the base model (or uses a custom path), which unlocks
the IQ low-bit quants iq2_xxs, iq2_m, iq3_xxs and iq4_xs. Merged export gains an
FP8 / NVFP4 compressed-tensors precision selector that runs llm-compressor for vLLM.

Backend threads imatrix_file through routes -> orchestrator -> worker -> export_gguf
(both the local save and the hub push), and maps the new compressed format_type
values onto the fp8/nvfp4 save_method, reporting the "<dir>-<suffix>" sibling output
directory. Frontend adds the imatrix Switch on the GGUF card and a merged precision
picker on the merged card, threaded through the export runtime store.

Depends on unslothai/unsloth#6706 (save.py imatrix_file and compressed-tensors
export) and unslothai/unsloth-zoo#839 (quantize_gguf imatrix flag).

* Studio export: guard imatrix/compressed against older unsloth builds and force imatrix for IQ quants

Addresses review feedback on the export wiring:
- GGUF: pass imatrix_file only when set, so a plain no-imatrix export (e.g. Q4_K_M) no
  longer fails with an unexpected-keyword error against an unsloth build that predates the
  imatrix_file parameter. When imatrix is requested but unsupported, return a clear
  upgrade message instead of a TypeError.
- Merged: gate FP8/NVFP4 compressed-tensors export on the installed unsloth actually
  supporting it, returning a clear message rather than a cryptic save_method failure.
- Frontend: IQ quants (iq2_xxs, iq2_m, iq3_xxs, iq4_xs) are imatrix-only, so force the
  imatrix on when one is selected and lock the toggle, instead of submitting an IQ quant
  with no imatrix that llama.cpp would reject.

Extends the backend tests for the new capability guards and the conditional kwarg wiring.

* Studio: upload compressed merged models to the Hub without recompressing

For an FP8/NVFP4 Hub export the model is already produced locally in the "<dir>-<suffix>"
output. Uploading it directly with HfApi.upload_folder (mirroring export_base_model) avoids
re-running the expensive compressed-tensors quantization a second time inside
push_to_hub_merged, which for NVFP4 also re-runs calibration and risks OOM. Falls back to
push_to_hub_merged when there is no local compressed output to reuse.
2026-06-30 03:41:02 -07:00
Michael Han
11469a60fe
(feat) Add project names to studio training runs (#6512)
* (feat) Add project names to studio training runs to avoid models being overwritten when doing similar training runs

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

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

* Update studio/frontend/src/features/export/export-page.tsx

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* Update studio/frontend/src/features/export/export-page.tsx

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* Update studio/frontend/src/features/export/export-page.tsx

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* better project name sanitization, removed duplicated project name normalization

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

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

* implement checkpoint scanning utilities and tests for base model inference

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

* Guard project_name against null and use leading important modifiers

* Fix/adjust training project names for PR #6512

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

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

* Fix/adjust training project names for PR #6512

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

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

* Address project-name review feedback

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

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

* Show project names in training recents

* Keep GGUF export directories source-specific

---------

Co-authored-by: NZ-Linix <nz-linix@outlook.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: NZ-Linix <linus.ordowski@outlook.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
2026-06-29 16:06:36 +02: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
Lee Jackson
1d4c4d4a76
Merge branch 'main' into image-generation 2026-06-26 18:22:14 +01: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
Daniel Han
80d3434d61
Studio: require signed capability tokens for /p preview links (#6666)
* Studio: require signed capability tokens for /p preview links

The public /p preview routes added in #6486 run model load and chat
generation as the admin user with no authentication. The only gate is the
preview ref, a deterministic outputs-root path (run or run/checkpoint) that
is guessable rather than secret. On a network-reachable Studio (--secure
tunnel or -H 0.0.0.0), an unauthenticated caller who guesses a ref can
consume GPU and probe a private fine-tuned checkpoint.

Make the share link an unguessable, revocable capability:

- Sign the canonical ref with a dedicated server-side secret (HMAC-SHA256,
  stored in app_secrets, independent of the JWT/login secret).
- Require a valid token on every /p chat, models, and page request before
  resolving a checkpoint or loading a model; missing or invalid tokens get a
  generic 404 so the surface never confirms a ref exists.
- Accept the token via ?k= (browser link and preview page) or
  Authorization: Bearer (OpenAI-compatible clients).
- Rotate the secret to revoke every outstanding link
  (POST /api/settings/preview-links/rotate).
- Clamp preview generation (max_tokens/max_completion_tokens <= 1024, n = 1)
  and set Referrer-Policy: no-referrer on the page so the token is not
  leaked via Referer.

Training history hands the authenticated owner the signed token, and the
copy-link button builds /p/{ref}?k={sig}.

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

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

* Studio: honor a lower caller token limit in the preview clamp

Codex review: when only the legacy max_tokens was sent, the clamp left
max_completion_tokens at the 1024 default, and _effective_max_tokens prefers
max_completion_tokens, so a request like max_tokens=16 could still generate up
to 1024 tokens. Derive one effective limit (max_completion_tokens wins, else the
legacy max_tokens) and pin both fields to it so a caller's lower limit is kept.

* Studio: add preview kill switch, rate limit, and revoke-links UI

Follow-ups to the /p preview capability work:

- Public-sharing kill switch: a persisted setting (default on) gates the public
  /p surface. When off, every preview request 404s even with a valid token, and
  the owner UI stops offering share links. GET/PUT /api/settings/preview-sharing;
  enforced in _verify_or_404.
- Per-IP rate limit on the preview chat route: a coarse in-process sliding-window
  limiter (20 req/min/IP) returns 429 + Retry-After before the GPU lock is taken.
  Client IP honors X-Forwarded-For only when UNSLOTH_STUDIO_TRUST_FORWARDED is
  set, matching the login limiter's trust model.
- Settings UI: a "Preview sharing" section with the public-sharing toggle and a
  "Revoke all preview links" button (confirm dialog) that rotates the secret.

Tests cover the kill switch (404 when off), the 429 path, the sliding window,
client-IP trust behavior, and the setting default.

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

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

* Studio: fix preview-fields sharing arg and refresh sigs after revoke

Codex review:
- P1: get_training_run_detail and update_training_run called _preview_fields
  with only output_dir after it gained a required sharing_on parameter, raising
  a 500 TypeError once get_run succeeded. Pass get_preview_sharing_enabled() at
  both sites; add a detail-endpoint regression test.
- P2: after rotating the preview secret from settings, the history grid still
  held stale preview_sig values, so a freshly copied link would 404. Emit
  emitTrainingRunsChanged() after a successful revoke so the grid refetches
  freshly signed refs.

* Studio: harden preview sharing controls (Codex review)

- Fail closed: a read failure on the preview-sharing kill switch now returns
  False instead of defaulting to enabled, so an unavailable settings DB can't
  reopen the public surface. A missing key still defaults to enabled.
- Per-IP rate limit behind the managed Cloudflare tunnel: client_ip now honors
  CF-Connecting-IP when the socket peer is loopback, so tunneled visitors are
  keyed by their real IP instead of collapsing onto the local cloudflared peer.
- GET /p no longer mints key/share_url when sharing is disabled; it returns
  sharing_enabled=false so clients don't distribute links that 404.
- Settings UI: toggling public sharing emits the training-runs-changed event so
  the history grid shows/hides Copy preview link without a manual refresh.

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

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

* Studio: harden preview rate limiter and IP keying (Opus review)

From a two-agent review of the PR:

- Rate limiter no longer evicts an active bucket when the table is full: a flood
  of distinct keys could otherwise cycle out a throttled bucket and reset its
  counter. Evict only aged-out buckets; if the table is full of live clients,
  fail closed (deny the new key) instead.
- client_ip keys on the rightmost (proxy-appended) X-Forwarded-For hop when the
  trust env is set; the leftmost is client-spoofable. Documented the
  append/overwrite-proxy assumption.
- _verify_or_404 checks the capability token before the kill-switch DB read, so
  unauthenticated /p spam can't be used as an unbounded settings-DB sink and the
  response is identical regardless of the sharing on/off state.

Tests: nested run/checkpoint happy path + wrong-ref rejection, the eviction
fail-closed behavior, and route-level coverage for the rotate / preview-sharing
settings endpoints.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-25 21:40:48 -07: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
efdc5ad9a2 transparent image rail; hide fine-tuned shortcut; custom folders filtered to diffusion GGUFs 2026-06-24 16:00:17 -03:00