Fix quantized-load LoRA bake, prequant family exclusions and outpaint canvas

Six review findings across the Images page and model scanning:

- The quantized (int8/fp8) load path can only attach LoRA adapters
  before quantization, but the frontend load request had no loras field,
  so every generation after such a load was rejected and each reload
  repeated it. Send the selection with the load.
- build_prequant_checkpoint passed no family to the scheme exclusions
  while recording the family in metadata, so a Qwen int8 artifact baked
  the short-M text-stream linears and was then rejected wholesale by the
  loader's family-keyed check.
- Registering a bare single-file checkpoint directory produced no On
  Device row even though the images loader can load it; only its parent
  worked. Admit that shape when nothing else matched.
- Unload left the Reapply target set, so the repair path was skipped and
  Reapply reloaded the ejected model. Clear it, as the video page does.
- Both FLUX.2 bases were trusted for training but not inference, so
  Deploy to Create rejected every FLUX.2 adapter.
- Outpaint allocated the grown canvas before downscaling, exceeding the
  browser canvas area cap on a large photo; an over-cap canvas is
  unusable, so Extend silently posted a fully transparent image and
  mask. Scale the source first.

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

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",

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",

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
 
 

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

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.

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
 

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 {

studio/frontend/src/features/images/images-page.tsx

@@ -813,6 +813,18 @@ function loadImage(src: string): Promise<HTMLImageElement> {
 // 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());