diff --git a/plans/controlnet-workflow.md b/plans/controlnet-workflow.md deleted file mode 100644 index 9d02a1a0ab..0000000000 --- a/plans/controlnet-workflow.md +++ /dev/null @@ -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 `, - `--control-image `, `--control-strength `, `--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. diff --git a/plans/diffusion-popularity-findings.md b/plans/diffusion-popularity-findings.md deleted file mode 100644 index 630ef93291..0000000000 --- a/plans/diffusion-popularity-findings.md +++ /dev/null @@ -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). diff --git a/plans/diffusion-workflows-pr-plan.md b/plans/diffusion-workflows-pr-plan.md deleted file mode 100644 index df07714cf4..0000000000 --- a/plans/diffusion-workflows-pr-plan.md +++ /dev/null @@ -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. diff --git a/plans/diffusion-workflows-studio.md b/plans/diffusion-workflows-studio.md deleted file mode 100644 index 126bc79129..0000000000 --- a/plans/diffusion-workflows-studio.md +++ /dev/null @@ -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. diff --git a/plans/wobbly-jumping-narwhal.md b/plans/wobbly-jumping-narwhal.md deleted file mode 100644 index 2261a2fe87..0000000000 --- a/plans/wobbly-jumping-narwhal.md +++ /dev/null @@ -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--bin-.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--` 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--bin-Darwin-macOS-arm64.zip`, `sd--bin-Darwin-macOS-x86_64.zip`, - `sd--bin-Linux-Ubuntu-24.04-x86_64.zip`, `sd--bin-Linux-Ubuntu-24.04-aarch64.zip`, - `sd--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). diff --git a/scripts/build_prequant_checkpoint.py b/scripts/build_prequant_checkpoint.py index c15f10a171..9766467351 100644 --- a/scripts/build_prequant_checkpoint.py +++ b/scripts/build_prequant_checkpoint.py @@ -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", diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 91a2acb699..8797f9b3d7 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -218,12 +218,18 @@ _TRUSTED_NON_GGUF_REPOS = frozenset( "stabilityai/stable-diffusion-xl-base-1.0", "stabilityai/sdxl-turbo", # Official vendor, safetensors-only base repos allowlisted as LoRA TRAINING bases - # (diffusion training loads the full pipeline from these). Same rule as above: no - # pickled weights, no remote code, exact-match lowercased. FLUX.1-dev is gated on - # the Hub (needs the user's token); the other two are open. + # (diffusion training loads the full pipeline from these) and as the official + # BF16 artifact behind each catalog group (model-catalog.ts). Same rule as above: + # no pickled weights, no remote code, exact-match lowercased. FLUX.1-dev/schnell/ + # Kontext are gated on the Hub (need the user's token); the Qwen and Z-Image repos + # are open. All verified as diffusers model_index pipelines. "black-forest-labs/flux.1-dev", + "black-forest-labs/flux.1-schnell", + "black-forest-labs/flux.1-kontext-dev", "tongyi-mai/z-image-turbo", "qwen/qwen-image", + "qwen/qwen-image-2512", + "qwen/qwen-image-edit-2511", # Krea 2: official vendor repos, safetensors-only, no remote code. Loaded # per-component via core/inference/diffusion_krea2.py (no GGUF variant yet). # Turbo is the inference model; Raw is the undistilled base Krea recommends @@ -300,6 +306,12 @@ class _LoadState: # Attention backend engaged via the diffusers dispatcher (e.g. "_native_cudnn"), or # None for the default SDPA. Set before compile; orthogonal to the weight quant. attention_backend: Optional[str] = None + # The caller's ORIGINAL attention request (None / "auto" left it to the backend, else + # an explicit alias like "native" / "sage" / "flash"). Carried so the deferred-speed + # engagement re-runs the SAME selection the load-time path did, instead of forcing the + # auto cuDNN upgrade -- otherwise an explicitly pinned backend (e.g. "native" to avoid + # cuDNN) is silently discarded when the 3rd generation engages the deferred profile. + attention_request: Optional[str] = None # Step cache engaged ("fbcache") or None. Opt-in, for many-step models. transformer_cache: Optional[str] = None # True when the cache decision was AUTO on a cache-capable transformer: generate() @@ -312,6 +324,13 @@ class _LoadState: # Shared eager monkey-patches (diffusion_eager_patches) installed for this load (any # non-off speed tier). Uninstalled on unload so a later `off` load is bit-identical. eager_patched: bool = False + # Deferred speed auto (dense models with speed_mode unset): the load stays fully + # eager/bit-identical, and generate() engages the `default` compile profile at the + # start of the 3rd generation this session -- repeated use is established by then, + # so the one-time compile warmup amortises. Cleared once engaged (or failed). + speed_deferred: bool = False + # Successful generations on this load; drives the deferred engagement above. + generation_count: int = 0 # Pre-warmed torch.compile cache context (diffusion_compile_cache.CacheContext) when a # compiled tier ran, else None. Carries the per-key inductor dir + bundle for save/restore. compile_cache_ctx: Any = None @@ -1258,6 +1277,19 @@ class DiffusionBackend: "(quantized transformer must be compiled; eager is ~30x slower)" ) effective_speed = SPEED_DEFAULT + # Deferred speed auto for dense models: the load stays eager (a one-off + # image should not pay the 25-60s compile warmup, and eager is the + # bit-identical reference), but a user starting their 3rd image in one + # session has revealed repeated use -- generate() then engages the + # `default` profile, where the warmup starts paying back. Only when + # the request left speed unset, nothing forced a compiled tier, and + # this device/family could actually compile. + speed_deferred = ( + speed_mode is None + and effective_speed == SPEED_OFF + and transformer_quant_engaged is None + and compile_eligible(target, is_gguf = False, family = fam) + ) # Opt-in speed optims run BEFORE placement (channels_last / compile # must precede CPU offload). Snapshot the process-wide backend flags # first so unload can restore them: TF32 / cudnn.benchmark are global, @@ -1425,12 +1457,15 @@ class DiffusionBackend: "compiled; eager torchao quant is ~30x slower than GGUF here", transformer_quant_engaged, ) - # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), - # also before placement so the offload hooks move the smaller weights. + # Quantise the dense companion text encoder(s) (opt-in fp8 / fp8_dynamic / + # int8 / nvfp4), also before placement so the offload hooks move the smaller + # weights. int8 needs a per-family keep-bf16 schedule, so pass the family. te_quant = quantize_text_encoders( pipe, target, mode = text_encoder_quant, + family = fam.name, + offload_active = plan.offload_policy != OFFLOAD_NONE, logger = logger, ) @@ -1451,10 +1486,13 @@ class DiffusionBackend: { "speed_mode": ( speed_mode, - effective_speed, + "deferred" if speed_deferred else effective_speed, "quantized transformer requires compile" if transformer_quant_engaged is not None and normalize_speed_mode(speed_mode) in (None, SPEED_OFF) + else "auto: exact eager for the first two images; " + "the compile profile engages on the 3rd" + if speed_deferred else "per-kind default" if speed_mode is None else "requested", @@ -1480,7 +1518,9 @@ class DiffusionBackend: "memory_mode": ( memory_mode, effective_policy, - "planned from measured free VRAM vs estimated footprint", + "everything fits on the GPU, no offload needed" + if effective_policy == OFFLOAD_NONE + else "planned from measured free VRAM vs estimated footprint", ), "transformer_cache": ( None if cache_auto else transformer_cache, @@ -1513,11 +1553,13 @@ class DiffusionBackend: text_encoder_quant = te_quant, transformer_quant = transformer_quant_engaged, attention_backend = attention_engaged, + attention_request = attention_backend, transformer_cache = cache_engaged, cache_auto = cache_may_toggle, cache_quant_active = cache_quant_active, cache_threshold = transformer_cache_threshold, eager_patched = eager_patched, + speed_deferred = speed_deferred, compile_cache_ctx = compile_ctx, hf_token = hf_token, resolved = resolved, @@ -2011,6 +2053,88 @@ class DiffusionBackend: except Exception: # noqa: BLE001 — reset is best-effort, never fail a generation pass + def _engage_deferred_speed(self, state: _LoadState) -> None: + """Engage the deferred `default` speed profile at the start of the 3rd + generation this session. + + The load left the pipe fully eager (bit-identical reference); by the 3rd + image repeated use is established, so pay the one-time compile now: eager + patches + attention auto upgrade + regional compile -- exactly what an + unset-speed GGUF load gets at load time. Runs under _generate_lock (the + caller), so no denoise can race the mutation. The flag is cleared FIRST so + a failure never retries per generation; unload cleans everything up via the + same state fields the load-time path uses (backend flags were snapshotted + at load, before any speed layer could mutate them).""" + object.__setattr__(state, "speed_deferred", False) + from .diffusion_eager_patches import install_compile_safe_patches + from .diffusion_arch_patches import install_arch_patches + + target = self._resolve_device_target(state.family) + install_compile_safe_patches() + install_arch_patches() + object.__setattr__(state, "eager_patched", True) + # Re-run the load-time selection with the caller's ORIGINAL request (not a bare + # None): an explicit backend must survive the deferred upgrade. auto still upgrades + # to cuDNN here (speed_active=True), but an explicit "native"/"sage"/"flash" is + # honored verbatim rather than silently replaced by the auto cuDNN choice. + attention_engaged = apply_attention_backend( + state.pipe, + select_attention_backend(target, state.attention_request, speed_active = True), + logger = logger, + ) + object.__setattr__(state, "attention_backend", attention_engaged) + gguf_transformer = state.kind == "gguf" and state.transformer_quant is None + if compile_eligible(target, is_gguf = gguf_transformer, family = state.family): + compile_ctx = compile_cache.begin( + family = state.family.name, + transformer = getattr(state.pipe, "transformer", None), + dtype = getattr(target, "dtype", None), + quant = state.transformer_quant, + attention_backend = attention_engaged, + compile_kwargs = { + # Mirrors the load-time fullgraph decision: an engaged or + # still-toggleable step cache OR an offload graph-breaks. + "fullgraph": state.transformer_cache is None + and not state.cache_auto + and state.offload_policy == OFFLOAD_NONE, + "dynamic": True, + "mode": "default", + }, + logger = logger, + ) + object.__setattr__(state, "compile_cache_ctx", compile_ctx) + speed_applied = apply_speed_optims( + state.pipe, + target, + is_gguf = gguf_transformer, + family = state.family, + speed_mode = SPEED_DEFAULT, + cache_active = state.transformer_cache is not None or state.cache_auto, + offload_active = state.offload_policy != OFFLOAD_NONE, + logger = logger, + ) + object.__setattr__(state, "speed_mode", SPEED_DEFAULT) + object.__setattr__(state, "speed_optims", tuple(k for k, v in speed_applied.items() if v)) + entry = (state.resolved or {}).get("speed_mode") + if isinstance(entry, dict): + entry["value"] = SPEED_DEFAULT + entry["reason"] = ( + "auto: compiled on the 3rd image this session " + "(repeated use amortises the one-time compile)" + ) + att = (state.resolved or {}).get("attention_backend") + if isinstance(att, dict) and att.get("source") == "auto": + att["value"] = attention_engaged or "native" + att["reason"] = ( + "cuDNN fused attention upgrade" if attention_engaged else "diffusers default" + ) + logger.info( + "diffusion.speed: deferred profile engaged on generation 3 " + "(optims=%s, attention=%s)", + ",".join(state.speed_optims) or "none", + attention_engaged or "native", + ) + def generate( self, *, @@ -2076,6 +2200,39 @@ class DiffusionBackend: seed = int(seed) generator.manual_seed(seed) + # Deferred speed auto: by the 3rd image in one session repeated use is + # established, so engage the compile profile now -- before the LoRA / + # workflow wiring, matching the load-time ordering. Best-effort: a + # failure logs, leaves the eager pipe running, and never retries + # (the helper clears the flag first). + # + # But NOT when this generation requests a LoRA: a compiled transformer rejects + # LoRA (supports_lora is False once compiled), and _apply_loras raises before its + # unchanged-selection no-op, so engaging compile here would permanently break every + # LoRA generation on this load. Compile and LoRA are mutually exclusive; keep the + # pipe eager and let compile defer to a later LoRA-free generation. + lora_requested = any(w != 0 for (_id, w) in (loras or [])) + # Also stay eager while adapters from a PRIOR generation are still attached: this + # request may clear them (lora_requested False), but _apply_loras runs AFTER the + # engage below, so compiling here would bake the resident adapter into the graph and + # the subsequent unload_lora_weights() (swallowed on a compiled pipe) would leave it + # active forever -- silent wrong output on every later LoRA-free generation. Deferring + # lets _apply_loras clear it on the still-eager pipe; compile engages a gen later. + loras_attached = bool(getattr(state.pipe, "_unsloth_loras", ())) + if ( + state.speed_deferred + and state.generation_count >= 2 + and not lora_requested + and not loras_attached + ): + try: + self._engage_deferred_speed(state) + except Exception as exc: # noqa: BLE001 — speed is best-effort + logger.warning( + "diffusion.speed: deferred engagement failed, staying eager: %s", + exc, + ) + # Apply/adjust LoRA adapters on the resident pipe (non-fused) before picking # the workflow pipe; from_pipe pipes share the transformer, so it propagates. self._apply_loras(state, loras, cancel) @@ -2393,6 +2550,9 @@ class DiffusionBackend: compile_cache.save(state.compile_cache_ctx, logger = logger) except Exception: # noqa: BLE001 — cache persistence is best-effort pass + # Count the finished generation (drives the deferred speed + # engagement above); a batch of N images is one generation. + object.__setattr__(state, "generation_count", state.generation_count + 1) # Return the PIL images (not yet encoded): the route embeds each # image's recipe and persists it via the gallery. return {"images": list(images), "seed": int(seed), "repo_id": state.repo_id} diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index 030b6da7f2..68f6307526 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/studio/backend/core/inference/diffusion_precision.py @@ -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) diff --git a/studio/backend/core/inference/diffusion_prequant.py b/studio/backend/core/inference/diffusion_prequant.py index 485c4a1847..903e323d0f 100644 --- a/studio/backend/core/inference/diffusion_prequant.py +++ b/studio/backend/core/inference/diffusion_prequant.py @@ -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") diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 8c2524dfda..00bc565e8c 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -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) diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index b990efc563..a4d7561b7f 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -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). diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 5d8e49a23f..fb182ea814 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -43,7 +43,14 @@ from typing import Any, Optional from loggers import get_logger from .diffusion_attention import apply_attention_backend, select_attention_backend -from .diffusion_cache import apply_step_cache, normalize_transformer_cache +from .diffusion_cache import ( + FBCACHE_MIN_STEPS, + TC_AUTO, + TC_FBCACHE, + apply_step_cache, + maybe_toggle_step_cache, + normalize_transformer_cache, +) from .diffusion_device import resolve_diffusion_device_target from .diffusion_memory import ( apply_memory_plan, @@ -65,11 +72,13 @@ from .diffusion_speed import ( ) from .diffusion_auto_policy import _QUANT_STEADY_FACTOR, build_resolved_record from .diffusion_transformer_quant import ( + TQ_AUTO, dense_transformer_supported, normalize_transformer_quant, quantize_transformer, select_transformer_quant_scheme, ) +from .diffusion_precision import normalize_te_quant, quantize_text_encoders from .video_families import ( VIDEO_CANCELLED_MSG, VIDEO_NOT_LOADED_MSG, @@ -213,11 +222,22 @@ class _VideoLoadState: backend_flags: Optional[dict] = None attention_backend: Optional[str] = None transformer_cache: Optional[str] = None + # True when the cache decision was AUTO on a cache-capable DiT: generate() then + # re-checks the actual step count and toggles FBCache across FBCACHE_MIN_STEPS. + # An explicit request (off / fbcache) is never toggled. + cache_auto: bool = False + # Inputs the generation-time toggle re-applies (quantised threshold + override). + cache_quant_active: bool = False + cache_threshold: Optional[float] = None # Dense transformer quant actually engaged ("int8" | "fp8" | "nvfp4" | "mxfp8") or # None. Mirrors the image backend's _LoadState.transformer_quant: on a pipeline-kind # load the dense DiT(s) can be torchao-quantised in place onto the low-precision # tensor cores; None means they run at their loaded (bf16) precision. transformer_quant: Optional[str] = None + # Text-encoder quant actually engaged ("fp8" | "fp8_dynamic" | "int8" | "nvfp4") or None. + # The companion text encoder (UMT5 / Gemma3 / Qwen2.5-VL) loads dense bf16 and is often the + # largest resident component; this shrinks it in place, mirroring the image backend. + text_encoder_quant: Optional[str] = None resolved: Optional[dict] = None @@ -322,6 +342,7 @@ class VideoBackend: family_override: Optional[str] = None, model_kind: Optional[str] = None, transformer_quant: Optional[str] = None, + text_encoder_quant: Optional[str] = None, ) -> VideoFamily: """Cheap, network-free validation shared by the route and the load path.""" kind = resolve_video_model_kind(gguf_filename, model_kind) @@ -381,6 +402,9 @@ class VideoBackend: # only on pipeline-kind loads (the dense DiT from the base repo); an ignored value # on a gguf/single_file load is left to the loader, matching the image backend. normalize_transformer_quant(transformer_quant) + # Reject a malformed text_encoder_quant the same way (applies to any load kind: the dense + # text encoder is resident for pipeline / gguf / single_file alike). + normalize_te_quant(text_encoder_quant) _ensure_mp4_encoder_available() return fam @@ -400,6 +424,7 @@ class VideoBackend: transformer_cache: Optional[str] = None, transformer_cache_threshold: Optional[float] = None, transformer_quant: Optional[str] = None, + text_encoder_quant: Optional[str] = None, model_kind: Optional[str] = None, ) -> dict[str, Any]: """Validate, then run the (slow) load on a daemon thread. Returns at once.""" @@ -411,6 +436,7 @@ class VideoBackend: family_override = family_override, model_kind = model_kind, transformer_quant = transformer_quant, + text_encoder_quant = text_encoder_quant, ) with self._lock: if self._loading is not None and self._loading.error is None: @@ -434,6 +460,7 @@ class VideoBackend: transformer_cache = transformer_cache, transformer_cache_threshold = transformer_cache_threshold, transformer_quant = transformer_quant, + text_encoder_quant = text_encoder_quant, model_kind = model_kind, _load_token = token, ), @@ -717,6 +744,19 @@ class VideoBackend: expected_bytes = int(expected) if expected else None, ) + def loading_repo_ids(self) -> tuple[str, ...]: + """Repo ids an in-flight background load is downloading (empty when idle). + + The delete-cached guard needs this: during a load ``status()["loaded"]`` is + still False, but deleting the target repo (or its companion base) would yank + blobs and snapshot files from under the download/assembly. Mirrors the image + backend's guard (DiffusionBackend.loading_repo_ids).""" + with self._lock: + loading = self._loading + if loading is None or loading.error is not None: + return () + return tuple(r for r in (loading.repo_id, loading.base_repo) if r) + # ── the load itself ────────────────────────────────────────────────────── def load_pipeline( @@ -733,6 +773,7 @@ class VideoBackend: transformer_cache: Optional[str] = None, transformer_cache_threshold: Optional[float] = None, transformer_quant: Optional[str] = None, + text_encoder_quant: Optional[str] = None, model_kind: Optional[str] = None, _load_token: Optional[int] = None, _base_local_dir: Optional[str] = None, @@ -747,6 +788,7 @@ class VideoBackend: family_override = family_override, model_kind = model_kind, transformer_quant = transformer_quant, + text_encoder_quant = text_encoder_quant, ) kind = resolve_video_model_kind(gguf_filename, model_kind) base = repo_id if kind == "pipeline" else resolve_video_base_repo(fam, base_repo) @@ -788,6 +830,20 @@ class VideoBackend: # stay quantised on disk and in memory, so only dense estimates scale. dtype_scale = 2.0 if device != "cpu" and dtype is torch.float32 else 1.0 + # Precision tri-state, mirroring the image backend: an UNSET request (or + # "auto") hands the decision to the hardware ladder -- on a dense-capable + # GPU the quantised DiT (int8 minimum, fp8 on data-center silicon) is + # faster at the same resident-or-better footprint. An explicit + # "none"/"off" pins dense bf16 and an explicit scheme pins that scheme. + # Only the pipeline kind can engage it (gguf/single_file checkpoints + # already carry their own precision), and the offload guard below still + # skips it when the plan moves the DiT. + if transformer_quant is None or str(transformer_quant).strip().lower() in ( + "", + "auto", + ): + transformer_quant = TQ_AUTO + # ── memory plan: family-table resident estimate + frames-aware headroom. device_memory = snapshot_device_memory(target) components = fam.bf16_components_gb @@ -993,12 +1049,34 @@ class VideoBackend: if quant_replanned and transformer_quant_engaged is None: plan = bf16_plan + # ── dense text-encoder quant (opt-in): the DiT arrives quantised in a GGUF, but the + # companion encoder (Gemma3 / UMT5 / Qwen2.5-VL) loads dense bf16 from the base repo and + # is often the largest resident component. Quantise it in place, mirroring the image + # backend (diffusion.py): applied for every kind (the encoder is dense regardless of how + # the DiT was sourced) and before placement so the offload hooks move the smaller weights. + # Best-effort: quantize_text_encoders leaves any encoder it can't cast dense. int8 needs a + # per-family keep-bf16 schedule, so the family name is passed. + text_encoder_quant_engaged = quantize_text_encoders( + pipe, + target, + mode = text_encoder_quant, + family = fam.name, + offload_active = plan.offload_policy != "none", + logger = logger, + ) + # ── optimisation layers, in the image backend's order: step cache FIRST # (compile keys its fullgraph decision off an active cache: FBCache hooks # graph-break, so compiling fullgraph before installing the cache crashes # the first cached generation), then attention, the speed profile, and # placement/offload last. - effective_speed = resolve_speed_mode(speed_mode, is_gguf = kind == "gguf") + # A clip denoise runs minutes, so even a dense (non-GGUF) load amortises the + # one-time regional compile within a single generation: unset resolves to the + # near-lossless `default` profile for every kind. Explicit values (incl. + # "off") are honored verbatim, and `max` is never an auto choice. + effective_speed = resolve_speed_mode( + speed_mode, is_gguf = kind == "gguf", dense_default = SPEED_DEFAULT + ) # A torchao-quantised DiT must be compiled (eager dynamic quant is ~30x slower and # would lose to the bf16 it replaced), so force at least the regional-compile # profile when quant engaged and the effective speed was off, matching diffusion.py. @@ -1014,18 +1092,51 @@ class VideoBackend: # (_run_load's error handler calls _rollback_precommit_globals with this # token). Registered BEFORE the first mutating call. self._precommit_globals = (_load_token, backend_flags) - # Run the step cache per expert so both denoisers cache; the engaged mode is - # identical across experts. + # Step cache tri-state, mirroring the image backend: unset / "auto" lets the + # step-count policy decide (engage when this model's DEFAULT schedule reaches + # FBCACHE_MIN_STEPS, re-checked against the actual step count per generation); + # explicit "off" / "fbcache" are pinned and never toggled. Run it per expert + # so both denoisers cache; the engaged mode is identical across experts. + cache_request = normalize_transformer_cache(transformer_cache) + cache_auto = transformer_cache is None or cache_request == TC_AUTO + # GGUF checkpoints and torchao-quantised DiTs both need the higher quantised + # threshold for the cache to still trigger over the quant noise. + cache_quant_active = kind == "gguf" or transformer_quant_engaged is not None + default_cache_steps: Optional[int] = None + if cache_auto: + default_cache_steps, _ = default_video_generation_params(gguf_filename, repo_id, base) + cache_request = TC_FBCACHE if default_cache_steps >= FBCACHE_MIN_STEPS else None cache_engaged = None for view in views: engaged = apply_step_cache( view, - mode = normalize_transformer_cache(transformer_cache), + mode = cache_request, threshold = transformer_cache_threshold, + quant_active = cache_quant_active, logger = logger, ) if view is pipe: cache_engaged = engaged + # The auto decision can flip at generation time, but only on a DiT that + # supports caching at all (a non-CacheMixin transformer can never engage). + cache_may_toggle = cache_auto and callable( + getattr(getattr(pipe, "transformer", None), "enable_cache", None) + ) + if cache_auto: + if cache_engaged: + cache_reason = ( + f"auto: {default_cache_steps}-step default schedule reaches " + f"{FBCACHE_MIN_STEPS}; re-checked per generation" + ) + elif cache_request is not None: + cache_reason = "auto: model does not support step caching" + else: + cache_reason = ( + f"auto: {default_cache_steps}-step default schedule is below " + f"{FBCACHE_MIN_STEPS}; re-checked per generation" + ) + else: + cache_reason = "requested" attention_engaged = None speed_optims: tuple = () for view in views: @@ -1048,7 +1159,10 @@ class VideoBackend: is_gguf = gguf_transformer, family = fam, speed_mode = effective_speed, - cache_active = cache_engaged is not None, + # An auto cache that could still engage mid-session also drops + # fullgraph: enabling FBCache under a fullgraph-compiled DiT would + # crash the first cached generation. + cache_active = cache_engaged is not None or cache_may_toggle, offload_active = plan.offload_policy != "none", ) if view is pipe: @@ -1089,7 +1203,9 @@ class VideoBackend: effective_speed, "quantized transformer requires compile" if transformer_quant_engaged is not None - else "GGUF video loads default to the near-lossless compile profile", + else "clip denoises amortise the one-time compile within a single run" + if speed_mode is None + else "requested", ), "attention_backend": ( attention_backend, @@ -1097,9 +1213,9 @@ class VideoBackend: "cuDNN fused attention on NVIDIA when a speed profile is active", ), "transformer_cache": ( - transformer_cache, + None if cache_auto else transformer_cache, cache_engaged or "off", - "step cache engages on many-step schedules only", + cache_reason, ), "transformer_quant": ( transformer_quant, @@ -1113,6 +1229,13 @@ class VideoBackend: else "not engaged (dense bf16 DiT loaded)" ), ), + "text_encoder_quant": ( + text_encoder_quant, + text_encoder_quant_engaged or "off", + "dense text encoder quantised in place" + if text_encoder_quant_engaged is not None + else "not engaged (dense bf16 text encoder loaded)", + ), } ) @@ -1141,7 +1264,11 @@ class VideoBackend: backend_flags = backend_flags, attention_backend = attention_engaged, transformer_cache = cache_engaged, + cache_auto = cache_may_toggle, + cache_quant_active = cache_quant_active, + cache_threshold = transformer_cache_threshold, transformer_quant = transformer_quant_engaged, + text_encoder_quant = text_encoder_quant_engaged, resolved = resolved, ) # Ownership of the globals transferred to _state / _teardown_state. @@ -1316,6 +1443,34 @@ class VideoBackend: # cancel and restore it afterwards. progress_ctx = _scheduler_step_progress(pipe, _on_scheduler_step) + # An AUTO cache decision is re-checked against the ACTUAL step count, + # mirroring the image backend: a many-step request gains FBCache even + # when the load's default schedule kept it off, and a few-step request + # drops it. Explicit choices never toggle. Runs per view so a dual-DiT + # MoE toggles both experts. + if state.cache_auto: + toggled = state.transformer_cache + for view in _views_for(pipe, fam): + toggled = maybe_toggle_step_cache( + view, + steps = steps, + quant_active = state.cache_quant_active, + threshold = state.cache_threshold, + logger = logger, + ) + if toggled != state.transformer_cache: + # _VideoLoadState is frozen (loads swap it as one unit); this + # tracks the pipe-level toggle that already happened so + # status() reports the true cache state. + object.__setattr__(state, "transformer_cache", toggled) + entry = (state.resolved or {}).get("transformer_cache") + if isinstance(entry, dict): + entry["value"] = toggled or "off" + entry["reason"] = ( + f"auto: {steps}-step generation " + + ("reaches" if toggled else "is below") + + f" {FBCACHE_MIN_STEPS}" + ) if state.transformer_cache: self._reset_step_cache(pipe) try: @@ -1464,6 +1619,7 @@ class VideoBackend: "attention_backend": None, "transformer_cache": None, "transformer_quant": None, + "text_encoder_quant": None, "has_audio": False, "defaults": None, "resolved": None, @@ -1488,6 +1644,7 @@ class VideoBackend: "attention_backend": state.attention_backend, "transformer_cache": state.transformer_cache, "transformer_quant": state.transformer_quant, + "text_encoder_quant": state.text_encoder_quant, "has_audio": fam.has_audio, "defaults": { "steps": default_steps, diff --git a/studio/backend/core/inference/video_gallery.py b/studio/backend/core/inference/video_gallery.py index 06753c5384..6a207f9780 100644 --- a/studio/backend/core/inference/video_gallery.py +++ b/studio/backend/core/inference/video_gallery.py @@ -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" diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 0d82a34f1e..34f323eaa8 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1743,12 +1743,13 @@ class DiffusionLoadRequest(BaseModel): "default (also regional torch.compile where eligible), " "max (also TF32 + fused QKV).", ) - text_encoder_quant: Optional[Literal["fp8", "nvfp4"]] = Field( + text_encoder_quant: Optional[Literal["fp8", "fp8_dynamic", "int8", "nvfp4"]] = Field( None, - description = "Quantise the companion text encoder(s): fp8 (~2x smaller, " - "CUDA cc>=8.9) or nvfp4 (~4x smaller, Blackwell sm_100+). A " - "memory-vs-quality tradeoff (shifts fine detail), not free; " - "pairs well with balanced mode.", + description = "Quantise the companion text encoder(s): fp8 (layerwise cast, ~2x smaller, " + "CUDA cc>=8.9), fp8_dynamic (torchao compute fp8 on the tensor cores, ~2x + faster, " + "cc>=8.9), int8 (torchao compute int8 with per-family keep-bf16 layers; falls back to " + "fp8 where no schedule exists; cc>=8.0), or nvfp4 (~4x smaller, Blackwell sm_100+). A " + "memory-vs-quality tradeoff (shifts fine detail), not free; pairs well with balanced mode.", ) transformer_quant: Optional[Literal["auto", "none", "off", "int8", "fp8", "nvfp4", "mxfp8"]] = ( Field( @@ -2341,6 +2342,16 @@ class VideoLoadRequest(BaseModel): "backend's transformer_quant field.", ) ) + text_encoder_quant: Optional[Literal["fp8", "fp8_dynamic", "int8", "nvfp4"]] = Field( + None, + description = "Quantise the dense companion text encoder (Gemma3 / UMT5 / Qwen2.5-VL), " + "which loads bf16 from the base repo regardless of how the DiT was sourced and is often " + "the largest resident component. fp8 = diffusers layerwise casting (memory only, cc >= " + "8.9); fp8_dynamic = torchao per-row fp8 COMPUTE on the tensor cores (cc >= 8.9); int8 = " + "torchao int8 COMPUTE with per-family keep-bf16 selection (cc >= 8.0; falls back to fp8 " + "for a family without a measured schedule); nvfp4 = torchao 4-bit weight-only (Blackwell " + "sm_100+). null keeps the encoder dense. Mirrors the image backend's field.", + ) @field_validator("attention_backend", mode = "before") @classmethod @@ -2506,6 +2517,12 @@ class VideoStatusResponse(BaseModel): "mxfp8 | null (null = the DiT(s) run at their loaded bf16 precision). For a dual-expert " "MoE family both experts share the reported scheme.", ) + text_encoder_quant: Optional[str] = Field( + None, + description = "Text-encoder quant engaged: fp8 | fp8_dynamic | int8 | nvfp4 | null " + "(null = the dense bf16 encoder is loaded). An int8 request without a per-family " + "keep-bf16 schedule is reported as the fp8 it fell back to.", + ) has_audio: bool = Field( False, description = "Whether the loaded family produces a synchronized audio track" ) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index f910c27a46..0feedd2504 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -32,6 +32,10 @@ class CachedModelRepo(BaseModel): # drop the value the handler sets, letting image-only repos pass the chat picker's # task gate. task: Optional[str] = None + # True when the snapshot is incomplete (a cancelled/partial download left only some + # weights). The picker must not treat a partial base repo as a usable download, or an + # On Device click routes to a fresh multi-GB re-download instead of the complete GGUF. + partial: Optional[bool] = None class CachedModelsResponse(BaseModel): @@ -3363,6 +3367,44 @@ def _repo_is_diffusers(repo_info) -> bool: return False +def _cached_repo_partial(repo_id: str, repo_cache_dir: Optional[Path] = None) -> bool: + """Whether the cached model snapshot is incomplete (cancelled/partial download). + Reuses the hub inventory scan's snapshot-partial detector (cancel marker, legacy + .incomplete blob, manifest walk -- cheapest first). ``repo_cache_dir`` scopes all three + signals to the specific snapshot being listed: without it the scan spans every HF cache + root, so a stale .incomplete copy in one root would flag a complete copy in another as + partial and hide it from the picker (the sibling inventory paths all scope the same way). + Best-effort: a detection error reports not-partial so a scan glitch never hides a + genuinely usable repo.""" + try: + from hub.utils.inventory_scan import is_snapshot_partial + return bool(is_snapshot_partial("model", repo_id, repo_cache_dir)) + except Exception: # noqa: BLE001 -- never fail the listing over a partial probe + return False + + +def _cached_repo_task(repo_info) -> Optional[str]: + """Pipeline task for a cached non-GGUF repo: 'text-to-video' for repos the + video backend can load as full pipelines (its trust list / family detector), + else 'text-to-image' for diffusers image repos, else None (chat). Without the + video tag, cached Lightricks / Wan / Hunyuan pipelines never surfaced in the + Video picker's On Device list -- everything diffusers was blanket-tagged + text-to-image.""" + repo_id = getattr(repo_info, "repo_id", "") or "" + try: + from core.inference.video import _is_trusted_video_repo + from core.inference.video_families import detect_video_family + + # Both gates: a detected video family (so unsloth image repos don't + # match) AND the load path's own trust rule (so an untrusted video repo + # isn't advertised as loadable). + if detect_video_family(repo_id) is not None and _is_trusted_video_repo(repo_id): + return _VIDEO_GEN_TASK + except Exception: + pass + return "text-to-image" if _repo_is_diffusers(repo_info) else None + + @router.get("/cached-models", response_model = CachedModelsResponse) async def list_cached_models( current_subject: str = Depends(get_current_subject), @@ -3405,12 +3447,22 @@ async def list_cached_models( ) key = repo_id.lower() existing = seen_lower.get(key) - if existing is None or total_size > existing["size_bytes"]: + is_partial = _cached_repo_partial(repo_id, Path(repo_info.repo_path)) + # Prefer the most COMPLETE snapshot, then the largest. The picker drops partial + # rows, so a partial copy in one cache root must not shadow a smaller COMPLETE + # copy in another (that would make a usable model vanish from On Device). + # Completeness wins outright; size only breaks ties among equal completeness. + if existing is None or (not is_partial, total_size) > ( + not bool(existing.get("partial")), + existing["size_bytes"], + ): row = { "repo_id": repo_id, "size_bytes": total_size, - "task": "text-to-image" if _repo_is_diffusers(repo_info) else None, + "task": _cached_repo_task(repo_info), } + if is_partial: + row["partial"] = True # Keep the newest timestamp across duplicate caches; # attach only when known so absent rows sort as oldest. lm = max(last_modified, (existing or {}).get("last_modified", 0.0)) @@ -3516,6 +3568,34 @@ async def delete_cached_model( except Exception: pass + # And refuse if the Video backend has this repo loaded or is downloading it: cached non-GGUF + # video repos now surface in the Video On-Device picker with the normal delete action, but the + # guards above only cover chat + the Images engine, so without this a loaded/loading Wan / LTX / + # Hunyuan pipeline could have its HF snapshot removed from under it. Mirror the Images guard. + try: + from core.inference.video import get_video_backend + + video_backend = get_video_backend() + video_status = video_backend.status() + if video_status.get("loaded") and video_status.get("repo_id"): + loaded_id = str(video_status["repo_id"]).lower() + if loaded_id == repo_id.lower() or loaded_id.startswith(repo_id.lower()): + raise HTTPException( + status_code = 400, + detail = "Unload the model before deleting", + ) + for lid in getattr(video_backend, "loading_repo_ids", tuple)(): + lid = str(lid).lower() + if lid == repo_id.lower() or lid.startswith(repo_id.lower()): + raise HTTPException( + status_code = 400, + detail = "A Video model load is using this repo; wait for it to finish", + ) + except HTTPException: + raise + except Exception: + pass + try: cache_scans = _all_hf_cache_scans() diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index 511f2413b4..9ccb3237e0 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -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 diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index 0f8f773e3e..fbc1471fd8 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -260,6 +260,45 @@ def test_list_cached_models_skips_non_suffix_repo_when_gguf_files_exist(monkeypa assert result["cached"] == [] +def test_list_cached_models_prefers_complete_over_larger_partial(monkeypatch, tmp_path): + # The same repo cached in two roots: a LARGER but PARTIAL copy must not shadow a SMALLER but + # COMPLETE one. The picker drops partial rows, so picking the partial winner by size alone would + # make a usable model vanish from On Device. Completeness wins over size. + complete = _repo( + "Org/Dup", + [_file("model.safetensors", 10_000)], + tmp_path / "root_a" / "models--Org--Dup", + ) + partial = _repo( + "Org/Dup", + [_file("model.safetensors", 15_000)], + tmp_path / "root_b" / "models--Org--Dup", + ) + + # The larger copy (root_b) is the partial one; the smaller (root_a) is complete. + monkeypatch.setattr( + models_route, + "_cached_repo_partial", + lambda repo_id, repo_cache_dir = None: "root_b" in str(repo_cache_dir), + ) + monkeypatch.setattr(models_route, "_cached_repo_task", lambda repo_info: None) + # List the partial (larger) FIRST, so the old size-only rule would have picked it. + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [partial, complete])], + ) + + result = asyncio.run(models_route.list_cached_models(current_subject = "test-user")) + + assert len(result["cached"]) == 1 + row = result["cached"][0] + assert row["repo_id"] == "Org/Dup" + # The COMPLETE (smaller) copy won: it is not flagged partial and carries its 10_000 size. + assert row.get("partial") is not True + assert row["size_bytes"] == 10_000 + + def test_list_cached_gguf_includes_mixed_repo_with_gguf_and_safetensors(monkeypatch, tmp_path): """Mixed repo still surfaces in cached-gguf as a GGUF download.""" mixed = _repo( @@ -812,3 +851,84 @@ def test_delete_cached_refuses_diffusion_loaded_repo(monkeypatch): except HTTPException as e: assert e.status_code == 400 assert "Unload the model before deleting" in e.detail + + +def test_delete_cached_refuses_video_loaded_repo(monkeypatch): + # The guard also refuses deleting a repo the Video backend has loaded: cached non-GGUF video + # repos now surface in the Video On-Device picker with the normal delete action, so without + # this a live Wan / LTX / Hunyuan pipeline's HF snapshot could be removed from under it. + from fastapi import HTTPException + + import core.inference.diffusion as diffusion_mod + import core.inference.video as video_mod + import routes.inference as routes_inference + + # Chat + Images backends report nothing loaded; only the Video backend holds the repo. + monkeypatch.setattr( + routes_inference, + "get_llama_cpp_backend", + lambda: SimpleNamespace(is_loaded = False, model_identifier = None), + ) + monkeypatch.setattr( + models_route, "get_inference_backend", lambda: SimpleNamespace(active_model_name = None) + ) + monkeypatch.setattr( + diffusion_mod, + "get_diffusion_backend", + lambda: SimpleNamespace( + status = lambda: {"loaded": False, "repo_id": None}, loading_repo_ids = lambda: () + ), + ) + monkeypatch.setattr( + video_mod, + "get_video_backend", + lambda: SimpleNamespace( + status = lambda: {"loaded": True, "repo_id": "Lightricks/LTX-2"}, + loading_repo_ids = lambda: (), + ), + ) + + try: + asyncio.run( + models_route.delete_cached_model( + repo_id = "Lightricks/LTX-2", variant = None, current_subject = "u" + ) + ) + assert False, "expected HTTPException refusing the delete" + except HTTPException as e: + assert e.status_code == 400 + assert "Unload the model before deleting" in e.detail + + +def test_cached_repo_partial_scopes_probe_to_snapshot_dir(monkeypatch): + # The partial probe must be scoped to the snapshot row being listed. Unscoped, the scan + # spans every HF cache root, so a stale .incomplete copy in one root would flag a complete + # copy living in another root as partial and hide the usable model from the picker. Verify + # _cached_repo_partial forwards the snapshot dir (matching the sibling inventory paths). + import hub.utils.inventory_scan as scan + + calls = [] + + def _fake( + repo_type, + repo_id, + repo_cache_dir = None, + ): + calls.append((repo_type, repo_id, repo_cache_dir)) + return False + + monkeypatch.setattr(scan, "is_snapshot_partial", _fake) + snapshot_dir = Path("/root_a/models--Org--Repo/snapshots/abc") + assert models_route._cached_repo_partial("Org/Repo", snapshot_dir) is False + assert calls == [("model", "Org/Repo", snapshot_dir)] + + # When that specific snapshot is partial, the row is flagged. + monkeypatch.setattr(scan, "is_snapshot_partial", lambda *a, **k: True) + assert models_route._cached_repo_partial("Org/Repo", snapshot_dir) is True + + # A probe error is swallowed (never hides a usable repo over a scan glitch). + def _boom(*a, **k): + raise RuntimeError("scan glitch") + + monkeypatch.setattr(scan, "is_snapshot_partial", _boom) + assert models_route._cached_repo_partial("Org/Repo", snapshot_dir) is False diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index bfe93e5b58..c1272de1af 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_precision.py b/studio/backend/tests/test_diffusion_precision.py index 072a43e5f8..2561352b76 100644 --- a/studio/backend/tests/test_diffusion_precision.py +++ b/studio/backend/tests/test_diffusion_precision.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_prequant.py b/studio/backend/tests/test_diffusion_prequant.py index a84ffcf24a..8b11cbd98f 100644 --- a/studio/backend/tests/test_diffusion_prequant.py +++ b/studio/backend/tests/test_diffusion_prequant.py @@ -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. diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index 9a18090e36..d2c53245ad 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -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 ──────────────────────────────────────────────────────────── diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index a621bd3874..b6e998c768 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -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 diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index d634ce7a41..b3d3f44be8 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -145,6 +145,9 @@ class _FakeWanDiT: def enable_cache(self, config) -> None: self.cache_config = config + def disable_cache(self) -> None: + self.cache_config = None + def set_attention_backend(self, backend) -> None: self.attention = backend @@ -801,6 +804,66 @@ def test_load_wan_ti2v_5b_pipeline(fake_runtime): assert _FakeWanPipelineSingle.last["repo"] == "Wan-AI/Wan2.2-TI2V-5B-Diffusers" +def test_video_dense_speed_defaults_to_compile_profile(fake_runtime): + # A clip denoise amortises the one-time compile within a single run, so an + # UNSET speed on a dense (pipeline) load resolves to `default` -- never `max`, + # never `off`. Explicit "off" is still honored verbatim. + backend = VideoBackend() + status = backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline") + assert status["speed_mode"] == "default" + assert status["resolved"]["speed_mode"]["source"] == "auto" + backend.unload() + status_off = backend.load_pipeline( + "Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline", speed_mode = "off" + ) + assert status_off["speed_mode"] == "off" + assert status_off["resolved"]["speed_mode"]["source"] == "explicit" + + +def test_video_step_cache_auto_from_default_schedule(fake_runtime, tmp_path): + # Unset step cache is AUTO, decided from the model's default schedule: Wan's + # 50-step default engages FBCache at load; the LTX distilled 8-step default + # keeps it off. Both are re-checked per generation (toggle test below). + backend = VideoBackend() + status = backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline") + assert status["transformer_cache"] == "fbcache" + assert status["resolved"]["transformer_cache"]["source"] == "auto" + backend.unload() + + (tmp_path / "ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf").write_bytes(b"w") + status2 = backend.load_pipeline( + str(tmp_path), + gguf_filename = "ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf", + base_repo = "Lightricks/LTX-2", + family_override = "ltx-2", + ) + assert status2["transformer_cache"] is None + assert status2["resolved"]["transformer_cache"]["source"] == "auto" + backend.unload() + + +def test_video_step_cache_auto_toggles_on_actual_steps(fake_runtime): + # The AUTO decision follows the ACTUAL step count of each generation: a + # few-step request drops the load-time cache, a many-step request restores + # it. An explicit "off" never toggles. + backend = VideoBackend() + backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline") + assert backend.status()["transformer_cache"] == "fbcache" + backend.generate(prompt = "a sloth", steps = 8) + assert backend.status()["transformer_cache"] is None + backend.generate(prompt = "a sloth", steps = 30) + assert backend.status()["transformer_cache"] == "fbcache" + backend.unload() + + backend.load_pipeline( + "Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline", transformer_cache = "off" + ) + assert backend.status()["transformer_cache"] is None + backend.generate(prompt = "a sloth", steps = 30) + assert backend.status()["transformer_cache"] is None + backend.unload() + + def test_wan_frame_snapping_4k_plus_1(fake_runtime): # Wan snaps num_frames to 4k+1 (temporal factor 4), unlike LTX-2's 8k+1. backend = VideoBackend() diff --git a/studio/backend/tests/test_video_gallery.py b/studio/backend/tests/test_video_gallery.py index 7248ec2c63..8df0bcff3b 100644 --- a/studio/backend/tests/test_video_gallery.py +++ b/studio/backend/tests/test_video_gallery.py @@ -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") diff --git a/studio/backend/tests/test_video_routes.py b/studio/backend/tests/test_video_routes.py index 784baf7897..99e52d75f6 100644 --- a/studio/backend/tests/test_video_routes.py +++ b/studio/backend/tests/test_video_routes.py @@ -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"] diff --git a/studio/frontend/package.json b/studio/frontend/package.json index a2eddecda3..e5b67d8035 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -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" }, diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 5df3ff4482..c2e389ca5f 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -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": diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index 65073258b9..c3f424bbc2 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -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={ ); diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.check.ts b/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.check.ts new file mode 100644 index 0000000000..8d9678d712 --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.check.ts @@ -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(); + 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"); diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.ts b/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.ts new file mode 100644 index 0000000000..20a273628a --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.ts @@ -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 => ({ + repoId, + format: "gguf", + loadKind: "gguf", + label: "GGUF", + keywords: ["gguf", "quantized"], + ...extra, +}); + +const bnb4bit = ( + repoId: string, + approxSizeGb: number, + extra?: Partial, +): 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 => ({ + 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; + /** canonical suffix-stripped key -> group */ + byKey: Map; + /** exact lowercased artifact id -> artifact */ + artifactById: Map; +} + +// 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(); + +function indexFor(catalog: CatalogGroup[]): CatalogIndex { + const cached = indexCache.get(catalog); + if (cached) return cached; + const byId = new Map(); + const byKey = new Map(); + const artifactById = new Map(); + 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 = { + 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]; +} diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index f13f454719..5af37ee192 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -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( + soleGguf?.repoId ?? null, + ); + if (soleGguf) { + return ( + + ); + } + return ( +
+ {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 ( +
+ + {open && ( + + )} +
+ ); + } + return ( + + ); + })} +
+ ); +} + // ── 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(); 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>(new Set()); + const [routingGroupId, setRoutingGroupId] = useState(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({
{ + 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 ( +
+
+
+ 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} + /> +
+ {routingGroupId === group.canonicalId ? ( + + + + ) : ( + + )} +
+ {expanded && ( + + )} +
+ ); + }; + + // 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 ( +
+
+
+ { + // 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} + /> +
+ {routingGroupId === group.canonicalId ? ( + + + + ) : ( + + )} +
+ {expanded && ( +
+ {rows.gguf.map(renderDownloadedGgufRow)} + {rows.models.map(renderDownloadedModelRow)} +
+ )} +
+ ); + })} + {ungroupedGguf.map(renderDownloadedGgufRow)} + {ungroupedModels.map(renderDownloadedModelRow)} + + ); + }; + return ( <>
@@ -2977,9 +3559,11 @@ export function HubModelPicker({ {sortedLmStudio.length > 0 ? "Unsloth" : "Downloaded"} {!downloadedCollapsed && - unslothCachedGguf.map(renderDownloadedGgufRow)} - {!downloadedCollapsed && - unslothCachedModelRows.map(renderDownloadedModelRow)} + renderCachedRows( + unslothCachedGguf, + unslothCachedModelRows, + "cached-group", + )} ) : null} @@ -2998,9 +3582,11 @@ export function HubModelPicker({ Other models {!otherModelsCollapsed && - otherCachedGguf.map(renderDownloadedGgufRow)} - {!otherModelsCollapsed && - otherCachedModelRows.map(renderDownloadedModelRow)} + renderCachedRows( + otherCachedGguf, + otherCachedModelRows, + "other-cached-group", + )}
) : 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 ( -
- handleModelClick(m.id)} - gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} - /> -
- ); - })} + {/* 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 ( +
+ handleModelClick(m.id)} + gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} + /> +
+ ); + })} {recommendedSearch.isLoading && recommendedRows.length === 0 ? (
@@ -3560,7 +4155,8 @@ export function HubModelPicker({
) : recommendedRows.length === 0 && - curatedSafetensorsRows.length === 0 ? ( + curatedSafetensorsRows.length === 0 && + !catalog?.length ? (
No models found.
@@ -3573,7 +4169,12 @@ export function HubModelPicker({ return (
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( diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 4699a031cd..acf3044077 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -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( diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 648f8a0c15..4ccb88797d 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -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 = { - "unsloth/Z-Image-Turbo-unsloth-bnb-4bit": { kind: "pipeline" }, - // Krea 2 Turbo: official vendor repo (bf16 pipeline), on the backend allowlist. - "krea/Krea-2-Turbo": { kind: "pipeline" }, - // Ideogram 4: official vendor pipelines, on the backend allowlist. No bf16 repo - // exists: -fp8 stores its two DiTs as raw float8 (highest precision; ~46 GB - // resident after the bf16 cast); -nf4-diffusers is the bnb-4bit export (~11 GB). - "ideogram-ai/ideogram-4-fp8": { kind: "pipeline" }, - "ideogram-ai/ideogram-4-nf4-diffusers": { kind: "pipeline" }, - "unsloth/Qwen-Image-2512-unsloth-bnb-4bit": { kind: "pipeline" }, - "unsloth/Qwen-Image-2512-FP8": { - kind: "single_file", - filename: "qwen-image-2512-fp8.safetensors", - }, - // SDXL is a U-Net family loaded as a whole pipeline (from_pretrained). These - // official base repos are on the backend's non-GGUF allowlist. - "stabilityai/sdxl-turbo": { kind: "pipeline" }, - "stabilityai/stable-diffusion-xl-base-1.0": { kind: "pipeline" }, -}; -// Curated non-GGUF picker entries (isGguf:false -> no quant expander, direct load). -const safetensors = (id: string, name: string, label: string): ModelOption => ({ - id, - name, - description: `Text-to-image · ${label}`, - isGguf: false, -}); - -const MODELS: ModelOption[] = [ - txt2img("unsloth/Z-Image-Turbo-GGUF", "Z-Image-Turbo"), - txt2img("unsloth/Z-Image-GGUF", "Z-Image"), - txt2img("unsloth/Qwen-Image-2512-GGUF", "Qwen-Image 2512"), - txt2img("unsloth/Qwen-Image-GGUF", "Qwen-Image"), - txt2img("unsloth/FLUX.1-schnell-GGUF", "FLUX.1 schnell"), - txt2img("unsloth/FLUX.1-dev-GGUF", "FLUX.1 dev"), - txt2img("unsloth/FLUX.2-klein-4B-GGUF", "FLUX.2 klein 4B"), - txt2img("unsloth/FLUX.2-klein-9B-GGUF", "FLUX.2 klein 9B"), - editGguf("unsloth/Qwen-Image-Edit-2511-GGUF", "Qwen-Image-Edit 2511"), - editGguf("unsloth/FLUX.1-Kontext-dev-GGUF", "FLUX.1 Kontext dev"), - safetensors( - "unsloth/Z-Image-Turbo-unsloth-bnb-4bit", - "Z-Image-Turbo (bnb-4bit)", - "Safetensors · bnb-4bit", - ), - safetensors("krea/Krea-2-Turbo", "Krea 2 Turbo", "Safetensors · bf16"), - safetensors("ideogram-ai/ideogram-4-fp8", "Ideogram 4 (FP8)", "Safetensors · fp8"), - safetensors( - "ideogram-ai/ideogram-4-nf4-diffusers", - "Ideogram 4 (bnb-4bit)", - "Safetensors · bnb-4bit", - ), - safetensors( - "unsloth/Qwen-Image-2512-unsloth-bnb-4bit", - "Qwen-Image 2512 (bnb-4bit)", - "Safetensors · bnb-4bit", - ), - safetensors( - "unsloth/Qwen-Image-2512-FP8", - "Qwen-Image 2512 (FP8)", - "Safetensors · fp8", - ), - safetensors("stabilityai/sdxl-turbo", "SDXL Turbo", "Safetensors · SDXL"), - safetensors( - "stabilityai/stable-diffusion-xl-base-1.0", - "SDXL Base 1.0", - "Safetensors · SDXL", - ), -]; +// Curated models come from the shared catalog: one canonical group per model, +// its artifacts (GGUF / FP8 / bnb-4bit / BF16) as data, and the load kind per +// artifact via loadSpecFor (replacing the old SAFETENSORS_MODELS table). The +// picker renders groups with a format second level and routes bare clicks to +// the best artifact for the device. +const MODELS: ModelOption[] = catalogToModelOptions(IMAGE_CATALOG); // Workflow tabs. `requires` is the backend workflow id (status.workflows) that must // be supported by the loaded model for the tab to enable; null = always available. @@ -321,23 +246,70 @@ const PAGE_SIZE = 50; // Export filename, e.g. Unsloth_20260624-143005_123.png. Batch siblings share // the seed + timestamp, so they get a "_" 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((resolve) => + canvas.toBlob(resolve, `image/${format}`, 0.95), + ); + if (!blob) throw new Error(`could not encode ${format}`); + const url = URL.createObjectURL(blob); + try { + saveBlobUrl(url, exportFilename(image, format)); + } finally { + // Give the click a tick to start before revoking. + setTimeout(() => URL.revokeObjectURL(url), 10_000); + } + } catch { + // Conversion failed (decode/encode); fall back to the original PNG bytes. + saveBlobUrl(src, exportFilename(image, "png")); + } +} + function formatTimestamp(epochSeconds: number): string { return new Date(epochSeconds * 1000).toLocaleString(); } @@ -489,6 +461,9 @@ function formatResolvedValue(key: string, value: string | boolean | null): strin if (value === null || value === "") return "Off"; if (typeof value === "boolean") return value ? "On" : "Off"; if (value === "_native_cudnn" || value.toLowerCase() === "cudnn") return "cuDNN"; + // Deferred speed auto: the dense pipe stays exact/eager and compiles on the + // 3rd image of the session (the tooltip carries the full reason). + if (value === "deferred") return "On from 3rd image"; return value.toUpperCase(); } @@ -1044,7 +1019,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { "auto", ); const [memoryMode, setMemoryMode] = useState<"auto" | "fast" | "balanced" | "low_vram">("auto"); - const [transformerCache, setTransformerCache] = useState<"off" | "fbcache">("off"); + const [transformerCache, setTransformerCache] = useState<"auto" | "off" | "fbcache">("auto"); const [cpuOffload, setCpuOffload] = useState(false); // The last load descriptor, so "Reapply" can reload the same model with new advanced // options without the user re-picking it from the dropdown. @@ -1527,7 +1502,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { transformer_quant: transformerQuant === "auto" ? undefined : transformerQuant, attention_backend: attentionBackend === "auto" ? undefined : attentionBackend, memory_mode: memoryMode === "auto" ? undefined : memoryMode, - transformer_cache: transformerCache === "off" ? undefined : transformerCache, + transformer_cache: transformerCache === "auto" ? undefined : transformerCache, }); } catch (err) { dismissLoadToast(); @@ -1576,8 +1551,8 @@ export function ImagesPage({ active = true }: { active?: boolean }) { // busy, while the backend rejects the second load with a 409. if (busy !== null) return; // Curated non-GGUF model: load as a full pipeline or single-file safetensors. - const spec = SAFETENSORS_MODELS[id]; - if (spec) { + const spec = loadSpecFor(id, IMAGE_CATALOG); + if (spec && spec.kind !== "gguf") { setQuant(null); const d = defaultsFor(id); setSteps(d.steps); @@ -1913,7 +1888,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { <> } value={speedMode} onValueChange={(v) => setSpeedMode(v as typeof speedMode)} @@ -1930,8 +1905,8 @@ export function ImagesPage({ active = true }: { active?: boolean }) { to GGUF (or nothing loaded) and otherwise show why it is unavailable. */} {!status?.loaded || status.model_kind === "gguf" ? ( } value={transformerQuant} onValueChange={(v) => setTransformerQuant(v as typeof transformerQuant)} @@ -1947,14 +1922,14 @@ export function ImagesPage({ active = true }: { active?: boolean }) { ) : (
- Dtype + Precision GGUF models only
)} } value={attentionBackend} onValueChange={(v) => setAttentionBackend(v as typeof attentionBackend)} @@ -1981,11 +1956,12 @@ export function ImagesPage({ active = true }: { active?: boolean }) { /> } value={transformerCache} onValueChange={(v) => setTransformerCache(v as typeof transformerCache)} options={[ + ["auto", "Auto"], ["off", "Off"], ["fbcache", "First-Block-Cache"], ]} @@ -2029,6 +2005,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { variant="ghost" className="!h-[34px]" task={IMAGE_GEN_TASKS} + catalog={IMAGE_CATALOG} open={active && selectorOpen} onOpenChange={(o) => setSelectorOpen(active && o)} /> @@ -2037,11 +2014,22 @@ export function ImagesPage({ active = true }: { active?: boolean }) { Create is the generation workspace; Train is the LoRA training workspace. */} setPageMode(v as "create" | "train")}> - - 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. */} + + + + Create + - - Train + + + + Train + @@ -2629,15 +2617,31 @@ export function ImagesPage({ active = true }: { active?: boolean }) { Size/seed live in the Recipe popover, so no separate chip here. */}
- + + + + + + void downloadImage(selectedSrc, selected, "png")} + > + PNG (original, keeps recipe) + + void downloadImage(selectedSrc, selected, "jpeg")} + > + JPEG (smaller) + + void downloadImage(selectedSrc, selected, "webp")} + > + WebP + + +
{lrScheduler !== "constant" && @@ -870,8 +870,7 @@ export function DiffusionTrainPanel({

- 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.

@@ -898,15 +897,10 @@ export function DiffusionTrainPanel({ ))}

- 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'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. )}

@@ -924,7 +918,7 @@ export function DiffusionTrainPanel({

- Mixed-precision autocast for the U-Net. bf16 suits modern GPUs. + How the math runs during training. bf16 is right for modern GPUs.

)} @@ -944,8 +938,7 @@ export function DiffusionTrainPanel({

- 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.

)} @@ -1084,8 +1077,8 @@ export function DiffusionTrainPanel({

- 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.

) : ( @@ -1113,8 +1106,8 @@ export function DiffusionTrainPanel({ )} {selectedDataset.caption_count === 0 && !gridOpen && (

- 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.

)} diff --git a/studio/frontend/src/features/video/api.ts b/studio/frontend/src/features/video/api.ts index d0159a7b42..28231197cf 100644 --- a/studio/frontend/src/features/video/api.ts +++ b/studio/frontend/src/features/video/api.ts @@ -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 { 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 { + const res = await authFetch( + `/api/inference/video/gallery/${id}/export?format=${format}`, + ); + if (!res.ok) throw new Error(await readFastApiError(res)); + return res.blob(); +} diff --git a/studio/frontend/src/features/video/video-page.tsx b/studio/frontend/src/features/video/video-page.tsx index c838b6b05b..991d0f9b23 100644 --- a/studio/frontend/src/features/video/video-page.tsx +++ b/studio/frontend/src/features/video/video-page.tsx @@ -14,6 +14,12 @@ import { import { HugeiconsIcon } from "@hugeicons/react"; import { Button } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; import { Input } from "@/components/ui/input"; import { Popover, @@ -33,6 +39,11 @@ import { Textarea } from "@/components/ui/textarea"; import { InfoHint } from "@/components/ui/info-hint"; import { ModelSelector } from "@/components/assistant-ui/model-selector"; import { VIDEO_GEN_TASKS } from "@/components/assistant-ui/model-selector/pickers"; +import { + VIDEO_CATALOG, + catalogToModelOptions, + loadSpecFor, +} from "@/components/assistant-ui/model-selector/model-catalog"; import type { ModelOption, ModelSelectorChangeMeta, @@ -51,6 +62,7 @@ import { cancelVideoGeneration, clearVideoGallery, deleteGalleryVideo, + fetchGalleryVideoExport, fetchGalleryVideoObjectUrl, generateVideo, getVideoGallery, @@ -61,60 +73,13 @@ import { unloadVideoModel, } from "./api"; -// How to load a curated non-GGUF (safetensors) video model. "pipeline" = a full diffusers -// repo (from_pretrained). The backend gates these to unsloth/* repos plus the official -// family base repos. Keyed by repo id so the load handler knows the kind. -type PipelineSpec = { kind: "pipeline"; filename?: string }; -const PIPELINE_MODELS: Record = { - "Lightricks/LTX-2": { kind: "pipeline" }, - // Wan2.2 diffusers base repos (no GGUF variant yet): loaded as full pipelines. TI2V-5B - // is a single-DiT 720p-class model; T2V-A14B is the dual-expert MoE. The backend gates - // these to the Wan-AI base repos (see _TRUSTED_NON_GGUF_VIDEO_REPOS). - "Wan-AI/Wan2.2-TI2V-5B-Diffusers": { kind: "pipeline" }, - "Wan-AI/Wan2.2-T2V-A14B-Diffusers": { kind: "pipeline" }, - // HunyuanVideo-1.5 community Diffusers repack (tencent's own repo is the original - // non-diffusers layout and cannot load as a pipeline). - "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v": { kind: "pipeline" }, -}; - -// A curated GGUF picker entry: isGguf true expands its .gguf files in the quant expander -// (like the image GGUF repos), and the backend resolves the pipeline + base repo from the id. -const ggufModel = (id: string, name: string): ModelOption => ({ - id, - name, - description: "Text-to-video · GGUF", - isGguf: true, -}); - -// A curated non-GGUF pipeline entry (isGguf false -> no quant expander, direct load). -const pipelineModel = (id: string, name: string, description: string): ModelOption => ({ - id, - name, - description, - isGguf: false, -}); - -// Curated text-to-video models the picker recommends. The chat ModelSelector also surfaces -// any other on-device video GGUF (via the VIDEO_GEN_TASKS filter). -const VIDEO_MODELS: ModelOption[] = [ - ggufModel("unsloth/LTX-2.3-GGUF", "LTX 2.3 distilled"), - pipelineModel("Lightricks/LTX-2", "LTX 2 (base, bf16)", "Text-to-video with audio · Safetensors"), - pipelineModel( - "Wan-AI/Wan2.2-TI2V-5B-Diffusers", - "Wan 2.2 TI2V 5B", - "Text-to-video 720p · Safetensors", - ), - pipelineModel( - "Wan-AI/Wan2.2-T2V-A14B-Diffusers", - "Wan 2.2 T2V A14B (MoE)", - "Text-to-video, dual-expert · Safetensors", - ), - pipelineModel( - "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v", - "HunyuanVideo 1.5 (480p)", - "Text-to-video 480p · Safetensors", - ), -]; +// Curated models come from the shared catalog: one canonical group per model +// with its artifacts as data (the HunyuanVideo group carries both the 480p and +// 720p repacks), and the load kind per artifact via loadSpecFor (replacing the +// old PIPELINE_MODELS table). The picker renders groups with a format second +// level -- which also finally surfaces LTX-2.3 in Recommended (its HF +// pipeline_tag is image-to-video, so the live text-to-video listing missed it). +const VIDEO_MODELS: ModelOption[] = catalogToModelOptions(VIDEO_CATALOG); // Per-model generation defaults (steps + guidance), matched by repo-id substring, most // specific first. The distilled model wants very few steps and no guidance; the full base @@ -175,23 +140,45 @@ const galleryCache: { const PAGE_SIZE = 50; // Export filename, e.g. Unsloth_video_20260624-143005_123.mp4. -function exportFilename(video: GalleryVideo): string { +type VideoExportFormat = "mp4" | "webm" | "gif"; + +function exportFilename(video: GalleryVideo, format: VideoExportFormat = "mp4"): string { const d = new Date(video.created_at); const p = (n: number) => String(n).padStart(2, "0"); const stamp = Number.isNaN(d.getTime()) ? "unknown" : `${d.getFullYear()}${p(d.getMonth() + 1)}${p(d.getDate())}` + `-${p(d.getHours())}${p(d.getMinutes())}${p(d.getSeconds())}`; - return `Unsloth_video_${stamp}_${video.seed}.mp4`; + return `Unsloth_video_${stamp}_${video.seed}.${format}`; } -function downloadVideo(src: string, video: GalleryVideo) { +function saveBlobUrl(href: string, filename: string) { const link = document.createElement("a"); - link.href = src; - link.download = exportFilename(video); + link.href = href; + link.download = filename; link.click(); } +// MP4 saves the already-fetched original bytes; WebM / GIF are transcoded by +// the backend on demand (501 with a readable reason when the codec is absent). +async function downloadVideo( + src: string, + video: GalleryVideo, + format: VideoExportFormat = "mp4", +) { + if (format === "mp4") { + saveBlobUrl(src, exportFilename(video, format)); + return; + } + const blob = await fetchGalleryVideoExport(video.id, format); + const url = URL.createObjectURL(blob); + try { + saveBlobUrl(url, exportFilename(video, format)); + } finally { + setTimeout(() => URL.revokeObjectURL(url), 10_000); + } +} + function formatTimestamp(iso: string): string { const d = new Date(iso); return Number.isNaN(d.getTime()) ? iso : d.toLocaleString(); @@ -530,7 +517,10 @@ export function VideoPage({ active = true }: { active?: boolean }) { const [attentionBackend, setAttentionBackend] = useState< "auto" | "native" | "cudnn" | "flash3" | "sage" >("auto"); - const [transformerCache, setTransformerCache] = useState<"off" | "fbcache">("off"); + const [transformerCache, setTransformerCache] = useState<"auto" | "off" | "fbcache">("auto"); + const [transformerQuant, setTransformerQuant] = useState< + "auto" | "none" | "fp8" | "int8" | "nvfp4" | "mxfp8" + >("auto"); // The last load descriptor, so "Reapply" can reload the same model with new advanced // options without the user re-picking it from the dropdown. const lastLoad = useRef<{ repoId: string; kind: "gguf" | "single_file" | "pipeline"; filename?: string } | null>( @@ -549,6 +539,12 @@ export function VideoPage({ active = true }: { active?: boolean }) { const [videos, setVideos] = useState(() => galleryCache.videos); const [hasMore, setHasMore] = useState(() => galleryCache.hasMore); const [selectedId, setSelectedId] = useState(() => 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>(() => Object.fromEntries(galleryCache.srcById), ); @@ -700,6 +696,29 @@ export function VideoPage({ active = true }: { active?: boolean }) { void loadGallery(); }, [loadGallery]); + // WebM/GIF go through a server-side transcode that can take a few seconds + // (and 501s with a readable reason when the codec is missing), so wrap the + // helper with progress + error toasts; MP4 saves instantly. + const handleDownload = useCallback( + async (src: string, video: GalleryVideo, format: "mp4" | "webm" | "gif") => { + if (format === "mp4") { + void downloadVideo(src, video, format); + return; + } + const toastId = toast.loading(`Converting to ${format.toUpperCase()}…`); + try { + await downloadVideo(src, video, format); + toast.dismiss(toastId); + } catch (err) { + toast.dismiss(toastId); + toast.error( + err instanceof Error ? err.message : `Failed to export ${format}`, + ); + } + }, + [], + ); + const handleDelete = useCallback(async (id: string) => { try { await deleteGalleryVideo(id); @@ -894,7 +913,8 @@ export function VideoPage({ active = true }: { active?: boolean }) { memory_mode: memoryMode === "auto" ? undefined : memoryMode, speed_mode: speedMode === "auto" ? undefined : speedMode, attention_backend: attentionBackend === "auto" ? undefined : attentionBackend, - transformer_cache: transformerCache === "off" ? undefined : transformerCache, + transformer_cache: transformerCache === "auto" ? undefined : transformerCache, + transformer_quant: transformerQuant === "auto" ? undefined : transformerQuant, }); } catch (err) { dismissLoadToast(); @@ -914,6 +934,7 @@ export function VideoPage({ active = true }: { active?: boolean }) { speedMode, attentionBackend, transformerCache, + transformerQuant, ], ); @@ -931,10 +952,15 @@ export function VideoPage({ active = true }: { active?: boolean }) { // Ignore picks while a load/generation/unload is in flight. if (busy !== null) return; // Curated non-GGUF model: load as a full pipeline. - const spec = PIPELINE_MODELS[id]; - if (spec) { + const spec = loadSpecFor(id, VIDEO_CATALOG); + if (spec && spec.kind !== "gguf") { setQuant(null); - const d = defaultsFor(id); + // The distilled variant lives in the single-file checkpoint name + // (ltx-2.3-...-distilled...), not the repo id, so include the filename when + // seeding defaults -- mirroring the GGUF branch below. Without it these + // distilled BF16/FP8 entries fall through to the generic LTX 40-step/CFG-4 + // defaults instead of the distilled 8-step/guidance-1 schedule. + const d = defaultsFor(spec.filename ? `${id}/${spec.filename}` : id); setSteps(d.steps); setGuidance(d.guidance); void handleLoad(id, { kind: spec.kind, filename: spec.filename }); @@ -1125,7 +1151,7 @@ export function VideoPage({ active = true }: { active?: boolean }) { /> } value={speedMode} onValueChange={(v) => setSpeedMode(v as typeof speedMode)} @@ -1137,9 +1163,36 @@ export function VideoPage({ active = true }: { active?: boolean }) { ["max", "Max"], ]} /> + {/* The dense transformer_quant fast path only engages on a full-pipeline load; a + GGUF / single-file checkpoint already carries its own precision, so gate the + control and otherwise show why it is unavailable. */} + {!status?.loaded || status.model_kind === "pipeline" ? ( + } + 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)"], + ]} + /> + ) : ( +
+ + Precision + + Full-pipeline models only +
+ )} } value={attentionBackend} onValueChange={(v) => setAttentionBackend(v as typeof attentionBackend)} @@ -1153,11 +1206,12 @@ export function VideoPage({ active = true }: { active?: boolean }) { /> } value={transformerCache} onValueChange={(v) => setTransformerCache(v as typeof transformerCache)} options={[ + ["auto", "Auto"], ["off", "Off"], ["fbcache", "First-Block-Cache"], ]} @@ -1191,6 +1245,7 @@ export function VideoPage({ active = true }: { active?: boolean }) { variant="ghost" className="!h-[34px]" task={VIDEO_GEN_TASKS} + catalog={VIDEO_CATALOG} open={active && selectorOpen} onOpenChange={(o) => setSelectorOpen(active && o)} /> @@ -1350,16 +1405,27 @@ export function VideoPage({ active = true }: { active?: boolean }) {
{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). */}