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

Merge remote-tracking branch 'origin/video-hunyuan-gate' into fold-integration

Browse source 
# Conflicts:
#	studio/backend/core/inference/video.py
#	studio/backend/routes/models.py
#	studio/backend/tests/test_cached_gguf_routes.py
#	studio/frontend/src/features/images/images-page.tsx
#	studio/frontend/src/features/images/train/diffusion-train-panel.tsx
This commit is contained in:
Daniel Han 2026-07-07 01:15:20 +00:00
commit de099eaecd
40 changed files with 4585 additions and 1007 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",

513
scripts/video_quality.py Normal file
View file

@ -0,0 +1,513 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Video quality-vs-cost harness for the Studio video backend.
The video analogue of scripts/diffusion_quality.py: hold the prompt + seed +
shape fixed, render one clip with a high-fidelity reference configuration
(default the family's BF16 artifact), then render the same clip with each
candidate configuration (a GGUF quant, a dense torchao quant, a speed profile,
a step cache) and measure how far the output drifts from the reference.
Per candidate it reports:
- mean PSNR / SSIM over evenly sampled frames (pixel + structural fidelity),
- a temporal-consistency deviation: the relative error between the reference's
and the candidate's frame-to-frame motion-energy series, which catches
flicker/juddering that per-frame SSIM alone can miss,
- black-frame and NaN collapse checks (the failure mode quant bugs actually
produce, per the image backend's qwen fp8 incident),
- an audio check for families that generate sound (LTX-2): RMS ratio vs the
reference and a silence trip-wire,
- wall time per generate and peak VRAM.
Verdict bands map the standing accuracy budget: a candidate that keeps mean
SSIM at or above 0.75 PASSes (a ~25 percent structural drift is acceptable for
a large speed/memory win), 0.50-0.75 WARNs, and below 0.50 or any black/NaN/
silence collapse FAILs regardless of how fast it is.
Runtime-budgeted: ONE short clip per candidate (default 33 frames at 480p-class
sizes) so a full family sweep stays in minutes, not hours. Metrics are pure
numpy; torch / diffusers / the backend load lazily so --help and --selftest run
on a host without them.
Examples:
# CPU metric sanity check (no GPU, no model):
python scripts/video_quality.py --selftest
# LTX-2.3 GGUF quants against the BF16 GGUF reference:
CUDA_VISIBLE_DEVICES=1 python scripts/video_quality.py \\
--model unsloth/LTX-2.3-GGUF --model-kind gguf \\
--reference "gguf_filename=distilled-1.1/ltx-2.3-22b-distilled-1.1-BF16.gguf" \\
--candidates "gguf_filename=distilled-1.1/ltx-2.3-22b-distilled-1.1-Q8_0.gguf" \\
"gguf_filename=distilled-1.1/ltx-2.3-22b-distilled-1.1-UD-Q4_K_M.gguf" \\
--steps 8 --guidance 1.0 --out-dir outputs/video_quality/ltx23
# Wan2.2-5B dense int8 + speed profiles against plain bf16:
CUDA_VISIBLE_DEVICES=1 python scripts/video_quality.py \\
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \\
--reference "" \\
--candidates "transformer_quant=int8" "speed_mode=max" \\
--steps 20 --out-dir outputs/video_quality/wan5b
"""
from __future__ import annotations
import argparse
import json
import math
import sys
import time
from pathlib import Path
from typing import Any, Optional
_BACKEND_ROOT = Path(__file__).resolve().parent.parent / "studio" / "backend"
if str(_BACKEND_ROOT) not in sys.path:
sys.path.insert(0, str(_BACKEND_ROOT))
DEFAULT_PROMPT = (
"a golden retriever puppy runs through shallow ocean waves at sunset, "
"splashing water, cinematic, camera tracking sideways"
)
# Finite PSNR (dB) an identical clip is capped to when averaged, matching
# scripts/diffusion_quality.py.
_PERFECT_MATCH_PSNR = 100.0
# ── frame metrics (pure numpy; frames are uint8 HxWx3 arrays) ────────────────
def _gray(frame: Any) -> Any:
import numpy as np
f = np.asarray(frame, dtype = np.float64)
return f @ np.array([0.299, 0.587, 0.114])
def frame_psnr(a: Any, b: Any) -> float:
import numpy as np
a64 = np.asarray(a, dtype = np.float64)
b64 = np.asarray(b, dtype = np.float64)
if a64.shape != b64.shape:
return 0.0
mse = float(((a64 - b64) ** 2).mean())
if mse == 0.0:
return math.inf
return 20.0 * math.log10(255.0) - 10.0 * math.log10(mse)
def _box_mean(x: Any, w: int) -> Any:
import numpy as np
r = w // 2
xp = np.pad(x, r, mode = "edge")
ii = np.cumsum(np.cumsum(xp, axis = 0), axis = 1)
ii = np.pad(ii, ((1, 0), (1, 0)), mode = "constant")
h, wd = x.shape
total = ii[w : h + w, w : wd + w] - ii[0:h, w : wd + w] - ii[w : h + w, 0:wd] + ii[0:h, 0:wd]
return total / float(w * w)
def frame_ssim(
a: Any,
b: Any,
window: int = 7,
) -> float:
"""Pure numpy box-window SSIM on luminance (Wang et al. constants); identical
math to scripts/diffusion_quality.py so image and video budgets compare."""
ga, gb = _gray(a), _gray(b)
if ga.shape != gb.shape:
return 0.0
c1, c2 = (0.01 * 255) ** 2, (0.03 * 255) ** 2
mu_a, mu_b = _box_mean(ga, window), _box_mean(gb, window)
mu_a2, mu_b2, mu_ab = mu_a * mu_a, mu_b * mu_b, mu_a * mu_b
var_a = _box_mean(ga * ga, window) - mu_a2
var_b = _box_mean(gb * gb, window) - mu_b2
cov_ab = _box_mean(ga * gb, window) - mu_ab
ssim_map = ((2 * mu_ab + c1) * (2 * cov_ab + c2)) / (
(mu_a2 + mu_b2 + c1) * (var_a + var_b + c2)
)
return float(ssim_map.mean())
def motion_energy(frames: Any) -> list[float]:
"""Mean absolute frame-to-frame luminance difference, one value per frame
transition. The temporal signature of the clip: flicker inflates it, frozen
or smeared motion deflates it."""
import numpy as np
grays = [_gray(f) for f in frames]
return [float(np.abs(grays[i + 1] - grays[i]).mean()) for i in range(len(grays) - 1)]
def temporal_deviation(ref_frames: Any, cand_frames: Any) -> float:
"""Relative L1 error between the two motion-energy series (0 = identical
temporal behaviour). Series lengths must match (same frame count)."""
ref_series = motion_energy(ref_frames)
cand_series = motion_energy(cand_frames)
if len(ref_series) != len(cand_series) or not ref_series:
return math.inf
denom = sum(abs(v) for v in ref_series) + 1e-6
return sum(abs(r - c) for r, c in zip(ref_series, cand_series)) / denom
def clip_metrics(
ref_frames: Any,
cand_frames: Any,
sample_count: int = 5,
) -> dict[str, Any]:
"""All frame metrics for one candidate clip vs the reference clip."""
import numpy as np
# The gate holds the requested shape (num_frames) fixed for reference and
# candidate alike, so both clips must decode to the same frame count. A
# shorter candidate is a truncated/corrupt render, not a valid one; comparing
# only the shared prefix would let good early frames mask the missing tail, so
# the mismatch is recorded and gated as FAIL (see verdict()) rather than
# silently dropped. No off-by-one is tolerated: nothing in the encode/decode
# path justifies one.
ref_count, cand_count = len(ref_frames), len(cand_frames)
frame_count_mismatch = ref_count != cand_count
n = min(ref_count, cand_count)
if n == 0:
# An empty/corrupt decode must gate as FAIL, not crash the whole run.
return {
"frames_compared": 0,
"ref_frame_count": ref_count,
"cand_frame_count": cand_count,
"frame_count_mismatch": frame_count_mismatch,
"psnr_mean": 0.0,
"ssim_mean": 0.0,
"temporal_deviation": math.inf,
"min_luma": 0.0,
"has_nan": True,
}
idx = sorted({int(round(i * (n - 1) / max(1, sample_count - 1))) for i in range(sample_count)})
psnrs = [min(frame_psnr(ref_frames[i], cand_frames[i]), _PERFECT_MATCH_PSNR) for i in idx]
ssims = [frame_ssim(ref_frames[i], cand_frames[i]) for i in idx]
lumas = [float(_gray(cand_frames[i]).mean() / 255.0) for i in idx]
has_nan = any(
bool(np.isnan(np.asarray(f, dtype = np.float64)).any()) for f in (cand_frames[i] for i in idx)
)
return {
"frames_compared": len(idx),
"ref_frame_count": ref_count,
"cand_frame_count": cand_count,
"frame_count_mismatch": frame_count_mismatch,
"psnr_mean": sum(psnrs) / len(psnrs),
"ssim_mean": sum(ssims) / len(ssims),
"temporal_deviation": temporal_deviation(ref_frames[:n], cand_frames[:n]),
"min_luma": min(lumas),
"has_nan": has_nan,
}
def audio_metrics(ref_audio: Optional[Any], cand_audio: Optional[Any]) -> dict[str, Any]:
"""RMS comparison for families with sound. None audio on both sides is fine;
losing the track (or emitting silence) when the reference has one is not."""
import numpy as np
def _rms(a: Any) -> Optional[float]:
if a is None:
return None
arr = np.asarray(a, dtype = np.float64)
return float(np.sqrt((arr**2).mean())) if arr.size else 0.0
ref_rms, cand_rms = _rms(ref_audio), _rms(cand_audio)
# NaN candidate audio compares False against any threshold, so call it out
# explicitly: a NaN track is a collapse, not a pass.
silent_collapse = (
ref_rms is not None
and ref_rms >= 1e-3
and (cand_rms is None or math.isnan(cand_rms) or cand_rms < 1e-4)
)
return {"ref_rms": ref_rms, "cand_rms": cand_rms, "silent_collapse": silent_collapse}
def verdict(metrics: dict[str, Any], audio: dict[str, Any]) -> str:
"""PASS / WARN / FAIL per the standing accuracy budget (~25 percent structural
drift acceptable, 50 percent or a collapse never)."""
if (
metrics["has_nan"]
or metrics.get("frame_count_mismatch")
or metrics["min_luma"] < 0.02
or audio.get("silent_collapse")
):
return "FAIL"
if metrics["ssim_mean"] < 0.50 or metrics["temporal_deviation"] > 1.0:
return "FAIL"
if metrics["ssim_mean"] < 0.75 or metrics["temporal_deviation"] > 0.5:
return "WARN"
return "PASS"
# ── mp4 decode (PyAV, same dependency the backend encodes with) ─────────────
def decode_mp4(mp4_bytes: bytes, workdir: Path, name: str) -> tuple[list[Any], Optional[Any]]:
"""Frames (uint8 arrays) + mono audio samples (float array or None) from bytes."""
import av
import numpy as np
path = workdir / f"{name}.mp4"
path.write_bytes(mp4_bytes)
container = av.open(str(path))
frames = [f.to_ndarray(format = "rgb24") for f in container.decode(container.streams.video[0])]
audio = None
if container.streams.audio:
container.close()
container = av.open(str(path))
chunks = [c.to_ndarray() for c in container.decode(container.streams.audio[0])]
if chunks:
audio = np.concatenate([c.reshape(c.shape[0], -1).mean(axis = 0) for c in chunks])
container.close()
return frames, audio
# ── configuration plumbing ───────────────────────────────────────────────────
def parse_spec(spec: str) -> dict[str, str]:
"""'k=v;k=v' (or space-free 'k=v,k=v') -> dict; empty string -> {} (pure base)."""
out: dict[str, str] = {}
for part in spec.replace(",", ";").split(";"):
part = part.strip()
if not part:
continue
if "=" not in part:
raise ValueError(f"Bad candidate spec fragment '{part}' (expected key=value)")
key, value = part.split("=", 1)
out[key.strip()] = value.strip()
return out
def spec_label(spec: dict[str, str]) -> str:
if not spec:
return "base"
return ",".join(
f"{k}={Path(v).name if k == 'gguf_filename' else v}" for k, v in sorted(spec.items())
)
def run_config(
backend: Any, args: Any, spec: dict[str, str], workdir: Path, name: str
) -> dict[str, Any]:
"""Load per spec, generate the fixed clip, unload. Returns frames/audio/cost."""
import torch
load_kwargs: dict[str, Any] = {
"gguf_filename": spec.get("gguf_filename"),
"model_kind": spec.get("model_kind", args.model_kind),
"memory_mode": spec.get("memory_mode"),
"speed_mode": spec.get("speed_mode"),
"attention_backend": spec.get("attention_backend"),
"transformer_cache": spec.get("transformer_cache"),
"transformer_quant": spec.get("transformer_quant"),
}
t0 = time.monotonic()
status = backend.load_pipeline(args.model, **load_kwargs)
load_s = time.monotonic() - t0
if torch.cuda.is_available():
torch.cuda.reset_peak_memory_stats()
t0 = time.monotonic()
result = backend.generate(
prompt = args.prompt,
width = args.width,
height = args.height,
num_frames = args.frames,
fps = args.fps,
steps = args.steps,
guidance = args.guidance,
seed = args.seed,
)
generate_s = time.monotonic() - t0
peak_gib = torch.cuda.max_memory_allocated() / 2**30 if torch.cuda.is_available() else 0.0
backend.unload()
frames, audio = decode_mp4(result["mp4_bytes"], workdir, name)
return {
"frames": frames,
"audio": audio,
"load_s": round(load_s, 1),
"generate_s": round(generate_s, 1),
"peak_vram_gib": round(peak_gib, 2),
"resolved": {
k: v
for k, v in status.items()
if k
in (
"speed_mode",
"attention_backend",
"transformer_cache",
"transformer_quant",
"offload_policy",
"model_kind",
)
},
}
def run_gate(args: Any) -> int:
from core.inference.video import get_video_backend
out_dir = Path(args.out_dir)
out_dir.mkdir(parents = True, exist_ok = True)
backend = get_video_backend()
print(f"reference: {args.reference or 'base'}", flush = True)
ref = run_config(backend, args, parse_spec(args.reference), out_dir, "reference")
print(
f" load {ref['load_s']}s, generate {ref['generate_s']}s, "
f"peak {ref['peak_vram_gib']} GiB",
flush = True,
)
rows = []
for spec_str in args.candidates:
spec = parse_spec(spec_str)
label = spec_label(spec)
print(f"candidate: {label}", flush = True)
cand = run_config(backend, args, spec, out_dir, label.replace("/", "_").replace("=", "-"))
metrics = clip_metrics(ref["frames"], cand["frames"], sample_count = args.sample_frames)
audio = audio_metrics(ref["audio"], cand["audio"])
row = {
"candidate": label,
**{
k: (round(v, 4) if isinstance(v, float) and math.isfinite(v) else v)
for k, v in metrics.items()
},
**{f"audio_{k}": v for k, v in audio.items()},
"load_s": cand["load_s"],
"generate_s": cand["generate_s"],
"ref_generate_s": ref["generate_s"],
"peak_vram_gib": cand["peak_vram_gib"],
"resolved": cand["resolved"],
"verdict": verdict(metrics, audio),
}
rows.append(row)
print(
f" ssim {row['ssim_mean']:.3f} | psnr {row['psnr_mean']:.1f} dB | "
f"temporal {row['temporal_deviation']:.3f} | luma>={row['min_luma']:.3f} | "
f"gen {row['generate_s']}s (ref {ref['generate_s']}s) | "
f"vram {row['peak_vram_gib']} GiB | {row['verdict']}",
flush = True,
)
report = {
"model": args.model,
"reference": args.reference or "base",
"prompt": args.prompt,
"shape": [args.width, args.height, args.frames, args.fps],
"steps": args.steps,
"guidance": args.guidance,
"seed": args.seed,
"reference_cost": {k: ref[k] for k in ("load_s", "generate_s", "peak_vram_gib")},
"candidates": rows,
}
(out_dir / "report.json").write_text(json.dumps(report, indent = 1))
print(f"report: {out_dir / 'report.json'}", flush = True)
return 0 if all(r["verdict"] != "FAIL" for r in rows) else 1
# ── selftest (CPU-only, synthetic clips, no torch/model) ────────────────────
def selftest() -> int:
import numpy as np
rng = np.random.default_rng(0)
h, w, n = 64, 96, 12
def make_clip(
offset = 0.0,
noise = 0.0,
black = False,
):
frames = []
for t in range(n):
x = np.linspace(0, 1, w)[None, :] + t * 0.05 + offset
base = (np.sin(x * 6.283) * 0.5 + 0.5) * 255.0
frame = np.repeat(base[..., None], 3, axis = 2) * np.ones((h, 1, 1))
if noise:
frame = frame + rng.normal(0, noise, frame.shape)
if black:
frame = frame * 0.0
frames.append(np.clip(frame, 0, 255).astype(np.uint8))
return frames
ref = make_clip()
ok = True
def check(cond, msg):
nonlocal ok
print(("PASS: " if cond else "FAIL: ") + msg)
ok = ok and cond
same = clip_metrics(ref, make_clip())
check(
same["ssim_mean"] > 0.99 and same["temporal_deviation"] < 0.01,
f"identical clip scores ~1 (ssim {same['ssim_mean']:.3f})",
)
check(verdict(same, {"silent_collapse": False}) == "PASS", "identical clip verdict PASS")
noisy = clip_metrics(ref, make_clip(noise = 12.0))
check(0.3 < noisy["ssim_mean"] < 0.99, f"noisy clip degrades ssim ({noisy['ssim_mean']:.3f})")
black = clip_metrics(ref, make_clip(black = True))
check(
verdict(black, {"silent_collapse": False}) == "FAIL",
f"black clip verdict FAIL (min_luma {black['min_luma']:.3f})",
)
shifted = clip_metrics(ref, make_clip(offset = 0.5))
check(shifted["ssim_mean"] < same["ssim_mean"], "content shift lowers ssim")
# A truncated render whose surviving prefix is pixel-identical must still FAIL
# on the frame-count mismatch alone, not PASS on the good early frames.
truncated = clip_metrics(ref, make_clip()[: n // 2])
check(
truncated["frame_count_mismatch"] is True
and verdict(truncated, {"silent_collapse": False}) == "FAIL",
f"truncated clip verdict FAIL ({truncated['cand_frame_count']}/{truncated['ref_frame_count']} frames)",
)
audio = audio_metrics(np.sin(np.linspace(0, 100, 16000)), np.zeros(16000))
check(audio["silent_collapse"] is True, "silent audio collapse detected")
audio_ok = audio_metrics(
np.sin(np.linspace(0, 100, 16000)), np.sin(np.linspace(0, 100, 16000)) * 0.8
)
check(audio_ok["silent_collapse"] is False, "attenuated audio is not a collapse")
print("VIDEO-QUALITY-SELFTEST", "PASS" if ok else "FAIL")
return 0 if ok else 1
def main() -> int:
parser = argparse.ArgumentParser(description = __doc__.split("\n")[0])
parser.add_argument("--selftest", action = "store_true", help = "CPU metric sanity check")
parser.add_argument("--model", help = "Repo id handed to the video backend")
parser.add_argument("--model-kind", default = None, help = "pipeline | gguf | single_file")
parser.add_argument(
"--reference", default = "", help = "Reference spec 'k=v;k=v' ('' = plain base load)"
)
parser.add_argument("--candidates", nargs = "+", default = [], help = "Candidate specs 'k=v;k=v'")
parser.add_argument("--prompt", default = DEFAULT_PROMPT)
parser.add_argument("--width", type = int, default = 768)
parser.add_argument("--height", type = int, default = 512)
parser.add_argument("--frames", type = int, default = 33)
parser.add_argument("--fps", type = int, default = 24)
parser.add_argument("--steps", type = int, default = None)
parser.add_argument("--guidance", type = float, default = None)
parser.add_argument("--seed", type = int, default = 7)
parser.add_argument("--sample-frames", type = int, default = 5)
parser.add_argument("--out-dir", default = "outputs/video_quality")
args = parser.parse_args()
if args.selftest:
return selftest()
if not args.model or not args.candidates:
parser.error("--model and --candidates are required (or use --selftest)")
return run_gate(args)
if __name__ == "__main__":
sys.exit(main())

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

@ -223,12 +223,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
@ -305,6 +311,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()
@ -317,6 +329,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
@ -1305,6 +1324,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,
@ -1472,12 +1504,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,
)
@ -1498,10 +1533,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",
@ -1527,7 +1565,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,
@ -1560,11 +1600,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,
@ -2080,6 +2122,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,
*,
@ -2145,6 +2269,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)
@ -2480,6 +2637,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).

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

@ -30,6 +30,7 @@ family's official base repos, or a local path the user explicitly picked.
from __future__ import annotations
import contextlib
import inspect
import os
import tempfile
@ -42,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,
@ -64,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,
@ -101,6 +111,11 @@ _TRUSTED_NON_GGUF_VIDEO_REPOS = frozenset(
# remote code, so allowed as full (pipeline-kind) loads like the LTX-2 bases.
"wan-ai/wan2.2-ti2v-5b-diffusers",
"wan-ai/wan2.2-t2v-a14b-diffusers",
# HunyuanVideo-1.5 community Diffusers repacks (tencent's own repo is the
# original non-diffusers layout: config.json, no model_index.json, so it
# cannot load through HunyuanVideo15Pipeline at all).
"hunyuanvideo-community/hunyuanvideo-1.5-diffusers-480p_t2v",
"hunyuanvideo-community/hunyuanvideo-1.5-diffusers-720p_t2v",
}
)
@ -171,6 +186,37 @@ def _picked_gguf_arch(repo_id: str, gguf_filename: str) -> Optional[str]:
return None
class _VideoGenerationCancelled(Exception):
"""Unwinds a denoise loop that has no cooperative interrupt (no step callback);
generate() maps it to the VIDEO_CANCELLED_MSG sentinel the routes 409 on."""
@contextlib.contextmanager
def _scheduler_step_progress(pipe: Any, on_step: Any):
"""Progress + cancellation for pipelines WITHOUT callback_on_step_end.
HunyuanVideo15Pipeline exposes no per-step callback, but every denoise step
makes exactly one ``scheduler.step`` call, so wrapping that method gives the
same per-step tick the callback path gets. ``on_step`` receives the 1-based
step count and may raise (_VideoGenerationCancelled) to abort the loop. The
original method is always restored, even when the pipeline raises.
"""
scheduler = pipe.scheduler
original = scheduler.step
count = {"n": 0}
def _step(*args: Any, **kwargs: Any) -> Any:
count["n"] += 1
on_step(count["n"])
return original(*args, **kwargs)
scheduler.step = _step
try:
yield
finally:
scheduler.step = original
def _detect_load_family(
repo_id: str, gguf_filename: Optional[str], family_override: Optional[str]
) -> Optional[VideoFamily]:
@ -227,11 +273,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
@ -336,6 +393,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)
@ -395,6 +453,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
@ -414,6 +475,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."""
@ -425,6 +487,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:
@ -448,6 +511,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,
),
@ -760,6 +824,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,
@ -774,6 +839,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)
@ -815,6 +881,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
@ -978,10 +1058,31 @@ class VideoBackend:
# the DiT dense. Quant must precede compile (dynamic quant is ~30x slower eager),
# so it runs before apply_speed_optims below -- same order as diffusion.py.
transformer_quant_engaged: Optional[str] = None
quant_skipped_for_offload = False
if (
kind == "pipeline"
and normalize_transformer_quant(transformer_quant) is not None
and dense_transformer_supported(target)
and plan.offload_policy != "none"
):
# Offload hooks move modules with Module.to(), which torchao quantized
# tensors reject (aten._has_compatible_shallow_copy_type is
# unimplemented) -- observed as a hard crash on the Wan2.2-A14B gate
# run, where the 114 GB dual DiT plans model offload. A dense DiT
# under offload beats a crashed one, so quant is skipped, surfaced in
# the resolved record, and the user can force it by pinning a
# resident memory mode.
logger.info(
"video.transformer_quant: skipped (offload policy '%s' moves the "
"DiT via Module.to(), unsupported for torchao quantized tensors); "
"pin a resident memory mode to combine quant with this model",
plan.offload_policy,
)
quant_skipped_for_offload = True
elif (
kind == "pipeline"
and normalize_transformer_quant(transformer_quant) is not None
and dense_transformer_supported(target)
):
engaged = []
for view in views:
@ -1016,12 +1117,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.
@ -1037,23 +1160,55 @@ 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,
# A quantized transformer's block residuals are larger, so it needs the
# higher FBCache trigger threshold to cache at all. Mirror the image path
# (diffusion.py): both an engaged transformer_quant AND a GGUF checkpoint
# (quantized weights) count as quant-active here.
quant_active = transformer_quant_engaged is not None or kind == "gguf",
# (quantized weights) count as quant-active here (cache_quant_active, L1172).
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:
@ -1076,7 +1231,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:
@ -1117,7 +1275,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,
@ -1125,16 +1285,28 @@ 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,
transformer_quant_engaged or "off",
"dense DiT(s) torchao-quantised onto the low-precision tensor cores"
if transformer_quant_engaged is not None
else "not engaged (dense bf16 DiT loaded)",
else (
"skipped: offload moves the DiT, unsupported for torchao "
"tensors; pin a resident memory mode to combine them"
if quant_skipped_for_offload
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)",
),
}
)
@ -1164,7 +1336,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.
@ -1269,12 +1445,19 @@ class VideoBackend:
kwargs: dict[str, Any] = {
"prompt": prompt,
"num_inference_steps": steps,
fam.cfg_kwarg: guidance,
"width": width,
"height": height,
"num_frames": frames,
"generator": generator,
}
if fam.guidance_via_guider:
# HunyuanVideo-1.5: __call__ has no guidance kwarg at all; the
# CFG scale is a plain attribute on the pipeline's guider
# component, set per request. Near-1 scales auto-disable CFG
# inside the guider itself (_is_cfg_enabled's is_close check).
pipe.guider.guidance_scale = float(guidance)
else:
kwargs[fam.cfg_kwarg] = guidance
if negative_prompt and "negative_prompt" in call_params:
kwargs["negative_prompt"] = negative_prompt
# LTX-2 takes frame_rate (it shapes the audio track length); other
@ -1304,25 +1487,83 @@ class VideoBackend:
"error": None,
}
def _on_step(p, step_index, timestep, callback_kwargs):
if cancel.is_set():
p._interrupt = True
return callback_kwargs
done = step_index + 1
def _tick(done: int) -> None:
elapsed = time.monotonic() - started
self._gen.update(
step = done,
eta_seconds = (elapsed / max(1, done)) * max(0, steps - done),
)
def _on_step(p, step_index, timestep, callback_kwargs):
if cancel.is_set():
p._interrupt = True
return callback_kwargs
_tick(step_index + 1)
return callback_kwargs
def _on_scheduler_step(done: int) -> None:
# No cooperative _interrupt here: without a callback the pipeline
# never checks it, so cancellation must unwind the denoise loop
# via an exception (mapped to the cancelled sentinel below).
if cancel.is_set():
raise _VideoGenerationCancelled()
_tick(done)
if "callback_on_step_end" in call_params:
kwargs["callback_on_step_end"] = _on_step
progress_ctx = contextlib.nullcontext()
else:
# HunyuanVideo-1.5 has no step callback; every scheduler.step
# call is exactly one denoise step, so wrap it for progress +
# 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)
with torch.inference_mode():
output = pipe(**kwargs)
try:
with torch.inference_mode(), progress_ctx:
output = pipe(**kwargs)
except _VideoGenerationCancelled:
# This cancel unwinds pipe.__call__ by exception (the scheduler
# wrapper has no cooperative _interrupt), skipping the pipeline's
# end-of-call maybe_free_model_hooks(); under model/group offload
# the currently-onloaded modules would otherwise stay on the GPU
# until the next request touches them.
free_hooks = getattr(pipe, "maybe_free_model_hooks", None)
if callable(free_hooks):
try:
free_hooks()
except Exception: # noqa: BLE001 -- cleanup is best-effort
pass
raise RuntimeError(VIDEO_CANCELLED_MSG) from None
if cancel.is_set():
raise RuntimeError(VIDEO_CANCELLED_MSG)
@ -1453,6 +1694,7 @@ class VideoBackend:
"attention_backend": None,
"transformer_cache": None,
"transformer_quant": None,
"text_encoder_quant": None,
"has_audio": False,
"defaults": None,
"resolved": None,
@ -1480,6 +1722,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,

73
studio/backend/core/inference/video_families.py
View file

@ -50,6 +50,12 @@ class VideoFamily:
transformer2_class: Optional[str] = None
is_moe: bool = False
cfg2_kwarg: Optional[str] = None
# HunyuanVideo-1.5 style guidance: the pipeline __call__ takes NO guidance
# kwarg at all; CFG lives on a ``guider`` component (ClassifierFreeGuidance)
# whose ``guidance_scale`` is a plain attribute set per request. When True,
# generate() writes the scale onto ``pipe.guider`` instead of passing
# ``cfg_kwarg`` (which the pipeline would reject as an unexpected argument).
guidance_via_guider: bool = False
# Generation defaults + shape constraints. ``frame_step`` is the temporal
# compression: a valid frame count is k * frame_step + 1 (the +1 is the
# anchor frame), so requests are snapped BEFORE latents are allocated.
@ -206,6 +212,70 @@ _FAMILIES: tuple[VideoFamily, ...] = (
# No gguf_repo: community GGUFs ship the two experts as separate files, and a
# single-file load covers only one (validate_load_request refuses it).
),
# HunyuanVideo-1.5 (diffusers >= 0.39): an 8.3B video DiT with a Qwen2.5-VL
# text encoder plus a ByT5 glyph encoder. Three quirks, all verified against
# the installed pipeline source (pipeline_hunyuan_video1_5.py):
# 1. __call__ takes NO guidance kwarg; CFG lives on the ``guider`` component
# (ClassifierFreeGuidance; the 480p t2v repo ships guidance_scale = 6.0),
# hence guidance_via_guider.
# 2. __call__ has NO callback_on_step_end; generate() falls back to the
# scheduler.step progress wrapper automatically (capability-detected).
# 3. The tencent/HunyuanVideo-1.5 repo is the ORIGINAL layout (config.json,
# no model_index.json); only the hunyuanvideo-community Diffusers repacks
# load through HunyuanVideo15Pipeline, so those are the trusted repos.
# The transformer declares _repeated_blocks and inherits CacheMixin, so the
# regional compile profile and First-Block-Cache both apply.
VideoFamily(
name = "hunyuanvideo-1.5",
pipeline_class = "HunyuanVideo15Pipeline",
transformer_class = "HunyuanVideo15Transformer3DModel",
base_repo = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
# No bare "hunyuanvideo" alias: it would also claim the incompatible 1.0
# repos (HunyuanVideoPipeline), which this family cannot load.
aliases = ("hunyuanvideo-1-5", "hunyuanvideo1.5", "hunyuanvideo1-5", "hv15"),
has_audio = False,
guidance_via_guider = True,
default_steps = 50,
default_guidance = 6.0,
# 121 frames at 24 fps is ~5s, the pipeline's own num_frames default.
default_num_frames = 121,
default_fps = 24,
# The HV15 VAE compresses 16x spatial / 4x temporal (vae config:
# spatial_compression_ratio 16, temporal_compression_ratio 4) with a
# patch-1 transformer, so sizes snap to /16 and frames to 4k+1.
frame_step = 4,
resolution_multiple = 16,
# 480p-class presets (the base repo is the 480p t2v variant): landscape,
# vertical, square.
resolution_presets = ((832, 480), (480, 832), (624, 624)),
# Disk shards are fp32 for the DiT (32.0 GB -> 16.6 bf16-resident) and the
# VAE (4.7 -> 2.4); the Qwen2.5-VL TE is stored bf16 (14.0) plus ByT5 0.8.
bf16_components_gb = (16.6, 14.8, 2.4),
),
# The 720p t2v repack: same architecture, pipeline quirks, guider config
# (guidance 6.0) and shard footprint as the 480p entry above; only the
# trained resolution class differs. Kept as its OWN family so a 720p load
# defaults to 720p-class sizes instead of silently rendering at 832x480.
# The repo-id alias is the full path segment, so it out-lengths (and thus
# outranks) the generic "hunyuanvideo-1.5" token for this repo only.
VideoFamily(
name = "hunyuanvideo-1.5-720p",
pipeline_class = "HunyuanVideo15Pipeline",
transformer_class = "HunyuanVideo15Transformer3DModel",
base_repo = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v",
aliases = ("hunyuanvideo-1.5-diffusers-720p_t2v", "hv15-720p"),
has_audio = False,
guidance_via_guider = True,
default_steps = 50,
default_guidance = 6.0,
default_num_frames = 121,
default_fps = 24,
frame_step = 4,
resolution_multiple = 16,
# 720p-class presets: landscape, vertical, square (all /16).
resolution_presets = ((1280, 720), (720, 1280), (960, 960)),
bf16_components_gb = (16.6, 14.8, 2.4),
),
)
@ -278,6 +348,9 @@ _VIDEO_GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = (
# and the base repo. A future distilled Wan GGUF is caught by the "distilled"
# row above (listed first), exactly as the LTX-2.3 distilled checkpoints are.
("wan", 50, 5.0),
# HunyuanVideo-1.5 runs the pipeline's 50 steps with the guider's shipped
# CFG 6.0 (guider_config.json in the community Diffusers repacks).
("hunyuanvideo", 50, 6.0),
)

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(
@ -2342,6 +2343,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
@ -2507,6 +2518,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"
)

63
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):
@ -3371,6 +3375,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),
@ -3413,12 +3455,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))
@ -3524,9 +3576,10 @@ async def delete_cached_model(
except Exception:
pass
# And refuse if the Video backend has this repo loaded: it shares the On-Device GGUF
# delete UI, so without this guard a loaded video GGUF could be removed from under a live
# pipeline -- the same invariant the three guards above enforce. Repo-level match.
# 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

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

73
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(
@ -863,3 +902,37 @@ def test_delete_cached_refuses_video_loaded_repo(monkeypatch):
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

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

@ -18,7 +18,7 @@ from core.inference.video import (
get_video_backend,
resolve_video_model_kind,
)
from core.inference.video_families import VIDEO_NOT_LOADED_MSG
from core.inference.video_families import VIDEO_CANCELLED_MSG, VIDEO_NOT_LOADED_MSG
class _FakeDtype:
@ -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
@ -278,6 +281,86 @@ class _FakeWanPipelineSingle:
return _FakeWanPipeMoE() if moe else _FakeWanPipeSingle()
# ── HunyuanVideo-1.5 fakes: the __call__ signature has NO guidance kwarg and NO
# callback_on_step_end (matching pipeline_hunyuan_video1_5.py in diffusers 0.39),
# a guider object carries the CFG scale, and the denoise loop drives
# scheduler.step -- so the guider write and the scheduler-wrap progress/cancel
# paths actually exercise.
class _FakeHV15Scheduler:
def __init__(self) -> None:
self.calls = 0
# Test hook fired from the ORIGINAL step (i.e. inside the wrapped call),
# letting a test cancel mid-denoise exactly as a user request would land.
self.on_step = None
def step(self, *args, **kwargs):
self.calls += 1
if self.on_step is not None:
self.on_step(self.calls)
return object()
class _FakeHV15Pipe:
def __init__(self) -> None:
self.vae = _FakeWanVae()
self.transformer = _FakeWanDiT()
self.scheduler = _FakeHV15Scheduler()
self.guider = types.SimpleNamespace(guidance_scale = 6.0)
self.components = {"transformer": self.transformer, "vae": self.vae}
self.moved_to = None
self.last_kwargs = None
self.hooks_freed = 0
def maybe_free_model_hooks(self):
self.hooks_freed += 1
def to(self, device):
self.moved_to = device
return self
def enable_vae_tiling(self) -> None:
self.vae.tiled = True
def __call__(
self,
*,
prompt = None,
negative_prompt = None,
height = None,
width = None,
num_frames = None,
num_inference_steps = None,
generator = None,
**kwargs,
):
self.last_kwargs = {
"prompt": prompt,
"negative_prompt": negative_prompt,
"num_inference_steps": num_inference_steps,
"width": width,
"height": height,
"num_frames": num_frames,
**kwargs,
}
for _ in range(int(num_inference_steps or 1)):
self.scheduler.step()
frames = [[object() for _ in range(int(num_frames or 1))]]
return types.SimpleNamespace(frames = frames, audio = None)
class _FakeHV15Pipeline:
last: dict = {}
instance = None
@classmethod
def from_pretrained(cls, repo, **kwargs):
_FakeHV15Pipeline.last = {"repo": repo, **kwargs}
_FakeHV15Pipeline.instance = _FakeHV15Pipe()
return _FakeHV15Pipeline.instance
@pytest.fixture
def fake_runtime(monkeypatch):
torch = types.ModuleType("torch")
@ -296,6 +379,8 @@ def fake_runtime(monkeypatch):
# Wan2.2: one pipeline class serves both families (it dispatches on the repo id).
diffusers.WanPipeline = _FakeWanPipelineSingle
diffusers.WanTransformer3DModel = _FakeTransformer
diffusers.HunyuanVideo15Pipeline = _FakeHV15Pipeline
diffusers.HunyuanVideo15Transformer3DModel = _FakeTransformer
diffusers.FirstBlockCacheConfig = lambda threshold = None: ("fbcache", threshold)
monkeypatch.setitem(sys.modules, "torch", torch)
@ -727,6 +812,52 @@ def test_generate_progress_and_cancel_idle(fake_runtime):
assert backend.cancel_generate() is False
def test_hv15_guider_and_scheduler_progress(fake_runtime):
# HunyuanVideo-1.5: no guidance kwarg (CFG set on the guider), no step
# callback (progress via the scheduler.step wrapper, restored afterwards).
backend = VideoBackend()
status = backend.load_pipeline(
"hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
model_kind = "pipeline",
)
assert status["family"] == "hunyuanvideo-1.5"
assert status["has_audio"] is False
assert status["defaults"]["frame_step"] == 4
pipe = _FakeHV15Pipeline.instance
result = backend.generate(
prompt = "a fox in the snow", steps = 4, guidance = 3.5, num_frames = 9, fps = 24
)
assert "guidance_scale" not in pipe.last_kwargs
assert "callback_on_step_end" not in pipe.last_kwargs
assert pipe.guider.guidance_scale == 3.5
# One wrapped tick per denoise step, then the original method back in place.
assert pipe.scheduler.calls == 4
assert pipe.scheduler.step.__func__ is _FakeHV15Scheduler.step
assert result["num_frames"] == 9 and result["has_audio"] is False
def test_hv15_cancel_unwinds_scheduler_loop(fake_runtime):
backend = VideoBackend()
backend.load_pipeline(
"hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
model_kind = "pipeline",
)
pipe = _FakeHV15Pipeline.instance
# Cancel lands during the FIRST real step; the next wrapped call must raise out
# of the denoise loop and generate() must surface the cancelled sentinel.
pipe.scheduler.on_step = lambda n: backend.cancel_generate() if n == 1 else None
with pytest.raises(RuntimeError, match = VIDEO_CANCELLED_MSG):
backend.generate(prompt = "a fox", steps = 4)
assert pipe.scheduler.calls == 1
# The wrapper must restore scheduler.step even on the exception path.
assert pipe.scheduler.step.__func__ is _FakeHV15Scheduler.step
# The exception unwound pipe.__call__ before its own end-of-call cleanup, so
# generate() must have freed the offload hooks itself (VRAM would otherwise
# stay onloaded until the next request).
assert pipe.hooks_freed == 1
def test_singleton():
assert get_video_backend() is get_video_backend()
@ -749,6 +880,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()
@ -933,6 +1124,57 @@ def test_wan_a14b_dense_quant_applies_to_both_dits(fake_runtime, monkeypatch):
assert status["transformer_quant"] == "int8"
def test_dense_quant_skipped_under_offload(fake_runtime, monkeypatch):
# Offload hooks move modules with Module.to(), which torchao quantized tensors
# reject (observed as a hard crash on the A14B gate run). When the memory plan
# resolves to any offload policy, quant must be SKIPPED, not attempted: the
# load succeeds dense and the resolved record explains why.
import core.inference.video as video_mod
monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True)
quantised = []
def _fake_quant(
view,
target,
*,
mode,
family,
logger = None,
):
quantised.append(view.transformer)
return "int8"
monkeypatch.setattr(video_mod, "quantize_transformer", _fake_quant)
# The CPU fake target never plans an offload, so force one at the plan seam
# (frozen dataclass -> dataclasses.replace) and stub the apply step, which
# would otherwise call offload hooks the fake pipe does not have.
import dataclasses
real_plan = video_mod.plan_diffusion_memory
monkeypatch.setattr(
video_mod,
"plan_diffusion_memory",
lambda **kwargs: dataclasses.replace(real_plan(**kwargs), offload_policy = "model"),
)
monkeypatch.setattr(
video_mod,
"apply_memory_plan",
lambda pipe, plan, device = None, logger = None: ("model", True),
)
backend = VideoBackend()
status = backend.load_pipeline(
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
model_kind = "pipeline",
transformer_quant = "int8",
)
assert status["offload_policy"] == "model"
assert quantised == []
assert status["transformer_quant"] is None
assert "offload moves the DiT" in status["resolved"]["transformer_quant"]["reason"]
def test_wan_a14b_partial_quant_fails_the_load(fake_runtime, monkeypatch):
# If the first expert quantises but the second does not, the pipe is left at
# mismatched precision with no way back (in-place mutation), so the load must
@ -1133,6 +1375,21 @@ def test_base_download_files_ltx23_keeps_only_shared_components():
)
def test_hv15_720p_repo_gets_720p_family_defaults():
# The 720p repack is trusted, but it must resolve its OWN family entry: the
# generic hunyuanvideo-1.5 entry would default generation to 832x480.
from core.inference.video_families import detect_video_family
fam = detect_video_family("hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v")
assert fam is not None and fam.name == "hunyuanvideo-1.5-720p"
assert fam.resolution_presets[0] == (1280, 720)
assert fam.base_repo.endswith("720p_t2v")
# The 480p repo keeps the original entry.
fam480 = detect_video_family("hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v")
assert fam480 is not None and fam480.name == "hunyuanvideo-1.5"
assert fam480.resolution_presets[0] == (832, 480)
def test_predownload_base_honors_cancel_between_files(monkeypatch):
# A warm-cache sweep returns each file instantly without consulting the event,
# so the loop must check it explicitly or an unload mid-predownload is ignored.

30
studio/backend/tests/test_video_families.py
View file

@ -146,7 +146,13 @@ def test_generation_defaults_distilled_vs_dev():
def test_supported_names():
assert supported_video_family_names() == ("ltx-2", "wan2.2-ti2v-5b", "wan2.2-t2v-a14b")
assert supported_video_family_names() == (
"ltx-2",
"wan2.2-ti2v-5b",
"wan2.2-t2v-a14b",
"hunyuanvideo-1.5",
"hunyuanvideo-1.5-720p",
)
def test_wan_snap_num_frames_4k_plus_1():
@ -259,3 +265,25 @@ def test_family_size_table_present():
# that would let auto planning under-reserve by ~50 GB.
assert text_encoder_gb > transformer_gb > 20.0
assert companions_gb > 0.0
def test_hv15_detection_and_flags():
fam = detect_video_family("hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v")
assert fam is not None and fam.name == "hunyuanvideo-1.5"
# CFG lives on the guider component (no guidance kwarg in __call__), and the
# HV15 VAE compresses 16x spatial / 4x temporal.
assert fam.guidance_via_guider is True
assert fam.frame_step == 4 and fam.resolution_multiple == 16
assert fam.has_audio is False
assert detect_video_family("x/y", override = "hv15") is fam
# The incompatible HunyuanVideo 1.0 repos must NOT be claimed: their
# model_index pins HunyuanVideoPipeline, which this family cannot load.
assert detect_video_family("hunyuanvideo-community/HunyuanVideo") is None
def test_hv15_generation_defaults():
# The community repacks ship a guider with guidance_scale 6.0 and the
# pipeline's own 50-step schedule.
assert default_video_generation_params(
None, "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
) == (50, 6.0)

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(

243
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.
@ -325,23 +250,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();
}
@ -493,6 +465,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();
}
@ -1580,8 +1555,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);
@ -1917,7 +1892,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)}
@ -1934,8 +1909,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)}
@ -1951,14 +1926,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)}
@ -1985,7 +1960,7 @@ 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). Auto enables it for many-step schedules and skips it for few-step distilled models; Off disables it entirely."
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)}
@ -2034,6 +2009,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)}
/>
@ -2042,11 +2018,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>
@ -2634,15 +2621,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.",
},
];
@ -865,7 +865,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.
How fast the model learns over time. Constant is fine for most runs.
</p>
</div>
{lrScheduler !== "constant" &&
@ -885,8 +885,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>
@ -913,15 +912,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>
@ -939,7 +933,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>
)}
@ -959,8 +953,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>
)}
@ -1099,8 +1092,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>
) : (
@ -1128,8 +1121,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();
}

241
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,52 +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" },
};
// 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",
),
];
// 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
@ -120,6 +93,9 @@ const MODEL_DEFAULTS: Array<{ match: string; steps: number; guidance: number }>
// Wan2.2 pipelines default to 50 steps at CFG 5.0 (WanPipeline defaults, verified in
// diffusers 0.39). The backend supplies the fps per family (24 for TI2V-5B, 16 for A14B).
{ match: "wan", steps: 50, guidance: 5 },
// HunyuanVideo-1.5 runs 50 steps; guidance 6 matches the guider the repo ships
// (the backend writes it onto the guider component, there is no pipeline kwarg).
{ match: "hunyuanvideo", steps: 50, guidance: 6 },
];
function defaultsFor(repoId: string): { steps: number; guidance: number } {
@ -164,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();
@ -519,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>(
@ -542,6 +543,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),
);
@ -713,6 +720,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);
@ -908,7 +938,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();
@ -928,6 +959,7 @@ export function VideoPage({ active = true }: { active?: boolean }) {
speedMode,
attentionBackend,
transformerCache,
transformerQuant,
],
);
@ -945,10 +977,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 });
@ -1141,7 +1178,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)}
@ -1153,9 +1190,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)}
@ -1169,11 +1233,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"],
]}
@ -1207,6 +1272,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)}
/>
@ -1366,16 +1432,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 && (
@ -1387,15 +1464,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"
@ -1488,8 +1581,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;