From 186f381bc66bdfceac4c4fb28cd387c7c6e07016 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 17:48:39 -0700 Subject: [PATCH] Studio: diffusion UX polish and stronger auto policies (images + video) (#6885) * Auto policies: deferred dense compile, video compile default, step cache and precision auto Image dense loads with speed unset no longer sit at plain off: the load stays bit-identical eager, and the 3rd generation in a session engages the default compile profile plus the cuDNN attention upgrade mid-session (a one-off image never pays the warmup, repeated use amortises it). Video dense loads resolve straight to the default profile since a clip denoise amortises the compile within a single run, and never to max. Video also gains the image backend's tri-state auto policies: unset step cache now decides from the default schedule and re-checks the actual step count per generation, and unset precision (transformer_quant) hands the decision to the hardware ladder instead of staying off. Memory badge reason now says plainly that everything fits when no offload is planned. * Rename Dtype to Precision, add the video Precision control, step cache Auto option The images Advanced panel's Dtype row is now Precision (same control, clearer name), and the video Advanced panel gains the matching Precision select wired to the load route's existing transformer_quant field, gated to full-pipeline loads the way the image control gates to GGUF. Step cache selects on both pages gain an explicit Auto option as the default (the previous Off default silently behaved as auto and never let anyone pin off), and the Speed and Attention tooltips now state the deferred dense compile and the SageAttention black-frame caveat. * Model catalog: canonical diffusion model groups with device-aware routing One canonical name per image/video model, its published artifacts (GGUF, FP8, bnb-4bit, official BF16) as data, and pure routing helpers: suffix-stripped canonical keys (owner-preserving; cross-owner merges only via explicit aliases), group/artifact lookups, a flat back-compat options shim, load-spec resolution replacing the pages' lookup tables, search matching over old ids and format tokens, the GGUF fit ladder extracted from the variant expander, and pickDefaultArtifact/pickDefaultQuant deciding what a bare group click loads (downloaded first, then the best quality that fits 70 percent of VRAM, GGUF as the safe fallback). Checked by npm run catalog:check, following the i18n:check pattern. * Picker: one canonical row per diffusion model with a format second level The Images and Video pickers now render the curated catalog as one row per model in Recommended: clicking loads the best artifact for the device (the routed GGUF quant, a prequant FP8/bnb-4bit that fits, or the official BF16), and a chevron opens the per-format list, with the GGUF row nesting the usual quant expander. Live HF listing rows that belong to a group are deduplicated, search collapses member repos into their group (old ids and format tokens still match), and the On Device sections group cached member repos under the same canonical name with the per-repo rows inside. Curated groups render from the catalog rather than the HF listing, which finally surfaces LTX-2.3 in the video Recommended list (its hub pipeline_tag is image-to-video, so the text-to-video listing always missed it) and exposes the HunyuanVideo 720p repack next to 480p. Backend: /cached-models now tags trusted video-family repos text-to-video instead of blanket text-to-image, and the pickers admit catalog-known non-unsloth repos On Device, so cached Lightricks/Wan/Hunyuan pipelines finally appear in the Video picker. Chat pickers pass no catalog and are unchanged. * Download formats, tab icons, plain-language train tips, 3-loop autoplay The image Download button becomes a menu: PNG saves the original bytes with the embedded recipe, JPEG and WebP re-encode client-side from the fetched blob (JPEG flattened onto white). The video Download button gains MP4 (original, keeps audio), WebM and GIF; the latter two transcode server-side from the stored MP4 via PyAV (VP9 realtime profile for WebM, ~12 fps adaptive palette for GIF) behind a new gallery export route that 501s with a readable message when a codec is missing. Generated clips no longer loop forever: the player replays a clip three times per selection, then pauses with controls up; a new generation or a refresh gets its own three plays. The Create/Train tabs reuse the sidebar's New Chat and Train icons (TestTubeOutlineIcon moved to a shared lib module), and every Train tab helper text is now one plain sentence. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep the Create/Train tab icon and label on one line TabsTrigger renders its children inside a plain inline span and the Tailwind preflight gives svg display:block, so the HugeiconsIcon forced the label onto a second line. Wrap icon plus label in their own inline flex row inside each trigger. * Strip -int8 and -nvfp4 prequant suffixes in the model catalog key canonicalKeyFor already lowercases before matching, so -GGUF/-FP8 in any case were covered; -int8 and -nvfp4 were not in the suffix table, so such repos rendered as standalone rows in Recommended and On Device instead of standardizing into their base-name group and routing through pickDefaultArtifact. Added both suffixes plus case-insensitivity and routing assertions to the catalog check. * Standardize non-catalog picker rows to their base model name The curated catalog already collapses its own groups, but hub listing rows and cached repos outside the catalog (ERNIE-Image, FLUX.2-klein, Qwen-Image-Edit-2509, FLUX.2-dev) still rendered raw ids with -GGUF / -FP8 style suffixes in Recommended and On Device. - model-catalog.ts: new stripArtifactSuffixesForDisplay, a case-preserving twin of canonicalKeyFor's stripping that keeps the owner prefix and original casing for display. - pickers.tsx: recommended hub rows and the downloaded GGUF/model rows pass their labels through it when a catalog is present, so only the diffusion pickers change; chat rows keep raw ids. Click targets keep the full repo id, and the format badge still shows the artifact kind. - Catalog check covers the new helper across GGUF/FP8/int8/nvfp4 in both cases plus no-op and suffix-only names. * Offer official BF16/FP8 artifacts per model group and fix gallery label clipping Model picker changes so groups are not limited to unsloth quant repos: - model-catalog.ts: each image group that has an official vendor pipeline now carries its BF16 (official) artifact as the top (highest quality) entry - Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image, Qwen/Qwen-Image-2512, Qwen/Qwen-Image-Edit-2511, black-forest-labs/FLUX.1-dev, FLUX.1-schnell and FLUX.1-Kontext-dev. The LTX-2.3 video group now lists Lightricks' own bf16 and fp8 distilled single-file checkpoints alongside the GGUF. Resident sizes are set from the actual weight totals (FLUX ships a duplicate single-file that from_pretrained ignores, so FLUX bf16 is ~32 GB not 54). The repos that used to be aliases are now real artifacts. - The router already prefers the highest-quality artifact that fits the 0.7 x GPU budget, so a datacenter GPU now defaults to official BF16 while consumer GPUs still route to the fitting quant or GGUF. That is why bnb-4bit was the Z-Image-Turbo default before: it was the only non-GGUF artifact and it was already downloaded. - diffusion.py: allowlist the four official image repos not previously trusted (qwen/qwen-image-2512, qwen/qwen-image-edit-2511, black-forest-labs/flux.1-schnell, flux.1-kontext-dev). All verified as safetensors-only diffusers model_index pipelines. The LTX-2.3 checkpoints are already on the video trust list. - catalog check: BF16-wins-on-datacenter, quant-wins-on-consumer, and the single-file load specs for the LTX-2.3 checkpoints. Also fixes the video gallery thumbnail caption: the leading duration was clipped by the rounded corner and selection border, so the strip now has enough left/bottom padding to clear the curve. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * video gallery: guard export transcode against a stream-less clip _transcode_webm and _transcode_gif indexed src.streams.video[0] before checking the stream list, so a container with no video stream raised a bare IndexError that the broad handlers then re-labeled as a missing libvpx or decoder. Raise an explicit RuntimeError naming the real cause in both the WebM and GIF paths. * Studio: honor explicit attention/format choices, fix distilled-LTX defaults and On Device catalog routing * Remove stray planning notes accidentally committed to the branch * video: add transformerQuant to the load callback deps handleLoad reads transformerQuant but omitted it from the useCallback dep array, so after the user changes only Precision and then selects a model or clicks Reapply, the memoized callback keeps the stale closure and loads the previous precision. The image page's equivalent callback already lists it. * model picker: honor the format filter when routing catalog clicks; add catalog rows to the roving list - routedArtifactFor now scopes a group's artifacts to the active format filter (the same matchesFormatFilter predicate the visibility check uses) before pickDefaultArtifact, so a group shown only because it owns a GGUF no longer routes a click to a large non-GGUF download. Covers both the Recommended and On Device grouped paths. - hubOptionKeys now includes the catalog-group, search-catalog-group, and grouped On Device row keys in exact render order, so arrow/Home/End roving reaches the catalog rows instead of giving them a duplicate missing id and skipping them. * model picker: don't treat a partial base cache as downloaded A partially-cached base repo (a cancelled download that left only some weights) was counted as downloaded, so an On Device click routed to a fresh multi-GB re-download instead of the complete GGUF. The picker's endpoint (/api/models/ cached-models) did not carry a partial flag at all, so a frontend-only guard could not see it. Surface partial from that endpoint by reusing the hub inventory scan's snapshot-partial detector, plumb it through CachedModelRepo (backend + frontend types), and skip partial base repos when building the downloaded set. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * model picker + diffusion: drop partial/unloadable cached rows, skip defer-compile before a LoRA gen - On Device (cached non-GGUF) rows filtered partial-download snapshots back in: sortedCachedModels gated on passesTaskGate + a groupForRepoId key match but, unlike downloadedSet, never checked c.partial, so an incomplete unsloth snapshot showed as a loadable On Device row (click errors or silently re-fetches multi-GB). It also admitted repos that only match the catalog by group KEY (a base / uncurated-quant sibling like Qwen/Qwen-Image-2512) which have no loadable artifact and dead-end at the trust gate. Add !c.partial and gate on artifactForRepoId (what loadSpecFor resolves) instead of groupForRepoId, so a cached row shows only when the backend can load it. - Deferred speed-auto engaged the compile profile on the 3rd generation BEFORE _apply_loras. A compiled transformer rejects LoRA (supports_lora is False) and _apply_loras raises before its unchanged-selection no-op, so once compile engaged every LoRA generation on that load failed permanently. Skip the deferral when a LoRA is requested (compile and LoRA are mutually exclusive) and let it engage on a later LoRA-free generation. * Scope the cached-model partial probe to the listed snapshot dir list_cached_models builds each row from the largest/complete copy across HF cache roots, but _cached_repo_partial probed is_snapshot_partial with no repo_cache_dir, so the scan spanned every root: a stale .incomplete copy in one root would flag a complete copy in another as partial and hide the usable model from the picker (the click then routes to a re-download). Forward the winning snapshot's repo_path so all three partial signals are scoped to that copy, matching the sibling inventory paths (models/dataset cache_inventory, local_inventory). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Do not auto-route to gated repos, prefer complete cached copies, defer compile past attached LoRA, scope group expand keys Four fixes: - pickDefaultArtifact's not-downloaded ladder returned the gated BF16 FLUX.1-dev / Kontext-dev before the open GGUF on a large GPU, so a bare group click routed to a repo the user may lack license/token access to. Add a gated flag and skip gated artifacts in the not-downloaded ladder (an already-downloaded gated artifact is still returned). - list_cached_models picked the largest duplicate cache copy and computed partial only on it, so a larger partial copy shadowed a smaller complete one; since partial rows are dropped from the picker the usable model vanished. Prefer completeness, then size. - the deferred-speed compile engaged on a no-LoRA generation while an adapter from a prior generation was still attached, baking it into the compiled graph (the later unload is swallowed on a compiled pipe); also defer while adapters remain attached. - routeGroupClick's GGUF fallback toggled the context-free canonicalId while the chevron toggles the context-scoped expandKey, leaving the format list un-collapsible in one context, dead in the other, and risking cross-context expansion; thread expandKey through. * Guard video pipeline repos from deletion, drop the always-failing LTX FP8 artifact, prefer 720p Hunyuan Three round-6 fixes: - cached non-GGUF video repos now surface in the Video On-Device picker with the normal delete action, but /delete-cached only guarded chat + the Images engine, so a loaded/loading Wan / LTX / Hunyuan pipeline could have its HF snapshot removed from under it. Add a VideoBackend loading_repo_ids accessor and a video loaded/loading guard mirroring the Images one. - the catalog advertised Lightricks/LTX-2.3-fp8 as loadable, but the LTX-2.3 loader refuses the official scaled-FP8 single file (.weight_scale/.input_scale) and points to GGUF/BF16, so a pick routed to a ~76 GB download that always fails on load. Remove the FP8 artifact. - pickDefaultArtifact only sorts by format, so the HunyuanVideo group's 480p (listed first) beat the 720p even on GPUs where 720p fits the budget. List 720p first so the fit loop prefers it and falls back to 480p only on smaller cards. * diffusion: add compute int8/fp8_dynamic text-encoder quant, wire into video Add two torchao compute text-encoder quant modes to the diffusion precision engine, alongside the existing layerwise fp8 and weight-only nvfp4: - int8: per-token activation + per-channel weight (torch._int_mm), with per-layer keep-bf16 selection. int8 degrades on large encoders unless the most quant-sensitive decoder blocks stay bf16, so it engages only for families with a measured keep-bf16 schedule (qwen-image / qwen-image-edit keep first+last 6, flux.2-dev keeps first 3); a family without one falls back to fp8. - fp8_dynamic: per-row fp8 compute (torch._scaled_mm), keeping the matmul in fp8 on the tensor cores instead of upcasting each forward like the layerwise fp8. The selective int8 caster reuses the committed transformer-quant factory (_make_quant_config / make_filter_fn / exclude_tokens_for_scheme) plus a small structural first/last-N block skip, so it depends only on committed APIs. Wire text-encoder quant into the video backend, which previously loaded the companion encoder (Gemma3 / UMT5 / Qwen2.5-VL) dense bf16 while quantising only the DiT. text_encoder_quant is plumbed through the load request, validation, the load chain, the resolved record, and status, mirroring the image backend; it applies for every load kind (the encoder is dense regardless of how the DiT was sourced). Widen the image and video load request Literals and add the video status field. Tests: int8 family-schedule routing and fp8 fallback, fp8_dynamic routing, hardware gates (int8 sm_80+, fp8_dynamic sm_89+), the structural block selection, the real int8 filter closure (keeps the first blocks plus the vision tower / lm_head / T5 wo dense), and the video route threading and 422 validation. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * text-encoder quant: skip the torchao modes under offload (both backends) quantize_text_encoders applied int8-with-schedule / fp8_dynamic / nvfp4 (all torchao) to the text encoder regardless of the offload policy. An offload placement then moves the quantized encoder with Module.to(), which torchao tensor subclasses reject (aten._has_compatible_shallow_copy_type is unimplemented) -- a hard crash, the same one the DiT path already skips torchao quant under offload to avoid. Add offload_active to quantize_text_encoders and skip the torchao modes when set; layerwise fp8 is not torchao and still streams under offload. Both the video and image loaders pass offload_active = (offload policy != none). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * diffusion: skip non-bf16 linears for scaled_mm quant schemes The fp8 / mxfp8 / nvfp4 schemes run on torch._scaled_mm and the fp4 / mx GEMMs, which assert a bfloat16 input weight. On a mixed-precision DiT that keeps some linears in fp32 for numerical stability (the Wan and Hunyuan video transformers do this), quantize_ hits the first fp32 linear, raises, and the best-effort wrapper swallows it to None, so the whole transformer stays dense with no error and no speedup or memory saving. Add a require_bf16 gate to make_filter_fn and pass it for the scaled_mm schemes in quantize_transformer (and the fp8_dynamic text-encoder caster). The gate skips non-bf16 linears so the scheme engages on the bf16 ones. int8 uses torch._int_mm, which quantizes fp32/fp16 weights fine, so it leaves the gate off and keeps its current coverage. Verified on Wan2.2-TI2V-5B: fp8 and mxfp8 now quantize 303 linears via the committed quantize_transformer path where they previously engaged 0. * prequant builder: mirror the scaled-mm bf16 gate offline The runtime DiT quantizer skips non-bf16 Linears for the scaled_mm schemes (fp8, nvfp4, mxfp8) so the scheme engages on a mixed-precision transformer instead of aborting on the first fp32 Linear. The offline prequant builder reused make_filter_fn without that gate, so building an fp8/nvfp4/mxfp8 checkpoint for a mixed-precision DiT (Wan, Hunyuan keep _keep_in_fp32_modules in fp32 even under torch_dtype=bf16) would hit the same fp32 Linear and abort, breaking the builder's stated offline == runtime, LPIPS-0 invariant. Thread require_bf16 = scheme in _SCALED_MM_SCHEMES through the builder, record it in the checkpoint metadata, and verify it on load (mirrors the existing exclude_name_tokens guard) so a future _SCALED_MM_SCHEMES change cannot silently load a checkpoint built under the old filter. * Keep nvfp4 fp32 linears quantised (bf16 gate is fp8/mxfp8 only) Verified on torchao 0.17 / B200: fp8 per-row asserts 'PerRow quantization only works for bfloat16 precision input weight' and mxfp8 asserts 'Only supporting bf16 out dtype', but NVFP4's high-precision conversion quantises an fp32 weight fine (forward included). So the bf16 skip-gate must be fp8/mxfp8 only, not all scaled_mm schemes -- otherwise nvfp4 leaves large fp32 projections dense, losing the intended memory/speed gain. Rename _SCALED_MM_SCHEMES -> _REQUIRE_BF16_SCHEMES = (fp8, mxfp8) and thread it through the runtime filter, the offline builder, and the loader require_bf16 verification (offline == runtime preserved). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- plans/controlnet-workflow.md | 113 --- plans/diffusion-popularity-findings.md | 113 --- plans/diffusion-workflows-pr-plan.md | 156 ---- plans/diffusion-workflows-studio.md | 125 --- plans/wobbly-jumping-narwhal.md | 141 ---- scripts/build_prequant_checkpoint.py | 20 +- studio/backend/core/inference/diffusion.py | 174 ++++- .../core/inference/diffusion_precision.py | 186 ++++- .../core/inference/diffusion_prequant.py | 19 + .../backend/core/inference/diffusion_speed.py | 18 +- .../inference/diffusion_transformer_quant.py | 42 +- studio/backend/core/inference/video.py | 175 ++++- .../backend/core/inference/video_gallery.py | 94 +++ studio/backend/models/inference.py | 27 +- studio/backend/routes/models.py | 84 +- studio/backend/routes/video.py | 32 + .../backend/tests/test_cached_gguf_routes.py | 120 +++ .../backend/tests/test_diffusion_backend.py | 203 +++++ .../backend/tests/test_diffusion_precision.py | 202 ++++- .../backend/tests/test_diffusion_prequant.py | 38 + studio/backend/tests/test_diffusion_speed.py | 4 + .../tests/test_diffusion_transformer_quant.py | 41 + studio/backend/tests/test_video_backend.py | 63 ++ studio/backend/tests/test_video_gallery.py | 40 + studio/backend/tests/test_video_routes.py | 54 ++ studio/frontend/package.json | 1 + .../frontend/src/components/app-sidebar.tsx | 9 +- .../assistant-ui/model-selector.tsx | 10 + .../model-selector/model-catalog.check.ts | 472 ++++++++++++ .../model-selector/model-catalog.ts | 677 ++++++++++++++++ .../assistant-ui/model-selector/pickers.tsx | 727 ++++++++++++++++-- .../src/features/chat/api/chat-api.ts | 3 + .../src/features/images/images-page.tsx | 248 +++--- .../images/train/diffusion-train-panel.tsx | 37 +- studio/frontend/src/features/video/api.ts | 18 + .../src/features/video/video-page.tsx | 246 ++++-- studio/frontend/src/lib/hugeicons-derived.ts | 14 + 37 files changed, 3747 insertions(+), 999 deletions(-) delete mode 100644 plans/controlnet-workflow.md delete mode 100644 plans/diffusion-popularity-findings.md delete mode 100644 plans/diffusion-workflows-pr-plan.md delete mode 100644 plans/diffusion-workflows-studio.md delete mode 100644 plans/wobbly-jumping-narwhal.md create mode 100644 studio/frontend/src/components/assistant-ui/model-selector/model-catalog.check.ts create mode 100644 studio/frontend/src/components/assistant-ui/model-selector/model-catalog.ts create mode 100644 studio/frontend/src/lib/hugeicons-derived.ts 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). */}