diff --git a/scripts/build_prequant_checkpoint.py b/scripts/build_prequant_checkpoint.py index 96f04b039a..f511867c5e 100644 --- a/scripts/build_prequant_checkpoint.py +++ b/scripts/build_prequant_checkpoint.py @@ -86,8 +86,13 @@ def main(argv = None) -> int: # Mirror the runtime path EXACTLY (offline == runtime, LPIPS-0 invariant): for int8 also skip # the M=1 AdaLN-modulation / conditioning-embedder projections, else the checkpoint 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) + # / fp4 / mx use scaled_mm (no M limit) -> exclude_tokens_for_scheme returns (). Pass the + # family: int8 also carries PER-FAMILY exclusions (Qwen-Image's unpadded text stream runs at + # M = prompt tokens, so a short prompt breaks _int_mm), and the loader validates the baked + # list against exclude_tokens_for_scheme(scheme, metadata["family"]) -- so building with + # family=None both bakes the crashing text-stream linears and yields an artifact the runtime + # then rejects (silently falling back to the dense quantise this script exists to avoid). + exclude_name_tokens = exclude_tokens_for_scheme(scheme, fam.name) # 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) keeps its _keep_in_fp32_modules in # fp32 even under torch_dtype=bf16, so quantising one raises inside quantize_ and aborts the diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 2fc5178ada..aafa6abe8f 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -298,6 +298,14 @@ _TRUSTED_NON_GGUF_REPOS = frozenset( # Krea's guidance-distilled FLUX.1-dev finetune: same arch/layout as dev (FluxPipeline, # CLIP+T5+ae), gated like dev. Detected as the flux.1 family via the "flux.1" token. "black-forest-labs/flux.1-krea-dev", + # FLUX.2: the LoRA TRAINING bases for the flux.2-dev / flux.2-klein families (dev is + # Hub-gated, klein-4B is open). Both ship a safetensors-only diffusers layout + # (model_index.json + transformer/text_encoder/vae, no pickled weights, no remote code), + # so a pipeline load is the same shape as FLUX.1-dev's. Trusted here too because the Train + # tab's "Deploy to Create" reloads the trained-on base as a pipeline: without these the + # deploy of every FLUX.2 adapter 400s on this gate. + "black-forest-labs/flux.2-dev", + "black-forest-labs/flux.2-klein-4b", "tongyi-mai/z-image-turbo", "qwen/qwen-image", "qwen/qwen-image-2512", diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 0f3a1206c3..1ffe7c6c64 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -1001,9 +1001,11 @@ def _restore_perf_flags(snap: Optional[dict]) -> None: # Official safetensors-only TRAINING bases trusted in addition to the inference allowlist -# (_TRUSTED_NON_GGUF_REPOS in core/inference/diffusion.py): the FLUX.2 bases train LoRAs but -# are not (yet) non-GGUF inference bases. Exact-match lowercased, same rules as the loader -# list: extend deliberately; never add pickled weights or remote code. +# (_TRUSTED_NON_GGUF_REPOS in core/inference/diffusion.py). The FLUX.2 bases are now in that +# list too (deploying a trained FLUX.2 adapter reloads the base as an inference pipeline), and +# are kept here so the training gate stays independent of a future edit to the loader list. +# Exact-match lowercased, same rules as the loader list: extend deliberately; never add pickled +# weights or remote code. _TRAIN_EXTRA_TRUSTED_REPOS = frozenset( { "black-forest-labs/flux.2-dev", diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 7add994e32..77d5930205 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -391,6 +391,29 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca ), ) + # A scan folder can also point directly at a BARE single-file checkpoint dir (one loose + # .safetensors / weight .bin, no config.json and no model_index.json): both root checks above + # reject it, yet the child loop admits exactly that shape via _has_non_gguf_weights and the + # Images/Video load path reinterprets such a directory through resolve_local_single_file. So + # registering the model folder itself returned no On Device row while registering its parent + # worked. Only when nothing else matched, so a models root holding a stray loose weight file + # next to real model subdirs still lists those children instead of collapsing to one row. + if not found and (limit is None or limit > 0) and _has_non_gguf_weights(models_dir): + try: + updated_at = models_dir.stat().st_mtime + except OSError: + updated_at = None + found.append( + LocalModelInfo( + id = str(models_dir), + display_name = models_dir.name, + path = str(models_dir), + source = "models_dir", + model_format = _dir_model_format(models_dir), + updated_at = updated_at, + ), + ) + return found diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index 58ee2389d9..53844f4ea2 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -163,6 +163,24 @@ def test_flux2_bases_pass_the_trusted_base_gate(): _assert_trusted_base_model("someone/random-flux2-finetune") +def test_every_train_base_is_deployable_as_an_inference_pipeline(): + # "Deploy to Create" reloads the trained-on base (or the family's deploy_base) through + # /images/load as a PIPELINE, which gates non-GGUF loads on _is_trusted_diffusion_repo. Any + # advertised training base that fails that gate makes Deploy 400 for every adapter trained on + # it -- which is what happened to both FLUX.2 families, trusted for training only. + from core.inference.diffusion import _is_trusted_diffusion_repo + from core.inference.diffusion_families import _FAMILIES + + for fam in _FAMILIES: + if not fam.trainable: + continue + for base in fam.train_base_repos: + deploy_base = fam.deploy_base_repo or base + assert _is_trusted_diffusion_repo(deploy_base), ( + f"{fam.name}: deploy base {deploy_base!r} is not loadable for inference" + ) + + def test_gated_access_requires_token(): assert "black-forest-labs/flux.1-dev" in _GATED_TRAIN_REPOS assert "black-forest-labs/flux.2-dev" in _GATED_TRAIN_REPOS diff --git a/studio/backend/tests/test_diffusion_prequant.py b/studio/backend/tests/test_diffusion_prequant.py index bceed8731f..52a00070b0 100644 --- a/studio/backend/tests/test_diffusion_prequant.py +++ b/studio/backend/tests/test_diffusion_prequant.py @@ -406,6 +406,28 @@ def test_load_exclude_tokens_match_ok(monkeypatch, tmp_path): assert _load(monkeypatch, tmp_path, ckpt, scheme = "int8") is not None +def test_load_exclude_tokens_need_the_recorded_family(monkeypatch, tmp_path): + # int8 carries PER-FAMILY exclusions (Qwen's unpadded text stream runs at M = prompt tokens, + # under _int_mm's M > 16 floor). An offline artifact that recorded the family but built its + # exclusion set with family=None baked those linears as int8, so the loader must reject it -- + # and accept only the family-aware set. Pins the offline builder + # (scripts/build_prequant_checkpoint.py) to exclude_tokens_for_scheme(scheme, fam.name). + from core.inference.diffusion_transformer_quant import exclude_tokens_for_scheme + + for family in ("qwen-image", "qwen-image-edit"): + family_less = _good_ckpt(scheme = "int8") + family_less["metadata"]["family"] = family + family_less["metadata"]["exclude_name_tokens"] = list(exclude_tokens_for_scheme("int8")) + assert _load(monkeypatch, tmp_path, family_less, scheme = "int8") is None + + family_aware = _good_ckpt(scheme = "int8") + family_aware["metadata"]["family"] = family + family_aware["metadata"]["exclude_name_tokens"] = list( + exclude_tokens_for_scheme("int8", family) + ) + assert _load(monkeypatch, tmp_path, family_aware, 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. diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index ebbc761215..a4f79a51ce 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -150,6 +150,32 @@ def test_scan_models_dir_surfaces_root_diffusers_pipeline(tmp_path): assert rows[0].model_format is None +def test_scan_models_dir_surfaces_root_single_file_checkpoint(tmp_path): + # A scan folder can also point DIRECTLY at a bare single-file checkpoint dir (one loose + # .safetensors, no config.json / model_index.json). The child loop admits exactly that shape + # and the images route reinterprets it via resolve_local_single_file, so the root must be + # surfaced too -- otherwise registering the model folder itself yields no On Device row while + # registering its parent works. + root = tmp_path / "qwen-image-2509" + _touch(root / "qwen-image-2509.safetensors") + + rows = models_route._scan_models_dir(root) + + assert [r.path for r in rows] == [str(root)] + assert rows[0].model_format is None + + +def test_scan_models_dir_root_weights_do_not_hide_child_models(tmp_path): + # A stray loose .safetensors at a models ROOT must not collapse the scan to a single row and + # hide the real model subdirs: the root fallback applies only when nothing else matched. + root = tmp_path / "models" + _touch(root / "stray.safetensors") + _touch(root / "llama" / "config.json") + _touch(root / "llama" / "model.safetensors") + + assert [Path(r.path).name for r in models_route._scan_models_dir(root)] == ["llama"] + + # ── Images picker task tag for local (non-GGUF) diffusers models ────────────── from models.models import LocalModelInfo # noqa: E402 diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index 9b858dceec..255050a6cb 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -84,6 +84,11 @@ export interface DiffusionLoadRequest { | "aiter"; memory_mode?: "auto" | "fast" | "balanced" | "low_vram"; transformer_cache?: "off" | "fbcache"; + // LoRA adapters to BAKE into a torchao int8/fp8 build: the backend can only attach them to + // the dense transformer BEFORE quantisation + compilation, so a quantized load that omits + // them rejects every generation with "reload the model with the adapter selection". Ignored + // by every other load kind (bf16 / bnb-4bit take adapters at generation time). + loras?: LoraSpecInput[]; } export interface DiffusionGenerateRequest { diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index bd8c39ea37..fbcf14c12a 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -813,6 +813,18 @@ function loadImage(src: string): Promise { // Which sides to grow when outpainting. type ExtendSides = { left: boolean; right: boolean; top: boolean; bottom: boolean }; +// Redraw an image/canvas at (w, h). Used to clamp an outpaint source (and the built pair) to a +// size the browser can actually back and the backend can decode. +function scaleToCanvas(source: CanvasImageSource, w: number, h: number): HTMLCanvasElement { + const dst = document.createElement("canvas"); + dst.width = w; + dst.height = h; + const dctx = dst.getContext("2d"); + if (!dctx) throw new Error("Could not scale the extended canvas"); + dctx.drawImage(source, 0, 0, w, h); + return dst; +} + // Build the (image, mask) pair for outpaint by reusing the inpaint backend: grow the canvas // by `pct` per dimension on the selected sides, edge-bleed the original pixels into the new // bands (so the VAE encodes plausible content), and mask the new bands white (= repaint) with @@ -822,9 +834,28 @@ async function buildOutpaint( sides: ExtendSides, pct: number, ): Promise<{ image: string; mask: string }> { - const img = await loadImage(src); - const w = img.naturalWidth; - const h = img.naturalHeight; + const source = await loadImage(src); + // Scale the SOURCE so the grown canvas fits MAX_SIDE, before allocating anything. Growing all + // four sides by 100% multiplies the source area by 9, and a canvas past the browser's limit is + // "unusable -- drawing commands will not work" (MDN): every drawImage silently no-ops, so + // Extend posted a fully transparent init_image + mask instead of the user's photo. Measured + // Chromium cap is 268,435,456 px of area (a 48MP phone photo grown on four sides is 439M); + // iOS caps a canvas at 4096x4096, which a 2048px source at 100% already passes. Pre-scaling + // also keeps the two full-size allocations (image + mask) from reaching a gigabyte on a large + // source. The backend rounds to /16, so exact dims aren't required. + const MAX_SIDE = 4096; + const grow = (a: boolean, b: boolean) => 1 + (a ? pct / 100 : 0) + (b ? pct / 100 : 0); + const fit = Math.min( + 1, + MAX_SIDE / + Math.max( + source.naturalWidth * grow(sides.left, sides.right), + source.naturalHeight * grow(sides.top, sides.bottom), + ), + ); + const w = fit < 1 ? Math.max(1, Math.floor(source.naturalWidth * fit)) : source.naturalWidth; + const h = fit < 1 ? Math.max(1, Math.floor(source.naturalHeight * fit)) : source.naturalHeight; + const img: CanvasImageSource = fit < 1 ? scaleToCanvas(source, w, h) : source; const px = Math.round((pct / 100) * w); const py = Math.round((pct / 100) * h); const l = sides.left ? px : 0; @@ -865,28 +896,17 @@ async function buildOutpaint( mctx.fillStyle = "#000000"; // ...except the kept original (inset by the seam overlap). mctx.fillRect(l + ol, t + ot, w - ol - or, h - ot - ob); - // The grown canvas can exceed the backend's 4096px-per-side decode limit (e.g. a 2048px - // source at 100% on both sides -> 6144px), which would 400 the load. Scale the built pair - // down proportionally to fit so Extend still returns an outpaint. The backend rounds to /16, - // so exact dims aren't required. - const MAX_SIDE = 4096; + // The source pre-scale above sizes the canvases to fit MAX_SIDE, but the per-side rounding + // (round(pct% * w) on each grown edge) can still overshoot it by a pixel or two, which the + // backend's 4096px-per-side decode limit would 400. Trim that slack here. const longest = Math.max(nw, nh); if (longest > MAX_SIDE) { const scale = MAX_SIDE / longest; const sw = Math.max(1, Math.round(nw * scale)); const sh = Math.max(1, Math.round(nh * scale)); - const scaleCanvas = (source: HTMLCanvasElement): HTMLCanvasElement => { - const dst = document.createElement("canvas"); - dst.width = sw; - dst.height = sh; - const dctx = dst.getContext("2d"); - if (!dctx) throw new Error("Could not scale the extended canvas"); - dctx.drawImage(source, 0, 0, sw, sh); - return dst; - }; return { - image: scaleCanvas(ic).toDataURL("image/png"), - mask: scaleCanvas(mc).toDataURL("image/png"), + image: scaleToCanvas(ic, sw, sh).toDataURL("image/png"), + mask: scaleToCanvas(mc, sw, sh).toDataURL("image/png"), }; } @@ -1652,6 +1672,21 @@ export function ImagesPage({ active = true }: { active?: boolean }) { // guard) leaves the previous model resident, so Reapply and the resident-default seeding // must keep pointing at it, not the failed pick. const prevLastLoad = lastLoad.current; + // A torchao int8/fp8 transformer (what the default GGUF fast path resolves to on any + // capable GPU) can only take adapters at LOAD time: they attach on the dense transformer + // before quantize_ + compile. Generation-time /images/generate then rejects a new adapter + // set with "reload the model with the adapter selection", so a reload that drops the + // selection leaves the advertised LoRA picker permanently unusable. Carry the same + // filtered nonzero list Generate sends, but only when reloading the SAME target (Reapply + // or re-picking the loaded model): a fresh pick of a different model can still hold a + // cross-family selection that the family-swap effect only clears after the load lands. + // Ignored by every other load kind (bf16 / bnb-4bit apply adapters at generation time). + const sameTarget = repoId === (prevLastLoad?.repoId ?? status?.repo_id ?? null); + const bakeLoras = sameTarget + ? loras + .map((l) => ({ id: l.id.trim(), weight: l.weight })) + .filter((l) => l.id && l.weight > 0) + : []; lastLoad.current = { repoId, kind: opts.kind, filename: opts.filename }; setCanReapply(true); // Carry the prior target so the async poll can restore it if the background load fails @@ -1674,6 +1709,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { attention_backend: attentionBackend === "auto" ? undefined : attentionBackend, memory_mode: memoryMode === "auto" ? undefined : memoryMode, transformer_cache: transformerCache === "auto" ? undefined : transformerCache, + loras: bakeLoras.length > 0 ? bakeLoras : undefined, }); } catch (err) { lastLoad.current = prevLastLoad; @@ -1698,6 +1734,8 @@ export function ImagesPage({ active = true }: { active?: boolean }) { attentionBackend, memoryMode, transformerCache, + loras, + status?.repo_id, ], ); @@ -1866,6 +1904,12 @@ export function ImagesPage({ active = true }: { active?: boolean }) { pollTimer.current = null; dismissLoadToast(); lastLoadSig.current = null; + // Drop the Reapply target with the model (like the Video page): the ejected pick is no + // longer resident, and leaving it set makes the resident-status seeding effect above skip + // repair (it is guarded on lastLoad.current === null) while "Reapply to loaded model" + // reloads the ejected model instead of whatever another client/session left loaded. + lastLoad.current = null; + setCanReapply(false); setBusy("unloading"); try { setStatus(await unloadDiffusionModel());