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

Studio: diffusion UX polish and stronger auto policies (images + video) (#6885)

Browse source 
* Auto policies: deferred dense compile, video compile default, step cache and precision auto

Image dense loads with speed unset no longer sit at plain off: the load stays
bit-identical eager, and the 3rd generation in a session engages the default
compile profile plus the cuDNN attention upgrade mid-session (a one-off image
never pays the warmup, repeated use amortises it). Video dense loads resolve
straight to the default profile since a clip denoise amortises the compile
within a single run, and never to max.

Video also gains the image backend's tri-state auto policies: unset step cache
now decides from the default schedule and re-checks the actual step count per
generation, and unset precision (transformer_quant) hands the decision to the
hardware ladder instead of staying off. Memory badge reason now says plainly
that everything fits when no offload is planned.

* Rename Dtype to Precision, add the video Precision control, step cache Auto option

The images Advanced panel's Dtype row is now Precision (same control, clearer
name), and the video Advanced panel gains the matching Precision select wired
to the load route's existing transformer_quant field, gated to full-pipeline
loads the way the image control gates to GGUF. Step cache selects on both
pages gain an explicit Auto option as the default (the previous Off default
silently behaved as auto and never let anyone pin off), and the Speed and
Attention tooltips now state the deferred dense compile and the SageAttention
black-frame caveat.

* Model catalog: canonical diffusion model groups with device-aware routing

One canonical name per image/video model, its published artifacts (GGUF, FP8,
bnb-4bit, official BF16) as data, and pure routing helpers: suffix-stripped
canonical keys (owner-preserving; cross-owner merges only via explicit
aliases), group/artifact lookups, a flat back-compat options shim, load-spec
resolution replacing the pages' lookup tables, search matching over old ids
and format tokens, the GGUF fit ladder extracted from the variant expander,
and pickDefaultArtifact/pickDefaultQuant deciding what a bare group click
loads (downloaded first, then the best quality that fits 70 percent of VRAM,
GGUF as the safe fallback). Checked by npm run catalog:check, following the
i18n:check pattern.

* Picker: one canonical row per diffusion model with a format second level

The Images and Video pickers now render the curated catalog as one row per
model in Recommended: clicking loads the best artifact for the device (the
routed GGUF quant, a prequant FP8/bnb-4bit that fits, or the official BF16),
and a chevron opens the per-format list, with the GGUF row nesting the usual
quant expander. Live HF listing rows that belong to a group are deduplicated,
search collapses member repos into their group (old ids and format tokens
still match), and the On Device sections group cached member repos under the
same canonical name with the per-repo rows inside. Curated groups render from
the catalog rather than the HF listing, which finally surfaces LTX-2.3 in the
video Recommended list (its hub pipeline_tag is image-to-video, so the
text-to-video listing always missed it) and exposes the HunyuanVideo 720p
repack next to 480p.

Backend: /cached-models now tags trusted video-family repos text-to-video
instead of blanket text-to-image, and the pickers admit catalog-known
non-unsloth repos On Device, so cached Lightricks/Wan/Hunyuan pipelines
finally appear in the Video picker. Chat pickers pass no catalog and are
unchanged.

* Download formats, tab icons, plain-language train tips, 3-loop autoplay

The image Download button becomes a menu: PNG saves the original bytes with
the embedded recipe, JPEG and WebP re-encode client-side from the fetched
blob (JPEG flattened onto white). The video Download button gains MP4
(original, keeps audio), WebM and GIF; the latter two transcode server-side
from the stored MP4 via PyAV (VP9 realtime profile for WebM, ~12 fps adaptive
palette for GIF) behind a new gallery export route that 501s with a readable
message when a codec is missing.

Generated clips no longer loop forever: the player replays a clip three times
per selection, then pauses with controls up; a new generation or a refresh
gets its own three plays. The Create/Train tabs reuse the sidebar's New Chat
and Train icons (TestTubeOutlineIcon moved to a shared lib module), and every
Train tab helper text is now one plain sentence.

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

* Keep the Create/Train tab icon and label on one line

TabsTrigger renders its children inside a plain inline span and the
Tailwind preflight gives svg display:block, so the HugeiconsIcon forced
the label onto a second line. Wrap icon plus label in their own
inline flex row inside each trigger.

* Strip -int8 and -nvfp4 prequant suffixes in the model catalog key

canonicalKeyFor already lowercases before matching, so -GGUF/-FP8 in any
case were covered; -int8 and -nvfp4 were not in the suffix table, so
such repos rendered as standalone rows in Recommended and On Device
instead of standardizing into their base-name group and routing through
pickDefaultArtifact. Added both suffixes plus case-insensitivity and
routing assertions to the catalog check.

* Standardize non-catalog picker rows to their base model name

The curated catalog already collapses its own groups, but hub listing
rows and cached repos outside the catalog (ERNIE-Image, FLUX.2-klein,
Qwen-Image-Edit-2509, FLUX.2-dev) still rendered raw ids with -GGUF /
-FP8 style suffixes in Recommended and On Device.

- model-catalog.ts: new stripArtifactSuffixesForDisplay, a
  case-preserving twin of canonicalKeyFor's stripping that keeps the
  owner prefix and original casing for display.
- pickers.tsx: recommended hub rows and the downloaded GGUF/model rows
  pass their labels through it when a catalog is present, so only the
  diffusion pickers change; chat rows keep raw ids. Click targets keep
  the full repo id, and the format badge still shows the artifact kind.
- Catalog check covers the new helper across GGUF/FP8/int8/nvfp4 in
  both cases plus no-op and suffix-only names.

* Offer official BF16/FP8 artifacts per model group and fix gallery label clipping

Model picker changes so groups are not limited to unsloth quant repos:

- model-catalog.ts: each image group that has an official vendor pipeline
  now carries its BF16 (official) artifact as the top (highest quality)
  entry - Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image, Qwen/Qwen-Image-2512,
  Qwen/Qwen-Image-Edit-2511, black-forest-labs/FLUX.1-dev, FLUX.1-schnell
  and FLUX.1-Kontext-dev. The LTX-2.3 video group now lists Lightricks'
  own bf16 and fp8 distilled single-file checkpoints alongside the GGUF.
  Resident sizes are set from the actual weight totals (FLUX ships a
  duplicate single-file that from_pretrained ignores, so FLUX bf16 is ~32
  GB not 54). The repos that used to be aliases are now real artifacts.
- The router already prefers the highest-quality artifact that fits the
  0.7 x GPU budget, so a datacenter GPU now defaults to official BF16
  while consumer GPUs still route to the fitting quant or GGUF. That is
  why bnb-4bit was the Z-Image-Turbo default before: it was the only
  non-GGUF artifact and it was already downloaded.
- diffusion.py: allowlist the four official image repos not previously
  trusted (qwen/qwen-image-2512, qwen/qwen-image-edit-2511,
  black-forest-labs/flux.1-schnell, flux.1-kontext-dev). All verified as
  safetensors-only diffusers model_index pipelines. The LTX-2.3
  checkpoints are already on the video trust list.
- catalog check: BF16-wins-on-datacenter, quant-wins-on-consumer, and the
  single-file load specs for the LTX-2.3 checkpoints.

Also fixes the video gallery thumbnail caption: the leading duration was
clipped by the rounded corner and selection border, so the strip now has
enough left/bottom padding to clear the curve.

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

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

* video gallery: guard export transcode against a stream-less clip

_transcode_webm and _transcode_gif indexed src.streams.video[0] before
checking the stream list, so a container with no video stream raised a bare
IndexError that the broad handlers then re-labeled as a missing libvpx or
decoder. Raise an explicit RuntimeError naming the real cause in both the
WebM and GIF paths.

* Studio: honor explicit attention/format choices, fix distilled-LTX defaults and On Device catalog routing

* Remove stray planning notes accidentally committed to the branch

* video: add transformerQuant to the load callback deps

handleLoad reads transformerQuant but omitted it from the useCallback dep array,
so after the user changes only Precision and then selects a model or clicks
Reapply, the memoized callback keeps the stale closure and loads the previous
precision. The image page's equivalent callback already lists it.

* model picker: honor the format filter when routing catalog clicks; add catalog rows to the roving list

- routedArtifactFor now scopes a group's artifacts to the active format filter
  (the same matchesFormatFilter predicate the visibility check uses) before
  pickDefaultArtifact, so a group shown only because it owns a GGUF no longer
  routes a click to a large non-GGUF download. Covers both the Recommended and
  On Device grouped paths.
- hubOptionKeys now includes the catalog-group, search-catalog-group, and grouped
  On Device row keys in exact render order, so arrow/Home/End roving reaches the
  catalog rows instead of giving them a duplicate missing id and skipping them.

* model picker: don't treat a partial base cache as downloaded

A partially-cached base repo (a cancelled download that left only some weights)
was counted as downloaded, so an On Device click routed to a fresh multi-GB
re-download instead of the complete GGUF. The picker's endpoint (/api/models/
cached-models) did not carry a partial flag at all, so a frontend-only guard
could not see it. Surface partial from that endpoint by reusing the hub inventory
scan's snapshot-partial detector, plumb it through CachedModelRepo (backend +
frontend types), and skip partial base repos when building the downloaded set.

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

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

* model picker + diffusion: drop partial/unloadable cached rows, skip defer-compile before a LoRA gen

- On Device (cached non-GGUF) rows filtered partial-download snapshots back in: sortedCachedModels
  gated on passesTaskGate + a groupForRepoId key match but, unlike downloadedSet, never checked
  c.partial, so an incomplete unsloth snapshot showed as a loadable On Device row (click errors or
  silently re-fetches multi-GB). It also admitted repos that only match the catalog by group KEY
  (a base / uncurated-quant sibling like Qwen/Qwen-Image-2512) which have no loadable artifact and
  dead-end at the trust gate. Add !c.partial and gate on artifactForRepoId (what loadSpecFor
  resolves) instead of groupForRepoId, so a cached row shows only when the backend can load it.

- Deferred speed-auto engaged the compile profile on the 3rd generation BEFORE _apply_loras. A
  compiled transformer rejects LoRA (supports_lora is False) and _apply_loras raises before its
  unchanged-selection no-op, so once compile engaged every LoRA generation on that load failed
  permanently. Skip the deferral when a LoRA is requested (compile and LoRA are mutually exclusive)
  and let it engage on a later LoRA-free generation.

* Scope the cached-model partial probe to the listed snapshot dir

list_cached_models builds each row from the largest/complete copy across HF cache
roots, but _cached_repo_partial probed is_snapshot_partial with no repo_cache_dir,
so the scan spanned every root: a stale .incomplete copy in one root would flag a
complete copy in another as partial and hide the usable model from the picker (the
click then routes to a re-download). Forward the winning snapshot's repo_path so all
three partial signals are scoped to that copy, matching the sibling inventory paths
(models/dataset cache_inventory, local_inventory).

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

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

* Do not auto-route to gated repos, prefer complete cached copies, defer compile past attached LoRA, scope group expand keys

Four fixes:
- pickDefaultArtifact's not-downloaded ladder returned the gated BF16 FLUX.1-dev / Kontext-dev
  before the open GGUF on a large GPU, so a bare group click routed to a repo the user may lack
  license/token access to. Add a gated flag and skip gated artifacts in the not-downloaded ladder
  (an already-downloaded gated artifact is still returned).
- list_cached_models picked the largest duplicate cache copy and computed partial only on it, so a
  larger partial copy shadowed a smaller complete one; since partial rows are dropped from the
  picker the usable model vanished. Prefer completeness, then size.
- the deferred-speed compile engaged on a no-LoRA generation while an adapter from a prior
  generation was still attached, baking it into the compiled graph (the later unload is swallowed
  on a compiled pipe); also defer while adapters remain attached.
- routeGroupClick's GGUF fallback toggled the context-free canonicalId while the chevron toggles
  the context-scoped expandKey, leaving the format list un-collapsible in one context, dead in the
  other, and risking cross-context expansion; thread expandKey through.

* Guard video pipeline repos from deletion, drop the always-failing LTX FP8 artifact, prefer 720p Hunyuan

Three round-6 fixes:
- cached non-GGUF video repos now surface in the Video On-Device picker with the normal delete
  action, but /delete-cached only guarded chat + the Images engine, so a loaded/loading Wan / LTX /
  Hunyuan pipeline could have its HF snapshot removed from under it. Add a VideoBackend
  loading_repo_ids accessor and a video loaded/loading guard mirroring the Images one.
- the catalog advertised Lightricks/LTX-2.3-fp8 as loadable, but the LTX-2.3 loader refuses the
  official scaled-FP8 single file (.weight_scale/.input_scale) and points to GGUF/BF16, so a pick
  routed to a ~76 GB download that always fails on load. Remove the FP8 artifact.
- pickDefaultArtifact only sorts by format, so the HunyuanVideo group's 480p (listed first) beat
  the 720p even on GPUs where 720p fits the budget. List 720p first so the fit loop prefers it and
  falls back to 480p only on smaller cards.

* diffusion: add compute int8/fp8_dynamic text-encoder quant, wire into video

Add two torchao compute text-encoder quant modes to the diffusion precision
engine, alongside the existing layerwise fp8 and weight-only nvfp4:

- int8: per-token activation + per-channel weight (torch._int_mm), with per-layer
  keep-bf16 selection. int8 degrades on large encoders unless the most
  quant-sensitive decoder blocks stay bf16, so it engages only for families with
  a measured keep-bf16 schedule (qwen-image / qwen-image-edit keep first+last 6,
  flux.2-dev keeps first 3); a family without one falls back to fp8.
- fp8_dynamic: per-row fp8 compute (torch._scaled_mm), keeping the matmul in fp8
  on the tensor cores instead of upcasting each forward like the layerwise fp8.

The selective int8 caster reuses the committed transformer-quant factory
(_make_quant_config / make_filter_fn / exclude_tokens_for_scheme) plus a small
structural first/last-N block skip, so it depends only on committed APIs.

Wire text-encoder quant into the video backend, which previously loaded the
companion encoder (Gemma3 / UMT5 / Qwen2.5-VL) dense bf16 while quantising only
the DiT. text_encoder_quant is plumbed through the load request, validation, the
load chain, the resolved record, and status, mirroring the image backend; it
applies for every load kind (the encoder is dense regardless of how the DiT was
sourced). Widen the image and video load request Literals and add the video
status field.

Tests: int8 family-schedule routing and fp8 fallback, fp8_dynamic routing,
hardware gates (int8 sm_80+, fp8_dynamic sm_89+), the structural block selection,
the real int8 filter closure (keeps the first blocks plus the vision tower /
lm_head / T5 wo dense), and the video route threading and 422 validation.

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

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

* text-encoder quant: skip the torchao modes under offload (both backends)

quantize_text_encoders applied int8-with-schedule / fp8_dynamic / nvfp4 (all torchao) to the
text encoder regardless of the offload policy. An offload placement then moves the quantized encoder
with Module.to(), which torchao tensor subclasses reject (aten._has_compatible_shallow_copy_type is
unimplemented) -- a hard crash, the same one the DiT path already skips torchao quant under offload to
avoid. Add offload_active to quantize_text_encoders and skip the torchao modes when set; layerwise fp8
is not torchao and still streams under offload. Both the video and image loaders pass
offload_active = (offload policy != none).

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

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

* diffusion: skip non-bf16 linears for scaled_mm quant schemes

The fp8 / mxfp8 / nvfp4 schemes run on torch._scaled_mm and the fp4 / mx GEMMs,
which assert a bfloat16 input weight. On a mixed-precision DiT that keeps some
linears in fp32 for numerical stability (the Wan and Hunyuan video transformers
do this), quantize_ hits the first fp32 linear, raises, and the best-effort
wrapper swallows it to None, so the whole transformer stays dense with no error
and no speedup or memory saving.

Add a require_bf16 gate to make_filter_fn and pass it for the scaled_mm schemes
in quantize_transformer (and the fp8_dynamic text-encoder caster). The gate
skips non-bf16 linears so the scheme engages on the bf16 ones. int8 uses
torch._int_mm, which quantizes fp32/fp16 weights fine, so it leaves the gate off
and keeps its current coverage.

Verified on Wan2.2-TI2V-5B: fp8 and mxfp8 now quantize 303 linears via the
committed quantize_transformer path where they previously engaged 0.

* prequant builder: mirror the scaled-mm bf16 gate offline

The runtime DiT quantizer skips non-bf16 Linears for the scaled_mm schemes (fp8,
nvfp4, mxfp8) so the scheme engages on a mixed-precision transformer instead of
aborting on the first fp32 Linear. The offline prequant builder reused make_filter_fn
without that gate, so building an fp8/nvfp4/mxfp8 checkpoint for a mixed-precision DiT
(Wan, Hunyuan keep _keep_in_fp32_modules in fp32 even under torch_dtype=bf16) would hit
the same fp32 Linear and abort, breaking the builder's stated offline == runtime,
LPIPS-0 invariant. Thread require_bf16 = scheme in _SCALED_MM_SCHEMES through the builder,
record it in the checkpoint metadata, and verify it on load (mirrors the existing
exclude_name_tokens guard) so a future _SCALED_MM_SCHEMES change cannot silently load a
checkpoint built under the old filter.

* Keep nvfp4 fp32 linears quantised (bf16 gate is fp8/mxfp8 only)

Verified on torchao 0.17 / B200: fp8 per-row asserts 'PerRow quantization only
works for bfloat16 precision input weight' and mxfp8 asserts 'Only supporting bf16
out dtype', but NVFP4's high-precision conversion quantises an fp32 weight fine
(forward included). So the bf16 skip-gate must be fp8/mxfp8 only, not all scaled_mm
schemes -- otherwise nvfp4 leaves large fp32 projections dense, losing the intended
memory/speed gain. Rename _SCALED_MM_SCHEMES -> _REQUIRE_BF16_SCHEMES = (fp8, mxfp8)
and thread it through the runtime filter, the offline builder, and the loader
require_bf16 verification (offline == runtime preserved).

* [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>
This commit is contained in:
Daniel Han 2026-07-06 17:48:39 -07:00 • committed by GitHub
commit 186f381bc6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
37 changed files with 3748 additions and 1000 deletions
Download patch file Download diff file Expand all files Collapse all files

113
plans/controlnet-workflow.md
View file

@ -1,113 +0,0 @@
# Plan: ControlNet for the Studio Images workflow (stacked on #6769)
## Context & research
ControlNet is the **#2 most-used "beyond text-to-image" diffusion workflow** after LoRA (now shipped
in #6771). It conditions generation on a spatial control map so the output follows a structure. Web
research (this session) on ComfyUI / Forge / A1111 usage:
- The dominant control types are **Depth, Canny (edges), and OpenPose (human pose)**.
- The biggest 2025 shift is toward **Union / all-in-one ControlNets** that bundle many control modes in
one model: **InstantX / Shakker-Labs `FLUX.1-dev-ControlNet-Union-Pro`** for FLUX, and **`xinsir/
controlnet-union-sdxl-1.0`** for SDXL. SDXL has no official ControlNet (community: xinsir, TheMistoAI,
BRIA). SD1.5 has the original lllyasviel set.
- Sources: comfyui-wiki.com ControlNet collections (flux-1 / sdxl), stable-diffusion-art.com ControlNet
ComfyUI, education.civitai.com ControlNet guide, stablediffusiontutorials.com Qwen-Image ControlNets.
**Both Studio backends can do ControlNet** (verified against the live tree):
- diffusers has `FluxControlNetPipeline` / `FluxControlNetModel`, `StableDiffusionXLControlNetPipeline` /
`ControlNetModel` / `ControlNetUnionModel`, `QwenImageControlNetPipeline`, `FluxControlNetInpaintPipeline`
(diffusers 0.38 in the studio venv). Also `FluxControlPipeline` (the Flux.1-Canny/Depth "Control"
in-model variants).
- native sd.cpp (`stable-diffusion.cpp`, the #6769 sd-cli base) has `--control-net <path>`,
`--control-image <path>`, `--control-strength <f>`, `--control-net-cpu`, and a built-in
`preprocess_canny` (examples/cli/main.cpp:704, examples/common/common.cpp:422+).
Studio has **no ControlNet wiring today**: `diffusion_families.py` has no controlnet fields and
`sd_cpp_args.py` has no `--control-*`. This adds it, mirroring the LoRA architecture (#6771) and the
existing img2img/inpaint/reference workflow patterns.
## Scope (this PR = diffusers ControlNet, Union-first)
Ship the highest-value slice first; keep it shippable and consistent with the shipped LoRA design.
- **In scope:** diffusers ControlNet for the families with the strongest ecosystems and pipeline support:
**FLUX** (FluxControlNetPipeline + Union Pro), **SDXL** (StableDiffusionXLControlNetPipeline + xinsir
Union), **Qwen-Image** (QwenImageControlNetPipeline). Single ControlNet per generation. A `control_type`
hint (canny / depth / pose / tile / passthrough). **Canny preprocessing built in** (cheap, cv2/PIL) plus
**passthrough** for user-supplied control maps (depth/pose maps made elsewhere, matching ComfyUI where
preprocessing is separate). Strength + guidance start/end. Discovery endpoint + family-gated picker.
- **Out of scope (follow-ups):** native sd.cpp ControlNet (`--control-net`, needs GGUF ControlNet assets +
family support probe); heavy preprocessors (Depth-Anything, OpenPose detector) as optional server-side
auto-preprocess; multi-ControlNet stacking; ControlNet + inpaint combo.
## Key facts (verified)
- diffusers ControlNet pipelines are built with `Pipeline.from_pipe(base_pipe, controlnet=cn_model)` (or
`from_pretrained(base, controlnet=...)`), so the resident base modules are reused with **no reload** --
same `from_pipe` machinery the img2img/inpaint/edit workflows already use (`diffusion.py` ~:1104-1130,
`_workflow_pipe`). The ControlNet model (`ControlNetModel` / `FluxControlNetModel` / `ControlNetUnionModel`)
is a small extra module loaded once and cached on `_LoadState`.
- ControlNet models are **family-specific** (SD1.5 CN != FLUX CN != SDXL CN). So discovery must be
**family-gated**, exactly like the LoRA picker's `supports_lora`/family filter.
- Generate-time contract mirrors reference/inpaint: a control image (b64) + params, threaded through
`routes/inference.py` into both backends (diffusers serves it; native rejects clearly until the
follow-up wires `--control-*`).
- Reuse: `diffusion_lora.py` discovery/resolve/family-gate patterns; the reference-image upload component +
the LoRA picker UI shape; `_workflow_pipe` from_pipe; `hf_hub_download_with_xet_fallback`.
## Approach
### Families (`core/inference/diffusion_families.py`)
- Add per-family ControlNet declaration: `controlnet_pipeline_class` (e.g. "FluxControlNetPipeline",
"StableDiffusionXLControlNetPipeline", "QwenImageControlNetPipeline"), `controlnet_model_class`
("FluxControlNetModel" / "ControlNetModel" / a union class), and a small curated list of recommended
ControlNet repos tagged by control type. Expose a `controlnet: bool` capability (like `reference`).
### Discovery (`core/inference/diffusion_controlnet.py`, new -- mirrors diffusion_lora.py)
- `list_controlnets(family)` = curated family-tagged repos + a local scan; `resolve_controlnet(id, family,
hf_token)` downloads via the xet-fallback helper; `preprocess_control(image, control_type)` (canny via
cv2/PIL; passthrough otherwise); `supports_controlnet(engine, family, model_kind, quant)` gate
(diffusers bf16 / bnb-4bit yes; GGUF-via-diffusers + torchao fp8/int8 dense = no, same rule as LoRA;
native = follow-up).
### Backend -- diffusers (`core/inference/diffusion.py`)
- A ControlNet manager parallel to `_apply_loras`: load the requested `ControlNetModel` once (cache on
`_LoadState`, reset on unload/model change), build the CN pipeline via `from_pipe(base, controlnet=...)`
in `_workflow_pipe`, and pass `control_image` + `controlnet_conditioning_scale` +
`control_guidance_start/end` at generate time. Never fuse; CN model stays bf16.
### Backend -- native (`core/inference/sd_cpp_backend.py`, sd_cpp_args.py) -- FOLLOW-UP
- Add `--control-net` / `--control-image` / `--control-strength` to the arg builder and a GGUF-ControlNet
resolve; gate to families sd.cpp supports. Deferred out of this PR.
### Routes + request models (`models/inference.py`, `routes/inference.py`)
- Add optional `controlnet: ControlNetSpec` to `DiffusionGenerateRequest` (`{id, image, control_type,
strength (0..2, default 1), guidance_start (0..1), guidance_end (0..1)}`); thread into `backend.generate`;
surface `supports_controlnet` in status; persist the chosen CN + type in gallery recipe metadata.
- New `GET /api/models/diffusion-controlnets?family=` (mirror the LoRA discovery route).
### Frontend (`features/images/images-page.tsx`, `api.ts`)
- A "ControlNet" control in the left rail (reuse the reference-image uploader + a control-type Select +
ControlNet-model Select gated by `supports_controlnet`/family + a strength SliderField). Show a small
preview of the preprocessed control map. Thread `controlnet` into `generateDiffusionImage`; omit when no
control image.
## Verification
- **Unit:** request validation (optional/empty unchanged; bad strength rejected; unsupported family/quant
rejected); discovery (family filter, resolve, canny preprocess shape); diffusers manager (loads CN once,
from_pipe built, scale threaded, reset on model change) with a fake pipe; routes (no-CN path unchanged).
- **Live smoke (critical):** on GPU 4, drive the real diffusers backend with a real family + Union CN and a
canny control image; same prompt/seed at strength 0 vs 0.8; assert (a) output DIFFERS from no-control and
(b) the strong-control output structurally follows the control map (edge-overlap / SSIM vs the control).
- **Playwright (`unsloth_studio_workflow`):** upload a control image, pick type + model + strength,
generate; capture screenshots/GIF against the live secure studio.
- Full `pytest studio/backend/tests/` green; frontend `vite build` clean; ruff clean.
## Delivery
- New branch `diffusion-controlnet` off `diffusion-image-workflows` (#6769 head) in an isolated worktree,
sibling to #6771 (LoRA) and #6772 (fp8 fix). PR on `unslothai/unsloth`, base = diffusion-image-workflows,
part of the single logical stack rooted on #6763 (continuation of #6658). Commit as Daniel Han; no AI/bot
mentions, no emojis, no em dashes.
- Follow-ups: native sd.cpp ControlNet; server-side Depth/OpenPose auto-preprocessors; multi-ControlNet;
ControlNet+inpaint.

113
plans/diffusion-popularity-findings.md
View file

@ -1,113 +0,0 @@
# Diffusion workflow popularity findings (HF download data)
Read-only HF metadata pull (`scripts/investigate_popularity.py`), to ground the Studio
Images scope against what people actually download. Downloads are HF's 30-day count and
all-time count; pulled 2026-06-30.
## Qwen-Image-Edit vs Qwen-Image-Layered (the explicit "determine popularity" question)
| Model | dl / 30d | dl all-time | likes | pipeline |
|---|---:|---:|---:|---|
| Qwen/Qwen-Image-Edit-2509 | 511,996 | 2,942,966 | 1,185 | image-to-image |
| Qwen/Qwen-Image-Edit-2511 | 162,185 | 1,088,015 | 1,087 | image-to-image |
| Qwen/Qwen-Image-Edit (base) | 70,728 | 1,161,044 | 2,440 | image-to-image |
| **Qwen-Image-Edit (all variants)** | **~745,000** | **~5,192,000** | - | - |
| Qwen/Qwen-Image-Layered | 51,303 | 234,785 | 1,112 | image-text-to-image |
| unsloth/Qwen-Image-Edit-2511-GGUF | 218,313 | - | - | image-to-image |
**Conclusion:** Qwen-Image-Edit is ~10-14x more downloaded than Layered (combined 745K/30d
vs 51K, 5.2M vs 235K all-time). Shipping Edit (2511 + the unsloth GGUF, which alone pulls
218K/30d) and rejecting/deferring Layered is the correct, data-backed call. Layered also
needs a dedicated pipeline (`additional_t_cond=True`) the standard QwenImagePipeline can't
drive, so it would be both niche AND extra engineering. Reject stands.
## ControlNet is niche on the modern (diffusers/FLUX/Qwen) stack
| Model | dl / 30d | dl all-time | likes |
|---|---:|---:|---:|
| InstantX/FLUX.1-dev-Controlnet-Canny | 2,891 | 136,727 | 194 |
| lllyasviel/ControlNet (SD1.5-era) | 0 | 14 | 3,820 |
| stabilityai/stable-diffusion-x4-upscaler | 10,040 | 2,976,405 | 725 |
**Conclusion:** ControlNet's large user base lives in the older SD1.5 / A1111 ecosystem, not
the diffusers/FLUX/Qwen stack Studio targets (the modern FLUX ControlNet is ~3K/30d). It is
NOT part of the "most popular ~80%" for current-gen models, so deferring it is justified by
the data, not just by effort. The dedicated x4 upscaler is also low 30-day (10K) though high
all-time; our generic hires-fix upscale (img2img re-detail) covers the use case for any
loaded family without an extra model.
## The shipped six cover the popular workflows
Top text-to-image (HF list, 30-day): SD1.5 (1.78M), SDXL (1.32M), FLUX.1-dev (1.09M),
dreamshaper-7 (1.03M), **Tongyi-MAI/Z-Image-Turbo (886K)**, sd-turbo (684K), SD3.5-medium
(606K), sdxl-turbo (598K), Qwen-Image-Lightning (483K). All are plain txt2img -> our Create
tab; the GGUF/bnb families + Z-Image cover the modern ones.
Top image-to-image (HF list, 30-day): Qwen-Image-Edit-2509 (512K) -> Edit tab; SDXL-refiner
(162K) -> Upscale/Transform; Kontext (150K) -> Edit tab.
So Create / Transform / Inpaint / Extend / Upscale / Edit map onto the head of both
distributions.
## SHIPPED: FLUX.2-klein image (reference) conditioning
> Status: IMPLEMENTED + verified live (2026-06-30). `flux.2-klein` now has `reference=True`;
> the backend exposes a "reference" workflow that passes the image to the loaded
> Flux2KleinPipeline directly (no from_pipe, no strength, output at the requested size); the
> frontend has a "Reference" tab. Verified with `scripts/verify_reference_http.py` on
> `unsloth/FLUX.2-klein-4B-GGUF` (Q4_K_M): a reference-conditioned 1024x1024 result is
> non-blank, correctly sized, and DIFFERS from the identical-seed plain txt2img.
>
> FLUX.2-klein ALSO gained inpaint (`Flux2KleinInpaintPipeline` via from_pipe; verified with
> `scripts/verify_klein_inpaint.py`). It does NOT get outpaint/extend: FLUX.2 scales any >1MP
> input down to ~1MP, so a padded outpaint canvas shrinks back. "outpaint" is now a distinct
> capability advertised only for size-preserving inpaint families (`inpaint_preserves_size`).
> Multi-reference is shipped too (the pipeline accepts a list; the Reference tab has add/remove
> slots, backend caps at 3 extra; verified with `scripts/verify_multiref_http.py`: two
> references differ from one at the same seed). The analysis that motivated the work follows.
The data surfaced this gap (now closed):
| Model | dl / 30d | pipeline |
|---|---:|---|
| black-forest-labs/FLUX.2-klein-4B | 470,482 | image-to-image (#2 overall) |
| black-forest-labs/FLUX.2-dev | 271,037 | image-to-image |
| **unsloth/FLUX.2-klein-4B-GGUF** | **243,307** | image-to-image |
| black-forest-labs/FLUX.2-klein-9B | 178,964 | image-to-image |
`flux.2-klein` is ALREADY a registered family in `diffusion_families.py` (txt2img only,
base `FLUX.2-klein-4B`, open repo). But `Flux2KleinPipeline.__call__` natively accepts an
`image` argument (verified in diffusers 0.38.0; params: image, prompt, height, width,
num_inference_steps, guidance_scale -- NOTE: no `strength`). FLUX.2 is a unified
text-to-image + reference/edit model: the SAME loaded pipe does both, depending on whether
`image` is passed. Today Studio exposes only txt2img for it, so the popular image-editing
mode of the #2 image-to-image model is unreachable.
### Why it's a separate PR, not a tail-of-session add
FLUX.2 reference conditioning is a DIFFERENT semantic from the shipped workflows:
- No `strength` (it is reference-conditioning, not a denoise blend like img2img).
- Output size comes from width/height (txt2img-style), not from the input image size, so the
image-conditioned width/height rule we added for img2img/inpaint/upscale does NOT apply.
- FLUX.2 supports MULTIPLE reference images; single-image is the common case but the UX
should not preclude multi-ref.
This needs: read the Flux2KleinPipeline source for exact `image` semantics (list vs single,
how it is resized/tiled, recommended guidance), decide the UX (a "Reference" workflow that is
available alongside Create for `reference=True` families, distinct from the strength-based
Transform tab), then verify on the open FLUX.2-klein-4B base (and the unsloth GGUF) with a
reference image before/after.
### Sketch (for the follow-up PR)
- `diffusion_families.py`: add `reference: bool = False`; set `reference=True` on flux.2-klein.
- `_family_workflows`: when `fam.reference`, expose `"reference"` (in addition to txt2img).
- `generate()`: a `reference` branch that passes `image` to `state.pipe` directly (no
from_pipe, no strength), with width/height = the requested size (NOT the input size).
- Frontend: a "Reference" tab (image dropzone + prompt), gated to `reference` families;
Create stays pure txt2img for the same model.
- Verify: load unsloth/FLUX.2-klein-4B-GGUF, pass a reference image, confirm the output is
conditioned on it and differs from a no-image run at the same seed.
## Net
The seven shipped workflows (create, transform, inpaint, extend, upscale, reference, edit)
cover the popular ~80% across both the txt2img and image-to-image distributions, including the
#1 image-to-image model (Qwen-Image-Edit) and the #2 (FLUX.2-klein, now via the reference tab).
ControlNet / SD1.5-era ControlNet remain deferred with data backing (niche on the modern stack).

156
plans/diffusion-workflows-pr-plan.md
View file

@ -1,156 +0,0 @@
# Stacked-PR plan: Studio diffusion workflows (Images redesign)
Branch tip: `diffusion-eager-and-compile-cache` (latest commit "Phase 16 review fixes").
Remote: `oobabooga/unsloth`. New PRs stack on top of the existing diffusion stack
(ultimately on top of unslothai/unsloth#6658), treated as one logical change.
Nothing here is committed yet (commit/push only on explicit request).
## CRITICAL: the working tree holds TWO uncommitted streams, and three core files INTERMINGLE them
A full `git status` / marker audit (branch `diffusion-eager-and-compile-cache`, tip "Phase 16
review fixes") shows the uncommitted tree is NOT a clean single feature. There are two streams:
A) **The eager/compile-cache phase** (the branch's own in-progress work, NOT this session's —
zero of this feature's markers). Purely-its files, safe to NOT touch in the workflow PRs:
- new modules: `diffusion_arch_patches.py`, `diffusion_compile_cache.py`,
`diffusion_eager_patches.py`, `diffusion_gguf_compile.py`, `diffusion_patch_backend.py`
- new tests: `test_diffusion_arch_patches.py`, `test_diffusion_compile_cache.py`,
`test_diffusion_eager_patches.py`, `test_diffusion_gguf_compile.py`
- modified: `diffusion_speed.py`, `test_diffusion_speed.py`, `conftest.py`,
`scripts/diffusion_bench.py`, and ~25 untracked `scripts/*bench*/*probe*/*orchestrator*`.
B) **The Images workflows feature** (this session): the workflow engine + frontend + installer.
**The two streams INTERMINGLE inside three shared files and are NOT separable by file:**
- `studio/backend/core/inference/diffusion.py` — this feature's workflow hunks are interleaved
with the eager/compile wiring (imports at L67-75; `install_arch_patches`/`compile_cache.begin`/
`.restore`/`.save` and the `eager_patched`/`compile_cache_ctx` state throughout
`load_pipeline`/`generate`/`unload`). A single `diffusion.py` cannot go into one PR without the
other stream's hunks.
- `studio/backend/models/inference.py` — this feature's `init_image`/`mask_image`/
`reference_images`/`upscale`/`model_kind` fields sit next to the pre-existing `speed_mode`/
`transformer_prequant_path` fields in the same request models.
- `studio/backend/tests/test_diffusion_backend.py` — this feature's workflow tests sit next to
the pre-existing `test_failed_load_rolls_back_eager_patches` (imports `diffusion_eager_patches`).
**Implication / options (USER DECIDES — it is their branch + their eager/compile work):**
- CLEANLY separable now (purely this feature, can be committed/PR'd on their own anytime):
frontend `images-page.tsx` + `api.ts` + `pickers.tsx`, and the sd.cpp installer
`install_sd_cpp_prebuilt.py` + `test_sd_cpp_install.py`. (These are PR 2 and PR 3 below.)
- The backend engine (PR 1) CANNOT be cleanly split from the eager/compile phase via files.
Realistic paths: (a) finalize + commit the eager/compile phase first, then this feature's
backend lands as a clean diff on top; or (b) commit both streams together as the branch's
next chunk (consistent with treating the stack as one logical change); or (c) a manual
`git add -p` hunk split of the three shared files (tedious, risks a non-compiling
intermediate). NOT auto-doable safely without the owner's intent for the eager/compile work.
## Proposed stack (3 PRs, bottom to top)
### PR 1 - Backend: diffusion workflow engine (safetensors + image-conditioned + editing)
Files:
- `studio/backend/core/inference/diffusion.py` (the feature hunks: three load "kinds"
gguf/single_file/pipeline; `_workflow_pipe` via `from_pipe(torch_dtype=None)`;
`_align_vae_dtype`; `generate()` routing for reference/img2img/inpaint/upscale/edit;
image-conditioned width/height from the input image (but reference + txt2img use the slider
size); `upscale` (hires fix) branch on the img2img pipe; `reference` (FLUX.2) branch that
passes the image(s) to the loaded pipe directly (no from_pipe, no strength) incl. multi-
reference (`reference_images` combined into a list, capped at 3 extra); branch ORDER
inpaint/upscale before reference so a mask/upscale request on a reference family still routes
right; `_family_workflows` (adds "upscale" wherever img2img is supported, "reference" for
reference families, "outpaint" only for size-preserving inpaint families); `kind` on state +
`model_kind` in status; `load_progress` double-count fix). NOTE: this file ALSO carries
pre-existing speed hunks if any landed here - review per-hunk and exclude non-feature hunks.
- `studio/backend/core/inference/diffusion_families.py` (img2img/inpaint pipeline slots;
`edit` flag + `reference` flag + `inpaint_preserves_size` flag; `qwen-image-edit` +
`flux.1-kontext` families; flux.2-klein gains reference + inpaint (no outpaint: FLUX.2
normalizes to ~1MP); `detect_family` longest-match + leftover-reject; `layered` reject).
- `studio/backend/core/inference/diffusion_engine_router.py` (model_kind -> diffusers for
non-gguf kinds).
- `studio/backend/core/inference/diffusion_memory.py` (`estimate_safetensors_dense_mib`).
- `studio/backend/core/inference/sd_cpp_backend.py` (model_kind passthrough; reject
img2img/inpaint on the native engine).
- `studio/backend/models/inference.py` (load request: optional gguf_filename, model_kind,
init/mask/strength, advanced knobs; status: workflows, model_kind).
- `studio/backend/routes/inference.py` (model_kind forwarding; ValueError -> 400;
exc_info logging).
- Tests: `test_diffusion_backend.py`, `test_diffusion_routes.py`.
Title: `Studio diffusion: safetensors + image-conditioned + instruction-editing workflows`
Summary: Adds non-GGUF safetensors loading (full bnb-4bit pipelines + single-file fp8,
gated to unsloth/*), the image-conditioned workflows (img2img, inpaint, outpaint via the
inpaint path) built with `Pipeline.from_pipe` for zero-extra-VRAM component reuse, and
instruction editing as its own family kind (Qwen-Image-Edit-2511 + FLUX.1-Kontext-dev).
Fixes two bugs: `from_pipe` defaulting to a float32 recast that crashed torchao-quantized
transformers, and image-conditioned calls forcing the slider size onto the input image.
### PR 2 - Frontend: redesigned Images page (workflow tabs + Advanced Options)
Files:
- `studio/frontend/src/features/images/images-page.tsx` (workflow tabs Create/Transform/
Inpaint/Extend/Upscale/Reference/Edit; capability gating + auto-switch; `MaskCanvas`;
`buildOutpaint`; Upscale tab with Scale + Detail-strength sliders; Reference tab (FLUX.2,
reference image + add/remove extra references, no strength); Advanced Options accordion gated
to GGUF for transformer-quant; spinner-overlap fix).
- `studio/frontend/src/features/images/api.ts` (request/status types incl. model_kind, upscale,
reference_images).
- `studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx` (curated
safetensors + edit GGUF rows; `SUPPORTED_EDIT_KEYWORDS` un-hide; layered hide).
Title: `Studio Images: workflow tabs (create/transform/inpaint/extend/upscale/reference/edit) + Advanced Options`
Summary: Redesigns the Images page around capability-gated workflow tabs with a brush mask
editor, client-side outpaint, a hires-fix upscale tab, a FLUX.2 reference tab, an instruction-edit
tab, and an Advanced Options panel (speed/quant/attention/memory/step-cache/offload), plus the
number-input spinner fix.
### PR 3 - sd.cpp prebuilt installer hardening
Files: `studio/install_sd_cpp_prebuilt.py`, `studio/backend/tests/test_sd_cpp_install.py`.
Title: `Studio sd.cpp: pin release + verify sha256 + mirror-ready source`
Summary: Pins the stable-diffusion.cpp release (was tracking `latest`), verifies each
download's sha256 against GitHub's published asset digest before extract/execute, adds a
download timeout, and makes the source repo configurable (`UNSLOTH_SD_CPP_REPO`) so a
future unslothai mirror needs no code change. Cleanly separable from the rest.
## Pre-PR review (done)
An independent 3-angle review (backend correctness, frontend/UX, security/robustness) ran over
the full session diff. No High findings; the load-gating to `unsloth/*` and the multi-reference
count caps were verified intact end-to-end. Fixes applied before the PRs:
- Frontend [Med]: multi-reference slots no longer renumber mid-edit (dropped the eager
`filter(Boolean)` in the per-slot onChange; empties dropped only at send).
- Backend [Med]: upscale now caps the ABSOLUTE output (longest side <= 2048), not just the
factor, so a large upload * 4x can't OOM.
- Backend/security [Med]: `_decode_b64_image` rejects images > 4096px/side (uniform guard for
init/mask/reference vs decompression-bomb / OOM inputs); base64 image fields capped at 32 MiB.
- Security [Low]: the native sd.cpp engine guard also rejects `reference_images`.
All covered by new tests (82 backend pass) and a post-fix e2e (all five workflows still pass).
## Post-deploy user feedback fixes (done)
From live use of the deployed studio:
- Backend [Med]: image-conditioned workflows passed the raw upload size to pipelines that
require multiples of 16 (Z-Image/Qwen/FLUX), so an odd upload (e.g. 186px) failed with
"Height must be divisible by 16". Added `_snap_to_multiple` and auto-resize init (and the
matched mask) to the nearest /16 for img2img/inpaint/extend/edit. Verified live: a 186x250
Transform and Inpaint now return 200 at 192x256. Tests added.
- Frontend [Med]: the Advanced options (FP8/INT8 quant, speed, attention, memory) were a
collapsed, muted accordion at the bottom of the left rail that users missed (HF screenshots
discussion #25). Moved them into a RIGHT-DOCKED panel mirroring Chat's settings panel.
Per follow-up (discussion #26): CLOSED by default, toggled by a SINGLE fixed top-bar button
using Chat's `LayoutAlignRightIcon` that stays in the exact same position in both states
(verified x/y identical open vs closed) and highlights when open. Controls extracted to a
render-local `advancedControls`; unused Accordion import removed.
## Constraints for execution (when authorized)
- Write as the user; no AI/bot mentions, no emojis, no em dashes; well-formatted bodies.
- `gh auth status` first. Push to `oobabooga/unsloth`, stack on the current branch.
- Keep PR 3 independent; PR 2 depends on PR 1 (frontend needs the backend contract).
- Re-run `pytest studio/backend/tests/test_diffusion_*.py test_sd_cpp_install.py` +
frontend `tsc`/`build` before each PR.
## Out of scope / follow-ups
Scope decisions are backed by HF download data in `plans/diffusion-popularity-findings.md`.
The seven shipped workflows are create, transform, inpaint, extend, upscale, reference, edit.
- Publish the unslothai/stable-diffusion.cpp mirror + macOS/Windows staging (#152).
- ControlNet / style-transfer: the goal's "most popular" set is covered by the seven shipped
workflows. ControlNet on the modern diffusers/FLUX/Qwen stack is niche by downloads
(~3K/30d), so deferring it is data-backed, not just an effort call.
- FLUX.2-klein inpaint and multi-reference are DONE (shipped). Outpaint is intentionally not
offered for FLUX.2 (it scales >1MP inputs to ~1MP). No further FLUX.2 follow-ups outstanding.

125
plans/diffusion-workflows-studio.md
View file

@ -1,125 +0,0 @@
# Plan: Unsloth Studio diffusion workflows + Images UI redesign
Stacked as NEW PRs on top of the existing 14-PR diffusion stack above unslothai/unsloth#6658
(treated as one logical base). Goal: cover ~80% of common real-world diffusion workflows,
across macOS/Linux/Windows/CPU, optimizing performance, accuracy, and memory.
## Current state (verified by recon)
- Backend is **text-to-image only** end-to-end. No image/mask/control plumbing in
`DiffusionGenerateRequest`, `/images/generate`, or `DiffusionBackend.generate()`.
- **diffusers 0.38.0 already imports every pipeline we need**: `*Img2ImgPipeline`,
`*InpaintPipeline`, `FluxFillPipeline`, `*ControlNetPipeline`/`ControlNetModel`,
`FluxKontextPipeline`, `QwenImageEditPipeline`/`QwenImageEditPlusPipeline`,
`StableDiffusion(Latent)UpscalePipeline`. No diffusers upgrade required.
- The native **sd.cpp engine already has dormant fields** (`init_img`, `strength`, `mask`,
`ref_images`) and a complete `upscale()` path — never wired to the request/route.
- **All advanced LOAD options are already wired** end-to-end (speed_mode, transformer_quant
fp8/int8/nvfp4/mxfp8, attention_backend, memory_mode, cpu_offload, transformer_cache,
vae_tiling status). The Advanced panel is mostly a FRONTEND surfacing job.
- Frontend has `Tabs` (`components/ui/tabs.tsx`) and `Accordion` ready. No dropzone and no
mask/brush canvas — both greenfield. `SliderField` is the customizer number input.
- `_EDIT_KEYWORDS = ("edit","kontext","inpaint","layered")` rejects edit/inpaint/Kontext/
Layered repos at family detection.
- sd.cpp binary is downloaded prebuilt from **upstream leejet/stable-diffusion.cpp** (no
unsloth mirror, no checksum/manifest/version pin, not wired into setup.sh). llama.cpp uses
an `unslothai/llama.cpp` mirror with manifest+sha256+version pin+source fallback.
- Chat "Images" pill = provider-side (OpenAI/Gemini) hosted tool, separate from local diffusion.
## Workflow popularity ranking (what to build for 80% coverage)
1. txt2img (keep, polished) — done
2. img2img / variations — **P0**
3. inpainting (mask edit) — **P0**
4. upscaling / hires fix — **P0**
5. ControlNet (Canny/Depth/Pose/Lineart/Tile) — **P0/P1**
6. outpainting (canvas extend) — **P1**
7. instruction image editing (Qwen-Image-Edit, FLUX Kontext) — **P1**
8. style transfer / reference — **P1** (via img2img/edit/control)
9. batch generation/edit/upscale — **P1**
10. LoRA/style packs — **P2**
## Editing-model decisions (researched)
- **Qwen-Image-Edit / Edit-2511**: popular, best-in-class clean targeted edits + multilingual
text. **Support** (instruction edit, mask-optional).
- **FLUX.1 Kontext**: popular, character-consistent in-context editing. **Support** (note:
Kontext-dev is non-commercial/gated — surface license, don't block local custom models).
- **Qwen-Image-Layered**: newer, niche (Photoshop RGBA layer decomposition). Needs a dedicated
pipeline (`additional_t_cond`) — **defer** (keep rejected for now; optional later behind a
layered-specific view). This already crashed the standard path (the earlier bug).
## UI design — workflow tabs (inside ImagesPage, route/nav unchanged)
`Tabs` across the top of the controls area. Combine related workflows:
- **Create** — txt2img (current behavior preserved)
- **Transform** — img2img + style transfer (upload + strength/denoise + presets)
- **Edit** — inpaint (mask brush/upload/invert/feather, masked-vs-whole) + instruction edit
(Qwen-Image-Edit / FLUX Kontext, mask-optional)
- **Extend** — outpainting (directional handles, aspect presets, overlap/feather)
- **Control** — ControlNet (one control slot first: Canny/Depth/Pose/Lineart/Tile + preview)
- **Enhance** — upscaling (ESRGAN/RealESRGAN + latent/tiled)
- **Advanced Options** — Accordion surfacing existing load knobs (speed/compile/attention/
quant fp8/int8/nvfp4/memory/offload/vae tiling/cache) with Auto defaults + resolved values.
Capability gating: a workflow/control is shown enabled only when the selected engine+family+
device+quant supports it; otherwise disabled with a plain-language "why".
## Backend architecture
- Extend `DiffusionGenerateRequest`: optional `workflow` (txt2img|img2img|inpaint|outpaint|
control|edit|upscale), `init_image` (b64), `mask_image` (b64), `control_image` (b64),
`strength`, `controlnet_conditioning_scale`, `control_start/end`, `upscale_factor`,
`ref_images`. Add an image-decode (b64→PIL) helper (none exists).
- `DiffusionFamily`: add optional pipeline-class slots (`img2img_pipeline_class`,
`inpaint_pipeline_class`, `edit_pipeline_class`, `controlnet_pipeline_class` + control repos).
Build the right pipeline around the already-loaded `transformer=` (reuse `_assemble_pipe`
shape); swap/cache pipeline class per workflow without reloading the transformer where
possible.
- `generate()` kwarg builder must branch: img2img/edit pipelines take `image=`/`strength=` and
reject `width/height`; inpaint adds `mask_image=`; control adds `control_image=`. Gate each
kwarg via `inspect.signature`.
- Capability resolver: maps engine+family+device+quant → supported workflows + reasons; echoed
in run metadata so the UI shows what actually ran.
- Memory planner must account for input/latent size, control models, VAE decode, upscale.
## PR breakdown (stacked, small, capability-gated)
- **PR-1 UI fixes + workflow shell**: fix number-input spinner overlap (DONE in tree), tab
scaffold (Create/Transform/Edit/Extend/Control/Enhance/Advanced), Advanced Options accordion
surfacing existing load knobs, capability banner, loading/empty/error states.
- **PR-2 Backend workflow contract + capability registry**: extend request/response, decode
helper, per-family pipeline slots, resolver. No new behavior yet beyond txt2img.
- **PR-3 img2img (Transform)**: backend + Transform tab + dropzone (adapt from
shared-composer `addFiles`/`PendingImageThumb`). Smoke test low vs high denoise.
- **PR-4 inpaint + instruction edit (Edit)**: mask canvas (greenfield), inpaint pipeline,
Qwen-Image-Edit/FLUX Kontext edit; relax `_EDIT_KEYWORDS` → route to edit family.
- **PR-5 outpaint (Extend)**: expanded-canvas inpaint, directional handles, feather/overlap.
- **PR-6 ControlNet (Control)**: one control slot + preprocessor preview + strength/start/end.
- **PR-7 upscaling (Enhance)**: wire dormant sd.cpp `upscale()` + diffusers upscale + `/images/upscale`.
- **PR-8 Advanced panel polish + FP8/INT8 verification matrix**.
- **PR-9 sd.cpp prebuilt packaging**: mirror to `unslothai/stable-diffusion.cpp`, manifest+
sha256+version pin+`--published-repo`+source fallback, wire into setup.sh (ref
install_llama_prebuilt.py).
- **PR-10 cross-platform staging validation** (danielhanchen staging repos, small GGUFs).
- **PR-11 Playwright tests + screenshots/GIFs per tab** (studio_test_kit / unsloth_studio_workflow).
- **PR-12 batch + multi-control + reproducibility polish** (later).
## Done so far
- Fixed the customizer number-input spinner overlap (`SliderField` in images-page.tsx): native
spinners covered the value on the narrow field; now fully suppressed (webkit inner+outer +
Firefox `appearance:textfield`) and field widened to `w-14`. Frontend rebuilt clean.
## Verification
- Playwright (studio_test_kit) per tab: screenshots + GIFs, capability gating, upload/mask,
progress/cancel/error, gallery.
- B200 functional: load + generate one image per workflow per representative family.
- FP8 + INT8 verified (build matrix: SDXL/FLUX/Qwen-Image/Qwen-Image-Edit/GGUF; measure
black-image/NaN rate, peak VRAM, time-to-first-image, prompt adherence, source preservation).
- Cross-platform staging (Linux CUDA/CPU, Windows CUDA/CPU, macOS MPS) with small GGUFs.
## Delivery
New branch(es) off the current tip; new stacked PRs. Commit/push only when asked.

141
plans/wobbly-jumping-narwhal.md
View file

@ -1,141 +0,0 @@
# Plan: Publish unslothai/stable-diffusion.cpp mirror + our own CPU/Apple prebuilts
## Context
The Unsloth Studio native diffusion engine downloads a prebuilt `sd-cli` / `sd-server`
(stable-diffusion.cpp) via `studio/install_sd_cpp_prebuilt.py`. Today it pulls from
**leejet/stable-diffusion.cpp** upstream releases. We want to own this like we own
**unslothai/llama.cpp**: a fork that builds and publishes OUR OWN prebuilt binaries on a
schedule, so we control reproducibility, integrity, and the macOS load floor.
**Why native is CPU/Apple-only.** On a GPU host, diffusers + our optimizations (regional
`torch.compile` ~2.2x, cuDNN/flash attention, FP8/INT8/NVFP4 quant, first-block-cache) is
faster than sd.cpp's CUDA path, which has none of those levers — so GPU hosts route to
diffusers. Native sd.cpp only wins where diffusers is weak: **CPU and Apple**. Therefore we
build native binaries ONLY for the platforms where native is actually the faster engine, and
skip CUDA/ROCm/Vulkan entirely (GPU = diffusers/torch). This also makes the CI far cheaper.
The Studio side is already prepared: `install_sd_cpp_prebuilt.py` reads `UNSLOTH_SD_CPP_REPO`
(repo override) + `UNSLOTH_SD_CPP_TAG` (pinned tag) and verifies the GitHub asset `digest`
(`_verify_sha256`). So the bulk of the work is the mirror repo + release CI; the Studio change
is a small default flip + resolver tweak.
## Coverage (user-confirmed): CPU / Apple ONLY
| Platform | Arch | Build | Runner | Notes |
|---|---|---|---|---|
| macOS | arm64 | Metal (`-DSD_METAL=ON`) | macos-26, `OSX_DEPLOYMENT_TARGET=14.0` | Apple fast path (diffusers/MPS weak) |
| macOS | x86_64 | CPU | macos-15-intel, `OSX_DEPLOYMENT_TARGET=13.3` | Intel Macs |
| Linux | x86_64 | CPU | ubuntu-22.04 (glibc 2.35) | **also covers WSL** (WSL = Linux x64) |
| Linux | aarch64 | CPU | ubuntu-24.04-arm | ARM servers |
| Windows | x86_64 | CPU | windows-2022 (MSVC+Ninja) | |
**Explicitly out of scope:** CUDA, ROCm, Vulkan native builds; GPU runners; cudart bundling;
per-gfx matrices. GPU stays on diffusers/torch.
## Reference pattern (verified this session)
`unslothai/llama.cpp` builds via `.github/workflows/unsloth-prebuilt.yml` (orchestrator) + six
per-accel children + `scripts/unsloth/` helpers (`assemble_metadata.py`, `package_bundle.py`,
`assert_macho_minos.sh`). Mechanisms to mirror: `resolve` (supply-chain aging — only build a
release public >=6h; stamp build-info + Unsloth fingerprint; upload ONE source artifact all
children extract) -> per-platform children (build from the source artifact, load-gate, package,
upload) -> `assemble` (fingerprint gate + manifest/sha256 index + coverage gate + **atomic
draft->publish**, no partial releases). Template files fetched to `workspace_81/temp/llamacpp_workflows/`.
## Key facts (verified)
- leejet builds both `sd-cli` and `sd-server` (`examples/cli`, `examples/server`) — the mirror
ships both (sd-server is used by PR #6768's persistent server).
- leejet naming: `sd-<tag>-bin-<Darwin-macOS-…-arm64 | Linux-Ubuntu-…-x86_64 | win-cpu-x64>.zip`.
leejet already ships macOS arm64, Linux x64 CPU, Windows CPU — we ADD macOS x86_64 and Linux
aarch64 (the gaps in our target set), and rebuild the rest under our own fingerprint/integrity.
- Studio resolver (`resolve_release_asset`, `install_sd_cpp_prebuilt.py:88`): filters to `.zip`;
macOS = darwin/macos + arch token; Linux = `linux` + arch + (no accel marker for auto/cpu);
Windows = `bin-win` + `avx2` else any. For a CPU-only mirror the resolver needs essentially NO
change — macOS x86_64 and Linux aarch64 already match by arch token; just confirm the Windows
CPU asset resolves (contains `bin-win`, falls back to the plain build).
## Design
### A. Mirror repo (fork of leejet/stable-diffusion.cpp)
Fork so upstream C++ stays intact; add only `.github/workflows/` + `scripts/unsloth/`. Adapt the
llama.cpp orchestrator, heavily simplified (no CUDA/ROCm/Vulkan, no PR-mix):
- **`resolve`**: pick the upstream leejet tag with the >=6h aging window; reuse leejet's
`master-<count>-<sha>` as the mirror tag (keeps `UNSLOTH_SD_CPP_TAG` comparable to upstream);
stamp a source tarball with build-info + the "Compiled by the Unsloth team" fingerprint; upload
the source artifact. Skip-if-already-published like llama.cpp.
- **Build children** (reusable `workflow_call`), each `cmake -DSD_BUILD_EXAMPLES=ON` (cli+server):
- `macos` (arm64 Metal + x64 CPU): pinned `CMAKE_OSX_DEPLOYMENT_TARGET`, `@loader_path` rpath,
load-gate via `assert_macho_minos.sh` (adapted for `sd-cli`/`sd-server`).
- `cpu-linux` (x64 + arm64) and `cpu-windows` (x64, MSVC+Ninja).
- **Asset naming = leejet-compatible**, all `.zip`:
`sd-<tag>-bin-Darwin-macOS-arm64.zip`, `sd-<tag>-bin-Darwin-macOS-x86_64.zip`,
`sd-<tag>-bin-Linux-Ubuntu-24.04-x86_64.zip`, `sd-<tag>-bin-Linux-Ubuntu-24.04-aarch64.zip`,
`sd-<tag>-bin-win-cpu-x64.zip`.
- **`assemble`**: fingerprint gate (verify the mark in every archive), generate
`sd-prebuilt-manifest.json` + `sd-prebuilt-sha256.json`, coverage gate (all 5 assets present),
atomic draft->publish. GitHub sets each asset `digest`, which the Studio already verifies.
- **Signing/notarization:** none. The Studio downloads via `urllib` (not a browser), so no macOS
quarantine xattr is set and Gatekeeper does not block CLI-run binaries (matches llama.cpp).
### B. Studio-side switch (PR on the diffusion stack, after the mirror's first green release)
Small, in `studio/install_sd_cpp_prebuilt.py` + its test:
1. `DEFAULT_REPO = "unslothai/stable-diffusion.cpp"`; `DEFAULT_TAG` = the mirror's first tag.
2. Confirm `resolve_release_asset` picks correctly for all 5 CPU/Apple hosts (add a Windows CPU
token only if the plain-`bin-win` fallback proves insufficient; likely no change needed).
Keep the leejet fallback (env override still points back upstream).
3. Extend `test_sd_cpp_install.py` `_ASSETS` to the mirror's 5-asset set; assert host->pick for
macOS arm64/x64, Linux x64/arm64, Windows x64; assert GPU hosts are unaffected (still diffusers).
## Critical files
- New (mirror repo): `.github/workflows/unsloth-sd-prebuilt.yml` (+ `-macos.yml`, `-cpu-linux.yml`,
`-cpu-windows.yml`), `scripts/unsloth/{assemble_metadata.py,package_bundle.py,assert_macho_minos.sh}`.
- Studio: `studio/install_sd_cpp_prebuilt.py`, `studio/backend/tests/test_sd_cpp_install.py`.
- Local templates to adapt: `workspace_81/temp/llamacpp_workflows/{unsloth-prebuilt.yml,unsloth-prebuilt-macos.yml,unsloth-prebuilt-cpu.yml}`.
## Sequencing (chicken-and-egg)
1. Build the mirror repo + CI; `publish=false` dry run to validate the 5-way matrix (~10-20 min,
no GPU runners so cheap).
2. First green **published** release with all 5 assets + manifest/sha256.
3. THEN the Studio PR flips `DEFAULT_REPO`/`DEFAULT_TAG` + resolver test (on the diffusion stack).
## Verification
- **Resolver unit tests** (hermetic): feed the mirror's 5 asset names to `resolve_release_asset`
for macOS arm64/x64, Linux x64/arm64, Windows x64 -> correct pick; and CUDA/GPU host -> still
routes to diffusers (native not selected).
- **CI dry run**: `publish=false` artifact-only run; inspect the 5 archives each contain `sd-cli`
(+ `sd-server`) and carry the fingerprint.
- **Live install smoke** (this Linux box): `UNSLOTH_SD_CPP_REPO=unslothai/stable-diffusion.cpp
python studio/install_sd_cpp_prebuilt.py --print-asset` then real `install()`, confirm
`sd-cli --version` + `sd-server` launch, and drive one native CPU generation via the Studio.
- **Integrity**: each published archive matches its manifest sha256 and the GitHub asset digest.
## Staging (user-confirmed): fork + push CI now
Execution order:
1. **Preflight permissions**: `gh auth status`; confirm the token can create/fork under the
`unslothai` org and enable Actions. If it CANNOT, stop and report (fall back to scaffold-only,
or a private fork under danielhanchen), rather than pushing somewhere unintended.
2. **Fork** leejet/stable-diffusion.cpp -> `unslothai/stable-diffusion.cpp` (clone into the
workspace to add files). Keep upstream C++ intact.
3. **Add CI + scripts** on a branch: `.github/workflows/unsloth-sd-prebuilt.yml` +
`-macos.yml`/`-cpu-linux.yml`/`-cpu-windows.yml`, `scripts/unsloth/*`. Commit as Daniel Han
(no AI/bot mentions, no emojis, no em dashes). `unset GH_TOKEN`/use `gh` creds for pushes that
touch `.github/workflows/*` (needs `workflow` scope).
4. **Dry run**: trigger the orchestrator with `publish=false` (artifact-only), confirm all 5
archives build + carry `sd-cli`/`sd-server` + the fingerprint. Iterate until green.
5. **First publish**: `publish=true` (or let the schedule run) -> a real release with the 5
assets + manifest/sha256.
6. **Studio PR** (section B) on the diffusion stack once the release tag exists.
## Follow-ups (not this task)
- Nightly schedule + auto-bump of the Studio `DEFAULT_TAG` (PR bot), like llama.cpp.
- Add GPU native builds later ONLY if a real need appears (today: GPU = diffusers/torch).

20
scripts/build_prequant_checkpoint.py
View file

@ -58,6 +58,7 @@ def main(argv = None) -> int:
FP8_GRANULARITY,
TQ_FP8,
TQ_SCHEMES,
_REQUIRE_BF16_SCHEMES,
_make_quant_config,
_resolve_fast_accum,
exclude_tokens_for_scheme,
@ -87,13 +88,24 @@ def main(argv = None) -> int:
# bakes them as int8 and crashes (torch._int_mm needs M>16) at the first denoise step on
# Flux / Qwen. fp8 / fp4 / mx use scaled_mm (no M limit) -> exclude_tokens_for_scheme returns ().
exclude_name_tokens = exclude_tokens_for_scheme(scheme)
# fp8 and mxfp8 assert a bf16 weight, so their filter must skip any non-bf16 Linear the
# transformer keeps: a mixed-precision DiT (Wan / Hunyuan) retains its _keep_in_fp32_modules in
# fp32 even under torch_dtype=bf16, so quantising one would raise inside quantize_ and abort the
# whole pass. nvfp4 quantises fp32 fine, so it is not gated. Runtime quantize_transformer gates
# this on scheme membership; mirror it here so the offline checkpoint quantises the exact same
# layer set (offline == runtime).
require_bf16 = scheme in _REQUIRE_BF16_SCHEMES
# fp8 bakes the accumulate mode into the saved kernels; record the resolved choice so the
# loader can refuse a checkpoint whose baked value contradicts an explicit runtime request.
fast_accum = _resolve_fast_accum(None) if scheme == TQ_FP8 else None
quantize_(
transformer,
_make_quant_config(scheme),
filter_fn = make_filter_fn(args.min_features, exclude_name_tokens = exclude_name_tokens),
filter_fn = make_filter_fn(
args.min_features,
exclude_name_tokens = exclude_name_tokens,
require_bf16 = require_bf16,
),
)
# Move the state dict to CPU for a portable, GPU-free artifact.
@ -107,9 +119,11 @@ def main(argv = None) -> int:
"scheme": scheme,
"min_features": args.min_features,
# The layers skipped for this scheme (int8's M=1 modulation projections; () for
# the scaled_mm schemes) and, for fp8, the baked accumulate mode. Both let the
# loader reject a checkpoint that would not match the runtime path.
# the scaled_mm schemes), whether non-bf16 Linears were skipped (the scaled_mm
# bf16 gate), and, for fp8, the baked accumulate mode. All let the loader reject a
# checkpoint that would not match the runtime path.
"exclude_name_tokens": list(exclude_name_tokens),
"require_bf16": require_bf16,
"fast_accum": fast_accum,
"torch_dtype": args.dtype,
"quant_backend": "torchao",

174
studio/backend/core/inference/diffusion.py
View file

@ -218,12 +218,18 @@ _TRUSTED_NON_GGUF_REPOS = frozenset(
"stabilityai/stable-diffusion-xl-base-1.0",
"stabilityai/sdxl-turbo",
# Official vendor, safetensors-only base repos allowlisted as LoRA TRAINING bases
# (diffusion training loads the full pipeline from these). Same rule as above: no
# pickled weights, no remote code, exact-match lowercased. FLUX.1-dev is gated on
# the Hub (needs the user's token); the other two are open.
# (diffusion training loads the full pipeline from these) and as the official
# BF16 artifact behind each catalog group (model-catalog.ts). Same rule as above:
# no pickled weights, no remote code, exact-match lowercased. FLUX.1-dev/schnell/
# Kontext are gated on the Hub (need the user's token); the Qwen and Z-Image repos
# are open. All verified as diffusers model_index pipelines.
"black-forest-labs/flux.1-dev",
"black-forest-labs/flux.1-schnell",
"black-forest-labs/flux.1-kontext-dev",
"tongyi-mai/z-image-turbo",
"qwen/qwen-image",
"qwen/qwen-image-2512",
"qwen/qwen-image-edit-2511",
# Krea 2: official vendor repos, safetensors-only, no remote code. Loaded
# per-component via core/inference/diffusion_krea2.py (no GGUF variant yet).
# Turbo is the inference model; Raw is the undistilled base Krea recommends
@ -300,6 +306,12 @@ class _LoadState:
# Attention backend engaged via the diffusers dispatcher (e.g. "_native_cudnn"), or
# None for the default SDPA. Set before compile; orthogonal to the weight quant.
attention_backend: Optional[str] = None
# The caller's ORIGINAL attention request (None / "auto" left it to the backend, else
# an explicit alias like "native" / "sage" / "flash"). Carried so the deferred-speed
# engagement re-runs the SAME selection the load-time path did, instead of forcing the
# auto cuDNN upgrade -- otherwise an explicitly pinned backend (e.g. "native" to avoid
# cuDNN) is silently discarded when the 3rd generation engages the deferred profile.
attention_request: Optional[str] = None
# Step cache engaged ("fbcache") or None. Opt-in, for many-step models.
transformer_cache: Optional[str] = None
# True when the cache decision was AUTO on a cache-capable transformer: generate()
@ -312,6 +324,13 @@ class _LoadState:
# Shared eager monkey-patches (diffusion_eager_patches) installed for this load (any
# non-off speed tier). Uninstalled on unload so a later `off` load is bit-identical.
eager_patched: bool = False
# Deferred speed auto (dense models with speed_mode unset): the load stays fully
# eager/bit-identical, and generate() engages the `default` compile profile at the
# start of the 3rd generation this session -- repeated use is established by then,
# so the one-time compile warmup amortises. Cleared once engaged (or failed).
speed_deferred: bool = False
# Successful generations on this load; drives the deferred engagement above.
generation_count: int = 0
# Pre-warmed torch.compile cache context (diffusion_compile_cache.CacheContext) when a
# compiled tier ran, else None. Carries the per-key inductor dir + bundle for save/restore.
compile_cache_ctx: Any = None
@ -1258,6 +1277,19 @@ class DiffusionBackend:
"(quantized transformer must be compiled; eager is ~30x slower)"
)
effective_speed = SPEED_DEFAULT
# Deferred speed auto for dense models: the load stays eager (a one-off
# image should not pay the 25-60s compile warmup, and eager is the
# bit-identical reference), but a user starting their 3rd image in one
# session has revealed repeated use -- generate() then engages the
# `default` profile, where the warmup starts paying back. Only when
# the request left speed unset, nothing forced a compiled tier, and
# this device/family could actually compile.
speed_deferred = (
speed_mode is None
and effective_speed == SPEED_OFF
and transformer_quant_engaged is None
and compile_eligible(target, is_gguf = False, family = fam)
)
# Opt-in speed optims run BEFORE placement (channels_last / compile
# must precede CPU offload). Snapshot the process-wide backend flags
# first so unload can restore them: TF32 / cudnn.benchmark are global,
@ -1425,12 +1457,15 @@ class DiffusionBackend:
"compiled; eager torchao quant is ~30x slower than GGUF here",
transformer_quant_engaged,
)
# Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4),
# also before placement so the offload hooks move the smaller weights.
# Quantise the dense companion text encoder(s) (opt-in fp8 / fp8_dynamic /
# int8 / nvfp4), also before placement so the offload hooks move the smaller
# weights. int8 needs a per-family keep-bf16 schedule, so pass the family.
te_quant = quantize_text_encoders(
pipe,
target,
mode = text_encoder_quant,
family = fam.name,
offload_active = plan.offload_policy != OFFLOAD_NONE,
logger = logger,
)
@ -1451,10 +1486,13 @@ class DiffusionBackend:
{
"speed_mode": (
speed_mode,
effective_speed,
"deferred" if speed_deferred else effective_speed,
"quantized transformer requires compile"
if transformer_quant_engaged is not None
and normalize_speed_mode(speed_mode) in (None, SPEED_OFF)
else "auto: exact eager for the first two images; "
"the compile profile engages on the 3rd"
if speed_deferred
else "per-kind default"
if speed_mode is None
else "requested",
@ -1480,7 +1518,9 @@ class DiffusionBackend:
"memory_mode": (
memory_mode,
effective_policy,
"planned from measured free VRAM vs estimated footprint",
"everything fits on the GPU, no offload needed"
if effective_policy == OFFLOAD_NONE
else "planned from measured free VRAM vs estimated footprint",
),
"transformer_cache": (
None if cache_auto else transformer_cache,
@ -1513,11 +1553,13 @@ class DiffusionBackend:
text_encoder_quant = te_quant,
transformer_quant = transformer_quant_engaged,
attention_backend = attention_engaged,
attention_request = attention_backend,
transformer_cache = cache_engaged,
cache_auto = cache_may_toggle,
cache_quant_active = cache_quant_active,
cache_threshold = transformer_cache_threshold,
eager_patched = eager_patched,
speed_deferred = speed_deferred,
compile_cache_ctx = compile_ctx,
hf_token = hf_token,
resolved = resolved,
@ -2011,6 +2053,88 @@ class DiffusionBackend:
except Exception: # noqa: BLE001 — reset is best-effort, never fail a generation
pass
def _engage_deferred_speed(self, state: _LoadState) -> None:
"""Engage the deferred `default` speed profile at the start of the 3rd
generation this session.
The load left the pipe fully eager (bit-identical reference); by the 3rd
image repeated use is established, so pay the one-time compile now: eager
patches + attention auto upgrade + regional compile -- exactly what an
unset-speed GGUF load gets at load time. Runs under _generate_lock (the
caller), so no denoise can race the mutation. The flag is cleared FIRST so
a failure never retries per generation; unload cleans everything up via the
same state fields the load-time path uses (backend flags were snapshotted
at load, before any speed layer could mutate them)."""
object.__setattr__(state, "speed_deferred", False)
from .diffusion_eager_patches import install_compile_safe_patches
from .diffusion_arch_patches import install_arch_patches
target = self._resolve_device_target(state.family)
install_compile_safe_patches()
install_arch_patches()
object.__setattr__(state, "eager_patched", True)
# Re-run the load-time selection with the caller's ORIGINAL request (not a bare
# None): an explicit backend must survive the deferred upgrade. auto still upgrades
# to cuDNN here (speed_active=True), but an explicit "native"/"sage"/"flash" is
# honored verbatim rather than silently replaced by the auto cuDNN choice.
attention_engaged = apply_attention_backend(
state.pipe,
select_attention_backend(target, state.attention_request, speed_active = True),
logger = logger,
)
object.__setattr__(state, "attention_backend", attention_engaged)
gguf_transformer = state.kind == "gguf" and state.transformer_quant is None
if compile_eligible(target, is_gguf = gguf_transformer, family = state.family):
compile_ctx = compile_cache.begin(
family = state.family.name,
transformer = getattr(state.pipe, "transformer", None),
dtype = getattr(target, "dtype", None),
quant = state.transformer_quant,
attention_backend = attention_engaged,
compile_kwargs = {
# Mirrors the load-time fullgraph decision: an engaged or
# still-toggleable step cache OR an offload graph-breaks.
"fullgraph": state.transformer_cache is None
and not state.cache_auto
and state.offload_policy == OFFLOAD_NONE,
"dynamic": True,
"mode": "default",
},
logger = logger,
)
object.__setattr__(state, "compile_cache_ctx", compile_ctx)
speed_applied = apply_speed_optims(
state.pipe,
target,
is_gguf = gguf_transformer,
family = state.family,
speed_mode = SPEED_DEFAULT,
cache_active = state.transformer_cache is not None or state.cache_auto,
offload_active = state.offload_policy != OFFLOAD_NONE,
logger = logger,
)
object.__setattr__(state, "speed_mode", SPEED_DEFAULT)
object.__setattr__(state, "speed_optims", tuple(k for k, v in speed_applied.items() if v))
entry = (state.resolved or {}).get("speed_mode")
if isinstance(entry, dict):
entry["value"] = SPEED_DEFAULT
entry["reason"] = (
"auto: compiled on the 3rd image this session "
"(repeated use amortises the one-time compile)"
)
att = (state.resolved or {}).get("attention_backend")
if isinstance(att, dict) and att.get("source") == "auto":
att["value"] = attention_engaged or "native"
att["reason"] = (
"cuDNN fused attention upgrade" if attention_engaged else "diffusers default"
)
logger.info(
"diffusion.speed: deferred profile engaged on generation 3 "
"(optims=%s, attention=%s)",
",".join(state.speed_optims) or "none",
attention_engaged or "native",
)
def generate(
self,
*,
@ -2076,6 +2200,39 @@ class DiffusionBackend:
seed = int(seed)
generator.manual_seed(seed)
# Deferred speed auto: by the 3rd image in one session repeated use is
# established, so engage the compile profile now -- before the LoRA /
# workflow wiring, matching the load-time ordering. Best-effort: a
# failure logs, leaves the eager pipe running, and never retries
# (the helper clears the flag first).
#
# But NOT when this generation requests a LoRA: a compiled transformer rejects
# LoRA (supports_lora is False once compiled), and _apply_loras raises before its
# unchanged-selection no-op, so engaging compile here would permanently break every
# LoRA generation on this load. Compile and LoRA are mutually exclusive; keep the
# pipe eager and let compile defer to a later LoRA-free generation.
lora_requested = any(w != 0 for (_id, w) in (loras or []))
# Also stay eager while adapters from a PRIOR generation are still attached: this
# request may clear them (lora_requested False), but _apply_loras runs AFTER the
# engage below, so compiling here would bake the resident adapter into the graph and
# the subsequent unload_lora_weights() (swallowed on a compiled pipe) would leave it
# active forever -- silent wrong output on every later LoRA-free generation. Deferring
# lets _apply_loras clear it on the still-eager pipe; compile engages a gen later.
loras_attached = bool(getattr(state.pipe, "_unsloth_loras", ()))
if (
state.speed_deferred
and state.generation_count >= 2
and not lora_requested
and not loras_attached
):
try:
self._engage_deferred_speed(state)
except Exception as exc: # noqa: BLE001 — speed is best-effort
logger.warning(
"diffusion.speed: deferred engagement failed, staying eager: %s",
exc,
)
# Apply/adjust LoRA adapters on the resident pipe (non-fused) before picking
# the workflow pipe; from_pipe pipes share the transformer, so it propagates.
self._apply_loras(state, loras, cancel)
@ -2393,6 +2550,9 @@ class DiffusionBackend:
compile_cache.save(state.compile_cache_ctx, logger = logger)
except Exception: # noqa: BLE001 — cache persistence is best-effort
pass
# Count the finished generation (drives the deferred speed
# engagement above); a batch of N images is one generation.
object.__setattr__(state, "generation_count", state.generation_count + 1)
# Return the PIL images (not yet encoded): the route embeds each
# image's recipe and persists it via the gallery.
return {"images": list(images), "seed": int(seed), "repo_id": state.repo_id}

186
studio/backend/core/inference/diffusion_precision.py
View file

@ -5,21 +5,29 @@
The transformer arrives quantised in the GGUF, but the companion text encoder loads
dense (bf16) from the base repo and is often the largest resident component (a Qwen3
/ T5-XXL / Mistral encoder runs to many GB). This shrinks it in place, with two
/ T5-XXL / Mistral encoder runs to many GB). This shrinks it in place, with four
backends:
fp8 - diffusers layerwise casting: 8-bit (e4m3) storage, upcast per layer to the
compute dtype. ~2x smaller. Works on any fp8-capable CUDA card (cc >= 8.9).
nvfp4 - torchao NVFP4 weight-only: 4-bit float with two-level microscaling, run on
Blackwell's (sm_100+) FP4 tensor cores. ~4x smaller and the lowest-VRAM
option, but a steeper quality cost than fp8.
fp8 - diffusers layerwise casting: 8-bit (e4m3) storage, upcast per layer to
the compute dtype. ~2x smaller. Works on any fp8-capable CUDA card (cc >= 8.9).
fp8_dynamic - torchao dynamic fp8 COMPUTE (per-row): keeps the matmul in fp8 on the
fp8 tensor cores (torch._scaled_mm) instead of upcasting each forward.
~2x smaller plus a tensor-core speedup; needs fp8-GEMM silicon (cc >= 8.9).
int8 - torchao dynamic int8 COMPUTE (per-token act + per-channel weight ->
torch._int_mm), with per-layer keep-bf16 selection. int8 degrades on large
encoders unless the most quant-sensitive decoder blocks stay bf16, so it is
applied only for families with a measured keep-bf16 schedule (else it falls
back to fp8). ~2x smaller; needs int8 tensor cores (cc >= 8.0).
nvfp4 - torchao NVFP4 weight-only: 4-bit float with two-level microscaling, run on
Blackwell's (sm_100+) FP4 tensor cores. ~4x smaller and the lowest-VRAM
option, but a steeper quality cost than fp8.
Both keep normalisations / embeddings full precision and are a memory-vs-quality
tradeoff, not free, so both are off by default. They pair especially well with
streamed (group) offload, where the text encoder stays resident -- this is where the
companion footprint dominates. Quantify the quality cost per model with the quality
harness (scripts/diffusion_quality.py). torch / diffusers / torchao are imported
lazily so the module stays importable in a no-torch runtime.
All keep normalisations / embeddings full precision and are a memory-vs-quality tradeoff,
not free, so all are off by default. They pair especially well with streamed (group)
offload, where the text encoder stays resident -- this is where the companion footprint
dominates. Quantify the quality cost per model with the quality harness
(scripts/diffusion_quality.py). torch / diffusers / torchao are imported lazily so the
module stays importable in a no-torch runtime.
"""
from __future__ import annotations
@ -28,11 +36,27 @@ from typing import Any, Optional
TE_QUANT_FP8 = "fp8"
TE_QUANT_NVFP4 = "nvfp4"
TE_QUANT_MODES = (TE_QUANT_FP8, TE_QUANT_NVFP4)
TE_QUANT_INT8 = "int8"
TE_QUANT_FP8_DYNAMIC = "fp8_dynamic"
TE_QUANT_MODES = (TE_QUANT_FP8, TE_QUANT_NVFP4, TE_QUANT_INT8, TE_QUANT_FP8_DYNAMIC)
# Pipeline attributes that hold a text encoder, in order.
_TEXT_ENCODER_ATTRS = ("text_encoder", "text_encoder_2", "text_encoder_3")
# int8 (torch._int_mm) degrades on large text encoders unless the most quant-sensitive decoder
# blocks stay bf16. Per-family (skip_first, skip_last) decoder blocks to keep dense, from measured
# hidden-state fidelity (mean per-token cosine vs the bf16 reference, at the layer each pipeline
# consumes): keeping the first blocks stops early-layer error seeding, keeping the last blocks
# protects the read layer. Families absent here have no int8 schedule that clears the bar, so an
# int8 request for them falls back to fp8.
# qwen-image (Qwen2.5-VL-7B): first+last 6 -> ~0.997 cosine (both ends needed; outlier-bound).
# flux.2-dev (Mistral-Small-24B): first 3 -> ~0.98 cosine (pure early-layer seeding).
_TE_INT8_SKIP: dict[str, tuple[int, int]] = {
"qwen-image": (6, 6),
"qwen-image-edit": (6, 6),
"flux.2-dev": (3, 0),
}
def normalize_te_quant(value: Optional[str]) -> Optional[str]:
"""Lower/strip a requested text-encoder quant; None / "" / "none" -> None.
@ -51,8 +75,9 @@ def normalize_te_quant(value: Optional[str]) -> Optional[str]:
def te_quant_supported(target: Any, mode: str) -> bool:
"""Whether ``mode`` is usable for ``target``: a CUDA device with a bf16 compute
dtype, plus fp8 dtype support (fp8) or Blackwell sm_100+ tensor cores (nvfp4)."""
"""Whether ``mode`` is usable for ``target``: a CUDA device with a bf16 compute dtype, plus
the tensor-core class each backend needs -- fp8 dtype (fp8 layerwise), fp8 GEMM sm_89+
(fp8_dynamic), int8 tensor cores sm_80+ (int8), or Blackwell sm_100+ (nvfp4)."""
if getattr(target, "device", None) != "cuda":
return False
try:
@ -62,6 +87,12 @@ def te_quant_supported(target: Any, mode: str) -> bool:
return False
if mode == TE_QUANT_FP8:
return hasattr(torch, "float8_e4m3fn")
if mode == TE_QUANT_FP8_DYNAMIC:
# Compute fp8 (torch._scaled_mm) needs fp8-GEMM silicon: Ada sm_89+ / Hopper / Blackwell.
return hasattr(torch, "float8_e4m3fn") and torch.cuda.get_device_capability() >= (8, 9)
if mode == TE_QUANT_INT8:
# int8 tensor cores (torch._int_mm) need Ampere sm_80+.
return torch.cuda.get_device_capability()[0] >= 8
if mode == TE_QUANT_NVFP4:
# NVFP4 tensor cores need Blackwell (compute capability major >= 10).
return torch.cuda.get_device_capability()[0] >= 10
@ -75,15 +106,49 @@ def quantize_text_encoders(
target: Any,
*,
mode: Optional[str],
family: Optional[str] = None,
offload_active: bool = False,
logger: Any = None,
) -> Optional[str]:
"""Quantise each present text encoder in place with ``mode`` (fp8 / nvfp4).
Returns the mode actually applied, or None when disabled, unsupported, or no
encoder was cast. Best-effort: any failure leaves the encoder dense."""
"""Quantise each present text encoder in place with ``mode`` (fp8 / fp8_dynamic / int8 / nvfp4).
Returns the mode actually applied, or None when disabled, unsupported, or no encoder was cast.
``int8`` needs a per-family keep-bf16 schedule (``_TE_INT8_SKIP``); a family without one falls
back to ``fp8``. When ``offload_active`` the torchao modes are skipped (their tensor subclasses
reject the ``Module.to()`` an offload hook uses); layerwise ``fp8`` still engages. Best-effort:
any failure leaves the encoder dense."""
mode = normalize_te_quant(mode)
if mode is None or not te_quant_supported(target, mode):
if mode is None:
return None
caster = _cast_fp8 if mode == TE_QUANT_FP8 else _cast_nvfp4
skip: Optional[tuple[int, int]] = None
if mode == TE_QUANT_INT8:
skip = _TE_INT8_SKIP.get((family or "").lower())
if skip is None:
_note(logger, f"int8 has no keep-bf16 schedule for family '{family}'; using fp8")
mode = TE_QUANT_FP8
# The torchao modes (int8 with a schedule, fp8_dynamic, nvfp4) produce tensor subclasses that
# reject Module.to(); an offload placement moves the encoder that way and hard-crashes -- the
# DiT path skips torchao quant under offload for exactly this reason. Layerwise fp8 is not
# torchao and streams fine, so it still engages. Skip the torchao modes under offload.
if offload_active and mode in (TE_QUANT_INT8, TE_QUANT_FP8_DYNAMIC, TE_QUANT_NVFP4):
_note(
logger,
f"text-encoder '{mode}' skipped under offload (torchao tensors reject Module.to()); "
"pin a resident memory mode or use fp8",
)
return None
if not te_quant_supported(target, mode):
return None
if mode == TE_QUANT_INT8:
first, last = skip # type: ignore[misc]
def caster(enc: Any, tgt: Any) -> None:
_cast_int8_selective(enc, tgt, first, last)
elif mode == TE_QUANT_FP8_DYNAMIC:
caster = _cast_fp8_dynamic
elif mode == TE_QUANT_NVFP4:
caster = _cast_nvfp4
else:
caster = _cast_fp8
cast: list[str] = []
for attr in _TEXT_ENCODER_ATTRS:
encoder = getattr(pipe, attr, None)
@ -97,6 +162,82 @@ def quantize_text_encoders(
return mode if cast else None
def _te_exclude_tokens(encoder: Any) -> tuple[str, ...]:
"""fqn tokens whose Linears stay bf16 in a torchao text-encoder quant: the VLM vision tower
and the unused lm_head (not used for prompt encoding), plus the encoder's own fp32-kept
modules (T5 ``wo``, which the gated feed-forward reads the dtype of and which explodes in
low precision)."""
tokens = ["visual", "vision_tower", "lm_head"]
tokens += [str(m).lower() for m in (getattr(encoder, "_keep_in_fp32_modules", None) or ())]
return tuple(dict.fromkeys(tokens))
def _keep_bf16_block_fqns(encoder: Any, skip_first: int, skip_last: int) -> set[str]:
"""FQNs of the decoder blocks to keep bf16: the first ``skip_first`` and last ``skip_last`` of
each top-level ``nn.ModuleList`` stack (a T5 ``encoder.block`` / a decoder ``...layers``).
Structural, so it needs no per-architecture table."""
import torch
keep: set[str] = set()
for name, module in encoder.named_modules():
if not isinstance(module, torch.nn.ModuleList):
continue
n = len(module)
if n <= skip_first + skip_last:
continue
for i in list(range(skip_first)) + list(range(n - skip_last, n)):
keep.add(f"{name}.{i}" if name else str(i))
return keep
def _cast_int8_selective(encoder: Any, target: Any, skip_first: int, skip_last: int) -> None:
# torchao dynamic int8 (per-token act + per-channel weight -> torch._int_mm) on the FLOP-heavy
# Linears, but keeping the first/last decoder blocks (and the vision tower / lm_head / T5 wo)
# in bf16. Reuses the committed transformer-quant factory so the config never drifts.
from torchao.quantization import quantize_
from .diffusion_transformer_quant import (
TQ_INT8,
DEFAULT_MIN_LINEAR_FEATURES,
_make_quant_config,
make_filter_fn,
exclude_tokens_for_scheme,
)
base = make_filter_fn(
DEFAULT_MIN_LINEAR_FEATURES,
exclude_tokens_for_scheme(TQ_INT8) + _te_exclude_tokens(encoder),
)
keep = _keep_bf16_block_fqns(encoder, skip_first, skip_last)
def filter_fn(module: Any, fqn: str = "") -> bool:
if not base(module, fqn):
return False
return not any(fqn == k or fqn.startswith(k + ".") for k in keep)
quantize_(encoder, _make_quant_config(TQ_INT8), filter_fn = filter_fn)
def _cast_fp8_dynamic(encoder: Any, target: Any) -> None:
# torchao dynamic fp8 COMPUTE, per-row (per-token activation + per-output-channel weight ->
# torch._scaled_mm on the fp8 tensor cores). Unlike the layerwise `fp8` backend this keeps the
# matmul in fp8 instead of upcasting each forward. fp8 is robust across encoder sizes, so no
# per-layer keep-bf16 is needed; only the vision tower / lm_head / T5 wo are excluded.
from torchao.quantization import quantize_
from .diffusion_transformer_quant import (
TQ_FP8,
DEFAULT_MIN_LINEAR_FEATURES,
_make_quant_config,
make_filter_fn,
)
# require_bf16: scaled_mm asserts a bf16 weight, so skip any stray non-bf16 Linear the encoder
# keeps (belt-and-suspenders over the named T5 wo exclusion) rather than aborting the pass.
filter_fn = make_filter_fn(
DEFAULT_MIN_LINEAR_FEATURES, _te_exclude_tokens(encoder), require_bf16 = True
)
quantize_(encoder, _make_quant_config(TQ_FP8), filter_fn = filter_fn)
def _cast_fp8(encoder: Any, target: Any) -> None:
import re
import torch
@ -155,3 +296,8 @@ def _cast_nvfp4(encoder: Any, target: Any) -> None:
def _warn(logger: Any, what: str, exc: Exception) -> None:
if logger is not None:
logger.warning("diffusion.precision: text-encoder quant (%s) failed: %s", what, exc)
def _note(logger: Any, msg: str) -> None:
if logger is not None:
logger.info("diffusion.precision: %s", msg)

19
studio/backend/core/inference/diffusion_prequant.py
View file

@ -329,6 +329,25 @@ def _validate_checkpoint(
),
)
return False
# require_bf16 (skip non-bf16 Linears) is the bf16-weight gate, pinned by the scheme (fp8 and
# mxfp8 assert a bf16 weight; nvfp4 / int8 quantise fp32 fine). Like exclude_name_tokens it is
# pinned by the scheme today, but recording and verifying it guards against a future
# _REQUIRE_BF16_SCHEMES change silently loading a checkpoint built under the old filter (it would
# carry a different quantised layer set). Absent (older artifact) is accepted since scheme already
# pins today's gate.
ckpt_require_bf16 = meta.get("require_bf16")
if ckpt_require_bf16 is not None:
from .diffusion_transformer_quant import _REQUIRE_BF16_SCHEMES
expected_require_bf16 = scheme in _REQUIRE_BF16_SCHEMES
if bool(ckpt_require_bf16) != expected_require_bf16:
_warn(
logger,
scheme,
ValueError(
f"checkpoint require_bf16 {bool(ckpt_require_bf16)!r} != {expected_require_bf16!r}"
),
)
return False
# fp8 fast-accum is baked into the saved kernels; only enforce when the caller forces it.
if fast_accum is not None:
ckpt_fa = meta.get("fast_accum")

18
studio/backend/core/inference/diffusion_speed.py
View file

@ -119,18 +119,26 @@ def normalize_speed_mode(value: Optional[str]) -> str:
return normalized
def resolve_speed_mode(value: Optional[str], *, is_gguf: bool) -> str:
def resolve_speed_mode(
value: Optional[str],
*,
is_gguf: bool,
dense_default: str = SPEED_OFF,
) -> str:
"""The effective speed mode when the caller leaves it UNSET (``None``).
A GGUF model defaults to ``default``: it compiles only the hot dequant op chain
(~70-80% of eager GGUF time) for ~1.24-1.64x at a small one-time compile and zero
extra VRAM -- a cheap, always-amortising win whose numeric perturbation sits well
below the quantisation noise floor (the dequant graph is unchanged, just
Inductor-fused). A dense (non-GGUF) model stays ``off`` / bit-identical, since there
compile would be the only source of drift. An explicit value -- including ``"off"``
-- is always honored verbatim."""
Inductor-fused). A dense (non-GGUF) model resolves to ``dense_default``: the image
backend keeps ``off`` at load (bit-identical first generations, with its own
deferred engagement after repeated use), while the video backend passes ``default``
-- a clip denoise runs long enough that the one-time compile always amortises
within a single generation. An explicit value -- including ``"off"`` -- is always
honored verbatim."""
if value is None:
return SPEED_DEFAULT if is_gguf else SPEED_OFF
return SPEED_DEFAULT if is_gguf else dense_default
return normalize_speed_mode(value)

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

@ -37,6 +37,16 @@ TQ_AUTO = "auto"
TQ_SCHEMES = (TQ_INT8, TQ_FP8, TQ_NVFP4, TQ_MXFP8)
TQ_MODES = (TQ_AUTO,) + TQ_SCHEMES
# Schemes whose torchao path asserts a bf16 weight, so their quantise filter must skip
# non-bf16 Linears (see make_filter_fn's require_bf16) rather than aborting the whole pass on a
# stray fp32 Linear (e.g. T5's fp32 `wo`). Verified on torchao 0.17 / B200: fp8 per-row asserts
# "PerRow quantization only works for bfloat16 precision input weight" and mxfp8 asserts
# "Only supporting bf16 out dtype", but NVFP4's high-precision conversion quantises an fp32
# weight fine (forward included) -- so nvfp4 is deliberately NOT gated here, keeping those fp32
# projections quantised instead of leaving them dense. int8 (torch._int_mm) also quantises
# fp32/fp16 weights fine, so it is not gated either.
_REQUIRE_BF16_SCHEMES = (TQ_FP8, TQ_MXFP8)
# The fp8 weight/activation granularity the runtime config uses (see _make_quant_config):
# per-ROW is REQUIRED for correctness on outlier-heavy DiTs. Stamped into a pre-quantized
# fp8 checkpoint's metadata at build time and required by the loader, so a stale checkpoint
@ -391,11 +401,23 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any:
raise ValueError(f"unknown transformer quant scheme '{scheme}'")
def make_filter_fn(min_features: int, exclude_name_tokens: tuple[str, ...] = ()):
def make_filter_fn(
min_features: int,
exclude_name_tokens: tuple[str, ...] = (),
*,
require_bf16: bool = False,
):
"""A torchao ``quantize_`` filter keeping only the FLOP-heavy linears: nn.Linear with both
in/out features >= ``min_features`` AND whose fully-qualified name contains none of
``exclude_name_tokens`` (used by int8 to skip the M=1 modulation / conditioning-embedder
projections that crash ``torch._int_mm``). Hides the (module, fqn) callback arity."""
projections that crash ``torch._int_mm``). Hides the (module, fqn) callback arity.
``require_bf16`` additionally skips any Linear whose weight is not bfloat16. The scaled_mm
schemes (fp8 / mxfp8 / nvfp4) assert a bf16 input weight, so a single non-bf16 Linear -- e.g.
the fp32 layers Wan / Hunyuan video DiTs keep for numerical stability -- otherwise raises inside
``quantize_`` and aborts the ENTIRE pass, leaving the module silently dense. Gating those layers
out lets the scheme engage on the bf16 linears and skip the fp32 ones. int8 (torch._int_mm)
tolerates non-bf16 weights, so it leaves this off and keeps quantising them."""
def filter_fn(module: Any, fqn: str = "") -> bool:
try:
@ -414,6 +436,11 @@ def make_filter_fn(min_features: int, exclude_name_tokens: tuple[str, ...] = ())
name = fqn.lower() if fqn else ""
if any(tok in name for tok in exclude_name_tokens):
return False
if require_bf16:
import torch
weight = getattr(module, "weight", None)
if weight is None or weight.dtype != torch.bfloat16:
return False
return True
return filter_fn
@ -446,12 +473,19 @@ def quantize_transformer(
from torchao.quantization import quantize_
# int8 (torch._int_mm, M>16) additionally skips the M=1 modulation / conditioning-embedder
# projections; fp8 / fp4 / mx (scaled_mm) have no such limit and quantise everything.
# projections; fp8 / fp4 / mx (scaled_mm) have no such limit and quantise everything -- but
# fp8 and mxfp8 assert a bf16 weight, so on a mixed-precision DiT (Wan / Hunyuan keep some
# fp32 linears) they must skip the non-bf16 ones or the whole pass raises and no-ops. nvfp4
# quantises fp32 weights fine, so it is not gated (see _REQUIRE_BF16_SCHEMES).
exclude = exclude_tokens_for_scheme(scheme)
quantize_(
transformer,
_make_quant_config(scheme, fast_accum = fast_accum),
filter_fn = make_filter_fn(min_features, exclude_name_tokens = exclude),
filter_fn = make_filter_fn(
min_features,
exclude_name_tokens = exclude,
require_bf16 = scheme in _REQUIRE_BF16_SCHEMES,
),
)
# Runtime-only marker (torchao tensors are not safetensors-serializable; this
# backend is inference-only, so this is purely diagnostic).

175
studio/backend/core/inference/video.py
View file

@ -43,7 +43,14 @@ from typing import Any, Optional
from loggers import get_logger
from .diffusion_attention import apply_attention_backend, select_attention_backend
from .diffusion_cache import apply_step_cache, normalize_transformer_cache
from .diffusion_cache import (
FBCACHE_MIN_STEPS,
TC_AUTO,
TC_FBCACHE,
apply_step_cache,
maybe_toggle_step_cache,
normalize_transformer_cache,
)
from .diffusion_device import resolve_diffusion_device_target
from .diffusion_memory import (
apply_memory_plan,
@ -65,11 +72,13 @@ from .diffusion_speed import (
)
from .diffusion_auto_policy import _QUANT_STEADY_FACTOR, build_resolved_record
from .diffusion_transformer_quant import (
TQ_AUTO,
dense_transformer_supported,
normalize_transformer_quant,
quantize_transformer,
select_transformer_quant_scheme,
)
from .diffusion_precision import normalize_te_quant, quantize_text_encoders
from .video_families import (
VIDEO_CANCELLED_MSG,
VIDEO_NOT_LOADED_MSG,
@ -213,11 +222,22 @@ class _VideoLoadState:
backend_flags: Optional[dict] = None
attention_backend: Optional[str] = None
transformer_cache: Optional[str] = None
# True when the cache decision was AUTO on a cache-capable DiT: generate() then
# re-checks the actual step count and toggles FBCache across FBCACHE_MIN_STEPS.
# An explicit request (off / fbcache) is never toggled.
cache_auto: bool = False
# Inputs the generation-time toggle re-applies (quantised threshold + override).
cache_quant_active: bool = False
cache_threshold: Optional[float] = None
# Dense transformer quant actually engaged ("int8" | "fp8" | "nvfp4" | "mxfp8") or
# None. Mirrors the image backend's _LoadState.transformer_quant: on a pipeline-kind
# load the dense DiT(s) can be torchao-quantised in place onto the low-precision
# tensor cores; None means they run at their loaded (bf16) precision.
transformer_quant: Optional[str] = None
# Text-encoder quant actually engaged ("fp8" | "fp8_dynamic" | "int8" | "nvfp4") or None.
# The companion text encoder (UMT5 / Gemma3 / Qwen2.5-VL) loads dense bf16 and is often the
# largest resident component; this shrinks it in place, mirroring the image backend.
text_encoder_quant: Optional[str] = None
resolved: Optional[dict] = None
@ -322,6 +342,7 @@ class VideoBackend:
family_override: Optional[str] = None,
model_kind: Optional[str] = None,
transformer_quant: Optional[str] = None,
text_encoder_quant: Optional[str] = None,
) -> VideoFamily:
"""Cheap, network-free validation shared by the route and the load path."""
kind = resolve_video_model_kind(gguf_filename, model_kind)
@ -381,6 +402,9 @@ class VideoBackend:
# only on pipeline-kind loads (the dense DiT from the base repo); an ignored value
# on a gguf/single_file load is left to the loader, matching the image backend.
normalize_transformer_quant(transformer_quant)
# Reject a malformed text_encoder_quant the same way (applies to any load kind: the dense
# text encoder is resident for pipeline / gguf / single_file alike).
normalize_te_quant(text_encoder_quant)
_ensure_mp4_encoder_available()
return fam
@ -400,6 +424,7 @@ class VideoBackend:
transformer_cache: Optional[str] = None,
transformer_cache_threshold: Optional[float] = None,
transformer_quant: Optional[str] = None,
text_encoder_quant: Optional[str] = None,
model_kind: Optional[str] = None,
) -> dict[str, Any]:
"""Validate, then run the (slow) load on a daemon thread. Returns at once."""
@ -411,6 +436,7 @@ class VideoBackend:
family_override = family_override,
model_kind = model_kind,
transformer_quant = transformer_quant,
text_encoder_quant = text_encoder_quant,
)
with self._lock:
if self._loading is not None and self._loading.error is None:
@ -434,6 +460,7 @@ class VideoBackend:
transformer_cache = transformer_cache,
transformer_cache_threshold = transformer_cache_threshold,
transformer_quant = transformer_quant,
text_encoder_quant = text_encoder_quant,
model_kind = model_kind,
_load_token = token,
),
@ -717,6 +744,19 @@ class VideoBackend:
expected_bytes = int(expected) if expected else None,
)
def loading_repo_ids(self) -> tuple[str, ...]:
"""Repo ids an in-flight background load is downloading (empty when idle).
The delete-cached guard needs this: during a load ``status()["loaded"]`` is
still False, but deleting the target repo (or its companion base) would yank
blobs and snapshot files from under the download/assembly. Mirrors the image
backend's guard (DiffusionBackend.loading_repo_ids)."""
with self._lock:
loading = self._loading
if loading is None or loading.error is not None:
return ()
return tuple(r for r in (loading.repo_id, loading.base_repo) if r)
# ── the load itself ──────────────────────────────────────────────────────
def load_pipeline(
@ -733,6 +773,7 @@ class VideoBackend:
transformer_cache: Optional[str] = None,
transformer_cache_threshold: Optional[float] = None,
transformer_quant: Optional[str] = None,
text_encoder_quant: Optional[str] = None,
model_kind: Optional[str] = None,
_load_token: Optional[int] = None,
_base_local_dir: Optional[str] = None,
@ -747,6 +788,7 @@ class VideoBackend:
family_override = family_override,
model_kind = model_kind,
transformer_quant = transformer_quant,
text_encoder_quant = text_encoder_quant,
)
kind = resolve_video_model_kind(gguf_filename, model_kind)
base = repo_id if kind == "pipeline" else resolve_video_base_repo(fam, base_repo)
@ -788,6 +830,20 @@ class VideoBackend:
# stay quantised on disk and in memory, so only dense estimates scale.
dtype_scale = 2.0 if device != "cpu" and dtype is torch.float32 else 1.0
# Precision tri-state, mirroring the image backend: an UNSET request (or
# "auto") hands the decision to the hardware ladder -- on a dense-capable
# GPU the quantised DiT (int8 minimum, fp8 on data-center silicon) is
# faster at the same resident-or-better footprint. An explicit
# "none"/"off" pins dense bf16 and an explicit scheme pins that scheme.
# Only the pipeline kind can engage it (gguf/single_file checkpoints
# already carry their own precision), and the offload guard below still
# skips it when the plan moves the DiT.
if transformer_quant is None or str(transformer_quant).strip().lower() in (
"",
"auto",
):
transformer_quant = TQ_AUTO
# ── memory plan: family-table resident estimate + frames-aware headroom.
device_memory = snapshot_device_memory(target)
components = fam.bf16_components_gb
@ -993,12 +1049,34 @@ class VideoBackend:
if quant_replanned and transformer_quant_engaged is None:
plan = bf16_plan
# ── dense text-encoder quant (opt-in): the DiT arrives quantised in a GGUF, but the
# companion encoder (Gemma3 / UMT5 / Qwen2.5-VL) loads dense bf16 from the base repo and
# is often the largest resident component. Quantise it in place, mirroring the image
# backend (diffusion.py): applied for every kind (the encoder is dense regardless of how
# the DiT was sourced) and before placement so the offload hooks move the smaller weights.
# Best-effort: quantize_text_encoders leaves any encoder it can't cast dense. int8 needs a
# per-family keep-bf16 schedule, so the family name is passed.
text_encoder_quant_engaged = quantize_text_encoders(
pipe,
target,
mode = text_encoder_quant,
family = fam.name,
offload_active = plan.offload_policy != "none",
logger = logger,
)
# ── optimisation layers, in the image backend's order: step cache FIRST
# (compile keys its fullgraph decision off an active cache: FBCache hooks
# graph-break, so compiling fullgraph before installing the cache crashes
# the first cached generation), then attention, the speed profile, and
# placement/offload last.
effective_speed = resolve_speed_mode(speed_mode, is_gguf = kind == "gguf")
# A clip denoise runs minutes, so even a dense (non-GGUF) load amortises the
# one-time regional compile within a single generation: unset resolves to the
# near-lossless `default` profile for every kind. Explicit values (incl.
# "off") are honored verbatim, and `max` is never an auto choice.
effective_speed = resolve_speed_mode(
speed_mode, is_gguf = kind == "gguf", dense_default = SPEED_DEFAULT
)
# A torchao-quantised DiT must be compiled (eager dynamic quant is ~30x slower and
# would lose to the bf16 it replaced), so force at least the regional-compile
# profile when quant engaged and the effective speed was off, matching diffusion.py.
@ -1014,18 +1092,51 @@ class VideoBackend:
# (_run_load's error handler calls _rollback_precommit_globals with this
# token). Registered BEFORE the first mutating call.
self._precommit_globals = (_load_token, backend_flags)
# Run the step cache per expert so both denoisers cache; the engaged mode is
# identical across experts.
# Step cache tri-state, mirroring the image backend: unset / "auto" lets the
# step-count policy decide (engage when this model's DEFAULT schedule reaches
# FBCACHE_MIN_STEPS, re-checked against the actual step count per generation);
# explicit "off" / "fbcache" are pinned and never toggled. Run it per expert
# so both denoisers cache; the engaged mode is identical across experts.
cache_request = normalize_transformer_cache(transformer_cache)
cache_auto = transformer_cache is None or cache_request == TC_AUTO
# GGUF checkpoints and torchao-quantised DiTs both need the higher quantised
# threshold for the cache to still trigger over the quant noise.
cache_quant_active = kind == "gguf" or transformer_quant_engaged is not None
default_cache_steps: Optional[int] = None
if cache_auto:
default_cache_steps, _ = default_video_generation_params(gguf_filename, repo_id, base)
cache_request = TC_FBCACHE if default_cache_steps >= FBCACHE_MIN_STEPS else None
cache_engaged = None
for view in views:
engaged = apply_step_cache(
view,
mode = normalize_transformer_cache(transformer_cache),
mode = cache_request,
threshold = transformer_cache_threshold,
quant_active = cache_quant_active,
logger = logger,
)
if view is pipe:
cache_engaged = engaged
# The auto decision can flip at generation time, but only on a DiT that
# supports caching at all (a non-CacheMixin transformer can never engage).
cache_may_toggle = cache_auto and callable(
getattr(getattr(pipe, "transformer", None), "enable_cache", None)
)
if cache_auto:
if cache_engaged:
cache_reason = (
f"auto: {default_cache_steps}-step default schedule reaches "
f"{FBCACHE_MIN_STEPS}; re-checked per generation"
)
elif cache_request is not None:
cache_reason = "auto: model does not support step caching"
else:
cache_reason = (
f"auto: {default_cache_steps}-step default schedule is below "
f"{FBCACHE_MIN_STEPS}; re-checked per generation"
)
else:
cache_reason = "requested"
attention_engaged = None
speed_optims: tuple = ()
for view in views:
@ -1048,7 +1159,10 @@ class VideoBackend:
is_gguf = gguf_transformer,
family = fam,
speed_mode = effective_speed,
cache_active = cache_engaged is not None,
# An auto cache that could still engage mid-session also drops
# fullgraph: enabling FBCache under a fullgraph-compiled DiT would
# crash the first cached generation.
cache_active = cache_engaged is not None or cache_may_toggle,
offload_active = plan.offload_policy != "none",
)
if view is pipe:
@ -1089,7 +1203,9 @@ class VideoBackend:
effective_speed,
"quantized transformer requires compile"
if transformer_quant_engaged is not None
else "GGUF video loads default to the near-lossless compile profile",
else "clip denoises amortise the one-time compile within a single run"
if speed_mode is None
else "requested",
),
"attention_backend": (
attention_backend,
@ -1097,9 +1213,9 @@ class VideoBackend:
"cuDNN fused attention on NVIDIA when a speed profile is active",
),
"transformer_cache": (
transformer_cache,
None if cache_auto else transformer_cache,
cache_engaged or "off",
"step cache engages on many-step schedules only",
cache_reason,
),
"transformer_quant": (
transformer_quant,
@ -1113,6 +1229,13 @@ class VideoBackend:
else "not engaged (dense bf16 DiT loaded)"
),
),
"text_encoder_quant": (
text_encoder_quant,
text_encoder_quant_engaged or "off",
"dense text encoder quantised in place"
if text_encoder_quant_engaged is not None
else "not engaged (dense bf16 text encoder loaded)",
),
}
)
@ -1141,7 +1264,11 @@ class VideoBackend:
backend_flags = backend_flags,
attention_backend = attention_engaged,
transformer_cache = cache_engaged,
cache_auto = cache_may_toggle,
cache_quant_active = cache_quant_active,
cache_threshold = transformer_cache_threshold,
transformer_quant = transformer_quant_engaged,
text_encoder_quant = text_encoder_quant_engaged,
resolved = resolved,
)
# Ownership of the globals transferred to _state / _teardown_state.
@ -1316,6 +1443,34 @@ class VideoBackend:
# cancel and restore it afterwards.
progress_ctx = _scheduler_step_progress(pipe, _on_scheduler_step)
# An AUTO cache decision is re-checked against the ACTUAL step count,
# mirroring the image backend: a many-step request gains FBCache even
# when the load's default schedule kept it off, and a few-step request
# drops it. Explicit choices never toggle. Runs per view so a dual-DiT
# MoE toggles both experts.
if state.cache_auto:
toggled = state.transformer_cache
for view in _views_for(pipe, fam):
toggled = maybe_toggle_step_cache(
view,
steps = steps,
quant_active = state.cache_quant_active,
threshold = state.cache_threshold,
logger = logger,
)
if toggled != state.transformer_cache:
# _VideoLoadState is frozen (loads swap it as one unit); this
# tracks the pipe-level toggle that already happened so
# status() reports the true cache state.
object.__setattr__(state, "transformer_cache", toggled)
entry = (state.resolved or {}).get("transformer_cache")
if isinstance(entry, dict):
entry["value"] = toggled or "off"
entry["reason"] = (
f"auto: {steps}-step generation "
+ ("reaches" if toggled else "is below")
+ f" {FBCACHE_MIN_STEPS}"
)
if state.transformer_cache:
self._reset_step_cache(pipe)
try:
@ -1464,6 +1619,7 @@ class VideoBackend:
"attention_backend": None,
"transformer_cache": None,
"transformer_quant": None,
"text_encoder_quant": None,
"has_audio": False,
"defaults": None,
"resolved": None,
@ -1488,6 +1644,7 @@ class VideoBackend:
"attention_backend": state.attention_backend,
"transformer_cache": state.transformer_cache,
"transformer_quant": state.transformer_quant,
"text_encoder_quant": state.text_encoder_quant,
"has_audio": fam.has_audio,
"defaults": {
"steps": default_steps,

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

@ -71,6 +71,100 @@ def video_path(video_id: str) -> Optional[Path]:
return path if path.is_file() else None
def transcode(video_id: str, fmt: str) -> Optional[bytes]:
"""Re-encode a stored MP4 for the Download menu: "webm" (VP9, web embeds)
or "gif" (shareable preview). Returns the encoded bytes, or None when the
id doesn't resolve. Raises RuntimeError when the codec/deps are missing so
the route can 501 with a clear message. MP4 downloads never come through
here -- the /file route streams the original bytes."""
path = video_path(video_id)
if path is None:
return None
normalized = fmt.strip().lower()
if normalized == "webm":
return _transcode_webm(path)
if normalized == "gif":
return _transcode_gif(path)
raise ValueError(f"Unsupported export format '{fmt}'. Use webm or gif.")
def _transcode_webm(path: Path) -> bytes:
import io
try:
import av
except Exception as exc: # noqa: BLE001 -- no PyAV -> no transcode
raise RuntimeError("WebM export needs the 'av' package (PyAV).") from exc
buf = io.BytesIO()
try:
with av.open(str(path)) as src, av.open(buf, "w", format = "webm") as dst:
if not src.streams.video:
raise RuntimeError("WebM export failed: the clip has no video stream.")
in_v = src.streams.video[0]
rate = in_v.average_rate or 24
out_v = dst.add_stream("libvpx-vp9", rate = rate)
out_v.width = in_v.codec_context.width
out_v.height = in_v.codec_context.height
out_v.pix_fmt = "yuv420p"
# Realtime-oriented settings: VP9's default "good" profile encodes a
# few frames per second; cpu-used 8 + row-mt is many times faster at
# a small quality cost, right for a download button.
out_v.options = {"deadline": "realtime", "cpu-used": "8", "row-mt": "1"}
for frame in src.decode(in_v):
for packet in out_v.encode(frame.reformat(format = "yuv420p")):
dst.mux(packet)
for packet in out_v.encode():
dst.mux(packet)
except RuntimeError:
raise
except Exception as exc: # noqa: BLE001 -- surface as "encoder unavailable"
raise RuntimeError(f"WebM export failed (libvpx-vp9 unavailable?): {exc}") from exc
# Audio is intentionally dropped: Opus muxing needs a 48 kHz resample chain
# and most exported clips are silent; the original MP4 keeps the audio.
return buf.getvalue()
def _transcode_gif(path: Path) -> bytes:
import io
try:
import av
from PIL import Image
except Exception as exc: # noqa: BLE001 -- missing deps -> no transcode
raise RuntimeError("GIF export needs the 'av' and 'Pillow' packages.") from exc
frames: list[Any] = []
try:
with av.open(str(path)) as src:
if not src.streams.video:
raise RuntimeError("GIF export failed: the clip has no video stream.")
in_v = src.streams.video[0]
rate = float(in_v.average_rate or 24)
# Full-rate GIFs are enormous and stutter in chat apps; ~12 fps is
# the sweet spot, achieved by skipping source frames.
step = max(1, round(rate / 12))
for i, frame in enumerate(src.decode(in_v)):
if i % step:
continue
frames.append(frame.to_image().convert("P", palette = Image.Palette.ADAPTIVE))
except RuntimeError:
raise
except Exception as exc: # noqa: BLE001 -- surface as "decoder unavailable"
raise RuntimeError(f"GIF export failed to decode the clip: {exc}") from exc
if not frames:
raise RuntimeError("GIF export decoded no frames.")
duration_ms = max(20, int(1000 * step / rate))
buf = io.BytesIO()
frames[0].save(
buf,
format = "GIF",
save_all = True,
append_images = frames[1:],
duration = duration_ms,
loop = 0,
)
return buf.getvalue()
def _sidecar_path(video_id: str) -> Path:
return gallery_dir() / f"{video_id}.json"

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

@ -1743,12 +1743,13 @@ class DiffusionLoadRequest(BaseModel):
"default (also regional torch.compile where eligible), "
"max (also TF32 + fused QKV).",
)
text_encoder_quant: Optional[Literal["fp8", "nvfp4"]] = Field(
text_encoder_quant: Optional[Literal["fp8", "fp8_dynamic", "int8", "nvfp4"]] = Field(
None,
description = "Quantise the companion text encoder(s): fp8 (~2x smaller, "
"CUDA cc>=8.9) or nvfp4 (~4x smaller, Blackwell sm_100+). A "
"memory-vs-quality tradeoff (shifts fine detail), not free; "
"pairs well with balanced mode.",
description = "Quantise the companion text encoder(s): fp8 (layerwise cast, ~2x smaller, "
"CUDA cc>=8.9), fp8_dynamic (torchao compute fp8 on the tensor cores, ~2x + faster, "
"cc>=8.9), int8 (torchao compute int8 with per-family keep-bf16 layers; falls back to "
"fp8 where no schedule exists; cc>=8.0), or nvfp4 (~4x smaller, Blackwell sm_100+). A "
"memory-vs-quality tradeoff (shifts fine detail), not free; pairs well with balanced mode.",
)
transformer_quant: Optional[Literal["auto", "none", "off", "int8", "fp8", "nvfp4", "mxfp8"]] = (
Field(
@ -2341,6 +2342,16 @@ class VideoLoadRequest(BaseModel):
"backend's transformer_quant field.",
)
)
text_encoder_quant: Optional[Literal["fp8", "fp8_dynamic", "int8", "nvfp4"]] = Field(
None,
description = "Quantise the dense companion text encoder (Gemma3 / UMT5 / Qwen2.5-VL), "
"which loads bf16 from the base repo regardless of how the DiT was sourced and is often "
"the largest resident component. fp8 = diffusers layerwise casting (memory only, cc >= "
"8.9); fp8_dynamic = torchao per-row fp8 COMPUTE on the tensor cores (cc >= 8.9); int8 = "
"torchao int8 COMPUTE with per-family keep-bf16 selection (cc >= 8.0; falls back to fp8 "
"for a family without a measured schedule); nvfp4 = torchao 4-bit weight-only (Blackwell "
"sm_100+). null keeps the encoder dense. Mirrors the image backend's field.",
)
@field_validator("attention_backend", mode = "before")
@classmethod
@ -2506,6 +2517,12 @@ class VideoStatusResponse(BaseModel):
"mxfp8 | null (null = the DiT(s) run at their loaded bf16 precision). For a dual-expert "
"MoE family both experts share the reported scheme.",
)
text_encoder_quant: Optional[str] = Field(
None,
description = "Text-encoder quant engaged: fp8 | fp8_dynamic | int8 | nvfp4 | null "
"(null = the dense bf16 encoder is loaded). An int8 request without a per-family "
"keep-bf16 schedule is reported as the fp8 it fell back to.",
)
has_audio: bool = Field(
False, description = "Whether the loaded family produces a synchronized audio track"
)

84
studio/backend/routes/models.py
View file

@ -32,6 +32,10 @@ class CachedModelRepo(BaseModel):
# drop the value the handler sets, letting image-only repos pass the chat picker's
# task gate.
task: Optional[str] = None
# True when the snapshot is incomplete (a cancelled/partial download left only some
# weights). The picker must not treat a partial base repo as a usable download, or an
# On Device click routes to a fresh multi-GB re-download instead of the complete GGUF.
partial: Optional[bool] = None
class CachedModelsResponse(BaseModel):
@ -3363,6 +3367,44 @@ def _repo_is_diffusers(repo_info) -> bool:
return False
def _cached_repo_partial(repo_id: str, repo_cache_dir: Optional[Path] = None) -> bool:
"""Whether the cached model snapshot is incomplete (cancelled/partial download).
Reuses the hub inventory scan's snapshot-partial detector (cancel marker, legacy
.incomplete blob, manifest walk -- cheapest first). ``repo_cache_dir`` scopes all three
signals to the specific snapshot being listed: without it the scan spans every HF cache
root, so a stale .incomplete copy in one root would flag a complete copy in another as
partial and hide it from the picker (the sibling inventory paths all scope the same way).
Best-effort: a detection error reports not-partial so a scan glitch never hides a
genuinely usable repo."""
try:
from hub.utils.inventory_scan import is_snapshot_partial
return bool(is_snapshot_partial("model", repo_id, repo_cache_dir))
except Exception: # noqa: BLE001 -- never fail the listing over a partial probe
return False
def _cached_repo_task(repo_info) -> Optional[str]:
"""Pipeline task for a cached non-GGUF repo: 'text-to-video' for repos the
video backend can load as full pipelines (its trust list / family detector),
else 'text-to-image' for diffusers image repos, else None (chat). Without the
video tag, cached Lightricks / Wan / Hunyuan pipelines never surfaced in the
Video picker's On Device list -- everything diffusers was blanket-tagged
text-to-image."""
repo_id = getattr(repo_info, "repo_id", "") or ""
try:
from core.inference.video import _is_trusted_video_repo
from core.inference.video_families import detect_video_family
# Both gates: a detected video family (so unsloth image repos don't
# match) AND the load path's own trust rule (so an untrusted video repo
# isn't advertised as loadable).
if detect_video_family(repo_id) is not None and _is_trusted_video_repo(repo_id):
return _VIDEO_GEN_TASK
except Exception:
pass
return "text-to-image" if _repo_is_diffusers(repo_info) else None
@router.get("/cached-models", response_model = CachedModelsResponse)
async def list_cached_models(
current_subject: str = Depends(get_current_subject),
@ -3405,12 +3447,22 @@ async def list_cached_models(
)
key = repo_id.lower()
existing = seen_lower.get(key)
if existing is None or total_size > existing["size_bytes"]:
is_partial = _cached_repo_partial(repo_id, Path(repo_info.repo_path))
# Prefer the most COMPLETE snapshot, then the largest. The picker drops partial
# rows, so a partial copy in one cache root must not shadow a smaller COMPLETE
# copy in another (that would make a usable model vanish from On Device).
# Completeness wins outright; size only breaks ties among equal completeness.
if existing is None or (not is_partial, total_size) > (
not bool(existing.get("partial")),
existing["size_bytes"],
):
row = {
"repo_id": repo_id,
"size_bytes": total_size,
"task": "text-to-image" if _repo_is_diffusers(repo_info) else None,
"task": _cached_repo_task(repo_info),
}
if is_partial:
row["partial"] = True
# Keep the newest timestamp across duplicate caches;
# attach only when known so absent rows sort as oldest.
lm = max(last_modified, (existing or {}).get("last_modified", 0.0))
@ -3516,6 +3568,34 @@ async def delete_cached_model(
except Exception:
pass
# And refuse if the Video backend has this repo loaded or is downloading it: cached non-GGUF
# video repos now surface in the Video On-Device picker with the normal delete action, but the
# guards above only cover chat + the Images engine, so without this a loaded/loading Wan / LTX /
# Hunyuan pipeline could have its HF snapshot removed from under it. Mirror the Images guard.
try:
from core.inference.video import get_video_backend
video_backend = get_video_backend()
video_status = video_backend.status()
if video_status.get("loaded") and video_status.get("repo_id"):
loaded_id = str(video_status["repo_id"]).lower()
if loaded_id == repo_id.lower() or loaded_id.startswith(repo_id.lower()):
raise HTTPException(
status_code = 400,
detail = "Unload the model before deleting",
)
for lid in getattr(video_backend, "loading_repo_ids", tuple)():
lid = str(lid).lower()
if lid == repo_id.lower() or lid.startswith(repo_id.lower()):
raise HTTPException(
status_code = 400,
detail = "A Video model load is using this repo; wait for it to finish",
)
except HTTPException:
raise
except Exception:
pass
try:
cache_scans = _all_hf_cache_scans()

32
studio/backend/routes/video.py
View file

@ -95,6 +95,7 @@ async def load_video_model(
family_override = request.family_override,
model_kind = request.model_kind,
transformer_quant = request.transformer_quant,
text_encoder_quant = request.text_encoder_quant,
)
# Refuse while training is running: a multi-GB video pipeline would compete
# with the training subprocess for VRAM. Mirrors the image-load guard.
@ -123,6 +124,7 @@ async def load_video_model(
transformer_cache = request.transformer_cache,
transformer_cache_threshold = request.transformer_cache_threshold,
transformer_quant = request.transformer_quant,
text_encoder_quant = request.text_encoder_quant,
model_kind = request.model_kind,
)
return VideoStatusResponse(**status_dict)
@ -288,6 +290,36 @@ async def get_gallery_video_file(
)
@router.get("/video/gallery/{video_id}/export")
async def export_gallery_video(
video_id: str,
format: str = "webm",
current_subject: str = Depends(get_current_subject),
):
"""Download-menu transcodes: WebM (VP9) or GIF, re-encoded on demand from the
stored MP4 (which the /file route serves verbatim). 501 with a clear message
when the codec/deps for the requested format are missing."""
from core.inference import video_gallery
fmt = format.strip().lower()
if fmt not in ("webm", "gif"):
raise HTTPException(status_code = 400, detail = "Unsupported format. Use webm or gif.")
try:
data = await asyncio.to_thread(video_gallery.transcode, video_id, fmt)
except RuntimeError as exc:
raise HTTPException(status_code = 501, detail = str(exc)) from exc
if data is None:
raise HTTPException(status_code = 404, detail = "Video not found.")
from fastapi.responses import Response
return Response(
content = data,
media_type = "video/webm" if fmt == "webm" else "image/gif",
# Transcodes are deterministic per id+format; let the browser cache them.
headers = {"Cache-Control": "private, max-age=31536000, immutable"},
)
@router.delete("/video/gallery/{video_id}")
async def delete_gallery_video(video_id: str, current_subject: str = Depends(get_current_subject)):
from core.inference import video_gallery

120
studio/backend/tests/test_cached_gguf_routes.py
View file

@ -260,6 +260,45 @@ def test_list_cached_models_skips_non_suffix_repo_when_gguf_files_exist(monkeypa
assert result["cached"] == []
def test_list_cached_models_prefers_complete_over_larger_partial(monkeypatch, tmp_path):
# The same repo cached in two roots: a LARGER but PARTIAL copy must not shadow a SMALLER but
# COMPLETE one. The picker drops partial rows, so picking the partial winner by size alone would
# make a usable model vanish from On Device. Completeness wins over size.
complete = _repo(
"Org/Dup",
[_file("model.safetensors", 10_000)],
tmp_path / "root_a" / "models--Org--Dup",
)
partial = _repo(
"Org/Dup",
[_file("model.safetensors", 15_000)],
tmp_path / "root_b" / "models--Org--Dup",
)
# The larger copy (root_b) is the partial one; the smaller (root_a) is complete.
monkeypatch.setattr(
models_route,
"_cached_repo_partial",
lambda repo_id, repo_cache_dir = None: "root_b" in str(repo_cache_dir),
)
monkeypatch.setattr(models_route, "_cached_repo_task", lambda repo_info: None)
# List the partial (larger) FIRST, so the old size-only rule would have picked it.
monkeypatch.setattr(
models_route,
"_all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [partial, complete])],
)
result = asyncio.run(models_route.list_cached_models(current_subject = "test-user"))
assert len(result["cached"]) == 1
row = result["cached"][0]
assert row["repo_id"] == "Org/Dup"
# The COMPLETE (smaller) copy won: it is not flagged partial and carries its 10_000 size.
assert row.get("partial") is not True
assert row["size_bytes"] == 10_000
def test_list_cached_gguf_includes_mixed_repo_with_gguf_and_safetensors(monkeypatch, tmp_path):
"""Mixed repo still surfaces in cached-gguf as a GGUF download."""
mixed = _repo(
@ -812,3 +851,84 @@ def test_delete_cached_refuses_diffusion_loaded_repo(monkeypatch):
except HTTPException as e:
assert e.status_code == 400
assert "Unload the model before deleting" in e.detail
def test_delete_cached_refuses_video_loaded_repo(monkeypatch):
# The guard also refuses deleting a repo the Video backend has loaded: cached non-GGUF video
# repos now surface in the Video On-Device picker with the normal delete action, so without
# this a live Wan / LTX / Hunyuan pipeline's HF snapshot could be removed from under it.
from fastapi import HTTPException
import core.inference.diffusion as diffusion_mod
import core.inference.video as video_mod
import routes.inference as routes_inference
# Chat + Images backends report nothing loaded; only the Video backend holds the repo.
monkeypatch.setattr(
routes_inference,
"get_llama_cpp_backend",
lambda: SimpleNamespace(is_loaded = False, model_identifier = None),
)
monkeypatch.setattr(
models_route, "get_inference_backend", lambda: SimpleNamespace(active_model_name = None)
)
monkeypatch.setattr(
diffusion_mod,
"get_diffusion_backend",
lambda: SimpleNamespace(
status = lambda: {"loaded": False, "repo_id": None}, loading_repo_ids = lambda: ()
),
)
monkeypatch.setattr(
video_mod,
"get_video_backend",
lambda: SimpleNamespace(
status = lambda: {"loaded": True, "repo_id": "Lightricks/LTX-2"},
loading_repo_ids = lambda: (),
),
)
try:
asyncio.run(
models_route.delete_cached_model(
repo_id = "Lightricks/LTX-2", variant = None, current_subject = "u"
)
)
assert False, "expected HTTPException refusing the delete"
except HTTPException as e:
assert e.status_code == 400
assert "Unload the model before deleting" in e.detail
def test_cached_repo_partial_scopes_probe_to_snapshot_dir(monkeypatch):
# The partial probe must be scoped to the snapshot row being listed. Unscoped, the scan
# spans every HF cache root, so a stale .incomplete copy in one root would flag a complete
# copy living in another root as partial and hide the usable model from the picker. Verify
# _cached_repo_partial forwards the snapshot dir (matching the sibling inventory paths).
import hub.utils.inventory_scan as scan
calls = []
def _fake(
repo_type,
repo_id,
repo_cache_dir = None,
):
calls.append((repo_type, repo_id, repo_cache_dir))
return False
monkeypatch.setattr(scan, "is_snapshot_partial", _fake)
snapshot_dir = Path("/root_a/models--Org--Repo/snapshots/abc")
assert models_route._cached_repo_partial("Org/Repo", snapshot_dir) is False
assert calls == [("model", "Org/Repo", snapshot_dir)]
# When that specific snapshot is partial, the row is flagged.
monkeypatch.setattr(scan, "is_snapshot_partial", lambda *a, **k: True)
assert models_route._cached_repo_partial("Org/Repo", snapshot_dir) is True
# A probe error is swallowed (never hides a usable repo over a scan glitch).
def _boom(*a, **k):
raise RuntimeError("scan glitch")
monkeypatch.setattr(scan, "is_snapshot_partial", _boom)
assert models_route._cached_repo_partial("Org/Repo", snapshot_dir) is False

203
studio/backend/tests/test_diffusion_backend.py
View file

@ -465,6 +465,209 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path):
assert backend.is_loaded is False
def test_dense_speed_auto_defers_compile_to_third_generation(fake_runtime, tmp_path, monkeypatch):
# Dense models with speed unset stay bit-identical eager for the first two
# generations; the 3rd engages the `default` profile mid-session (repeated
# use amortises the one-time compile), upgrading attention alongside it.
from core.inference import diffusion as dmod
monkeypatch.setattr(dmod, "compile_eligible", lambda *a, **k: True)
monkeypatch.setattr(
dmod,
"apply_speed_optims",
lambda pipe, target, **k: {"compiled": k.get("speed_mode") == "default"},
)
monkeypatch.setattr(dmod, "apply_attention_backend", lambda pipe, backend, logger = None: backend)
monkeypatch.setattr(
dmod,
"select_attention_backend",
lambda target, requested, speed_active = False: ("_native_cudnn" if speed_active else None),
)
monkeypatch.setattr(dmod.compile_cache, "begin", lambda **k: None)
(tmp_path / "model.safetensors").write_bytes(b"weights")
backend = DiffusionBackend()
status = backend.load_pipeline(
str(tmp_path),
gguf_filename = "model.safetensors",
base_repo = "base/repo",
family_override = "qwen-image",
)
assert status["speed_mode"] == "off"
assert status["resolved"]["speed_mode"]["value"] == "deferred"
assert status["resolved"]["speed_mode"]["source"] == "auto"
backend.generate(prompt = "one")
backend.generate(prompt = "two")
assert backend.status()["speed_mode"] == "off" # first two stay exact eager
backend.generate(prompt = "three")
status3 = backend.status()
assert status3["speed_mode"] == "default"
assert "compiled" in status3["speed_optims"]
assert status3["attention_backend"] == "_native_cudnn"
assert status3["resolved"]["speed_mode"]["value"] == "default"
# An explicit "off" is pinned: no deferral, still eager after 3 generations.
backend.unload()
status_off = backend.load_pipeline(
str(tmp_path),
gguf_filename = "model.safetensors",
base_repo = "base/repo",
family_override = "qwen-image",
speed_mode = "off",
)
assert status_off["resolved"]["speed_mode"]["value"] == "off"
for p in ("a", "b", "c"):
backend.generate(prompt = p)
assert backend.status()["speed_mode"] == "off"
backend.unload()
def test_deferred_speed_skips_when_lora_requested(fake_runtime, tmp_path, monkeypatch):
# A compiled transformer rejects LoRA (supports_lora is False once compiled), and _apply_loras
# raises before its unchanged-selection no-op, so engaging the deferred compile on a generation
# that requests a LoRA would permanently break every LoRA generation on this load. The deferral
# must skip while a LoRA is requested and engage only on a later LoRA-free generation.
from core.inference import diffusion as dmod
monkeypatch.setattr(dmod, "compile_eligible", lambda *a, **k: True)
engaged: list = []
def fake_engage(self, state):
engaged.append(state.generation_count)
state.speed_deferred = False # mirror the real helper: engage once, then clear
monkeypatch.setattr(DiffusionBackend, "_engage_deferred_speed", fake_engage)
# LoRA loading is covered elsewhere; stub it so this test needs no adapter file.
monkeypatch.setattr(DiffusionBackend, "_apply_loras", lambda self, state, loras, cancel: None)
(tmp_path / "model.safetensors").write_bytes(b"weights")
backend = DiffusionBackend()
backend.load_pipeline(
str(tmp_path),
gguf_filename = "model.safetensors",
base_repo = "base/repo",
family_override = "qwen-image",
)
backend.generate(prompt = "one")
backend.generate(prompt = "two")
# 3rd generation requests a LoRA: the deferral must be skipped (pipe stays eager, LoRA-capable).
backend.generate(prompt = "three", loras = [("adapter", 1.0)])
assert engaged == []
# 4th generation without a LoRA: the deferral now engages (the guard is LoRA-specific, not off).
backend.generate(prompt = "four")
assert len(engaged) == 1
def test_deferred_speed_skips_while_adapter_attached(fake_runtime, tmp_path, monkeypatch):
# Even a generation that requests NO LoRA must defer the compile while an adapter from a PRIOR
# generation is still attached: _apply_loras runs AFTER the engage, so compiling here would bake
# the resident adapter into the graph and the subsequent unload (swallowed on a compiled pipe)
# would leave it active forever -- silent wrong output. Defer until _apply_loras clears it.
from core.inference import diffusion as dmod
monkeypatch.setattr(dmod, "compile_eligible", lambda *a, **k: True)
engaged: list = []
def fake_engage(self, state):
engaged.append(state.generation_count)
state.speed_deferred = False
monkeypatch.setattr(DiffusionBackend, "_engage_deferred_speed", fake_engage)
# Track the attached set on the pipe, mirroring the real _apply_loras marker (_unsloth_loras).
def fake_apply(self, state, loras, cancel):
specs = [(i, w) for (i, w) in (loras or []) if w != 0]
state.pipe._unsloth_loras = tuple(specs)
monkeypatch.setattr(DiffusionBackend, "_apply_loras", fake_apply)
(tmp_path / "model.safetensors").write_bytes(b"weights")
backend = DiffusionBackend()
backend.load_pipeline(
str(tmp_path),
gguf_filename = "model.safetensors",
base_repo = "base/repo",
family_override = "qwen-image",
)
# Gens 1-2 attach an adapter, so it is still resident going into gen 3.
backend.generate(prompt = "one", loras = [("adapter", 1.0)])
backend.generate(prompt = "two", loras = [("adapter", 1.0)])
# Gen 3 requests NO LoRA but the adapter is still attached -> defer (no compile-with-adapter).
backend.generate(prompt = "three")
assert engaged == []
# Gen 3's _apply_loras([]) cleared the adapter; gen 4 is genuinely LoRA-free -> engage.
backend.generate(prompt = "four")
assert len(engaged) == 1
def test_deferred_speed_preserves_explicit_attention(fake_runtime, tmp_path, monkeypatch):
# A dense model loaded with Speed left on Auto but Attention explicitly pinned
# (e.g. "native" to avoid cuDNN) must KEEP that choice when the 3rd generation
# engages the deferred `default` profile. The auto cuDNN upgrade only applies when
# attention was left on auto -- never when the caller pinned a backend.
from core.inference import diffusion as dmod
monkeypatch.setattr(dmod, "compile_eligible", lambda *a, **k: True)
monkeypatch.setattr(
dmod,
"apply_speed_optims",
lambda pipe, target, **k: {"compiled": k.get("speed_mode") == "default"},
)
monkeypatch.setattr(dmod, "apply_attention_backend", lambda pipe, backend, logger = None: backend)
# A select mock that -- unlike a bare "auto -> cuDNN" stub -- HONORS an explicit
# request: "native" stays on the default (None) even under a speed profile, and only
# a left-unset ("auto"/None) request upgrades to cuDNN when speed is active.
def fake_select(
target,
requested,
speed_active = False,
):
if requested in (None, "", "auto"):
return "_native_cudnn" if speed_active else None
if str(requested).lower() in ("native", "sdpa"):
return None
return requested
monkeypatch.setattr(dmod, "select_attention_backend", fake_select)
monkeypatch.setattr(dmod.compile_cache, "begin", lambda **k: None)
(tmp_path / "model.safetensors").write_bytes(b"weights")
backend = DiffusionBackend()
backend.load_pipeline(
str(tmp_path),
gguf_filename = "model.safetensors",
base_repo = "base/repo",
family_override = "qwen-image",
attention_backend = "native",
)
backend.generate(prompt = "one")
backend.generate(prompt = "two")
backend.generate(prompt = "three") # deferred profile engages here
status = backend.status()
assert status["speed_mode"] == "default" # the compile profile still engaged
assert "compiled" in status["speed_optims"]
# The pinned "native" survived: NOT silently upgraded to cuDNN.
assert status["attention_backend"] is None
assert status["resolved"]["attention_backend"]["value"] == "native"
assert status["resolved"]["attention_backend"]["source"] == "explicit"
# Control: with attention left on auto, the same 3rd-generation deferral DOES upgrade
# to cuDNN -- so the assertion above is not vacuously passing.
backend.unload()
backend.load_pipeline(
str(tmp_path),
gguf_filename = "model.safetensors",
base_repo = "base/repo",
family_override = "qwen-image",
)
for p in ("a", "b", "c"):
backend.generate(prompt = p)
assert backend.status()["attention_backend"] == "_native_cudnn"
backend.unload()
def _tiny_png_b64() -> str:
import base64
import io

202
studio/backend/tests/test_diffusion_precision.py
View file

@ -14,9 +14,14 @@ import types
import pytest
import core.inference.diffusion_precision as dp
from core.inference.diffusion_precision import (
TE_QUANT_FP8,
TE_QUANT_FP8_DYNAMIC,
TE_QUANT_INT8,
TE_QUANT_NVFP4,
_cast_int8_selective,
_keep_bf16_block_fqns,
normalize_te_quant,
quantize_text_encoders,
te_quant_supported,
@ -44,8 +49,12 @@ def _stub_torch(
if with_fp8:
torch.float8_e4m3fn = "float8_e4m3fn"
# _cast_fp8 skips nn.Embedding tables (skip_modules_classes) to keep prompt
# tokens full precision, so the stub torch must expose torch.nn.Embedding.
torch.nn = types.SimpleNamespace(Embedding = type("Embedding", (), {}))
# tokens full precision, and _keep_bf16_block_fqns walks for nn.ModuleList block
# stacks, so the stub torch must expose both.
torch.nn = types.SimpleNamespace(
Embedding = type("Embedding", (), {}),
ModuleList = type("ModuleList", (list,), {}),
)
torch.cuda = types.SimpleNamespace(get_device_capability = lambda *a: cc)
monkeypatch.setitem(sys.modules, "torch", torch)
return torch
@ -77,6 +86,9 @@ def test_normalize_te_quant():
assert normalize_te_quant("none") is None
assert normalize_te_quant("FP8") == TE_QUANT_FP8
assert normalize_te_quant("NVFP4") == TE_QUANT_NVFP4
assert normalize_te_quant("int8") == TE_QUANT_INT8
# Hyphens fold to underscores so "fp8-dynamic" is accepted.
assert normalize_te_quant("FP8-Dynamic") == TE_QUANT_FP8_DYNAMIC
with pytest.raises(ValueError):
normalize_te_quant("int2")
@ -99,6 +111,31 @@ def test_nvfp4_supported_requires_blackwell(monkeypatch):
assert te_quant_supported(_target(), TE_QUANT_NVFP4) is False
def test_int8_supported_requires_sm80(monkeypatch):
# int8 tensor cores (torch._int_mm) need Ampere sm_80+.
_stub_torch(monkeypatch, cc = (8, 0))
assert te_quant_supported(_target(), TE_QUANT_INT8) is True
_stub_torch(monkeypatch, cc = (7, 5))
assert te_quant_supported(_target(), TE_QUANT_INT8) is False
# Still needs CUDA + bf16 like every mode.
_stub_torch(monkeypatch, cc = (8, 0))
assert te_quant_supported(_target(device = "cpu"), TE_QUANT_INT8) is False
def test_fp8_dynamic_supported_requires_sm89_and_fp8(monkeypatch):
# Compute fp8 (torch._scaled_mm) needs fp8-GEMM silicon: Ada sm_89+ / Hopper / Blackwell.
_stub_torch(monkeypatch, cc = (8, 9))
assert te_quant_supported(_target(), TE_QUANT_FP8_DYNAMIC) is True
_stub_torch(monkeypatch, cc = (9, 0))
assert te_quant_supported(_target(), TE_QUANT_FP8_DYNAMIC) is True
# Ampere (8.0) has int8 but not fp8 GEMM.
_stub_torch(monkeypatch, cc = (8, 0))
assert te_quant_supported(_target(), TE_QUANT_FP8_DYNAMIC) is False
# No fp8 dtype at all -> unsupported regardless of arch.
_stub_torch(monkeypatch, with_fp8 = False, cc = (9, 0))
assert te_quant_supported(_target(), TE_QUANT_FP8_DYNAMIC) is False
# ── apply ─────────────────────────────────────────────────────────────────────
@ -155,3 +192,164 @@ def test_quantize_tolerates_caster_failure(monkeypatch):
pipe = types.SimpleNamespace(text_encoder = object())
# The only encoder fails to cast -> nothing applied -> None.
assert quantize_text_encoders(pipe, _target(), mode = "fp8") is None
# ── int8 (selective) + fp8_dynamic routing ─────────────────────────────────────
def test_quantize_int8_uses_family_keep_bf16_schedule(monkeypatch):
# int8 for a family with a measured schedule routes to the selective caster with
# that family's (skip_first, skip_last); qwen-image keeps first+last 6 blocks bf16.
_stub_torch(monkeypatch, cc = (10, 0))
calls: list = []
monkeypatch.setattr(
dp, "_cast_int8_selective", lambda enc, tgt, first, last: calls.append((enc, first, last))
)
te = object()
pipe = types.SimpleNamespace(text_encoder = te)
mode = quantize_text_encoders(pipe, _target(), mode = "int8", family = "qwen-image")
assert mode == TE_QUANT_INT8
assert calls == [(te, 6, 6)]
def test_quantize_int8_unknown_family_falls_back_to_fp8(monkeypatch):
# A family without an int8 keep-bf16 schedule falls back to layerwise fp8 (logged),
# never silently running full int8 that would degrade the encoder.
_stub_torch(monkeypatch, cc = (10, 0))
int8_calls: list = []
fp8_calls: list = []
monkeypatch.setattr(dp, "_cast_int8_selective", lambda *a: int8_calls.append(a))
monkeypatch.setattr(dp, "_cast_fp8", lambda enc, tgt: fp8_calls.append(enc))
te = object()
pipe = types.SimpleNamespace(text_encoder = te)
mode = quantize_text_encoders(pipe, _target(), mode = "int8", family = "wan-umt5")
assert mode == TE_QUANT_FP8
assert int8_calls == [] and fp8_calls == [te]
def test_quantize_fp8_dynamic_uses_compute_caster(monkeypatch):
# fp8_dynamic routes to the torchao per-row compute caster (not the layerwise one)
# and needs no per-family schedule.
_stub_torch(monkeypatch, cc = (9, 0))
calls: list = []
monkeypatch.setattr(dp, "_cast_fp8_dynamic", lambda enc, tgt: calls.append(enc))
te = object()
pipe = types.SimpleNamespace(text_encoder = te)
mode = quantize_text_encoders(pipe, _target(), mode = "fp8_dynamic")
assert mode == TE_QUANT_FP8_DYNAMIC
assert calls == [te]
def test_quantize_int8_unsupported_hw_is_noop(monkeypatch):
# int8 on pre-Ampere silicon (no int8 tensor cores) applies nothing.
_stub_torch(monkeypatch, cc = (7, 5))
monkeypatch.setattr(dp, "_cast_int8_selective", lambda *a: pytest.fail("must not cast"))
pipe = types.SimpleNamespace(text_encoder = object())
assert quantize_text_encoders(pipe, _target(), mode = "int8", family = "qwen-image") is None
def test_quantize_te_skips_torchao_modes_under_offload(monkeypatch):
# The torchao modes (int8-with-schedule / fp8_dynamic / nvfp4) produce tensor subclasses that
# reject Module.to(), which an offload hook uses, so they must be skipped under offload (the DiT
# path skips torchao quant for the same reason). Hardware supports every mode here, so a None
# result proves the offload skip, not a capability gate; the casters fail if wrongly invoked.
_stub_torch(monkeypatch, cc = (10, 0))
monkeypatch.setattr(
dp, "_cast_fp8_dynamic", lambda *a: pytest.fail("torchao caster must not run")
)
monkeypatch.setattr(dp, "_cast_nvfp4", lambda *a: pytest.fail("torchao caster must not run"))
monkeypatch.setattr(
dp, "_cast_int8_selective", lambda *a: pytest.fail("torchao caster must not run")
)
pipe = types.SimpleNamespace(text_encoder = object())
assert quantize_text_encoders(pipe, _target(), mode = "fp8_dynamic", offload_active = True) is None
assert quantize_text_encoders(pipe, _target(), mode = "nvfp4", offload_active = True) is None
assert (
quantize_text_encoders(
pipe, _target(), mode = "int8", family = "qwen-image", offload_active = True
)
is None
)
# Layerwise fp8 is not torchao and streams fine under offload, so it still engages.
fp8_calls: list = []
monkeypatch.setattr(dp, "_cast_fp8", lambda enc, tgt: fp8_calls.append(enc))
assert quantize_text_encoders(pipe, _target(), mode = "fp8", offload_active = True) == TE_QUANT_FP8
assert len(fp8_calls) == 1
# ── block selection + real int8 filter closure ─────────────────────────────────
def test_keep_bf16_block_fqns_selects_first_and_last(monkeypatch):
torch = _stub_torch(monkeypatch)
module_list = torch.nn.ModuleList
layers = module_list([object() for _ in range(10)])
# A short stack (<= skip_first + skip_last) contributes nothing (keeping it all would
# leave no interior to quantise).
short = module_list([object() for _ in range(4)])
enc = types.SimpleNamespace()
enc.named_modules = lambda: [("", enc), ("model.layers", layers), ("aux.blocks", short)]
keep = _keep_bf16_block_fqns(enc, 3, 2)
assert keep == {
"model.layers.0",
"model.layers.1",
"model.layers.2",
"model.layers.8",
"model.layers.9",
}
def _stub_transformer_quant(monkeypatch, captured):
# Reuse the committed factory's names but record what the int8 caster hands quantize_().
dtq = types.ModuleType("core.inference.diffusion_transformer_quant")
dtq.TQ_INT8 = "int8"
dtq.TQ_FP8 = "fp8"
dtq.DEFAULT_MIN_LINEAR_FEATURES = 512
dtq._make_quant_config = lambda scheme, *a, **k: f"cfg:{scheme}"
dtq.exclude_tokens_for_scheme = lambda scheme: ("modulation",)
def _make_filter_fn(min_features, exclude_name_tokens = ()):
def _f(module, fqn = ""):
return not any(tok in fqn for tok in exclude_name_tokens)
return _f
dtq.make_filter_fn = _make_filter_fn
monkeypatch.setitem(sys.modules, "core.inference.diffusion_transformer_quant", dtq)
tq = types.ModuleType("torchao.quantization")
def _quantize_(
module,
config,
filter_fn = None,
):
captured["config"] = config
captured["filter_fn"] = filter_fn
tq.quantize_ = _quantize_
monkeypatch.setitem(sys.modules, "torchao.quantization", tq)
def test_int8_filter_keeps_blocks_and_towers_dense(monkeypatch):
# The real selective closure: interior Linears quantise, but the kept first blocks,
# the vision tower, lm_head, and the encoder's fp32-kept modules (T5 "wo") stay bf16.
torch = _stub_torch(monkeypatch)
captured: dict = {}
_stub_transformer_quant(monkeypatch, captured)
layers = torch.nn.ModuleList([object() for _ in range(8)])
enc = types.SimpleNamespace(_keep_in_fp32_modules = ["wo"])
enc.named_modules = lambda: [("model.layers", layers)]
_cast_int8_selective(enc, _target(), 3, 0)
assert captured["config"] == "cfg:int8"
ff = captured["filter_fn"]
# Kept first-3 decoder blocks stay bf16.
assert ff(object(), "model.layers.0.self_attn.q_proj") is False
assert ff(object(), "model.layers.2.mlp.gate_proj") is False
# An interior block is quantised.
assert ff(object(), "model.layers.5.self_attn.q_proj") is True
# Vision tower / lm_head / T5 wo are excluded by the shared token filter.
assert ff(object(), "visual.blocks.0.attn.qkv") is False
assert ff(object(), "lm_head") is False
assert ff(object(), "model.decoder.wo") is False

38
studio/backend/tests/test_diffusion_prequant.py
View file

@ -286,6 +286,44 @@ def test_load_exclude_tokens_match_ok(monkeypatch, tmp_path):
assert _load(monkeypatch, tmp_path, ckpt, scheme = "int8") is not None
def test_load_require_bf16_mismatch_is_none(monkeypatch, tmp_path):
# An fp8 (scaled_mm) checkpoint built WITHOUT the bf16 gate quantised a different layer set
# than the runtime filter now produces, so it must be rejected rather than loaded.
ckpt = _good_ckpt(scheme = "fp8")
ckpt["metadata"]["require_bf16"] = False
assert _load(monkeypatch, tmp_path, ckpt, scheme = "fp8") is None
def test_load_require_bf16_match_ok(monkeypatch, tmp_path):
ckpt = _good_ckpt(scheme = "fp8")
ckpt["metadata"]["require_bf16"] = True
assert _load(monkeypatch, tmp_path, ckpt, scheme = "fp8") is not None
def test_load_require_bf16_int8_true_is_none(monkeypatch, tmp_path):
# int8 (torch._int_mm) tolerates non-bf16 weights, so it never sets the gate; a checkpoint
# claiming it did contradicts the runtime filter and must be rejected.
ckpt = _good_ckpt(scheme = "int8")
ckpt["metadata"]["require_bf16"] = True
assert _load(monkeypatch, tmp_path, ckpt, scheme = "int8") is None
def test_load_require_bf16_nvfp4_false_ok(monkeypatch, tmp_path):
# nvfp4 quantises fp32 weights fine, so the runtime filter does NOT set the bf16 gate; a
# checkpoint built the same way (require_bf16=False) matches and loads.
ckpt = _good_ckpt(scheme = "nvfp4")
ckpt["metadata"]["require_bf16"] = False
assert _load(monkeypatch, tmp_path, ckpt, scheme = "nvfp4") is not None
def test_load_require_bf16_nvfp4_true_is_none(monkeypatch, tmp_path):
# An nvfp4 checkpoint claiming the bf16 gate contradicts the runtime filter (nvfp4 is not gated),
# so it quantised a different layer set and must be rejected.
ckpt = _good_ckpt(scheme = "nvfp4")
ckpt["metadata"]["require_bf16"] = True
assert _load(monkeypatch, tmp_path, ckpt, scheme = "nvfp4") is None
def test_resolve_checkpoint_path_expands_user(monkeypatch, tmp_path):
# The allowlist gate expands ~, so the existence check must too, or a "~/..." checkpoint
# that passed the gate is silently skipped.

4
studio/backend/tests/test_diffusion_speed.py
View file

@ -91,6 +91,10 @@ def test_resolve_speed_mode_gguf_auto_default():
assert resolve_speed_mode("off", is_gguf = True) == SPEED_OFF
assert resolve_speed_mode("max", is_gguf = True) == SPEED_MAX
assert resolve_speed_mode("max", is_gguf = False) == SPEED_MAX
# The video backend passes a dense default of `default` (clips amortise the
# compile within one run); it must not affect GGUF or explicit values.
assert resolve_speed_mode(None, is_gguf = False, dense_default = SPEED_DEFAULT) == SPEED_DEFAULT
assert resolve_speed_mode("off", is_gguf = False, dense_default = SPEED_DEFAULT) == SPEED_OFF
# ── compile gating ────────────────────────────────────────────────────────────

41
studio/backend/tests/test_diffusion_transformer_quant.py
View file

@ -325,6 +325,47 @@ def test_make_filter_fn(monkeypatch):
assert keep(types.SimpleNamespace(), "no_attrs") is False
def test_require_bf16_schemes_excludes_nvfp4():
# fp8 and mxfp8 assert a bf16 weight (torchao 0.17 / B200: "PerRow quantization only works for
# bfloat16 ..." and "Only supporting bf16 out dtype ..."), so they gate on it; nvfp4 quantises an
# fp32 weight fine, so it is NOT gated (leaving its large fp32 projections quantised, not dense).
from core.inference.diffusion_transformer_quant import (
_REQUIRE_BF16_SCHEMES,
TQ_FP8,
TQ_MXFP8,
TQ_NVFP4,
TQ_INT8,
)
assert TQ_FP8 in _REQUIRE_BF16_SCHEMES
assert TQ_MXFP8 in _REQUIRE_BF16_SCHEMES
assert TQ_NVFP4 not in _REQUIRE_BF16_SCHEMES
assert TQ_INT8 not in _REQUIRE_BF16_SCHEMES
def test_make_filter_fn_require_bf16_skips_non_bf16(monkeypatch):
# fp8 / mxfp8 assert a bf16 weight, so require_bf16 must skip a fp32 Linear (which Wan / Hunyuan
# video DiTs keep) while keeping the bf16 ones -- otherwise a single fp32 layer raises inside
# quantize_ and no-ops the whole pass. int8 and nvfp4 leave it off (they quantise fp32 fine).
torch = types.ModuleType("torch")
torch.bfloat16, torch.float32 = "bf16", "fp32"
class _Lin:
def __init__(self, i, o, dtype):
self.in_features, self.out_features = i, o
self.weight = types.SimpleNamespace(dtype = dtype)
torch.nn = types.SimpleNamespace(Linear = _Lin)
monkeypatch.setitem(sys.modules, "torch", torch)
gated = make_filter_fn(512, require_bf16 = True)
assert gated(_Lin(1024, 4096, torch.bfloat16), "blocks.0.attn.to_q") is True
assert gated(_Lin(1024, 4096, torch.float32), "blocks.0.attn.to_q") is False # fp32 -> skip
assert gated(types.SimpleNamespace(in_features = 1024, out_features = 4096), "no_weight") is False
# int8 (require_bf16 off, the default) still quantises the fp32 linear.
assert make_filter_fn(512)(_Lin(1024, 4096, torch.float32), "blocks.0.attn.to_q") is True
def test_make_filter_fn_int8_excludes_modulation_and_embedders(monkeypatch):
# The int8 path skips the large M=1 AdaLN modulation / conditioning-embedder projections
# (they crash torch._int_mm's M>16), while keeping the attention / FFN compute layers and

63
studio/backend/tests/test_video_backend.py
View file

@ -145,6 +145,9 @@ class _FakeWanDiT:
def enable_cache(self, config) -> None:
self.cache_config = config
def disable_cache(self) -> None:
self.cache_config = None
def set_attention_backend(self, backend) -> None:
self.attention = backend
@ -801,6 +804,66 @@ def test_load_wan_ti2v_5b_pipeline(fake_runtime):
assert _FakeWanPipelineSingle.last["repo"] == "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
def test_video_dense_speed_defaults_to_compile_profile(fake_runtime):
# A clip denoise amortises the one-time compile within a single run, so an
# UNSET speed on a dense (pipeline) load resolves to `default` -- never `max`,
# never `off`. Explicit "off" is still honored verbatim.
backend = VideoBackend()
status = backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline")
assert status["speed_mode"] == "default"
assert status["resolved"]["speed_mode"]["source"] == "auto"
backend.unload()
status_off = backend.load_pipeline(
"Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline", speed_mode = "off"
)
assert status_off["speed_mode"] == "off"
assert status_off["resolved"]["speed_mode"]["source"] == "explicit"
def test_video_step_cache_auto_from_default_schedule(fake_runtime, tmp_path):
# Unset step cache is AUTO, decided from the model's default schedule: Wan's
# 50-step default engages FBCache at load; the LTX distilled 8-step default
# keeps it off. Both are re-checked per generation (toggle test below).
backend = VideoBackend()
status = backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline")
assert status["transformer_cache"] == "fbcache"
assert status["resolved"]["transformer_cache"]["source"] == "auto"
backend.unload()
(tmp_path / "ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf").write_bytes(b"w")
status2 = backend.load_pipeline(
str(tmp_path),
gguf_filename = "ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf",
base_repo = "Lightricks/LTX-2",
family_override = "ltx-2",
)
assert status2["transformer_cache"] is None
assert status2["resolved"]["transformer_cache"]["source"] == "auto"
backend.unload()
def test_video_step_cache_auto_toggles_on_actual_steps(fake_runtime):
# The AUTO decision follows the ACTUAL step count of each generation: a
# few-step request drops the load-time cache, a many-step request restores
# it. An explicit "off" never toggles.
backend = VideoBackend()
backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline")
assert backend.status()["transformer_cache"] == "fbcache"
backend.generate(prompt = "a sloth", steps = 8)
assert backend.status()["transformer_cache"] is None
backend.generate(prompt = "a sloth", steps = 30)
assert backend.status()["transformer_cache"] == "fbcache"
backend.unload()
backend.load_pipeline(
"Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline", transformer_cache = "off"
)
assert backend.status()["transformer_cache"] is None
backend.generate(prompt = "a sloth", steps = 30)
assert backend.status()["transformer_cache"] is None
backend.unload()
def test_wan_frame_snapping_4k_plus_1(fake_runtime):
# Wan snaps num_frames to 4k+1 (temporal factor 4), unlike LTX-2's 8k+1.
backend = VideoBackend()

40
studio/backend/tests/test_video_gallery.py
View file

@ -165,3 +165,43 @@ def test_list_skips_corrupt_sidecar():
gallery.save(_mp4(), _meta(prompt = "ours"))
listed = gallery.list_videos()
assert [r["prompt"] for r in listed] == ["ours"]
def _real_mp4_bytes() -> bytes:
# A real (tiny) MP4 for the transcode tests: 8 frames of flat color at
# 32x32, encoded with mpeg4 (bundled in every PyAV build, unlike libx264).
av = pytest.importorskip("av")
np = pytest.importorskip("numpy")
import io
buf = io.BytesIO()
with av.open(buf, "w", format = "mp4") as out:
stream = out.add_stream("mpeg4", rate = 8)
stream.width = 32
stream.height = 32
stream.pix_fmt = "yuv420p"
for i in range(8):
frame = av.VideoFrame.from_ndarray(
np.full((32, 32, 3), i * 30, dtype = np.uint8), format = "rgb24"
)
for packet in stream.encode(frame):
out.mux(packet)
for packet in stream.encode():
out.mux(packet)
return buf.getvalue()
def test_transcode_gif_and_webm_produce_real_containers():
record = gallery.save(_real_mp4_bytes(), _meta())
gif = gallery.transcode(record["id"], "gif")
assert gif is not None and gif.startswith(b"GIF8")
webm = gallery.transcode(record["id"], "webm")
# EBML magic: WebM is a Matroska container.
assert webm is not None and webm[:4] == b"\x1a\x45\xdf\xa3"
def test_transcode_unknown_id_and_bad_format():
assert gallery.transcode("does-not-exist", "gif") is None
record = gallery.save(_real_mp4_bytes(), _meta())
with pytest.raises(ValueError):
gallery.transcode(record["id"], "avi")

54
studio/backend/tests/test_video_routes.py
View file

@ -53,6 +53,7 @@ def _unloaded_status():
"attention_backend": None,
"transformer_cache": None,
"transformer_quant": None,
"text_encoder_quant": None,
"has_audio": False,
"defaults": None,
"resolved": None,
@ -73,6 +74,7 @@ class _FakeBackend:
family_override = None,
model_kind = None,
transformer_quant = None,
text_encoder_quant = None,
):
# Mirror the real backend's cheap validation so the route's
# validate-before-evict ordering is exercised.
@ -289,6 +291,35 @@ def test_load_rejects_bad_transformer_quant_422(client):
assert resp.status_code == 422
def test_load_threads_text_encoder_quant(client):
# The load-time text_encoder_quant field reaches the backend (the video path now
# quantises the dense companion encoder, not just the DiT).
resp = client.post(
"/api/inference/video/load",
json = {
"model_path": "unsloth/LTX-2.3-GGUF",
"gguf_filename": "q.gguf",
"text_encoder_quant": "fp8",
},
)
assert resp.status_code == 200
kwargs = video_module.get_video_backend().last_load_kwargs
assert kwargs.get("text_encoder_quant") == "fp8"
def test_load_rejects_bad_text_encoder_quant_422(client):
# text_encoder_quant is a Literal, so an unknown scheme is a 422 at request validation.
resp = client.post(
"/api/inference/video/load",
json = {
"model_path": "unsloth/LTX-2.3-GGUF",
"gguf_filename": "q.gguf",
"text_encoder_quant": "bogus",
},
)
assert resp.status_code == 422
def test_load_progress_route(client):
idle = client.get("/api/inference/video/load-progress")
assert idle.status_code == 200 and idle.json()["phase"] is None
@ -497,3 +528,26 @@ def test_routes_require_auth():
app.include_router(video_router, prefix = "/api/inference")
unauth = TestClient(app)
assert unauth.get("/api/inference/video/status").status_code in (401, 403)
def test_export_endpoint_validation(client, monkeypatch):
# Unknown format is a 400 before any work happens.
resp = client.get("/api/inference/video/gallery/x/export?format=avi")
assert resp.status_code == 400
# Unknown id is a 404.
resp = client.get("/api/inference/video/gallery/does-not-exist/export?format=gif")
assert resp.status_code == 404
# A missing codec surfaces as 501 with the transcoder's message.
def _boom(video_id, fmt):
raise RuntimeError("WebM export needs the 'av' package (PyAV).")
monkeypatch.setattr(gallery_module, "transcode", _boom)
client.post(
"/api/inference/video/load",
json = {"model_path": "unsloth/LTX-2.3-GGUF", "gguf_filename": "q.gguf"},
)
video = client.post("/api/inference/video/generate", json = {"prompt": "a"}).json()["video"]
resp = client.get(f"/api/inference/video/gallery/{video['id']}/export?format=webm")
assert resp.status_code == 501
assert "PyAV" in resp.json()["detail"]

1
studio/frontend/package.json
View file

@ -13,6 +13,7 @@
"preview": "vite preview",
"typecheck": "tsc -b --pretty false",
"i18n:check": "node --experimental-strip-types --no-warnings src/i18n/check-parity.ts",
"catalog:check": "node --experimental-strip-types --no-warnings src/components/assistant-ui/model-selector/model-catalog.check.ts",
"biome:check": "biome check",
"biome:fix": "biome check --write"
},

9
studio/frontend/src/components/app-sidebar.tsx
View file

@ -79,10 +79,10 @@ import {
LayoutAlignLeftIcon,
Settings02Icon,
Sun03Icon,
TestTube01Icon,
Video01Icon,
ZapIcon,
} from "@hugeicons/core-free-icons";
import { TestTubeOutlineIcon } from "@/lib/hugeicons-derived";
import {
exportConversationRawJsonl,
exportConversationCsv,
@ -169,13 +169,6 @@ function getTourId(pathname: string): string | null {
return null;
}
// TestTube01Icon's last 2 paths are interior bubbles; slice to the first
// 3 (outline + cap + liquid line) to drop them. Original export untouched.
const TestTubeOutlineIcon = TestTube01Icon.slice(
0,
3,
) as typeof TestTube01Icon;
function runStatusDotClass(status: TrainingRunSummary["status"]): string {
switch (status) {
case "running":

10
studio/frontend/src/components/assistant-ui/model-selector.tsx
View file

@ -36,6 +36,7 @@ import {
import { Input } from "../ui/input";
import type { HfTaskFilter } from "@/features/hub/hooks/use-hub-model-search";
import { HubModelPicker, hasDownloadedModels } from "./model-selector/pickers";
import type { CatalogGroup } from "./model-selector/model-catalog";
import { PillTabs } from "./model-selector/pill-tabs";
import {
buildSourceTabs,
@ -140,6 +141,10 @@ interface ModelSelectorProps {
showCloudIndicator?: boolean;
/** Restrict the Hub tab to a pipeline task (e.g. text-to-image). */
task?: HfTaskFilter;
/** Canonical model groups (Images / Video pages): collapses a model's
* artifact repos into one row with a format second level and device-aware
* routing. Undefined (chat) changes nothing. */
catalog?: CatalogGroup[];
}
function ModelSelectorTrigger({
@ -321,6 +326,7 @@ function ModelSelectorContent({
className,
dataTour,
task,
catalog,
}: {
open: boolean;
models: ModelOption[];
@ -337,6 +343,7 @@ function ModelSelectorContent({
className?: string;
dataTour?: string;
task?: HfTaskFilter;
catalog?: CatalogGroup[];
}) {
const hasSelection = Boolean(value);
const chatOnly = usePlatformStore((s) => s.isChatOnly());
@ -507,6 +514,7 @@ function ModelSelectorContent({
section={effectiveHubSection}
onEject={hasSelection && onEject ? onEject : undefined}
task={task}
catalog={catalog}
sectionToggle={
<PillTabs
ariaLabel="Hub section"
@ -587,6 +595,7 @@ export function ModelSelector({
contentDataTour,
showCloudIndicator = false,
task,
catalog,
}: ModelSelectorProps) {
const [uncontrolledOpen, setUncontrolledOpen] = useState(false);
const open = controlledOpen ?? uncontrolledOpen;
@ -713,6 +722,7 @@ export function ModelSelector({
className={contentClassName}
dataTour={contentDataTour}
task={task}
catalog={catalog}
/>
</Popover>
);

472
studio/frontend/src/components/assistant-ui/model-selector/model-catalog.check.ts Normal file
View file

@ -0,0 +1,472 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
// Assertions over the model catalog: canonical keys, alias resolution, catalog
// integrity, the artifact/quant routing ladders, and search matching. Follows
// the i18n:check pattern (plain node:assert, no test framework).
//
// Run: npm run catalog:check
import assert from "node:assert/strict";
import {
IMAGE_CATALOG,
VIDEO_CATALOG,
canonicalKeyFor,
catalogToModelOptions,
classifyGgufFit,
groupForRepoId,
groupMatchesQuery,
loadSpecFor,
pickDefaultArtifact,
pickDefaultQuant,
stripArtifactSuffixesForDisplay,
} from "./model-catalog.ts";
// ── canonicalKeyFor: suffix stripping, owner preserved ─────────────────────────
assert.equal(canonicalKeyFor("unsloth/Qwen-Image-2512-GGUF"), "unsloth/qwen-image-2512");
assert.equal(canonicalKeyFor("unsloth/Qwen-Image-2512-FP8"), "unsloth/qwen-image-2512");
assert.equal(
canonicalKeyFor("unsloth/Qwen-Image-2512-unsloth-bnb-4bit"),
"unsloth/qwen-image-2512",
);
assert.equal(
canonicalKeyFor("ideogram-ai/ideogram-4-nf4-diffusers"),
"ideogram-ai/ideogram-4",
);
assert.equal(canonicalKeyFor("Wan-AI/Wan2.2-TI2V-5B-Diffusers"), "wan-ai/wan2.2-ti2v-5b");
assert.equal(canonicalKeyFor("lightricks/ltx-2.3-fp8"), "lightricks/ltx-2.3");
// Prequant suffixes strip regardless of case: -GGUF/-FP8/-int8/-nvfp4 all route
// to the base name.
assert.equal(canonicalKeyFor("unsloth/Qwen-Image-2512-int8"), "unsloth/qwen-image-2512");
assert.equal(canonicalKeyFor("unsloth/Qwen-Image-2512-INT8"), "unsloth/qwen-image-2512");
assert.equal(canonicalKeyFor("unsloth/Qwen-Image-2512-nvfp4"), "unsloth/qwen-image-2512");
assert.equal(canonicalKeyFor("unsloth/Qwen-Image-2512-NVFP4"), "unsloth/qwen-image-2512");
assert.equal(canonicalKeyFor("unsloth/qwen-image-2512-gguf"), "unsloth/qwen-image-2512");
assert.equal(canonicalKeyFor("unsloth/qwen-image-2512-fp8"), "unsloth/qwen-image-2512");
// ── stripArtifactSuffixesForDisplay: case-preserving base name for row labels ──
assert.equal(
stripArtifactSuffixesForDisplay("unsloth/ERNIE-Image-Turbo-GGUF"),
"unsloth/ERNIE-Image-Turbo",
);
assert.equal(
stripArtifactSuffixesForDisplay("unsloth/FLUX.2-klein-base-9B-GGUF"),
"unsloth/FLUX.2-klein-base-9B",
);
assert.equal(
stripArtifactSuffixesForDisplay("unsloth/Qwen-Image-2512-FP8"),
"unsloth/Qwen-Image-2512",
);
assert.equal(
stripArtifactSuffixesForDisplay("unsloth/Some-Model-int8"),
"unsloth/Some-Model",
);
assert.equal(
stripArtifactSuffixesForDisplay("unsloth/Some-Model-NVFP4"),
"unsloth/Some-Model",
);
// Non-suffixed names and suffix-only names come back unchanged, casing intact.
assert.equal(
stripArtifactSuffixesForDisplay("krea/Krea-2-Turbo"),
"krea/Krea-2-Turbo",
);
assert.equal(stripArtifactSuffixesForDisplay("someone/FP8"), "someone/FP8");
// Non-suffixed ids come back unchanged (lowercased).
assert.equal(canonicalKeyFor("krea/Krea-2-Turbo"), "krea/krea-2-turbo");
// Stripping never merges owners.
assert.notEqual(
canonicalKeyFor("Qwen/Qwen-Image-2512"),
canonicalKeyFor("unsloth/Qwen-Image-2512"),
);
// Stripping never empties a name that IS a suffix-looking token.
assert.equal(canonicalKeyFor("someone/fp8"), "someone/fp8");
// ── groupForRepoId: artifacts, aliases, canonical keys, unknowns ───────────────
const qwen2512 = groupForRepoId("unsloth/Qwen-Image-2512-GGUF", IMAGE_CATALOG);
assert.ok(qwen2512);
assert.equal(qwen2512.canonicalId, "unsloth/Qwen-Image-2512");
// Every artifact of the group resolves to the same group.
for (const artifact of qwen2512.artifacts) {
assert.equal(groupForRepoId(artifact.repoId, IMAGE_CATALOG), qwen2512);
}
// Cross-owner aliases resolve only because they are declared.
assert.equal(groupForRepoId("Qwen/Qwen-Image-2512", IMAGE_CATALOG), qwen2512);
// Undeclared prequant variants (any case) still route to the base group via the
// stripped key, so Recommended and On Device standardize them to the base name.
assert.equal(groupForRepoId("unsloth/Qwen-Image-2512-INT8", IMAGE_CATALOG), qwen2512);
assert.equal(groupForRepoId("unsloth/Qwen-Image-2512-NVFP4", IMAGE_CATALOG), qwen2512);
assert.equal(
groupForRepoId("Tongyi-MAI/Z-Image-Turbo", IMAGE_CATALOG)?.canonicalId,
"unsloth/Z-Image-Turbo",
);
// A sibling artifact of an aliased owner groups via the alias' stripped key.
assert.equal(groupForRepoId("Qwen/Qwen-Image-2512-FP8", IMAGE_CATALOG), qwen2512);
// Unknown repos pass through ungrouped.
assert.equal(groupForRepoId("someone/some-model-GGUF", IMAGE_CATALOG), null);
assert.equal(groupForRepoId("unsloth/Llama-3.3-70B-GGUF", VIDEO_CATALOG), null);
// Video: the Lightricks 2.3 checkpoints group under the unsloth 2.3 release.
const ltx23 = groupForRepoId("unsloth/LTX-2.3-GGUF", VIDEO_CATALOG);
assert.ok(ltx23);
assert.equal(groupForRepoId("lightricks/ltx-2.3", VIDEO_CATALOG), ltx23);
assert.equal(groupForRepoId("lightricks/ltx-2.3-fp8", VIDEO_CATALOG), ltx23);
// ...but the LTX-2.0 base stays its own group (different model).
assert.notEqual(groupForRepoId("Lightricks/LTX-2", VIDEO_CATALOG), ltx23);
// SDXL Turbo and Base stay separate groups (different checkpoints).
assert.notEqual(
groupForRepoId("stabilityai/sdxl-turbo", IMAGE_CATALOG),
groupForRepoId("stabilityai/stable-diffusion-xl-base-1.0", IMAGE_CATALOG),
);
// Both HunyuanVideo resolutions land in one group.
assert.equal(
groupForRepoId(
"hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
VIDEO_CATALOG,
),
groupForRepoId(
"hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v",
VIDEO_CATALOG,
),
);
// ── catalog integrity: unique ids, artifacts resolve to exactly one group ──────
for (const catalog of [IMAGE_CATALOG, VIDEO_CATALOG]) {
const seen = new Set<string>();
for (const group of catalog) {
for (const artifact of group.artifacts) {
const lowered = artifact.repoId.toLowerCase();
assert.ok(!seen.has(lowered), `duplicate artifact id: ${artifact.repoId}`);
seen.add(lowered);
assert.equal(
groupForRepoId(artifact.repoId, catalog),
group,
`artifact ${artifact.repoId} resolves to a different group`,
);
if (artifact.loadKind === "single_file") {
assert.ok(artifact.filename, `single_file ${artifact.repoId} needs a filename`);
}
}
for (const alias of group.aliases ?? []) {
assert.equal(
groupForRepoId(alias, catalog),
group,
`alias ${alias} resolves to a different group`,
);
}
}
}
// ── loadSpecFor reproduces the old page lookup tables exactly ──────────────────
const OLD_SAFETENSORS_MODELS: Record<
string,
{ kind: "pipeline" | "single_file"; filename?: string }
> = {
"unsloth/Z-Image-Turbo-unsloth-bnb-4bit": { kind: "pipeline" },
"krea/Krea-2-Turbo": { kind: "pipeline" },
"ideogram-ai/ideogram-4-fp8": { kind: "pipeline" },
"ideogram-ai/ideogram-4-nf4-diffusers": { kind: "pipeline" },
"unsloth/Qwen-Image-2512-unsloth-bnb-4bit": { kind: "pipeline" },
"unsloth/Qwen-Image-2512-FP8": {
kind: "single_file",
filename: "qwen-image-2512-fp8.safetensors",
},
"stabilityai/sdxl-turbo": { kind: "pipeline" },
"stabilityai/stable-diffusion-xl-base-1.0": { kind: "pipeline" },
};
for (const [id, spec] of Object.entries(OLD_SAFETENSORS_MODELS)) {
const got = loadSpecFor(id, IMAGE_CATALOG);
assert.ok(got, `missing image load spec for ${id}`);
assert.equal(got.kind, spec.kind, id);
assert.equal(got.filename, spec.filename, id);
}
const OLD_PIPELINE_MODELS = [
"Lightricks/LTX-2",
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
"hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
];
for (const id of OLD_PIPELINE_MODELS) {
const got = loadSpecFor(id, VIDEO_CATALOG);
assert.ok(got, `missing video load spec for ${id}`);
assert.equal(got.kind, "pipeline", id);
}
// GGUF artifacts report the gguf kind; unknown ids report null.
assert.equal(loadSpecFor("unsloth/Z-Image-Turbo-GGUF", IMAGE_CATALOG)?.kind, "gguf");
assert.equal(loadSpecFor("someone/unknown", IMAGE_CATALOG), null);
// Every old curated id is still present as an option (backwards compat).
const imageOptionIds = new Set(catalogToModelOptions(IMAGE_CATALOG).map((o) => o.id));
for (const id of [
"unsloth/Z-Image-Turbo-GGUF",
"unsloth/Z-Image-GGUF",
"unsloth/Qwen-Image-2512-GGUF",
"unsloth/Qwen-Image-GGUF",
"unsloth/FLUX.1-schnell-GGUF",
"unsloth/FLUX.1-dev-GGUF",
"unsloth/FLUX.2-klein-4B-GGUF",
"unsloth/FLUX.2-klein-9B-GGUF",
"unsloth/Qwen-Image-Edit-2511-GGUF",
"unsloth/FLUX.1-Kontext-dev-GGUF",
...Object.keys(OLD_SAFETENSORS_MODELS),
]) {
assert.ok(imageOptionIds.has(id), `image option missing: ${id}`);
}
const videoOptionIds = new Set(catalogToModelOptions(VIDEO_CATALOG).map((o) => o.id));
for (const id of ["unsloth/LTX-2.3-GGUF", ...OLD_PIPELINE_MODELS]) {
assert.ok(videoOptionIds.has(id), `video option missing: ${id}`);
}
// ── classifyGgufFit ────────────────────────────────────────────────────────────
const GB = 1024 ** 3;
assert.equal(classifyGgufFit(10 * GB, { gpuGb: 24, systemRamGb: 64 }), "fits");
assert.equal(classifyGgufFit(20 * GB, { gpuGb: 24, systemRamGb: 64 }), "tight");
assert.equal(classifyGgufFit(100 * GB, { gpuGb: 24, systemRamGb: 64 }), "oom");
// Unknown device: never scare with OOM.
assert.equal(classifyGgufFit(100 * GB, { gpuGb: 0, systemRamGb: 0 }), "fits");
// Unified-memory host (no GPU budget): fit-or-oom against RAM.
assert.equal(classifyGgufFit(30 * GB, { gpuGb: 0, systemRamGb: 64 }), "fits");
assert.equal(classifyGgufFit(60 * GB, { gpuGb: 0, systemRamGb: 64 }), "oom");
// ── pickDefaultQuant ───────────────────────────────────────────────────────────
const variants = [
{ quant: "Q4_K_M", filename: "m-Q4_K_M.gguf", size_bytes: 12 * GB },
{ quant: "Q8_0", filename: "m-Q8_0.gguf", size_bytes: 22 * GB },
{ quant: "BF16", filename: "m-BF16.gguf", size_bytes: 40 * GB },
];
const budget24 = { gpuGb: 24, systemRamGb: 64 };
// Repo default kept when it is not OOM.
assert.equal(pickDefaultQuant(variants, "Q4_K_M", budget24)?.quant, "Q4_K_M");
// Downloaded non-OOM quant beats the undownloaded default.
assert.equal(
pickDefaultQuant(
[variants[0], { ...variants[1], downloaded: true }, variants[2]],
"Q4_K_M",
budget24,
)?.quant,
"Q8_0",
);
// OOM default falls to the largest non-OOM quant (Q8_0 runs tight via RAM offload).
assert.equal(
pickDefaultQuant(variants, "BF16", { gpuGb: 24, systemRamGb: 16 })?.quant,
"Q8_0",
);
// Without RAM to offload into, the tight tier disappears and Q4_K_M wins.
assert.equal(
pickDefaultQuant(variants, "BF16", { gpuGb: 24, systemRamGb: 0 })?.quant,
"Q4_K_M",
);
// All OOM: smallest wins (closest to running).
assert.equal(
pickDefaultQuant(variants, "BF16", { gpuGb: 4, systemRamGb: 4 })?.quant,
"Q4_K_M",
);
// No budget knowledge: trust the repo default (expander parity).
assert.equal(
pickDefaultQuant(variants, "Q8_0", { gpuGb: 0, systemRamGb: 0 })?.quant,
"Q8_0",
);
assert.equal(pickDefaultQuant([], "Q4_K_M", budget24), null);
// ── pickDefaultArtifact ────────────────────────────────────────────────────────
const notDownloaded = () => false;
const qwenGroup = qwen2512;
// 8 GB consumer GPU: nothing prequant fits (fp8 24 GB, bnb 14 GB) -> GGUF.
assert.equal(
pickDefaultArtifact(qwenGroup, { gpuGb: 8, systemRamGb: 32, isDownloaded: notDownloaded })
.format,
"gguf",
);
// 24 GB: bnb-4bit (14 GB) fits the 16.8 GB budget, fp8 (24 GB) does not.
assert.equal(
pickDefaultArtifact(qwenGroup, { gpuGb: 24, systemRamGb: 64, isDownloaded: notDownloaded })
.format,
"bnb-4bit",
);
// 48 GB: fp8 fits -> highest quality that fits wins.
assert.equal(
pickDefaultArtifact(qwenGroup, { gpuGb: 48, systemRamGb: 64, isDownloaded: notDownloaded })
.format,
"fp8",
);
// Unknown device: GGUF (the backend plans offload itself).
assert.equal(
pickDefaultArtifact(qwenGroup, { gpuGb: 0, systemRamGb: 0, isDownloaded: notDownloaded })
.format,
"gguf",
);
// Downloaded-first: a downloaded bnb-4bit beats everything undownloaded.
assert.equal(
pickDefaultArtifact(qwenGroup, {
gpuGb: 48,
systemRamGb: 64,
isDownloaded: (id) => id === "unsloth/Qwen-Image-2512-unsloth-bnb-4bit",
}).format,
"bnb-4bit",
);
// A downloaded GGUF wins over undownloaded prequants even on a big GPU.
assert.equal(
pickDefaultArtifact(qwenGroup, {
gpuGb: 80,
systemRamGb: 128,
isDownloaded: (id) => id === "unsloth/Qwen-Image-2512-GGUF",
}).format,
"gguf",
);
// Ideogram on 24 GB: fp8 (46 GB) too big -> bnb-4bit (11 GB).
const ideogram = groupForRepoId("ideogram-ai/ideogram-4-fp8", IMAGE_CATALOG);
assert.ok(ideogram);
assert.equal(
pickDefaultArtifact(ideogram, { gpuGb: 24, systemRamGb: 64, isDownloaded: notDownloaded })
.repoId,
"ideogram-ai/ideogram-4-nf4-diffusers",
);
// A gated BF16 artifact (FLUX.1-dev) is NOT auto-routed when undownloaded, even on a big GPU that
// fits it: the download would fail without license/token access, so route to the open GGUF.
const fluxDevRoute = groupForRepoId("unsloth/FLUX.1-dev", IMAGE_CATALOG);
assert.ok(fluxDevRoute);
assert.equal(
pickDefaultArtifact(fluxDevRoute, { gpuGb: 80, systemRamGb: 128, isDownloaded: notDownloaded })
.format,
"gguf",
);
// But an already-downloaded gated BF16 (the user clearly has access) is still returned.
assert.equal(
pickDefaultArtifact(fluxDevRoute, {
gpuGb: 80,
systemRamGb: 128,
isDownloaded: (id) => id === "black-forest-labs/FLUX.1-dev",
}).repoId,
"black-forest-labs/FLUX.1-dev",
);
// FLUX.1-schnell is Apache-2.0 (not gated): its BF16 IS auto-routed on a GPU that fits it.
const fluxSchnellRoute = groupForRepoId("unsloth/FLUX.1-schnell", IMAGE_CATALOG);
assert.ok(fluxSchnellRoute);
assert.equal(
pickDefaultArtifact(fluxSchnellRoute, { gpuGb: 80, systemRamGb: 128, isDownloaded: notDownloaded })
.format,
"bf16",
);
// HunyuanVideo on 80 GB: the highest-quality artifact that FITS (720p, 52 GB <= budget 56) wins.
const hunyuan = groupForRepoId(
"hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
VIDEO_CATALOG,
);
assert.ok(hunyuan);
assert.equal(
pickDefaultArtifact(hunyuan, { gpuGb: 80, systemRamGb: 128, isDownloaded: notDownloaded })
.label,
"BF16 - 720p",
);
// Same-format artifacts keep declaration order, so 720p is listed first and the fit loop returns
// it when it fits; a smaller card (budget 42) skips 720p (52 > 42) and falls back to 480p (40).
assert.equal(
pickDefaultArtifact(hunyuan, { gpuGb: 60, systemRamGb: 128, isDownloaded: notDownloaded })
.label,
"BF16 - 480p",
);
// ── official BF16 artifacts (added so groups are not unsloth-quant-only) ────────
// Qwen-Image-2512 BF16 (54 GB) does not fit a 24/48 GB budget (bnb-4bit/fp8 win
// there, asserted above) but on an 80 GB datacenter GPU (budget 56) the official
// BF16 is the highest-quality artifact that fits and wins.
assert.equal(
pickDefaultArtifact(qwenGroup, { gpuGb: 80, systemRamGb: 128, isDownloaded: notDownloaded })
.format,
"bf16",
);
assert.equal(
pickDefaultArtifact(qwenGroup, { gpuGb: 80, systemRamGb: 128, isDownloaded: notDownloaded })
.repoId,
"Qwen/Qwen-Image-2512",
);
// Z-Image-Turbo BF16 (30 GB): does not fit 24 GB (bnb-4bit wins) but fits a
// 48 GB GPU (budget 33.6) where the official BF16 wins.
const zturbo = groupForRepoId("unsloth/Z-Image-Turbo", IMAGE_CATALOG);
assert.ok(zturbo);
assert.equal(
pickDefaultArtifact(zturbo, { gpuGb: 24, systemRamGb: 64, isDownloaded: notDownloaded })
.format,
"bnb-4bit",
);
assert.equal(
pickDefaultArtifact(zturbo, { gpuGb: 48, systemRamGb: 64, isDownloaded: notDownloaded })
.format,
"bf16",
);
// FLUX.1-dev BF16 (32 GB) fits a 48 GB GPU, but it is GATED: a bare click routes to the open
// GGUF unless the BF16 is already downloaded (see the gated-routing checks above). Small GPU -> GGUF.
const fluxDev = groupForRepoId("black-forest-labs/FLUX.1-dev", IMAGE_CATALOG);
assert.ok(fluxDev);
assert.equal(fluxDev.canonicalId, "unsloth/FLUX.1-dev");
assert.equal(
pickDefaultArtifact(fluxDev, { gpuGb: 48, systemRamGb: 64, isDownloaded: notDownloaded })
.format,
"gguf",
);
assert.equal(
pickDefaultArtifact(fluxDev, {
gpuGb: 48,
systemRamGb: 64,
isDownloaded: (id) => id === "black-forest-labs/FLUX.1-dev",
}).format,
"bf16",
);
assert.equal(
pickDefaultArtifact(fluxDev, { gpuGb: 24, systemRamGb: 64, isDownloaded: notDownloaded })
.format,
"gguf",
);
// LTX-2.3 video carries the official BF16 single-file checkpoint (no FP8: its scaled-fp8 file is
// refused by the loader), which keeps the ~50 GB Gemma3 encoder resident, so a consumer or 80 GB
// GPU routes to GGUF; only a B200-class budget picks the official BF16.
const ltxGroup = groupForRepoId("unsloth/LTX-2.3", VIDEO_CATALOG);
assert.ok(ltxGroup);
assert.equal(
pickDefaultArtifact(ltxGroup, { gpuGb: 24, systemRamGb: 64, isDownloaded: notDownloaded })
.format,
"gguf",
);
assert.equal(
pickDefaultArtifact(ltxGroup, { gpuGb: 80, systemRamGb: 128, isDownloaded: notDownloaded })
.format,
"gguf",
);
assert.equal(
pickDefaultArtifact(ltxGroup, { gpuGb: 192, systemRamGb: 256, isDownloaded: notDownloaded })
.format,
"bf16",
);
// The LTX-2.3 official checkpoints load as single-file against the family base.
assert.equal(loadSpecFor("Lightricks/LTX-2.3", VIDEO_CATALOG)?.kind, "single_file");
assert.equal(
loadSpecFor("Lightricks/LTX-2.3", VIDEO_CATALOG)?.filename,
"ltx-2.3-22b-distilled.safetensors",
);
// The official image BF16 pipelines load via from_pretrained (pipeline kind).
assert.equal(loadSpecFor("Tongyi-MAI/Z-Image-Turbo", IMAGE_CATALOG)?.kind, "pipeline");
assert.equal(loadSpecFor("Qwen/Qwen-Image-2512", IMAGE_CATALOG)?.kind, "pipeline");
// ── groupMatchesQuery ──────────────────────────────────────────────────────────
assert.ok(groupMatchesQuery(qwenGroup, "qwen"));
assert.ok(groupMatchesQuery(qwenGroup, "2512"));
assert.ok(groupMatchesQuery(qwenGroup, "gguf"));
assert.ok(groupMatchesQuery(qwenGroup, "fp8"));
assert.ok(groupMatchesQuery(qwenGroup, "4bit"));
assert.ok(groupMatchesQuery(qwenGroup, "q4_k_m"));
assert.ok(groupMatchesQuery(qwenGroup, "unsloth/qwen-image-2512-fp8"));
assert.ok(!groupMatchesQuery(qwenGroup, "mlx"));
assert.ok(!groupMatchesQuery(qwenGroup, "ideogram"));
assert.ok(groupMatchesQuery(ltx23, "ltx"));
assert.ok(groupMatchesQuery(ltx23, "lightricks/ltx-2.3"));
console.log("model-catalog check: all assertions passed");

677
studio/frontend/src/components/assistant-ui/model-selector/model-catalog.ts Normal file
View file

@ -0,0 +1,677 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
// One canonical name per diffusion model, with its published artifacts (GGUF
// quants, prequant FP8 / bnb-4bit repos, official BF16 pipelines) as a second
// level, plus a deterministic router that picks the best artifact for the
// device. Pure helpers -- no React/DOM deps so they are easy to test (see
// model-catalog.check.ts, run via `npm run catalog:check`).
import type { ModelOption } from "./types";
export type ArtifactFormat = "gguf" | "fp8" | "bnb-4bit" | "bf16";
export type LoadKind = "gguf" | "single_file" | "pipeline";
export interface ModelArtifact {
/** Exact artifact repo id (the pre-grouping id -- stays loadable/searchable). */
repoId: string;
format: ArtifactFormat;
loadKind: LoadKind;
/** single_file loads name their exact checkpoint inside the repo. */
filename?: string;
/** Second-level row label ("GGUF", "FP8", "BF16 (official)", "BF16 - 720p"). */
label: string;
/** Curated resident-size estimate for routing. Omitted = unknown: never
* auto-picked unless already downloaded. GGUF artifacts omit it too -- their
* per-quant ladder self-fits via pickDefaultQuant. */
approxSizeGb?: number;
/** Extra search tokens beyond the id/label ("4bit", "nf4", ...). */
keywords?: readonly string[];
/** Gated on the Hub (license acceptance + token needed to download). A bare group
* click must not auto-route to it when it isn't already downloaded -- the download
* would fail for a user without access -- so the not-downloaded ladder skips it and
* falls through to an open artifact (e.g. the GGUF). An already-downloaded gated
* artifact is still returned (the user clearly has access). */
gated?: boolean;
}
export interface CatalogGroup {
/** Canonical display id, owner spelled once ("unsloth/Qwen-Image-2512"). */
canonicalId: string;
displayName: string;
/** Row meta line ("Text-to-image", "Image editing", "Text-to-video with audio"). */
description: string;
scope: "image" | "video";
/** Descending quality order: bf16, fp8, bnb-4bit, gguf. The router walks it. */
artifacts: ModelArtifact[];
/** Cross-owner ids that resolve to this group. Suffix stripping never merges
* two owners on its own, so arbitrary cached repos cannot be mis-grouped. */
aliases?: readonly string[];
}
// ── artifact constructors (keep the data tables terse) ─────────────────────────
const gguf = (repoId: string, extra?: Partial<ModelArtifact>): ModelArtifact => ({
repoId,
format: "gguf",
loadKind: "gguf",
label: "GGUF",
keywords: ["gguf", "quantized"],
...extra,
});
const bnb4bit = (
repoId: string,
approxSizeGb: number,
extra?: Partial<ModelArtifact>,
): ModelArtifact => ({
repoId,
format: "bnb-4bit",
loadKind: "pipeline",
label: "bnb-4bit",
approxSizeGb,
keywords: ["4bit", "bnb", "nf4", "bitsandbytes"],
...extra,
});
const fp8Single = (
repoId: string,
filename: string,
approxSizeGb: number,
): ModelArtifact => ({
repoId,
format: "fp8",
loadKind: "single_file",
filename,
label: "FP8",
approxSizeGb,
keywords: ["fp8", "float8"],
});
const fp8Pipeline = (repoId: string, approxSizeGb: number): ModelArtifact => ({
repoId,
format: "fp8",
loadKind: "pipeline",
label: "FP8",
approxSizeGb,
keywords: ["fp8", "float8"],
});
const bf16Pipeline = (
repoId: string,
approxSizeGb?: number,
extra?: Partial<ModelArtifact>,
): ModelArtifact => ({
repoId,
format: "bf16",
loadKind: "pipeline",
label: "BF16 (official)",
approxSizeGb,
keywords: ["bf16", "safetensors", "full precision"],
...extra,
});
// A bf16 single-file DiT checkpoint (e.g. Lightricks' distilled LTX-2.3): loads
// via from_single_file against the family base repo for the VAE / text encoder,
// same load path as the fp8 single-file checkpoints.
const bf16Single = (
repoId: string,
filename: string,
approxSizeGb: number,
): ModelArtifact => ({
repoId,
format: "bf16",
loadKind: "single_file",
filename,
label: "BF16 (official)",
approxSizeGb,
keywords: ["bf16", "safetensors", "full precision"],
});
// ── curated catalogs ────────────────────────────────────────────────────────────
// Sizes are steady resident estimates (GB) used only for routing; a missing size
// means "never auto-pick unless downloaded". GGUF entries carry no size -- the
// quant ladder (pickDefaultQuant) sizes the individual .gguf files.
export const IMAGE_CATALOG: CatalogGroup[] = [
{
canonicalId: "unsloth/Z-Image-Turbo",
displayName: "Z-Image-Turbo",
description: "Text-to-image",
scope: "image",
artifacts: [
bf16Pipeline("Tongyi-MAI/Z-Image-Turbo", 30),
bnb4bit("unsloth/Z-Image-Turbo-unsloth-bnb-4bit", 8),
gguf("unsloth/Z-Image-Turbo-GGUF"),
],
},
{
canonicalId: "unsloth/Z-Image",
displayName: "Z-Image",
description: "Text-to-image",
scope: "image",
artifacts: [gguf("unsloth/Z-Image-GGUF")],
},
{
canonicalId: "unsloth/Qwen-Image-2512",
displayName: "Qwen-Image 2512",
description: "Text-to-image",
scope: "image",
artifacts: [
bf16Pipeline("Qwen/Qwen-Image-2512", 54),
fp8Single(
"unsloth/Qwen-Image-2512-FP8",
"qwen-image-2512-fp8.safetensors",
24,
),
bnb4bit("unsloth/Qwen-Image-2512-unsloth-bnb-4bit", 14),
gguf("unsloth/Qwen-Image-2512-GGUF"),
],
},
{
canonicalId: "unsloth/Qwen-Image",
displayName: "Qwen-Image",
description: "Text-to-image",
scope: "image",
artifacts: [
bf16Pipeline("Qwen/Qwen-Image", 54),
gguf("unsloth/Qwen-Image-GGUF"),
],
},
{
canonicalId: "unsloth/FLUX.1-schnell",
displayName: "FLUX.1 schnell",
description: "Text-to-image",
scope: "image",
artifacts: [
bf16Pipeline("black-forest-labs/FLUX.1-schnell", 32),
gguf("unsloth/FLUX.1-schnell-GGUF"),
],
},
{
canonicalId: "unsloth/FLUX.1-dev",
displayName: "FLUX.1 dev",
description: "Text-to-image",
scope: "image",
artifacts: [
// FLUX.1-dev is gated (license acceptance + token); FLUX.1-schnell above is Apache-2.0.
bf16Pipeline("black-forest-labs/FLUX.1-dev", 32, { gated: true }),
gguf("unsloth/FLUX.1-dev-GGUF"),
],
},
{
canonicalId: "unsloth/FLUX.2-klein-4B",
displayName: "FLUX.2 klein 4B",
description: "Text-to-image",
scope: "image",
artifacts: [gguf("unsloth/FLUX.2-klein-4B-GGUF")],
},
{
canonicalId: "unsloth/FLUX.2-klein-9B",
displayName: "FLUX.2 klein 9B",
description: "Text-to-image",
scope: "image",
artifacts: [gguf("unsloth/FLUX.2-klein-9B-GGUF")],
},
{
canonicalId: "unsloth/Qwen-Image-Edit-2511",
displayName: "Qwen-Image-Edit 2511",
description: "Image editing",
scope: "image",
artifacts: [
bf16Pipeline("Qwen/Qwen-Image-Edit-2511", 54),
gguf("unsloth/Qwen-Image-Edit-2511-GGUF"),
],
},
{
canonicalId: "unsloth/FLUX.1-Kontext-dev",
displayName: "FLUX.1 Kontext dev",
description: "Image editing",
scope: "image",
artifacts: [
// FLUX.1-Kontext-dev is gated on the Hub (license acceptance + token).
bf16Pipeline("black-forest-labs/FLUX.1-Kontext-dev", 32, { gated: true }),
gguf("unsloth/FLUX.1-Kontext-dev-GGUF"),
],
},
{
canonicalId: "krea/Krea-2-Turbo",
displayName: "Krea 2 Turbo",
description: "Text-to-image",
scope: "image",
artifacts: [bf16Pipeline("krea/Krea-2-Turbo", 18)],
},
{
// No bf16 repo exists for Ideogram 4: -fp8 stores its two DiTs as raw
// float8 (~46 GB resident after the bf16 cast); -nf4-diffusers is the
// bnb-4bit export (~11 GB).
canonicalId: "ideogram-ai/ideogram-4",
displayName: "Ideogram 4",
description: "Text-to-image",
scope: "image",
artifacts: [
fp8Pipeline("ideogram-ai/ideogram-4-fp8", 46),
bnb4bit("ideogram-ai/ideogram-4-nf4-diffusers", 11),
],
},
// SDXL Turbo and Base are different checkpoints with different step/guidance
// defaults -- two groups, not two formats of one model.
{
canonicalId: "stabilityai/sdxl-turbo",
displayName: "SDXL Turbo",
description: "Text-to-image",
scope: "image",
artifacts: [bf16Pipeline("stabilityai/sdxl-turbo", 8, { label: "Safetensors" })],
},
{
canonicalId: "stabilityai/stable-diffusion-xl-base-1.0",
displayName: "SDXL Base 1.0",
description: "Text-to-image",
scope: "image",
artifacts: [
bf16Pipeline("stabilityai/stable-diffusion-xl-base-1.0", 8, {
label: "Safetensors",
}),
],
},
];
export const VIDEO_CATALOG: CatalogGroup[] = [
{
// The distilled 2.3 release: Lightricks' own bf16/fp8 single-file DiT
// checkpoints (loaded against the LTX-2 base for the VAE / Gemma3 text
// encoder, both repos already on the backend trust list) plus the GGUF
// quants. The single-file checkpoints keep the ~50 GB Gemma3-27B encoder in
// bf16, so their resident footprint is datacenter-scale; consumer GPUs route
// to GGUF, which offloads.
canonicalId: "unsloth/LTX-2.3",
displayName: "LTX 2.3 distilled",
description: "Text-to-video with audio",
scope: "video",
artifacts: [
bf16Single(
"Lightricks/LTX-2.3",
"ltx-2.3-22b-distilled.safetensors",
90,
),
// No FP8 artifact: the LTX-2.3 loader refuses the official scaled-FP8 single file (it carries
// .weight_scale/.input_scale tensors) and points users to GGUF/BF16, so advertising it would
// route a bare click or manual pick to a ~76 GB download that always fails on load.
gguf("unsloth/LTX-2.3-GGUF"),
],
},
{
canonicalId: "Lightricks/LTX-2",
displayName: "LTX 2 (base)",
description: "Text-to-video with audio",
scope: "video",
artifacts: [bf16Pipeline("Lightricks/LTX-2", 90)],
},
{
canonicalId: "Wan-AI/Wan2.2-TI2V-5B",
displayName: "Wan 2.2 TI2V 5B",
description: "Text-to-video 720p",
scope: "video",
artifacts: [bf16Pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", 30)],
},
{
canonicalId: "Wan-AI/Wan2.2-T2V-A14B",
displayName: "Wan 2.2 T2V A14B (MoE)",
description: "Text-to-video, dual-expert",
scope: "video",
artifacts: [bf16Pipeline("Wan-AI/Wan2.2-T2V-A14B-Diffusers", 114)],
},
{
canonicalId: "hunyuanvideo-community/HunyuanVideo-1.5",
displayName: "HunyuanVideo 1.5",
description: "Text-to-video",
scope: "video",
artifacts: [
// Highest-quality first: pickDefaultArtifact only sorts by FORMAT, so among these two bf16
// artifacts it keeps catalog order and the fit loop returns the FIRST that fits the budget.
// The 720p (52 GB) must precede the 480p (40 GB) so a bare click on a GPU where 720p fits
// (e.g. 80 GB, 0.7*budget=56) picks 720p, falling back to 480p only on smaller cards.
bf16Pipeline("hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v", 52, {
label: "BF16 - 720p",
keywords: ["bf16", "720p"],
}),
bf16Pipeline("hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v", 40, {
label: "BF16 - 480p",
keywords: ["bf16", "480p"],
}),
],
},
];
// ── canonical keys and lookups ──────────────────────────────────────────────────
// Artifact/format suffixes stripped (repeatedly, longest-first) off the NAME part
// of a repo id to reach its generic key. Owner is preserved: cross-owner merges
// happen only through the explicit alias tables above.
const ARTIFACT_SUFFIXES = [
"-unsloth-bnb-4bit",
"-nf4-diffusers",
"-bnb-4bit",
"-bnb4bit",
"-fp8-dynamic",
"-safetensors",
"-diffusers",
"-nvfp4",
"-gguf",
"-int8",
"-4bit",
"-nf4",
"-fp8",
"-bf16",
] as const;
/** Owner-preserving generic key: lowercase, artifact suffixes stripped off the
* name part. "unsloth/Qwen-Image-2512-GGUF" -> "unsloth/qwen-image-2512". */
export function canonicalKeyFor(repoId: string): string {
const lowered = repoId.trim().toLowerCase();
const slash = lowered.indexOf("/");
const owner = slash >= 0 ? lowered.slice(0, slash + 1) : "";
let name = slash >= 0 ? lowered.slice(slash + 1) : lowered;
let stripped = true;
while (stripped) {
stripped = false;
for (const suffix of ARTIFACT_SUFFIXES) {
if (name.endsWith(suffix) && name.length > suffix.length) {
name = name.slice(0, -suffix.length);
stripped = true;
}
}
}
return owner + name;
}
/** Case-preserving display name: the artifact suffixes stripped off the name part
* while keeping the original casing and the owner prefix. Used by the diffusion
* pickers so rows OUTSIDE the curated catalog (arbitrary hub or cached repos)
* still read as their base model name ("ERNIE-Image-Turbo-GGUF" ->
* "ERNIE-Image-Turbo"); the format badge next to the row carries the artifact
* kind. The id used for loading is never touched. */
export function stripArtifactSuffixesForDisplay(repoId: string): string {
const trimmed = repoId.trim();
const slash = trimmed.indexOf("/");
const owner = slash >= 0 ? trimmed.slice(0, slash + 1) : "";
let name = slash >= 0 ? trimmed.slice(slash + 1) : trimmed;
let stripped = true;
while (stripped) {
stripped = false;
const lowered = name.toLowerCase();
for (const suffix of ARTIFACT_SUFFIXES) {
if (lowered.endsWith(suffix) && name.length > suffix.length) {
name = name.slice(0, -suffix.length);
stripped = true;
break;
}
}
}
return owner + name;
}
interface CatalogIndex {
/** exact lowercased artifact/alias/canonical id -> group */
byId: Map<string, CatalogGroup>;
/** canonical suffix-stripped key -> group */
byKey: Map<string, CatalogGroup>;
/** exact lowercased artifact id -> artifact */
artifactById: Map<string, ModelArtifact>;
}
// Rebuilt only when a new catalog array identity shows up (the curated arrays
// are module constants, so in practice this builds twice: images + video).
const indexCache = new WeakMap<CatalogGroup[], CatalogIndex>();
function indexFor(catalog: CatalogGroup[]): CatalogIndex {
const cached = indexCache.get(catalog);
if (cached) return cached;
const byId = new Map<string, CatalogGroup>();
const byKey = new Map<string, CatalogGroup>();
const artifactById = new Map<string, ModelArtifact>();
for (const group of catalog) {
byId.set(group.canonicalId.toLowerCase(), group);
byKey.set(canonicalKeyFor(group.canonicalId), group);
for (const alias of group.aliases ?? []) {
byId.set(alias.toLowerCase(), group);
// An alias also claims its own suffix-stripped key so sibling artifacts
// of the aliased owner group correctly (Qwen/Qwen-Image-2512-FP8 etc.).
byKey.set(canonicalKeyFor(alias), group);
}
for (const artifact of group.artifacts) {
byId.set(artifact.repoId.toLowerCase(), group);
byKey.set(canonicalKeyFor(artifact.repoId), group);
artifactById.set(artifact.repoId.toLowerCase(), artifact);
}
}
const built = { byId, byKey, artifactById };
indexCache.set(catalog, built);
return built;
}
/** The group a repo id belongs to, or null for unknown repos (callers render
* those ungrouped, exactly as before the catalog existed). */
export function groupForRepoId(
repoId: string,
catalog: CatalogGroup[],
): CatalogGroup | null {
const index = indexFor(catalog);
const lowered = repoId.trim().toLowerCase();
return index.byId.get(lowered) ?? index.byKey.get(canonicalKeyFor(lowered)) ?? null;
}
/** The exact curated artifact for a repo id (null when the repo only matches a
* group by key/alias -- e.g. a cached quant repo we know but did not curate). */
export function artifactForRepoId(
repoId: string,
catalog: CatalogGroup[],
): { group: CatalogGroup; artifact: ModelArtifact } | null {
const index = indexFor(catalog);
const artifact = index.artifactById.get(repoId.trim().toLowerCase());
if (!artifact) return null;
const group = index.byId.get(repoId.trim().toLowerCase());
return group ? { group, artifact } : null;
}
/** Back-compat: the flat ModelOption list the ModelSelector's `models` prop
* expects, one option per ARTIFACT (old ids keep working everywhere). */
export function catalogToModelOptions(catalog: CatalogGroup[]): ModelOption[] {
const options: ModelOption[] = [];
for (const group of catalog) {
for (const artifact of group.artifacts) {
options.push({
id: artifact.repoId,
name:
group.artifacts.length > 1
? `${group.displayName} (${artifact.label})`
: group.displayName,
description: `${group.description} - ${artifact.label}`,
isGguf: artifact.format === "gguf",
});
}
}
return options;
}
/** How to load a curated artifact: replaces the pages' SAFETENSORS_MODELS /
* PIPELINE_MODELS lookup tables. Null for unknown ids (GGUF picks carry their
* own variant metadata; local paths and hub GGUFs resolve elsewhere). */
export function loadSpecFor(
repoId: string,
catalog: CatalogGroup[],
): { kind: LoadKind; filename?: string } | null {
const hit = artifactForRepoId(repoId, catalog);
if (!hit) return null;
return { kind: hit.artifact.loadKind, filename: hit.artifact.filename };
}
// Quant-class tokens that should match every GGUF artifact ("q4" finds the
// group whose GGUF repo publishes Q4_K_M, etc.).
const GGUF_QUANT_TOKENS = [
"q2",
"q3",
"q4",
"q5",
"q6",
"q8",
"q4_k_m",
"q5_k_m",
"q6_k",
"q8_0",
"bf16",
"f16",
] as const;
/** Whether a (lowercased, trimmed) query matches the group: canonical id,
* display name, any artifact id, any label/keyword, or a quant-class token. */
export function groupMatchesQuery(group: CatalogGroup, query: string): boolean {
const q = query.trim().toLowerCase();
if (!q) return true;
if (group.canonicalId.toLowerCase().includes(q)) return true;
if (group.displayName.toLowerCase().includes(q)) return true;
if (group.description.toLowerCase().includes(q)) return true;
for (const alias of group.aliases ?? []) {
if (alias.toLowerCase().includes(q)) return true;
}
for (const artifact of group.artifacts) {
if (artifact.repoId.toLowerCase().includes(q)) return true;
if (artifact.label.toLowerCase().includes(q)) return true;
for (const keyword of artifact.keywords ?? []) {
if (keyword.includes(q) || q.includes(keyword)) return true;
}
if (artifact.format === "gguf" && GGUF_QUANT_TOKENS.some((t) => q === t)) {
return true;
}
}
return false;
}
// ── device fit + routing ─────────────────────────────────────────────────────────
export interface DeviceBudget {
/** Total GPU memory in GB (0/undefined = unknown or none). */
gpuGb: number;
/** Available system RAM in GB (for the GGUF offload tier). */
systemRamGb: number;
}
/** GGUF fit classification matching llama-server's _select_gpus logic:
* fits = model <= 0.7 * GPU; tight = fits with 0.7 * RAM offload; oom = neither.
* Extracted from GgufVariantExpander so the badge and the router agree. */
export function classifyGgufFit(
sizeBytes: number,
budget: DeviceBudget,
): "fits" | "tight" | "oom" {
const gpuBudgetGb = (budget.gpuGb || 0) * 0.7;
const totalBudgetGb = gpuBudgetGb + (budget.systemRamGb || 0) * 0.7;
if (totalBudgetGb <= 0) return "fits";
const gb = sizeBytes / 1024 ** 3;
if (gb <= 0 || gb <= gpuBudgetGb) return "fits";
if (gpuBudgetGb <= 0) return gb <= totalBudgetGb ? "fits" : "oom";
if (gb <= totalBudgetGb) return "tight";
return "oom";
}
export interface QuantVariant {
quant: string;
filename: string;
size_bytes: number;
downloaded?: boolean;
}
/** The quant a bare group/repo click should load. Preference order:
* largest downloaded non-OOM quant, the repo default when non-OOM, the largest
* fitting quant, then the smallest overall (closest to running). Mirrors the
* expander's effectiveRecommended, extended to prefer what is already on disk. */
export function pickDefaultQuant(
variants: QuantVariant[],
defaultVariant: string | null,
budget: DeviceBudget,
): QuantVariant | null {
if (!variants || variants.length === 0) return null;
const totalBudgetGb =
(budget.gpuGb || 0) * 0.7 + (budget.systemRamGb || 0) * 0.7;
const downloadedFitting = variants
.filter((v) => v.downloaded && classifyGgufFit(v.size_bytes, budget) !== "oom")
.sort((a, b) => b.size_bytes - a.size_bytes);
if (downloadedFitting.length > 0) return downloadedFitting[0];
const byQuant = (quant: string | null) =>
quant ? (variants.find((v) => v.quant === quant) ?? null) : null;
// No budget knowledge at all: trust the repo default.
if (totalBudgetGb <= 0) return byQuant(defaultVariant) ?? variants[0];
const defaultV = byQuant(defaultVariant);
if (defaultV && classifyGgufFit(defaultV.size_bytes, budget) !== "oom") {
return defaultV;
}
const fitting = variants
.filter((v) => classifyGgufFit(v.size_bytes, budget) !== "oom")
.sort((a, b) => b.size_bytes - a.size_bytes);
if (fitting.length > 0) return fitting[0];
const smallest = [...variants].sort((a, b) => a.size_bytes - b.size_bytes);
return smallest[0] ?? null;
}
export interface RoutingInput extends DeviceBudget {
/** Whether an artifact repo already has weights on disk. */
isDownloaded: (repoId: string) => boolean;
}
const FORMAT_QUALITY: Record<ArtifactFormat, number> = {
bf16: 0,
fp8: 1,
"bnb-4bit": 2,
gguf: 3,
};
function fitsResident(artifact: ModelArtifact, gpuGb: number): boolean {
if (artifact.approxSizeGb === undefined) return false;
return artifact.approxSizeGb <= gpuGb * 0.7;
}
/** The artifact a bare group click loads. Deterministic ladder:
* 1. Downloaded first: the highest-quality downloaded artifact that fits the
* 0.7 * GPU budget; else a downloaded GGUF (its quant ladder self-fits);
* else the smallest-footprint downloaded artifact.
* 2. No budget known: the GGUF artifact when the group has one (the backend
* GGUF path plans offload itself and the Precision auto ladder still
* upgrades capable GPUs), else the first artifact.
* 3. Best sized artifact that fits resident, walking descending quality
* (BF16 official, FP8, bnb-4bit). Unknown sizes never auto-picked.
* 4. Fallback: GGUF, else the smallest-footprint artifact. */
export function pickDefaultArtifact(
group: CatalogGroup,
input: RoutingInput,
): ModelArtifact {
const artifacts = [...group.artifacts].sort(
(a, b) => FORMAT_QUALITY[a.format] - FORMAT_QUALITY[b.format],
);
const ggufArtifact = artifacts.find((a) => a.format === "gguf") ?? null;
const downloaded = artifacts.filter((a) => input.isDownloaded(a.repoId));
if (downloaded.length > 0) {
const fitting = downloaded.find(
(a) => a.format !== "gguf" && fitsResident(a, input.gpuGb),
);
if (fitting) return fitting;
const downloadedGguf = downloaded.find((a) => a.format === "gguf");
if (downloadedGguf) return downloadedGguf;
return downloaded.sort(
(a, b) => (a.approxSizeGb ?? Infinity) - (b.approxSizeGb ?? Infinity),
)[0];
}
if (!input.gpuGb || input.gpuGb <= 0) {
return ggufArtifact ?? artifacts[0];
}
for (const artifact of artifacts) {
// Skip a gated, NOT-downloaded artifact: auto-routing to it would fail the download for a
// user without license/token access, so fall through to an open artifact (the GGUF below).
// The downloaded branch above still returns a gated artifact the user already fetched.
if (artifact.format !== "gguf" && !artifact.gated && fitsResident(artifact, input.gpuGb)) {
return artifact;
}
}
if (ggufArtifact) return ggufArtifact;
return artifacts.sort(
(a, b) => (a.approxSizeGb ?? Infinity) - (b.approxSizeGb ?? Infinity),
)[0];
}

727
studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
View file

@ -112,6 +112,16 @@ import {
paramsFromId,
} from "./recommended-fit";
import { parseMetaTokens, splitRepoLabel } from "./row-meta";
import {
type CatalogGroup,
type ModelArtifact,
artifactForRepoId,
groupForRepoId,
groupMatchesQuery,
pickDefaultArtifact,
pickDefaultQuant,
stripArtifactSuffixesForDisplay,
} from "./model-catalog";
import type {
DeletedModelRef,
ExternalModelOption,
@ -1026,6 +1036,149 @@ function GgufVariantExpander({
);
}
// ── Catalog group second level: one row per artifact format ────────────────────
/** The format list under an expanded catalog group row: one row per published
* artifact (BF16 / FP8 / bnb-4bit / GGUF). Non-GGUF rows load directly; the
* GGUF row nests the existing quant expander. A single-GGUF group renders the
* quant expander directly (identical to today's GGUF repo rows). */
function ArtifactFormatList({
group,
recommendedArtifactId,
isRepoDownloaded,
onSelect,
gpuGb,
systemRamGb,
hfToken,
parentOptionKey,
onDevice = false,
}: {
group: CatalogGroup;
/** What pickDefaultArtifact would route a bare group click to (badged). */
recommendedArtifactId?: string;
isRepoDownloaded: (repoId: string) => boolean;
onSelect: (id: string, meta: ModelSelectorChangeMeta) => void;
gpuGb?: number;
systemRamGb?: number;
hfToken?: string;
parentOptionKey?: string;
onDevice?: boolean;
}) {
const ggufArtifacts = group.artifacts.filter((a) => a.format === "gguf");
const soleGguf =
group.artifacts.length === 1 && ggufArtifacts.length === 1
? ggufArtifacts[0]
: null;
const [openGguf, setOpenGguf] = useState<string | null>(
soleGguf?.repoId ?? null,
);
if (soleGguf) {
return (
<GgufVariantExpander
repoId={soleGguf.repoId}
onDevice={onDevice}
onSelect={onSelect}
hfToken={hfToken}
parentOptionKey={parentOptionKey}
gpuGb={gpuGb}
systemRamGb={systemRamGb}
/>
);
}
return (
<div className="ml-3 my-1 border-l-2 border-accent/50 pl-2">
{group.artifacts.map((artifact) => {
const downloaded = isRepoDownloaded(artifact.repoId);
const isRecommended = artifact.repoId === recommendedArtifactId;
const sizeLabel =
artifact.approxSizeGb !== undefined
? `~${artifact.approxSizeGb} GB`
: null;
if (artifact.format === "gguf") {
const open = openGguf === artifact.repoId;
return (
<div key={artifact.repoId}>
<button
type="button"
onClick={() =>
setOpenGguf((prev) =>
prev === artifact.repoId ? null : artifact.repoId,
)
}
className="flex w-full items-center justify-between gap-2 rounded px-2 py-1.5 text-left text-xs transition-colors hover:bg-accent/50"
>
<span className="flex items-center gap-1.5 font-mono">
{artifact.label}
{open ? (
<ChevronDownIcon className="size-3 text-muted-foreground" />
) : (
<ChevronRightIcon className="size-3 text-muted-foreground" />
)}
</span>
<span className="flex shrink-0 items-center gap-1.5">
{isRecommended && (
<span className="rounded-sm bg-primary/10 px-1 py-px text-[9px] font-medium uppercase tracking-wider text-primary">
recommended
</span>
)}
{downloaded && (
<span className="rounded-sm bg-green-500/10 px-1 py-px text-[9px] font-medium uppercase tracking-wider text-green-600 dark:text-green-500">
downloaded
</span>
)}
</span>
</button>
{open && (
<GgufVariantExpander
repoId={artifact.repoId}
onDevice={onDevice}
onSelect={onSelect}
hfToken={hfToken}
gpuGb={gpuGb}
systemRamGb={systemRamGb}
/>
)}
</div>
);
}
return (
<button
key={artifact.repoId}
type="button"
onClick={() =>
onSelect(artifact.repoId, {
source: "hub",
isLora: false,
isDownloaded: downloaded,
})
}
className="flex w-full items-center justify-between gap-2 rounded px-2 py-1.5 text-left text-xs transition-colors hover:bg-accent/50"
>
<span className="font-mono">{artifact.label}</span>
<span className="flex shrink-0 items-center gap-1.5">
{sizeLabel && (
<span className="text-[10px] text-muted-foreground">
{sizeLabel}
</span>
)}
{isRecommended && (
<span className="rounded-sm bg-primary/10 px-1 py-px text-[9px] font-medium uppercase tracking-wider text-primary">
recommended
</span>
)}
{downloaded && (
<span className="rounded-sm bg-green-500/10 px-1 py-px text-[9px] font-medium uppercase tracking-wider text-green-600 dark:text-green-500">
downloaded
</span>
)}
</span>
</button>
);
})}
</div>
);
}
// ── Detect GGUF repos by naming convention or hub tag ────────────────────
function hasGgufSuffix(id: string): boolean {
@ -1291,6 +1444,23 @@ function localModelMatchesFormat(
);
}
/** Whether a curated catalog group offers any artifact matching the format toggle.
* A group spans several formats, so it stays visible under a filter only when at
* least one of its artifacts qualifies -- otherwise a GGUF/Safetensors/MLX filter
* would still list groups whose click loads a different format (and MLX, which no
* catalog artifact provides, would show every group). */
function catalogGroupMatchesFormat(
group: CatalogGroup,
filter: FormatFilter,
): boolean {
return (
filter === "all" ||
group.artifacts.some((a) =>
matchesFormatFilter(a.repoId, a.format === "gguf", filter),
)
);
}
export function HubModelPicker({
models,
loraModels = [],
@ -1305,6 +1475,7 @@ export function HubModelPicker({
sectionToggle,
onEject,
task,
catalog,
}: {
models: ModelOption[];
/** Fine-tuned models, shown as a section in the On Device view. */
@ -1327,6 +1498,11 @@ export function HubModelPicker({
/** Restrict Hub results to a pipeline task (e.g. text-to-image for the
* Images page). Undefined = all tasks (chat default). */
task?: HfTaskFilter;
/** Canonical model groups (Images / Video pages): Recommended and On Device
* collapse a model's artifact repos (GGUF / FP8 / bnb-4bit / BF16) into one
* row with a format second level, and a bare row click auto-routes to the
* best artifact for this device. Undefined (chat) changes nothing. */
catalog?: CatalogGroup[];
}) {
const gpu = useGpuInfo();
// The currently-loaded/running model id. We read params.checkpoint from the
@ -1738,7 +1914,10 @@ export function HubModelPicker({
const downloadedSet = useMemo(() => {
const s = new Set<string>();
for (const c of cachedGguf) s.add(c.repo_id.toLowerCase());
for (const c of cachedModels) s.add(c.repo_id.toLowerCase());
// Skip partial (cancelled/incomplete) base repos: a partial snapshot has only
// some weights, so treating it as downloaded routes an On Device click to a fresh
// multi-GB re-download instead of loading the complete GGUF.
for (const c of cachedModels) if (!c.partial) s.add(c.repo_id.toLowerCase());
return s;
}, [cachedGguf, cachedModels]);
@ -1813,6 +1992,9 @@ export function HubModelPicker({
// rows regardless of the format dropdown (mirrors hfIds and the empty
// Recommended view); selecting a non-GGUF row is a silent no-op.
if (task) rows = rows.filter((r) => r.isGguf);
// A catalog group already renders its member repos as one canonical row;
// drop them from the live listing so they don't appear twice.
if (catalog) rows = rows.filter((r) => !groupForRepoId(r.id, catalog));
// With no explicit format, show the device-recommended formats (GGUF, plus
// MLX on Mac). When the user picks a format, honor it instead so Safetensors
// is not dropped by the recommendation default.
@ -1838,6 +2020,7 @@ export function HubModelPicker({
gpu,
isChatSupported,
task,
catalog,
]);
// Curated non-GGUF (safetensors) models for the Images picker. The HF listing +
@ -1855,6 +2038,108 @@ export function HubModelPicker({
return models.filter((m) => m.isGguf === false);
}, [models, task]);
// Catalog grouping (Images / Video pages): expansion + router state. A bare
// group-row click loads the best artifact for this device (downloaded first,
// then quality-that-fits); the chevron reveals the per-format second level.
const [expandedGroups, setExpandedGroups] = useState<Set<string>>(new Set());
const [routingGroupId, setRoutingGroupId] = useState<string | null>(null);
const toggleGroupExpanded = useCallback((key: string) => {
setExpandedGroups((prev) => {
const next = new Set(prev);
if (next.has(key)) next.delete(key);
else next.add(key);
return next;
});
}, []);
const isRepoDownloaded = useCallback(
(repoId: string) => downloadedSet.has(repoId.toLowerCase()),
[downloadedSet],
);
const deviceBudget = useMemo(
() => ({
gpuGb: gpu.available ? gpu.memoryTotalGb : 0,
systemRamGb: gpu.systemRamAvailableGb || 0,
}),
[gpu],
);
const routedArtifactFor = useCallback(
(group: CatalogGroup): ModelArtifact => {
// Honor the format filter when routing a bare group click. A group is only
// visible under a GGUF/Safetensors/MLX filter because at least one of its
// artifacts matches (catalogGroupMatchesFormat), so restrict the routing
// candidates to those same artifacts before the ladder picks. Otherwise a
// GGUF-filtered group could route to a large BF16/FP8 download the filter
// never surfaced (pickDefaultArtifact prefers a fitting non-GGUF).
const scoped =
formatFilter === "all"
? group
: (() => {
const artifacts = group.artifacts.filter((a) =>
matchesFormatFilter(a.repoId, a.format === "gguf", formatFilter),
);
return artifacts.length > 0 ? { ...group, artifacts } : group;
})();
return pickDefaultArtifact(scoped, {
...deviceBudget,
isDownloaded: isRepoDownloaded,
});
},
[deviceBudget, isRepoDownloaded, formatFilter],
);
const routeGroupClick = useCallback(
async (group: CatalogGroup, expandKey: string) => {
const artifact = routedArtifactFor(group);
if (artifact.format !== "gguf") {
onSelect(artifact.repoId, {
source: "hub",
isLora: false,
isDownloaded: isRepoDownloaded(artifact.repoId),
});
return;
}
// GGUF route: resolve the quant list, then load the ladder's pick. A
// failed fetch falls back to opening the format list instead -- toggle the
// caller's CONTEXT-SCOPED expandKey (not the context-free canonicalId), so the
// chevron (which toggles expandKey) can still collapse it and the same group in
// another list is not expanded too.
setRoutingGroupId(group.canonicalId);
try {
const res = normalizeGgufVariantsResponse(
await listGgufVariants(artifact.repoId, hfToken || undefined),
);
const quant = pickDefaultQuant(
res.variants,
res.defaultVariant,
deviceBudget,
);
if (quant) {
onSelect(artifact.repoId, {
source: "hub",
isLora: false,
ggufVariant: quant.quant,
ggufFilename: quant.filename,
isDownloaded: quant.downloaded,
expectedBytes: quant.size_bytes,
});
} else {
toggleGroupExpanded(expandKey);
}
} catch {
toggleGroupExpanded(expandKey);
} finally {
setRoutingGroupId(null);
}
},
[
routedArtifactFor,
onSelect,
isRepoDownloaded,
hfToken,
deviceBudget,
toggleGroupExpanded,
],
);
// Per-row meta + VRAM badge from the recommended listing's own metadata.
const recommendedMeta = useMemo(() => {
const map = new Map<
@ -1957,13 +2242,22 @@ export function HubModelPicker({
sortCachedRepos(
cachedModels.filter(
(c) =>
// A partially-downloaded snapshot is not on-device: listing it as loadable
// errors or triggers a silent multi-GB re-fetch on click (mirrors downloadedSet).
!c.partial &&
passesTaskGate(c.task, c.repo_id, task) &&
(!task || isUnslothRepoId(c.repo_id)),
// Diffusion pickers: unsloth repos plus any repo the backend can actually LOAD.
// Gate on a curated ARTIFACT (artifactForRepoId, what loadSpecFor resolves), not a
// group-key match: a base / uncurated-quant sibling (Qwen/Qwen-Image-2512) matches
// the group by key but has no loadable artifact and dead-ends at the trust gate.
(!task ||
isUnslothRepoId(c.repo_id) ||
(catalog ? artifactForRepoId(c.repo_id, catalog) !== null : false)),
),
downloadedSort,
loadTimes,
),
[cachedModels, downloadedSort, loadTimes, task],
[cachedModels, downloadedSort, loadTimes, task, catalog],
);
// Each local section's search is scoped to its own models (matched by name).
const localQuery = normalizeForSearch(debouncedQuery.trim());
@ -2125,11 +2419,24 @@ export function HubModelPicker({
}, [results, recommendedSearch.results]);
// Recommended models that match the current search query
// Catalog groups matching a typed query (old ids, format tokens, quant names
// all match); rendered as canonical rows above the remaining search results.
const matchedCatalogGroups = useMemo(() => {
if (!catalog || !showHfSection) return [];
return catalog.filter(
(g) =>
groupMatchesQuery(g, debouncedQuery.trim()) &&
catalogGroupMatchesFormat(g, formatFilter),
);
}, [catalog, showHfSection, debouncedQuery, formatFilter]);
const filteredRecommendedIds = useMemo(() => {
if (!showHfSection) return [];
const q = normalizeForSearch(debouncedQuery.trim());
return recommendedIds
.filter((id) => normalizeForSearch(id).includes(q))
// Member repos of a catalog group collapse into the group row above.
.filter((id) => !catalog || !groupForRepoId(id, catalog))
.filter((id) =>
matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter),
)
@ -2158,6 +2465,7 @@ export function HubModelPicker({
downloadedSet,
recommendedParamCountById,
gpu,
catalog,
]);
const recommendedSet = useMemo(
@ -2212,6 +2520,66 @@ export function HubModelPicker({
const hubOptionKeys = useMemo(() => {
const keys: string[] = [];
// Roving keys for an On Device cached section, mirroring renderCachedRows'
// render order exactly so arrow/Home/End nav matches the visual order.
// Without a catalog the rows are flat; with one, catalog members collapse
// under a canonical group row (whose key must lead), and the per-repo child
// rows only render (and only join the roving list) while the group is
// expanded. Ungrouped GGUF rows then ungrouped model rows follow.
const groupedCachedKeys = (
ggufRows: { repo_id: string }[],
modelRows: { repo_id: string }[],
keyPrefix: string,
): string[] => {
if (!catalog) {
return [
...ggufRows.map((c) =>
makeModelOptionKey("downloaded-gguf", c.repo_id),
),
...modelRows.map((c) =>
makeModelOptionKey("downloaded-model", c.repo_id),
),
];
}
const grouped = new Map<
CatalogGroup,
{ gguf: string[]; models: string[] }
>();
const ungroupedGguf: string[] = [];
const ungroupedModels: string[] = [];
for (const c of ggufRows) {
const key = makeModelOptionKey("downloaded-gguf", c.repo_id);
const group = groupForRepoId(c.repo_id, catalog);
if (group) {
const entry = grouped.get(group) ?? { gguf: [], models: [] };
entry.gguf.push(key);
grouped.set(group, entry);
} else {
ungroupedGguf.push(key);
}
}
for (const c of modelRows) {
const key = makeModelOptionKey("downloaded-model", c.repo_id);
const group = groupForRepoId(c.repo_id, catalog);
if (group) {
const entry = grouped.get(group) ?? { gguf: [], models: [] };
entry.models.push(key);
grouped.set(group, entry);
} else {
ungroupedModels.push(key);
}
}
const out: string[] = [];
for (const [group, rows] of grouped.entries()) {
out.push(makeModelOptionKey(keyPrefix, group.canonicalId));
if (expandedGroups.has(`${keyPrefix}:${group.canonicalId}`)) {
out.push(...rows.gguf, ...rows.models);
}
}
out.push(...ungroupedGguf, ...ungroupedModels);
return out;
};
// Downloaded (Unsloth) rows (query-filtered) on the On Device tab only.
if (
section === "downloaded" &&
@ -2220,19 +2588,21 @@ export function HubModelPicker({
(unslothCachedGguf.length > 0 || unslothCachedModelRows.length > 0)
) {
keys.push(
...unslothCachedGguf.map((model) =>
makeModelOptionKey("downloaded-gguf", model.repo_id),
),
);
keys.push(
...unslothCachedModelRows.map((model) =>
makeModelOptionKey("downloaded-model", model.repo_id),
...groupedCachedKeys(
unslothCachedGguf,
unslothCachedModelRows,
"cached-group",
),
);
}
// Unsloth-tab search keys (curated matches + HF unsloth results).
// Unsloth-tab search keys (curated catalog matches + curated/HF results).
if (showHfSection && section === "recommended") {
keys.push(
...matchedCatalogGroups.map((g) =>
makeModelOptionKey("search-catalog-group", g.canonicalId),
),
);
keys.push(
...filteredRecommendedIds.map((id) =>
makeModelOptionKey("search-recommended", id),
@ -2250,13 +2620,10 @@ export function HubModelPicker({
(otherCachedGguf.length > 0 || otherCachedModelRows.length > 0)
) {
keys.push(
...otherCachedGguf.map((model) =>
makeModelOptionKey("downloaded-gguf", model.repo_id),
),
);
keys.push(
...otherCachedModelRows.map((model) =>
makeModelOptionKey("downloaded-model", model.repo_id),
...groupedCachedKeys(
otherCachedGguf,
otherCachedModelRows,
"other-cached-group",
),
);
}
@ -2292,14 +2659,25 @@ export function HubModelPicker({
}
if (section === "recommended") {
// Curated safetensors rows render ABOVE the recommended rows (and call
// getOptionProps), so their keys must lead here or they fall back to the
// duplicate ...-option-missing id and drop out of arrow-key navigation.
keys.push(
...curatedSafetensorsRows.map((m) =>
makeModelOptionKey("curated-safetensors", m.id),
),
);
// Curated rows render ABOVE the recommended rows (and call getOptionProps),
// so their keys must lead here or they fall back to the duplicate
// ...-option-missing id and drop out of arrow-key navigation. With a
// catalog (Images / Video) those are the canonical catalog-group rows,
// gated by the same format filter as the render; without one they are the
// flat curated safetensors rows.
if (catalog) {
keys.push(
...catalog
.filter((g) => catalogGroupMatchesFormat(g, formatFilter))
.map((g) => makeModelOptionKey("catalog-group", g.canonicalId)),
);
} else {
keys.push(
...curatedSafetensorsRows.map((m) =>
makeModelOptionKey("curated-safetensors", m.id),
),
);
}
keys.push(
...recommendedRows.map((r) => makeModelOptionKey("recommended", r.id)),
);
@ -2308,15 +2686,19 @@ export function HubModelPicker({
return keys;
}, [
cachedReady,
catalog,
chatOnly,
curatedSafetensorsRows,
sortedCustomFolderModels,
customFoldersCollapsed,
downloadedCollapsed,
expandedGroups,
fineTunedRows,
fineTunedCollapsed,
filteredRecommendedIds,
formatFilter,
hfIds,
matchedCatalogGroups,
sortedLmStudio,
lmStudioCollapsed,
recommendedRows,
@ -2629,7 +3011,11 @@ export function HubModelPicker({
<div className={downloadedRowShellClassName(isSelected)}>
<div className="min-w-0 flex-1">
<ModelRow
label={c.repo_id}
label={
catalog
? stripArtifactSuffixesForDisplay(c.repo_id)
: c.repo_id
}
meta="GGUF"
showVision={c.has_vision ?? visionByRepo[c.repo_id]}
selected={isSelected}
@ -2685,7 +3071,9 @@ export function HubModelPicker({
>
<div className="min-w-0 flex-1">
<ModelRow
label={c.repo_id}
label={
catalog ? stripArtifactSuffixesForDisplay(c.repo_id) : c.repo_id
}
meta={`${isMlxId(c.repo_id) ? "MLX" : "Safetensors"} · ${formatBytes(
c.size_bytes,
)}`}
@ -2724,6 +3112,200 @@ export function HubModelPicker({
);
};
// One canonical row per catalog group (Recommended): click loads the routed
// best artifact for this device; the chevron opens the format second level.
const renderCatalogGroupRow = (group: CatalogGroup, keyPrefix: string) => {
const optionKey = makeModelOptionKey(keyPrefix, group.canonicalId);
const expandKey = `${keyPrefix}:${group.canonicalId}`;
const expanded = expandedGroups.has(expandKey);
const anyDownloaded = group.artifacts.some((a) => isRepoDownloaded(a.repoId));
const selected = group.artifacts.some((a) => a.repoId === value);
const routed = routedArtifactFor(group);
const formats = group.artifacts.map((a) => a.label).join(" / ");
return (
<div key={expandKey}>
<div className={downloadedRowShellClassName(selected)}>
<div className="min-w-0 flex-1">
<ModelRow
label={group.canonicalId}
hideOwner={true}
downloaded={anyDownloaded}
meta={
group.artifacts.length > 1
? `${group.description} · ${formats}`
: (group.description ?? formats)
}
tooltipText={`Loads ${routed.label} (best for this device). Open the chevron to pick a format.`}
selected={selected}
optionProps={hubModelList.getOptionProps(optionKey, selected)}
onClick={() => void routeGroupClick(group, expandKey)}
vramStatus={null}
gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
className={downloadedRowButtonClassName}
/>
</div>
{routingGroupId === group.canonicalId ? (
<span className="mr-1 flex h-6 w-[26px] shrink-0 items-center justify-center">
<Spinner className="size-3 text-muted-foreground" />
</span>
) : (
<button
type="button"
aria-label={expanded ? "Hide formats" : "Show formats"}
onClick={() => toggleGroupExpanded(expandKey)}
className="mr-1 flex h-6 w-[26px] shrink-0 items-center justify-center rounded text-muted-foreground/60 transition-colors hover:text-foreground"
>
{expanded ? (
<ChevronDownIcon className="size-3.5" />
) : (
<ChevronRightIcon className="size-3.5" />
)}
</button>
)}
</div>
{expanded && (
<ArtifactFormatList
group={group}
recommendedArtifactId={routed.repoId}
isRepoDownloaded={isRepoDownloaded}
onSelect={onSelect}
gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
systemRamGb={gpu.systemRamAvailableGb || undefined}
hfToken={hfToken || undefined}
parentOptionKey={optionKey}
/>
)}
</div>
);
};
// On Device sections: collapse cached member repos of a catalog group under
// one canonical row (click = load the best on-disk artifact, chevron = show
// the per-repo rows with their usual quant expanders / delete actions).
// Unknown repos render exactly as before.
const renderCachedRows = (
ggufRows: typeof visibleCachedGguf,
modelRows: typeof visibleCachedModelRows,
keyPrefix: string,
) => {
if (!catalog) {
return (
<>
{ggufRows.map(renderDownloadedGgufRow)}
{modelRows.map(renderDownloadedModelRow)}
</>
);
}
const grouped = new Map<
CatalogGroup,
{ gguf: typeof ggufRows; models: typeof modelRows }
>();
const ungroupedGguf: typeof ggufRows = [];
const ungroupedModels: typeof modelRows = [];
for (const c of ggufRows) {
const group = groupForRepoId(c.repo_id, catalog);
if (group) {
const entry = grouped.get(group) ?? { gguf: [], models: [] };
entry.gguf.push(c);
grouped.set(group, entry);
} else {
ungroupedGguf.push(c);
}
}
for (const c of modelRows) {
const group = groupForRepoId(c.repo_id, catalog);
if (group) {
const entry = grouped.get(group) ?? { gguf: [], models: [] };
entry.models.push(c);
grouped.set(group, entry);
} else {
ungroupedModels.push(c);
}
}
return (
<>
{[...grouped.entries()].map(([group, rows]) => {
const expandKey = `${keyPrefix}:${group.canonicalId}`;
const optionKey = makeModelOptionKey(keyPrefix, group.canonicalId);
const expanded = expandedGroups.has(expandKey);
const memberCount = rows.gguf.length + rows.models.length;
const selected =
rows.gguf.some((c) => c.repo_id === value) ||
rows.models.some((c) => c.repo_id === value);
return (
<div key={expandKey}>
<div className={downloadedRowShellClassName(selected)}>
<div className="min-w-0 flex-1">
<ModelRow
label={group.canonicalId}
meta={`${group.description} · ${memberCount} format${
memberCount === 1 ? "" : "s"
} on disk`}
selected={selected}
optionProps={hubModelList.getOptionProps(optionKey, selected)}
onClick={() => {
// routeGroupClick picks the best CURATED artifact; when one is on
// disk pickDefaultArtifact returns it, so keep that path. But this
// group can appear in On Device solely because a cached member
// matched by key/alias (a sibling prequant that is not a curated
// artifact). In that case the routed artifact is NOT downloaded, so
// load an actual on-disk member instead of downloading a different
// artifact -- the On Device row must "load the best on-disk artifact".
if (!isRepoDownloaded(routedArtifactFor(group).repoId)) {
const cachedModel = rows.models[0];
if (cachedModel) {
onSelect(cachedModel.repo_id, {
source: "hub",
isLora: false,
isDownloaded: true,
});
return;
}
if (rows.gguf.length > 0) {
toggleGroupExpanded(expandKey);
return;
}
}
void routeGroupClick(group, expandKey);
}}
vramStatus={null}
className={downloadedRowButtonClassName}
/>
</div>
{routingGroupId === group.canonicalId ? (
<span className="mr-1 flex h-6 w-[26px] shrink-0 items-center justify-center">
<Spinner className="size-3 text-muted-foreground" />
</span>
) : (
<button
type="button"
aria-label={expanded ? "Hide formats" : "Show formats"}
onClick={() => toggleGroupExpanded(expandKey)}
className="mr-1 flex h-6 w-[26px] shrink-0 items-center justify-center rounded text-muted-foreground/60 transition-colors hover:text-foreground"
>
{expanded ? (
<ChevronDownIcon className="size-3.5" />
) : (
<ChevronRightIcon className="size-3.5" />
)}
</button>
)}
</div>
{expanded && (
<div className="ml-3 border-l-2 border-accent/50 pl-1">
{rows.gguf.map(renderDownloadedGgufRow)}
{rows.models.map(renderDownloadedModelRow)}
</div>
)}
</div>
);
})}
{ungroupedGguf.map(renderDownloadedGgufRow)}
{ungroupedModels.map(renderDownloadedModelRow)}
</>
);
};
return (
<>
<div className="relative space-y-2">
@ -2977,9 +3559,11 @@ export function HubModelPicker({
{sortedLmStudio.length > 0 ? "Unsloth" : "Downloaded"}
</ListLabel>
{!downloadedCollapsed &&
unslothCachedGguf.map(renderDownloadedGgufRow)}
{!downloadedCollapsed &&
unslothCachedModelRows.map(renderDownloadedModelRow)}
renderCachedRows(
unslothCachedGguf,
unslothCachedModelRows,
"cached-group",
)}
</>
) : null}
@ -2998,9 +3582,11 @@ export function HubModelPicker({
Other models
</ListLabel>
{!otherModelsCollapsed &&
otherCachedGguf.map(renderDownloadedGgufRow)}
{!otherModelsCollapsed &&
otherCachedModelRows.map(renderDownloadedModelRow)}
renderCachedRows(
otherCachedGguf,
otherCachedModelRows,
"other-cached-group",
)}
</div>
) : null}
@ -3527,30 +4113,39 @@ export function HubModelPicker({
{showRecommendedSection ? (
<>
{/* Curated safetensors models (full diffusers pipelines / single-file
fp8). Shown above the GGUF rows; clicking loads directly (no quant
expander), the same path as a non-GGUF Recommended row. */}
{curatedSafetensorsRows.map((m) => {
const optionKey = makeModelOptionKey("curated-safetensors", m.id);
return (
<div key={m.id}>
<ModelRow
label={m.id}
hideOwner={true}
downloaded={downloadedSet.has(m.id.toLowerCase())}
capabilities={capsById.get(m.id)}
meta={m.description ?? "Safetensors"}
selected={value === m.id}
optionProps={hubModelList.getOptionProps(
optionKey,
value === m.id,
)}
onClick={() => handleModelClick(m.id)}
gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
/>
</div>
);
})}
{/* Curated models. With a catalog (Images / Video), one canonical
row per model with its formats as a second level -- rendered
unconditionally from the catalog, never dependent on the HF
listing's task tags. Without one (legacy), the flat curated
safetensors rows. */}
{catalog
? catalog
.filter((g) => catalogGroupMatchesFormat(g, formatFilter))
.map((g) => renderCatalogGroupRow(g, "catalog-group"))
: curatedSafetensorsRows.map((m) => {
const optionKey = makeModelOptionKey(
"curated-safetensors",
m.id,
);
return (
<div key={m.id}>
<ModelRow
label={m.id}
hideOwner={true}
downloaded={downloadedSet.has(m.id.toLowerCase())}
capabilities={capsById.get(m.id)}
meta={m.description ?? "Safetensors"}
selected={value === m.id}
optionProps={hubModelList.getOptionProps(
optionKey,
value === m.id,
)}
onClick={() => handleModelClick(m.id)}
gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
/>
</div>
);
})}
{recommendedSearch.isLoading &&
recommendedRows.length === 0 ? (
<div className="flex items-center gap-2 px-5 py-3">
@ -3560,7 +4155,8 @@ export function HubModelPicker({
</span>
</div>
) : recommendedRows.length === 0 &&
curatedSafetensorsRows.length === 0 ? (
curatedSafetensorsRows.length === 0 &&
!catalog?.length ? (
<div className="px-2.5 py-2 text-xs text-muted-foreground">
No models found.
</div>
@ -3573,7 +4169,12 @@ export function HubModelPicker({
return (
<div key={id}>
<ModelRow
label={id}
// Diffusion pickers standardize hub rows to the base
// model name; the GGUF/format badge carries the
// artifact kind. The id used on click is untouched.
label={
catalog ? stripArtifactSuffixesForDisplay(id) : id
}
hideOwner={true}
downloaded={downloadedSet.has(id.toLowerCase())}
capabilities={capsById.get(id)}
@ -3647,8 +4248,12 @@ export function HubModelPicker({
{showHfSection &&
section === "recommended" &&
filteredRecommendedIds.length > 0 ? (
(filteredRecommendedIds.length > 0 ||
matchedCatalogGroups.length > 0) ? (
<>
{matchedCatalogGroups.map((g) =>
renderCatalogGroupRow(g, "search-catalog-group"),
)}
{filteredRecommendedIds.map((id) => {
const vram = recommendedVramMap.get(id);
const optionKey = makeModelOptionKey(

3
studio/frontend/src/features/chat/api/chat-api.ts
View file

@ -323,6 +323,9 @@ export interface CachedModelRepo {
/** HF pipeline task: "text-to-image" for a cached diffusers pipeline repo
* (model_index.json present), so the chat picker can hide it. Absent = chat. */
task?: string | null;
/** True when the snapshot is incomplete (a cancelled/partial download). Such a
* repo must not count as downloaded, or a click re-downloads the full weights. */
partial?: boolean;
}
export async function listCachedModels(

248
studio/frontend/src/features/images/images-page.tsx
View file

@ -10,11 +10,19 @@ import {
ImageAdd02Icon,
InformationCircleIcon,
LayoutAlignRightIcon,
PencilEdit02Icon,
Settings02Icon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { TestTubeOutlineIcon } from "@/lib/hugeicons-derived";
import { Button } from "@/components/ui/button";
import {
DropdownMenu,
DropdownMenuContent,
DropdownMenuItem,
DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import { Input } from "@/components/ui/input";
import {
Popover,
@ -36,6 +44,11 @@ import { Textarea } from "@/components/ui/textarea";
import { InfoHint } from "@/components/ui/info-hint";
import { ModelSelector } from "@/components/assistant-ui/model-selector";
import { IMAGE_GEN_TASKS } from "@/components/assistant-ui/model-selector/pickers";
import {
IMAGE_CATALOG,
catalogToModelOptions,
loadSpecFor,
} from "@/components/assistant-ui/model-selector/model-catalog";
import type {
ModelOption,
ModelSelectorChangeMeta,
@ -69,100 +82,12 @@ import {
} from "./api";
import { DiffusionTrainPanel } from "./train/diffusion-train-panel";
// Curated diffusion GGUFs the picker recommends. The backend resolves each one's
// pipeline + base diffusers repo from its repo id, so the rail just lists them;
// the chat ModelSelector also surfaces any other on-device image GGUF.
const txt2img = (id: string, name: string): ModelOption => ({
id,
name,
description: "Text-to-image · GGUF",
isGguf: true,
});
// Instruction-editing GGUF (Qwen-Image-Edit). Same single-file GGUF flow as txt2img
// (the picker expands quant variants); the backend resolves it to the edit family, which
// exposes only the "edit" workflow.
const editGguf = (id: string, name: string): ModelOption => ({
id,
name,
description: "Image editing · GGUF",
isGguf: true,
});
// How to load a curated non-GGUF (safetensors) model. "pipeline" = a full diffusers
// repo (from_pretrained, embedded bnb-4bit quant auto-applied); "single_file" = a
// single safetensors transformer (e.g. fp8) assembled onto its base repo. The backend
// gates these to unsloth/* repos plus a short allowlist of official base repos (SDXL).
// Keyed by repo id so the load handler knows the kind (and, for single_file, the exact
// filename).
type SafetensorsSpec = { kind: "pipeline" | "single_file"; filename?: string };
const SAFETENSORS_MODELS: Record<string, SafetensorsSpec> = {
"unsloth/Z-Image-Turbo-unsloth-bnb-4bit": { kind: "pipeline" },
// Krea 2 Turbo: official vendor repo (bf16 pipeline), on the backend allowlist.
"krea/Krea-2-Turbo": { kind: "pipeline" },
// Ideogram 4: official vendor pipelines, on the backend allowlist. No bf16 repo
// exists: -fp8 stores its two DiTs as raw float8 (highest precision; ~46 GB
// resident after the bf16 cast); -nf4-diffusers is the bnb-4bit export (~11 GB).
"ideogram-ai/ideogram-4-fp8": { kind: "pipeline" },
"ideogram-ai/ideogram-4-nf4-diffusers": { kind: "pipeline" },
"unsloth/Qwen-Image-2512-unsloth-bnb-4bit": { kind: "pipeline" },
"unsloth/Qwen-Image-2512-FP8": {
kind: "single_file",
filename: "qwen-image-2512-fp8.safetensors",
},
// SDXL is a U-Net family loaded as a whole pipeline (from_pretrained). These
// official base repos are on the backend's non-GGUF allowlist.
"stabilityai/sdxl-turbo": { kind: "pipeline" },
"stabilityai/stable-diffusion-xl-base-1.0": { kind: "pipeline" },
};
// Curated non-GGUF picker entries (isGguf:false -> no quant expander, direct load).
const safetensors = (id: string, name: string, label: string): ModelOption => ({
id,
name,
description: `Text-to-image · ${label}`,
isGguf: false,
});
const MODELS: ModelOption[] = [
txt2img("unsloth/Z-Image-Turbo-GGUF", "Z-Image-Turbo"),
txt2img("unsloth/Z-Image-GGUF", "Z-Image"),
txt2img("unsloth/Qwen-Image-2512-GGUF", "Qwen-Image 2512"),
txt2img("unsloth/Qwen-Image-GGUF", "Qwen-Image"),
txt2img("unsloth/FLUX.1-schnell-GGUF", "FLUX.1 schnell"),
txt2img("unsloth/FLUX.1-dev-GGUF", "FLUX.1 dev"),
txt2img("unsloth/FLUX.2-klein-4B-GGUF", "FLUX.2 klein 4B"),
txt2img("unsloth/FLUX.2-klein-9B-GGUF", "FLUX.2 klein 9B"),
editGguf("unsloth/Qwen-Image-Edit-2511-GGUF", "Qwen-Image-Edit 2511"),
editGguf("unsloth/FLUX.1-Kontext-dev-GGUF", "FLUX.1 Kontext dev"),
safetensors(
"unsloth/Z-Image-Turbo-unsloth-bnb-4bit",
"Z-Image-Turbo (bnb-4bit)",
"Safetensors · bnb-4bit",
),
safetensors("krea/Krea-2-Turbo", "Krea 2 Turbo", "Safetensors · bf16"),
safetensors("ideogram-ai/ideogram-4-fp8", "Ideogram 4 (FP8)", "Safetensors · fp8"),
safetensors(
"ideogram-ai/ideogram-4-nf4-diffusers",
"Ideogram 4 (bnb-4bit)",
"Safetensors · bnb-4bit",
),
safetensors(
"unsloth/Qwen-Image-2512-unsloth-bnb-4bit",
"Qwen-Image 2512 (bnb-4bit)",
"Safetensors · bnb-4bit",
),
safetensors(
"unsloth/Qwen-Image-2512-FP8",
"Qwen-Image 2512 (FP8)",
"Safetensors · fp8",
),
safetensors("stabilityai/sdxl-turbo", "SDXL Turbo", "Safetensors · SDXL"),
safetensors(
"stabilityai/stable-diffusion-xl-base-1.0",
"SDXL Base 1.0",
"Safetensors · SDXL",
),
];
// Curated models come from the shared catalog: one canonical group per model,
// its artifacts (GGUF / FP8 / bnb-4bit / BF16) as data, and the load kind per
// artifact via loadSpecFor (replacing the old SAFETENSORS_MODELS table). The
// picker renders groups with a format second level and routes bare clicks to
// the best artifact for the device.
const MODELS: ModelOption[] = catalogToModelOptions(IMAGE_CATALOG);
// Workflow tabs. `requires` is the backend workflow id (status.workflows) that must
// be supported by the loaded model for the tab to enable; null = always available.
@ -321,23 +246,70 @@ const PAGE_SIZE = 50;
// Export filename, e.g. Unsloth_20260624-143005_123.png. Batch siblings share
// the seed + timestamp, so they get a "_<n>" suffix past the first one.
function exportFilename(image: GalleryImage): string {
type ImageExportFormat = "png" | "jpeg" | "webp";
function exportFilename(image: GalleryImage, format: ImageExportFormat = "png"): string {
const d = new Date(image.created_at * 1000);
const p = (n: number) => String(n).padStart(2, "0");
const stamp =
`${d.getFullYear()}${p(d.getMonth() + 1)}${p(d.getDate())}` +
`-${p(d.getHours())}${p(d.getMinutes())}${p(d.getSeconds())}`;
const suffix = image.batch_index > 0 ? `_${image.batch_index}` : "";
return `Unsloth_${stamp}_${image.seed}${suffix}.png`;
const ext = format === "jpeg" ? "jpg" : format;
return `Unsloth_${stamp}_${image.seed}${suffix}.${ext}`;
}
function downloadImage(src: string, image: GalleryImage) {
function saveBlobUrl(href: string, filename: string) {
const link = document.createElement("a");
link.href = src;
link.download = exportFilename(image);
link.href = href;
link.download = filename;
link.click();
}
// PNG saves the stored bytes verbatim (keeps the embedded recipe metadata);
// JPEG / WebP re-encode client-side from the already-fetched object URL. JPEG
// has no alpha, so it is flattened onto white first.
async function downloadImage(
src: string,
image: GalleryImage,
format: ImageExportFormat = "png",
) {
if (format === "png") {
saveBlobUrl(src, exportFilename(image, format));
return;
}
try {
const el = new Image();
el.decoding = "async";
el.src = src;
await el.decode();
const canvas = document.createElement("canvas");
canvas.width = el.naturalWidth;
canvas.height = el.naturalHeight;
const ctx = canvas.getContext("2d");
if (!ctx) throw new Error("canvas 2d context unavailable");
if (format === "jpeg") {
ctx.fillStyle = "#ffffff";
ctx.fillRect(0, 0, canvas.width, canvas.height);
}
ctx.drawImage(el, 0, 0);
const blob = await new Promise<Blob | null>((resolve) =>
canvas.toBlob(resolve, `image/${format}`, 0.95),
);
if (!blob) throw new Error(`could not encode ${format}`);
const url = URL.createObjectURL(blob);
try {
saveBlobUrl(url, exportFilename(image, format));
} finally {
// Give the click a tick to start before revoking.
setTimeout(() => URL.revokeObjectURL(url), 10_000);
}
} catch {
// Conversion failed (decode/encode); fall back to the original PNG bytes.
saveBlobUrl(src, exportFilename(image, "png"));
}
}
function formatTimestamp(epochSeconds: number): string {
return new Date(epochSeconds * 1000).toLocaleString();
}
@ -489,6 +461,9 @@ function formatResolvedValue(key: string, value: string | boolean | null): strin
if (value === null || value === "") return "Off";
if (typeof value === "boolean") return value ? "On" : "Off";
if (value === "_native_cudnn" || value.toLowerCase() === "cudnn") return "cuDNN";
// Deferred speed auto: the dense pipe stays exact/eager and compiles on the
// 3rd image of the session (the tooltip carries the full reason).
if (value === "deferred") return "On from 3rd image";
return value.toUpperCase();
}
@ -1044,7 +1019,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
"auto",
);
const [memoryMode, setMemoryMode] = useState<"auto" | "fast" | "balanced" | "low_vram">("auto");
const [transformerCache, setTransformerCache] = useState<"off" | "fbcache">("off");
const [transformerCache, setTransformerCache] = useState<"auto" | "off" | "fbcache">("auto");
const [cpuOffload, setCpuOffload] = useState(false);
// The last load descriptor, so "Reapply" can reload the same model with new advanced
// options without the user re-picking it from the dropdown.
@ -1527,7 +1502,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
transformer_quant: transformerQuant === "auto" ? undefined : transformerQuant,
attention_backend: attentionBackend === "auto" ? undefined : attentionBackend,
memory_mode: memoryMode === "auto" ? undefined : memoryMode,
transformer_cache: transformerCache === "off" ? undefined : transformerCache,
transformer_cache: transformerCache === "auto" ? undefined : transformerCache,
});
} catch (err) {
dismissLoadToast();
@ -1576,8 +1551,8 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
// busy, while the backend rejects the second load with a 409.
if (busy !== null) return;
// Curated non-GGUF model: load as a full pipeline or single-file safetensors.
const spec = SAFETENSORS_MODELS[id];
if (spec) {
const spec = loadSpecFor(id, IMAGE_CATALOG);
if (spec && spec.kind !== "gguf") {
setQuant(null);
const d = defaultsFor(id);
setSteps(d.steps);
@ -1913,7 +1888,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
<>
<AdvancedSelect
label="Speed"
hint="Auto picks per model (GGUF compiles, dense stays eager). eager = fused kernels, no compile. default/max add torch.compile (max also TF32 + fused QKV)."
hint="Auto picks per model: GGUF compiles at load; a dense model keeps the first two images exact and eager, then compiles from the 3rd (~2x from there). eager = fused kernels, no compile. default/max add torch.compile (max also TF32 + fused QKV)."
badge={<ResolvedBadge status={status} controlKey="speed_mode" />}
value={speedMode}
onValueChange={(v) => setSpeedMode(v as typeof speedMode)}
@ -1930,8 +1905,8 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
to GGUF (or nothing loaded) and otherwise show why it is unavailable. */}
{!status?.loaded || status.model_kind === "gguf" ? (
<AdvancedSelect
label="Dtype"
hint="Transformer compute dtype. Auto picks the fastest precision the hardware supports (at least INT8 on a capable GPU; FP8 on data-center cards) by loading the FULL base model and quantising its transformer onto low-precision tensor cores, and falls back to running the GGUF as-is when the device, VRAM or disk can't take it. Off always runs the GGUF as-is."
label="Precision"
hint="How the model computes. Auto picks the fastest precision the hardware supports (at least INT8 on a capable GPU; FP8 on data-center cards) by loading the FULL base model and quantising its transformer onto low-precision tensor cores, and falls back to running the GGUF as-is when the device, VRAM or disk can't take it. Off always runs the GGUF as-is."
badge={<ResolvedBadge status={status} controlKey="transformer_quant" />}
value={transformerQuant}
onValueChange={(v) => setTransformerQuant(v as typeof transformerQuant)}
@ -1947,14 +1922,14 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
) : (
<div className="flex items-center justify-between gap-2">
<span className="flex items-center gap-1 text-xs font-medium text-muted-foreground">
Dtype
Precision
</span>
<span className="text-xs text-muted-foreground/60">GGUF models only</span>
</div>
)}
<AdvancedSelect
label="Attention"
hint="Attention kernel. Auto upgrades to cuDNN fused attention on NVIDIA when a speed profile is active. sage is INT8 attention (small quality cost)."
hint="Attention kernel. Auto upgrades to cuDNN fused attention on NVIDIA when a speed profile is active. sage is INT8 attention: fast (10-40%) but can black-frame some families (Qwen, Wan), so it never engages automatically."
badge={<ResolvedBadge status={status} controlKey="attention_backend" />}
value={attentionBackend}
onValueChange={(v) => setAttentionBackend(v as typeof attentionBackend)}
@ -1981,11 +1956,12 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
/>
<AdvancedSelect
label="Step cache"
hint="First-Block-Cache reuses the transformer tail across steps for many-step models (~1.4x). Leave off for few-step distilled models."
hint="First-Block-Cache reuses the transformer tail across steps for many-step models (~1.4x). Auto turns it on at 20+ steps and off for few-step distilled models, re-checked per image."
badge={<ResolvedBadge status={status} controlKey="transformer_cache" />}
value={transformerCache}
onValueChange={(v) => setTransformerCache(v as typeof transformerCache)}
options={[
["auto", "Auto"],
["off", "Off"],
["fbcache", "First-Block-Cache"],
]}
@ -2029,6 +2005,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
variant="ghost"
className="!h-[34px]"
task={IMAGE_GEN_TASKS}
catalog={IMAGE_CATALOG}
open={active && selectorOpen}
onOpenChange={(o) => setSelectorOpen(active && o)}
/>
@ -2037,11 +2014,22 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
Create is the generation workspace; Train is the LoRA training workspace. */}
<Tabs value={pageMode} onValueChange={(v) => setPageMode(v as "create" | "train")}>
<TabsList className="h-[34px]">
<TabsTrigger value="create" className="w-[64px]">
Create
{/* Same icons as the sidebar's New Chat / Train entries, so the
two workspaces read as the same actions everywhere. */}
{/* TabsTrigger renders children inside a plain inline span (and preflight
makes svg display:block), so the icon and label need their own flex
row to stay on one line. */}
<TabsTrigger value="create" className="w-[84px]">
<span className="flex items-center gap-1.5">
<HugeiconsIcon icon={PencilEdit02Icon} className="size-3.5" />
Create
</span>
</TabsTrigger>
<TabsTrigger value="train" className="w-[64px]">
Train
<TabsTrigger value="train" className="w-[84px]">
<span className="flex items-center gap-1.5">
<HugeiconsIcon icon={TestTubeOutlineIcon} className="size-3.5" />
Train
</span>
</TabsTrigger>
</TabsList>
</Tabs>
@ -2629,15 +2617,31 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
Size/seed live in the Recipe popover, so no separate chip here. */}
<div className="absolute bottom-4 right-4 flex items-center gap-0.5 rounded-xl bg-background/80 p-1 shadow-lg ring-1 ring-border backdrop-blur">
<RecipePopover image={selected} onRestore={restoreSettings} active={active} />
<Button
size="sm"
variant="ghost"
className="gap-1.5"
onClick={() => downloadImage(selectedSrc, selected)}
>
<HugeiconsIcon icon={Download01Icon} className="size-4" />
Download
</Button>
<DropdownMenu>
<DropdownMenuTrigger asChild={true}>
<Button size="sm" variant="ghost" className="gap-1.5">
<HugeiconsIcon icon={Download01Icon} className="size-4" />
Download
</Button>
</DropdownMenuTrigger>
<DropdownMenuContent align="end">
<DropdownMenuItem
onClick={() => void downloadImage(selectedSrc, selected, "png")}
>
PNG (original, keeps recipe)
</DropdownMenuItem>
<DropdownMenuItem
onClick={() => void downloadImage(selectedSrc, selected, "jpeg")}
>
JPEG (smaller)
</DropdownMenuItem>
<DropdownMenuItem
onClick={() => void downloadImage(selectedSrc, selected, "webp")}
>
WebP
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>
<Button
size="sm"
variant="ghost"

37
studio/frontend/src/features/images/train/diffusion-train-panel.tsx
View file

@ -64,7 +64,7 @@ const FAMILY_PRESETS: FamilyPreset[] = [
label: "FLUX.1-dev (12B)",
base_repos: ["black-forest-labs/FLUX.1-dev"],
defaults: { rank: 16, lr: 0.0001, resolution: 512 },
vram_note: "Gated repo - accept the license on Hugging Face and add your HF token. QLoRA (4-bit).",
vram_note: "Needs a free Hugging Face license: accept it on the model page, then add your token.",
gated: true,
},
{
@ -72,21 +72,21 @@ const FAMILY_PRESETS: FamilyPreset[] = [
label: "Qwen-Image (20B)",
base_repos: ["unsloth/Qwen-Image-2512-unsloth-bnb-4bit", "Qwen/Qwen-Image"],
defaults: { rank: 16, lr: 0.00005, resolution: 512 },
vram_note: "Largest model - QLoRA (4-bit) on a big GPU. Start at 512px, batch 1.",
vram_note: "The biggest model: needs a large GPU. Start at 512px.",
},
{
name: "z-image",
label: "Z-Image-Turbo (6B)",
base_repos: ["unsloth/Z-Image-Turbo-unsloth-bnb-4bit", "Tongyi-MAI/Z-Image-Turbo"],
defaults: { rank: 16, lr: 0.0001, resolution: 768 },
vram_note: "Lightest and fastest to train. bf16 only (fp16 is unstable for this family).",
vram_note: "The smallest and fastest to train. A great first pick.",
},
{
name: "sdxl",
label: "SDXL (U-Net)",
base_repos: ["stabilityai/stable-diffusion-xl-base-1.0", "stabilityai/sdxl-turbo"],
defaults: { rank: 16, lr: 0.0001, resolution: 1024 },
vram_note: "The classic text-to-image base. Trains comfortably at 1024px.",
vram_note: "The classic. Trains comfortably at 1024px.",
},
];
@ -850,7 +850,7 @@ export function DiffusionTrainPanel({
<option value="linear">Linear decay</option>
</select>
<p className="text-[11px] leading-snug text-muted-foreground">
How the learning rate evolves over the run (shown live in the LR chart).
How fast the model learns over time. Constant is fine for most runs.
</p>
</div>
{lrScheduler !== "constant" &&
@ -870,8 +870,7 @@ export function DiffusionTrainPanel({
<option value="off">Off (faster steps)</option>
</select>
<p className="text-[11px] leading-snug text-muted-foreground">
Recomputes activations in the backward pass: a large VRAM saving for a modest
per-step slowdown.
Saves a lot of GPU memory in exchange for slightly slower steps.
</p>
</div>
@ -898,15 +897,10 @@ export function DiffusionTrainPanel({
))}
</select>
<p className="text-[11px] leading-snug text-muted-foreground">
How the frozen base weights are quantised. nf4 (4-bit) uses the least VRAM;
bf16 is fastest but needs the most. Auto picks this family&apos;s recommended
mode.
How the base model is stored while training. Auto picks the best fit for
your GPU.
{basePrequantized && (
<>
{" "}
This base is already 4-bit quantised, so only nf4/auto apply; pick a dense
(bf16) base repo for the other modes.
</>
<> This base is already 4-bit, so only nf4/auto apply.</>
)}
</p>
</div>
@ -924,7 +918,7 @@ export function DiffusionTrainPanel({
<option value="no">fp32 (no mixed)</option>
</select>
<p className="text-[11px] leading-snug text-muted-foreground">
Mixed-precision autocast for the U-Net. bf16 suits modern GPUs.
How the math runs during training. bf16 is right for modern GPUs.
</p>
</div>
)}
@ -944,8 +938,7 @@ export function DiffusionTrainPanel({
<option value="off">Off</option>
</select>
<p className="text-[11px] leading-snug text-muted-foreground">
torch.compile the transformer blocks. Adds a one-time warmup, then speeds up
each step.
Warms up once at the start, then every training step runs faster.
</p>
</div>
)}
@ -1084,8 +1077,8 @@ export function DiffusionTrainPanel({
</Button>
</div>
<p className="text-[11px] text-muted-foreground">
10-50 images work well. Optional captions: a .txt per image or a
metadata.jsonl; without them the trigger prompt below captions every image.
10-50 images are plenty. Captions are optional: without them, the trigger
prompt below describes every image.
</p>
</div>
) : (
@ -1113,8 +1106,8 @@ export function DiffusionTrainPanel({
)}
{selectedDataset.caption_count === 0 && !gridOpen && (
<p className="text-[11px] text-muted-foreground">
No caption files - the trigger prompt below captions every image, or
open Review captions to label them.
No captions yet: the trigger prompt below will describe every image,
or open Review captions to write your own.
</p>
)}
</>

18
studio/frontend/src/features/video/api.ts
View file

@ -46,6 +46,8 @@ export interface VideoStatus {
speed_optims: string[];
attention_backend?: string | null;
transformer_cache?: string | null;
// Dense DiT precision actually engaged ("int8" | "fp8" | ...) or null for bf16.
transformer_quant?: string | null;
// Whether the loaded family produces a synchronized audio track.
has_audio: boolean;
// Per-family generation defaults + shape constraints; null when unloaded.
@ -102,6 +104,9 @@ export interface VideoLoadRequest {
| "aiter";
transformer_cache?: "off" | "fbcache";
transformer_cache_threshold?: number;
// Dense DiT precision on full-pipeline loads (omit for the hardware-ladder auto;
// "none" pins plain bf16). GGUF / single-file checkpoints carry their own precision.
transformer_quant?: "none" | "fp8" | "int8" | "nvfp4" | "mxfp8";
}
export interface VideoGenerateRequest {
@ -224,3 +229,16 @@ export async function fetchGalleryVideoObjectUrl(url: string): Promise<string> {
if (!res.ok) throw new Error(await readFastApiError(res));
return URL.createObjectURL(await res.blob());
}
/** Server-side transcode for the Download menu (WebM / GIF). The backend 501s
* with a readable message when the codec for that format is unavailable. */
export async function fetchGalleryVideoExport(
id: string,
format: "webm" | "gif",
): Promise<Blob> {
const res = await authFetch(
`/api/inference/video/gallery/${id}/export?format=${format}`,
);
if (!res.ok) throw new Error(await readFastApiError(res));
return res.blob();
}

246
studio/frontend/src/features/video/video-page.tsx
View file

@ -14,6 +14,12 @@ import {
import { HugeiconsIcon } from "@hugeicons/react";
import { Button } from "@/components/ui/button";
import {
DropdownMenu,
DropdownMenuContent,
DropdownMenuItem,
DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import { Input } from "@/components/ui/input";
import {
Popover,
@ -33,6 +39,11 @@ import { Textarea } from "@/components/ui/textarea";
import { InfoHint } from "@/components/ui/info-hint";
import { ModelSelector } from "@/components/assistant-ui/model-selector";
import { VIDEO_GEN_TASKS } from "@/components/assistant-ui/model-selector/pickers";
import {
VIDEO_CATALOG,
catalogToModelOptions,
loadSpecFor,
} from "@/components/assistant-ui/model-selector/model-catalog";
import type {
ModelOption,
ModelSelectorChangeMeta,
@ -51,6 +62,7 @@ import {
cancelVideoGeneration,
clearVideoGallery,
deleteGalleryVideo,
fetchGalleryVideoExport,
fetchGalleryVideoObjectUrl,
generateVideo,
getVideoGallery,
@ -61,60 +73,13 @@ import {
unloadVideoModel,
} from "./api";
// How to load a curated non-GGUF (safetensors) video model. "pipeline" = a full diffusers
// repo (from_pretrained). The backend gates these to unsloth/* repos plus the official
// family base repos. Keyed by repo id so the load handler knows the kind.
type PipelineSpec = { kind: "pipeline"; filename?: string };
const PIPELINE_MODELS: Record<string, PipelineSpec> = {
"Lightricks/LTX-2": { kind: "pipeline" },
// Wan2.2 diffusers base repos (no GGUF variant yet): loaded as full pipelines. TI2V-5B
// is a single-DiT 720p-class model; T2V-A14B is the dual-expert MoE. The backend gates
// these to the Wan-AI base repos (see _TRUSTED_NON_GGUF_VIDEO_REPOS).
"Wan-AI/Wan2.2-TI2V-5B-Diffusers": { kind: "pipeline" },
"Wan-AI/Wan2.2-T2V-A14B-Diffusers": { kind: "pipeline" },
// HunyuanVideo-1.5 community Diffusers repack (tencent's own repo is the original
// non-diffusers layout and cannot load as a pipeline).
"hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v": { kind: "pipeline" },
};
// A curated GGUF picker entry: isGguf true expands its .gguf files in the quant expander
// (like the image GGUF repos), and the backend resolves the pipeline + base repo from the id.
const ggufModel = (id: string, name: string): ModelOption => ({
id,
name,
description: "Text-to-video · GGUF",
isGguf: true,
});
// A curated non-GGUF pipeline entry (isGguf false -> no quant expander, direct load).
const pipelineModel = (id: string, name: string, description: string): ModelOption => ({
id,
name,
description,
isGguf: false,
});
// Curated text-to-video models the picker recommends. The chat ModelSelector also surfaces
// any other on-device video GGUF (via the VIDEO_GEN_TASKS filter).
const VIDEO_MODELS: ModelOption[] = [
ggufModel("unsloth/LTX-2.3-GGUF", "LTX 2.3 distilled"),
pipelineModel("Lightricks/LTX-2", "LTX 2 (base, bf16)", "Text-to-video with audio · Safetensors"),
pipelineModel(
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
"Wan 2.2 TI2V 5B",
"Text-to-video 720p · Safetensors",
),
pipelineModel(
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
"Wan 2.2 T2V A14B (MoE)",
"Text-to-video, dual-expert · Safetensors",
),
pipelineModel(
"hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
"HunyuanVideo 1.5 (480p)",
"Text-to-video 480p · Safetensors",
),
];
// Curated models come from the shared catalog: one canonical group per model
// with its artifacts as data (the HunyuanVideo group carries both the 480p and
// 720p repacks), and the load kind per artifact via loadSpecFor (replacing the
// old PIPELINE_MODELS table). The picker renders groups with a format second
// level -- which also finally surfaces LTX-2.3 in Recommended (its HF
// pipeline_tag is image-to-video, so the live text-to-video listing missed it).
const VIDEO_MODELS: ModelOption[] = catalogToModelOptions(VIDEO_CATALOG);
// Per-model generation defaults (steps + guidance), matched by repo-id substring, most
// specific first. The distilled model wants very few steps and no guidance; the full base
@ -175,23 +140,45 @@ const galleryCache: {
const PAGE_SIZE = 50;
// Export filename, e.g. Unsloth_video_20260624-143005_123.mp4.
function exportFilename(video: GalleryVideo): string {
type VideoExportFormat = "mp4" | "webm" | "gif";
function exportFilename(video: GalleryVideo, format: VideoExportFormat = "mp4"): string {
const d = new Date(video.created_at);
const p = (n: number) => String(n).padStart(2, "0");
const stamp = Number.isNaN(d.getTime())
? "unknown"
: `${d.getFullYear()}${p(d.getMonth() + 1)}${p(d.getDate())}` +
`-${p(d.getHours())}${p(d.getMinutes())}${p(d.getSeconds())}`;
return `Unsloth_video_${stamp}_${video.seed}.mp4`;
return `Unsloth_video_${stamp}_${video.seed}.${format}`;
}
function downloadVideo(src: string, video: GalleryVideo) {
function saveBlobUrl(href: string, filename: string) {
const link = document.createElement("a");
link.href = src;
link.download = exportFilename(video);
link.href = href;
link.download = filename;
link.click();
}
// MP4 saves the already-fetched original bytes; WebM / GIF are transcoded by
// the backend on demand (501 with a readable reason when the codec is absent).
async function downloadVideo(
src: string,
video: GalleryVideo,
format: VideoExportFormat = "mp4",
) {
if (format === "mp4") {
saveBlobUrl(src, exportFilename(video, format));
return;
}
const blob = await fetchGalleryVideoExport(video.id, format);
const url = URL.createObjectURL(blob);
try {
saveBlobUrl(url, exportFilename(video, format));
} finally {
setTimeout(() => URL.revokeObjectURL(url), 10_000);
}
}
function formatTimestamp(iso: string): string {
const d = new Date(iso);
return Number.isNaN(d.getTime()) ? iso : d.toLocaleString();
@ -530,7 +517,10 @@ export function VideoPage({ active = true }: { active?: boolean }) {
const [attentionBackend, setAttentionBackend] = useState<
"auto" | "native" | "cudnn" | "flash3" | "sage"
>("auto");
const [transformerCache, setTransformerCache] = useState<"off" | "fbcache">("off");
const [transformerCache, setTransformerCache] = useState<"auto" | "off" | "fbcache">("auto");
const [transformerQuant, setTransformerQuant] = useState<
"auto" | "none" | "fp8" | "int8" | "nvfp4" | "mxfp8"
>("auto");
// The last load descriptor, so "Reapply" can reload the same model with new advanced
// options without the user re-picking it from the dropdown.
const lastLoad = useRef<{ repoId: string; kind: "gguf" | "single_file" | "pipeline"; filename?: string } | null>(
@ -549,6 +539,12 @@ export function VideoPage({ active = true }: { active?: boolean }) {
const [videos, setVideos] = useState<GalleryVideo[]>(() => galleryCache.videos);
const [hasMore, setHasMore] = useState(() => galleryCache.hasMore);
const [selectedId, setSelectedId] = useState<string | null>(() => galleryCache.selectedId);
// Autoplay replays per selected clip (3 total plays, then pause). Reset on
// every selection change so a new generation or pick gets its own 3 plays.
const playCountRef = useRef(0);
useEffect(() => {
playCountRef.current = 0;
}, [selectedId]);
const [srcById, setSrcById] = useState<Record<string, string>>(() =>
Object.fromEntries(galleryCache.srcById),
);
@ -700,6 +696,29 @@ export function VideoPage({ active = true }: { active?: boolean }) {
void loadGallery();
}, [loadGallery]);
// WebM/GIF go through a server-side transcode that can take a few seconds
// (and 501s with a readable reason when the codec is missing), so wrap the
// helper with progress + error toasts; MP4 saves instantly.
const handleDownload = useCallback(
async (src: string, video: GalleryVideo, format: "mp4" | "webm" | "gif") => {
if (format === "mp4") {
void downloadVideo(src, video, format);
return;
}
const toastId = toast.loading(`Converting to ${format.toUpperCase()}…`);
try {
await downloadVideo(src, video, format);
toast.dismiss(toastId);
} catch (err) {
toast.dismiss(toastId);
toast.error(
err instanceof Error ? err.message : `Failed to export ${format}`,
);
}
},
[],
);
const handleDelete = useCallback(async (id: string) => {
try {
await deleteGalleryVideo(id);
@ -894,7 +913,8 @@ export function VideoPage({ active = true }: { active?: boolean }) {
memory_mode: memoryMode === "auto" ? undefined : memoryMode,
speed_mode: speedMode === "auto" ? undefined : speedMode,
attention_backend: attentionBackend === "auto" ? undefined : attentionBackend,
transformer_cache: transformerCache === "off" ? undefined : transformerCache,
transformer_cache: transformerCache === "auto" ? undefined : transformerCache,
transformer_quant: transformerQuant === "auto" ? undefined : transformerQuant,
});
} catch (err) {
dismissLoadToast();
@ -914,6 +934,7 @@ export function VideoPage({ active = true }: { active?: boolean }) {
speedMode,
attentionBackend,
transformerCache,
transformerQuant,
],
);
@ -931,10 +952,15 @@ export function VideoPage({ active = true }: { active?: boolean }) {
// Ignore picks while a load/generation/unload is in flight.
if (busy !== null) return;
// Curated non-GGUF model: load as a full pipeline.
const spec = PIPELINE_MODELS[id];
if (spec) {
const spec = loadSpecFor(id, VIDEO_CATALOG);
if (spec && spec.kind !== "gguf") {
setQuant(null);
const d = defaultsFor(id);
// The distilled variant lives in the single-file checkpoint name
// (ltx-2.3-...-distilled...), not the repo id, so include the filename when
// seeding defaults -- mirroring the GGUF branch below. Without it these
// distilled BF16/FP8 entries fall through to the generic LTX 40-step/CFG-4
// defaults instead of the distilled 8-step/guidance-1 schedule.
const d = defaultsFor(spec.filename ? `${id}/${spec.filename}` : id);
setSteps(d.steps);
setGuidance(d.guidance);
void handleLoad(id, { kind: spec.kind, filename: spec.filename });
@ -1125,7 +1151,7 @@ export function VideoPage({ active = true }: { active?: boolean }) {
/>
<AdvancedSelect
label="Speed"
hint="Auto picks per model (GGUF compiles near-losslessly, dense stays eager). eager = fused kernels, no compile. default/max add torch.compile (max also TF32 + fused QKV)."
hint="Auto compiles every model at load: a clip takes minutes to denoise, so the one-time compile always pays for itself within a single run. eager = fused kernels, no compile. max adds TF32 + fused QKV."
badge={<ResolvedBadge status={status} controlKey="speed_mode" />}
value={speedMode}
onValueChange={(v) => setSpeedMode(v as typeof speedMode)}
@ -1137,9 +1163,36 @@ export function VideoPage({ active = true }: { active?: boolean }) {
["max", "Max"],
]}
/>
{/* The dense transformer_quant fast path only engages on a full-pipeline load; a
GGUF / single-file checkpoint already carries its own precision, so gate the
control and otherwise show why it is unavailable. */}
{!status?.loaded || status.model_kind === "pipeline" ? (
<AdvancedSelect
label="Precision"
hint="How the model computes. Auto picks the fastest precision the hardware supports (at least INT8 on a capable GPU; FP8 on data-center cards) by quantising the transformer onto low-precision tensor cores, and keeps plain bf16 when the device or memory plan can't take it. Off always runs bf16."
badge={<ResolvedBadge status={status} controlKey="transformer_quant" />}
value={transformerQuant}
onValueChange={(v) => setTransformerQuant(v as typeof transformerQuant)}
options={[
["auto", "Auto (fastest for GPU)"],
["none", "Off (bf16)"],
["fp8", "FP8"],
["int8", "INT8"],
["nvfp4", "NVFP4 (Blackwell)"],
["mxfp8", "MXFP8 (Blackwell)"],
]}
/>
) : (
<div className="flex items-center justify-between gap-2">
<span className="flex items-center gap-1 text-xs font-medium text-muted-foreground">
Precision
</span>
<span className="text-xs text-muted-foreground/60">Full-pipeline models only</span>
</div>
)}
<AdvancedSelect
label="Attention"
hint="Attention kernel. Auto upgrades to cuDNN fused attention on NVIDIA when a speed profile is active. sage is INT8 attention (small quality cost)."
hint="Attention kernel. Auto upgrades to cuDNN fused attention on NVIDIA when a speed profile is active. sage is INT8 attention: fast (10-40%) but can black-frame some families (Qwen, Wan), so it never engages automatically."
badge={<ResolvedBadge status={status} controlKey="attention_backend" />}
value={attentionBackend}
onValueChange={(v) => setAttentionBackend(v as typeof attentionBackend)}
@ -1153,11 +1206,12 @@ export function VideoPage({ active = true }: { active?: boolean }) {
/>
<AdvancedSelect
label="Step cache"
hint="First-Block-Cache reuses the transformer tail across steps for many-step models. Leave off for few-step distilled models."
hint="First-Block-Cache reuses the transformer tail across steps for many-step models. Auto turns it on at 20+ steps and off for few-step distilled models, re-checked per clip."
badge={<ResolvedBadge status={status} controlKey="transformer_cache" />}
value={transformerCache}
onValueChange={(v) => setTransformerCache(v as typeof transformerCache)}
options={[
["auto", "Auto"],
["off", "Off"],
["fbcache", "First-Block-Cache"],
]}
@ -1191,6 +1245,7 @@ export function VideoPage({ active = true }: { active?: boolean }) {
variant="ghost"
className="!h-[34px]"
task={VIDEO_GEN_TASKS}
catalog={VIDEO_CATALOG}
open={active && selectorOpen}
onOpenChange={(o) => setSelectorOpen(active && o)}
/>
@ -1350,16 +1405,27 @@ export function VideoPage({ active = true }: { active?: boolean }) {
<div className="relative flex flex-1 items-center justify-center overflow-auto p-6">
{selected && selectedSrc ? (
<>
{/* The first video element in the app. autoPlay + loop + muted + playsInline
so it plays inline without a gesture; controls let the user scrub/unmute. */}
{/* The first video element in the app. autoPlay + muted + playsInline so
it plays inline without a gesture; controls let the user scrub/unmute.
Instead of a bare `loop`, onEnded replays up to 3 total plays then
pauses -- an endlessly looping clip is distracting once you've seen
it. The counter resets per selection (`key` remounts the element). */}
<video
key={selected.id}
src={selectedSrc}
controls
autoPlay
loop
muted
playsInline
onPlay={() => {
playCountRef.current += 1;
}}
onEnded={(e) => {
if (playCountRef.current < 3) {
e.currentTarget.currentTime = 0;
void e.currentTarget.play();
}
}}
className="max-h-full max-w-full rounded-xl object-contain shadow-sm"
/>
{selected.has_audio && (
@ -1371,15 +1437,31 @@ export function VideoPage({ active = true }: { active?: boolean }) {
{/* Actions grouped in one glass toolbar so they stay legible over any clip. */}
<div className="absolute bottom-4 right-4 flex items-center gap-0.5 rounded-xl bg-background/80 p-1 shadow-lg ring-1 ring-border backdrop-blur">
<RecipePopover video={selected} onRestore={restoreSettings} active={active} />
<Button
size="sm"
variant="ghost"
className="gap-1.5"
onClick={() => downloadVideo(selectedSrc, selected)}
>
<HugeiconsIcon icon={Download01Icon} className="size-4" />
Download
</Button>
<DropdownMenu>
<DropdownMenuTrigger asChild={true}>
<Button size="sm" variant="ghost" className="gap-1.5">
<HugeiconsIcon icon={Download01Icon} className="size-4" />
Download
</Button>
</DropdownMenuTrigger>
<DropdownMenuContent align="end">
<DropdownMenuItem
onClick={() => void handleDownload(selectedSrc, selected, "mp4")}
>
MP4 (original{selected.has_audio ? ", keeps audio" : ""})
</DropdownMenuItem>
<DropdownMenuItem
onClick={() => void handleDownload(selectedSrc, selected, "webm")}
>
WebM (web embeds)
</DropdownMenuItem>
<DropdownMenuItem
onClick={() => void handleDownload(selectedSrc, selected, "gif")}
>
GIF (preview, no audio)
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>
<Button
size="sm"
variant="ghost"
@ -1472,8 +1554,10 @@ export function VideoPage({ active = true }: { active?: boolean }) {
<Spinner className="size-4 text-muted-foreground" />
</span>
)}
{/* A terse caption strip so cards read at a glance. */}
<span className="relative z-10 truncate bg-gradient-to-t from-black/70 to-transparent px-1 pb-0.5 pt-2 text-left text-[9px] font-medium text-white">
{/* A terse caption strip so cards read at a glance. Left/bottom
padding clears the rounded-lg corner and the selection border
so the leading "5.0s" is never clipped by the curve. */}
<span className="relative z-10 truncate bg-gradient-to-t from-black/70 to-transparent px-2 pb-1 pt-2 text-left text-[9px] font-medium leading-none text-white">
{clipMeta(video)}
</span>
{/* Selection marker on a non-focusable overlay. */}

14
studio/frontend/src/lib/hugeicons-derived.ts Normal file
View file

@ -0,0 +1,14 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
// Derived HugeIcons shared between the sidebar and page tabs, so the same
// visual language ("Train" = test tube, "Create" = pencil) appears everywhere.
import { TestTube01Icon } from "@hugeicons/core-free-icons";
// TestTube01Icon's last 2 paths are interior bubbles; slice to the first
// 3 (outline + cap + liquid line) to drop them. Original export untouched.
export const TestTubeOutlineIcon = TestTube01Icon.slice(
0,
3,
) as typeof TestTube01Icon;