Studio diffusion: fix static compile shape registration and prequant path validation

Register the dims the forward actually compiled with: image-conditioned workflows (img2img, inpaint, upscale, edit) run at the input image's size, not the slider's, so recording the slider values marked never-compiled shapes as covered and warm restarts kept paying compile for the real one. Validate a request-supplied transformer_prequant_path (existence plus the UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH allowlist) before treating prequant as available at the resident-fit re-check: an unusable path skipped the dense fit check up front and then fell back to materializing dense bf16 after the previous pipeline was evicted, recreating the post-eviction OOM path. Shared as usable_prequant_source, also used by the auto-policy planner.
2026-07-11 16:50:19 +00:00 · 2026-07-11 16:50:19 +00:00 · e1fa4fec04
commit e1fa4fec04
parent 316c8b5a81
5 changed files with 216 additions and 12 deletions
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@ -95,6 +95,7 @@ from .diffusion_precision import normalize_te_quant, quantize_text_encoders
 from .diffusion_prequant import (
    load_prequantized_transformer,
    resolve_prequant_source,
+    usable_prequant_source,
 )
 from .diffusion_auto_policy import (
    build_resolved_record,
@ -265,6 +266,24 @@ def _clamp_max_side(img: Any, max_side: int) -> Any:
    return img.resize((nw, nh), Image.LANCZOS)


+def _compile_shape_dims(
+    workflow: str, init_pil: Any, width: int, height: int
+) -> tuple[int, int]:
+    """The (width, height) a generation's forward ACTUALLY runs at, for static
+    compile-cache shape registration.
+
+    txt2img / reference / controlnet generate at the requested slider size, but the
+    image-conditioned workflows (img2img / inpaint / upscale / edit) derive the output
+    from the (resized/snapped) input image -- registering the slider values there would
+    mark a shape covered that was never compiled, so the truly-used shape never
+    re-dirties the bundle and warm restarts keep paying its compile. Mirrors the
+    width/height kwarg derivation in generate()."""
+    if workflow in ("txt2img", "reference", "controlnet") or init_pil is None:
+        return int(width), int(height)
+    iw, ih = init_pil.size
+    return int(iw), int(ih)
+
+
 # A small allowlist of well-known official base repos that may load as a full
 # (non-GGUF) pipeline even though they are not under ``unsloth/``. These are
 # safetensors-only checkpoints from their original publisher (no pickle, no remote
@ -1378,8 +1397,14 @@ class DiffusionBackend:
                            transformer_quant,  # normalized above
                            family = getattr(fam, "name", None),
                        )
+                        # usable_prequant_source (not resolve_): a request-supplied local
+                        # path that is missing or outside the allowlist must NOT count as a
+                        # prequant source here, or it would skip the dense-fit re-check and
+                        # _load_dense_quant_pipeline's refusal would fall back to
+                        # materialising the dense bf16 transformer AFTER the eviction --
+                        # exactly the OOM this re-check exists to prevent.
                        prequant = (
-                            resolve_prequant_source(
+                            usable_prequant_source(
                                fam, scheme, path_override = transformer_prequant_path
                            )
                            if scheme is not None
@ -2892,9 +2917,16 @@ class DiffusionBackend:
                # with the enriched set. Idempotent + best-effort -- never fails a
                # generation.
                try:
+                    # Register the dims the forward ACTUALLY compiled with: the
+                    # image-conditioned workflows run at the input image's size, not the
+                    # slider's (see _compile_shape_dims), and a mis-registered slider
+                    # shape would keep the truly-used shape out of the saved bundle.
+                    reg_width, reg_height = _compile_shape_dims(
+                        workflow, init_pil, width, height
+                    )
                    compile_cache.register_shape(
                        state.compile_cache_ctx,
-                        (int(width), int(height), int(batch_size)),
+                        (reg_width, reg_height, int(batch_size)),
                        static = "compiled" in (state.speed_optims or ())
                        and compiled_shapes_are_static(state.pipe, state.speed_mode),
                    )
--- a/studio/backend/core/inference/diffusion_auto_policy.py
+++ b/studio/backend/core/inference/diffusion_auto_policy.py
@ -188,15 +188,14 @@ def resolve_dense_quant_candidate(
        return None
    prequant_available = False
    try:
-        from .diffusion_prequant import local_prequant_path_ready, resolve_prequant_source
+        from .diffusion_prequant import usable_prequant_source

-        src = resolve_prequant_source(fam, scheme, path_override = prequant_path)
-        # A request-supplied local path override is only usable if the loader will accept it
-        # (allowlisted AND present); otherwise load_prequantized_transformer refuses it and
-        # rebuilds dense after the resident pipe is unloaded -- the evict-then-OOM this
-        # small-plan prefetch exists to avoid. Hosted-repo sources keep the existing signal.
-        if src is not None and src.kind == "path" and not local_prequant_path_ready(src.location):
-            src = None
+        # usable_ (not resolve_): a request-supplied local path override counts only when
+        # the loader will accept it (allowlisted AND present); otherwise
+        # load_prequantized_transformer refuses it and rebuilds dense after the resident
+        # pipe is unloaded -- the evict-then-OOM this small-plan prefetch exists to avoid.
+        # Hosted-repo sources keep the existing signal.
+        src = usable_prequant_source(fam, scheme, path_override = prequant_path)
        prequant_available = src is not None
    except Exception:  # noqa: BLE001 -- prequant probing must never sink the candidate
        prequant_available = False
--- a/studio/backend/core/inference/diffusion_prequant.py
+++ b/studio/backend/core/inference/diffusion_prequant.py
@ -139,6 +139,26 @@ def resolve_prequant_source(
    return None


+def usable_prequant_source(
+    fam: Any,
+    scheme: str,
+    *,
+    path_override: Optional[str] = None,
+) -> Optional[PrequantSource]:
+    """``resolve_prequant_source``, but a request-supplied local path counts only when
+    the loader would actually accept it: inside the ``UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH``
+    allowlist AND present on disk. A path that fails either check resolves to None so
+    memory planning falls back to the dense-fit checks up front -- otherwise
+    ``load_prequantized_transformer`` refuses the path only AFTER the resident pipeline
+    was evicted, and the dense bf16 fallback then materialises under a plan that never
+    budgeted for it (the evict-then-OOM those checks exist to prevent). Hosted-repo
+    sources are unaffected."""
+    src = resolve_prequant_source(fam, scheme, path_override = path_override)
+    if src is not None and src.kind == "path" and not local_prequant_path_ready(src.location):
+        return None
+    return src
+
+
 def load_prequantized_transformer(
    transformer_cls: Any,
    base: str,
--- a/studio/backend/tests/test_diffusion_backend.py
+++ b/studio/backend/tests/test_diffusion_backend.py
@ -1101,6 +1101,59 @@ def test_image_conditioned_passes_image_size_not_slider(fake_runtime, tmp_path):
    assert _SizePipe.last == {"width": 96, "height": 64}


+def test_compile_shape_dims_follow_workflow():
+    """_compile_shape_dims mirrors generate()'s width/height derivation: slider size for
+    txt2img / reference / controlnet, the input image's size for the image-conditioned
+    workflows (whose forward runs at init_pil.size, whatever the sliders say)."""
+    from PIL import Image
+
+    from core.inference.diffusion import _compile_shape_dims
+
+    img = Image.new("RGB", (96, 64), (10, 20, 30))
+    assert _compile_shape_dims("txt2img", None, 1024, 512) == (1024, 512)
+    # reference generates at the slider size even though an init image is present.
+    assert _compile_shape_dims("reference", img, 1024, 512) == (1024, 512)
+    assert _compile_shape_dims("controlnet", None, 768, 768) == (768, 768)
+    for wf in ("img2img", "inpaint", "upscale", "edit"):
+        assert _compile_shape_dims(wf, img, 1024, 512) == (96, 64)
+
+
+def test_register_shape_uses_actual_forward_dims(fake_runtime, tmp_path, monkeypatch):
+    """The static compile-cache manifest must record the dims the forward ACTUALLY ran
+    at: an image-conditioned generate derives its output size from the input image, so
+    registering the slider values would mark a never-compiled shape as covered while the
+    truly-used shape never re-dirties/saves the bundle (warm restarts keep paying its
+    compile)."""
+    from core.inference import diffusion as diff
+
+    registered: list = []
+    monkeypatch.setattr(
+        diff.compile_cache,
+        "register_shape",
+        lambda ctx, shape, *, static: registered.append(tuple(shape)),
+    )
+    monkeypatch.setattr(diff.compile_cache, "save", lambda ctx, *, logger = None: True)
+    (tmp_path / "model.gguf").write_bytes(b"x")
+    backend = DiffusionBackend()
+    backend.load_pipeline(
+        str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image"
+    )
+    # txt2img registers the requested slider size.
+    backend.generate(prompt = "x", steps = 4, width = 1024, height = 512, seed = 1)
+    assert registered[-1] == (1024, 512, 1)
+    # img2img runs at the INPUT image's 64x64; the 1024x512 slider must not be recorded.
+    backend.generate(
+        prompt = "x",
+        steps = 4,
+        width = 1024,
+        height = 512,
+        seed = 1,
+        init_image = _tiny_png_b64(),
+        strength = 0.5,
+    )
+    assert registered[-1] == (64, 64, 1)
+
+
 def test_edit_family_uses_own_pipeline_and_requires_image(fake_runtime, tmp_path):
    """An instruction-editing family (Qwen-Image-Edit) exposes only the 'edit' workflow,
    runs the image through its OWN loaded pipeline (no from_pipe), and rejects a call with
@ -2495,7 +2548,9 @@ def test_dense_quant_prequant_skips_dense_refit(fake_runtime, tmp_path, monkeypa
    monkeypatch.setattr(
        dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8"
    )
-    monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: "prequant/path")
+    # usable_ (not resolve_): the re-check site only honours a source the loader would
+    # actually accept, so the fake must present a USABLE one (e.g. a hosted repo).
+    monkeypatch.setattr(dmod, "usable_prequant_source", lambda fam, scheme, **kw: "prequant/path")
    # Large dense shards cached: if the re-check ran, it would wrongly decline the fast path.
    monkeypatch.setattr(
        DiffusionBackend,
@ -2534,6 +2589,58 @@ def test_dense_quant_prequant_skips_dense_refit(fake_runtime, tmp_path, monkeypa
    assert attempted == [True]  # fast path still attempted (with the prequant)


+def test_dense_quant_unusable_prequant_path_runs_dense_refit(fake_runtime, tmp_path, monkeypatch):
+    # A request-supplied transformer_prequant_path the loader will refuse (missing, or
+    # outside UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH) resolves to NO usable prequant source,
+    # so the dense-transformer fit re-check MUST run: with real device budgets it is what
+    # declines the fast path up front instead of evicting the resident pipeline and
+    # OOMing in the dense bf16 fallback _load_dense_quant_pipeline falls into.
+    from core.inference import diffusion as dmod
+
+    backend = DiffusionBackend()
+    _force_cuda_target(backend, monkeypatch)
+    monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True)
+    monkeypatch.setattr(
+        dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8"
+    )
+    # The REAL usable_prequant_source refuses a non-allowlisted path (unit-tested in
+    # test_diffusion_prequant.py); returning None here pins that outcome at this site.
+    monkeypatch.setattr(dmod, "usable_prequant_source", lambda fam, scheme, **kw: None)
+    monkeypatch.setattr(
+        DiffusionBackend,
+        "_dense_transformer_resident_bytes",
+        staticmethod(lambda base: 999 * 1024**3),
+    )
+    dense_refit_ran = []
+    orig_plan = DiffusionBackend._plan_memory
+
+    def spy_plan(
+        self,
+        *a,
+        transformer_resident_override_mib = None,
+        **k,
+    ):
+        if transformer_resident_override_mib is not None:
+            dense_refit_ran.append(True)
+        return orig_plan(self, *a, **k)
+
+    monkeypatch.setattr(DiffusionBackend, "_plan_memory", spy_plan)
+    monkeypatch.setattr(
+        DiffusionBackend, "_load_dense_quant_pipeline", lambda self, *a, **k: (None, None)
+    )
+    (tmp_path / "m.gguf").write_bytes(b"x")
+    backend.load_pipeline(
+        str(tmp_path),
+        gguf_filename = "m.gguf",
+        family_override = "z-image",
+        transformer_quant = "fp8",
+        transformer_prequant_path = str(tmp_path / "not-allowlisted.pt"),
+    )
+    # Unusable path -> no prequant shortcut -> the dense fit re-check ran.
+    assert dense_refit_ran == [True]
+    assert backend.status()["loaded"] is True
+
+
 def test_transformer_quant_unsupported_scheme_skips_dense_download(
    fake_runtime, tmp_path, monkeypatch
 ):
--- a/studio/backend/tests/test_diffusion_prequant.py
+++ b/studio/backend/tests/test_diffusion_prequant.py
@ -75,6 +75,52 @@ def test_local_prequant_path_ready(tmp_path, monkeypatch):
    assert pq.local_prequant_path_ready(str(ckpt)) is False


+# ── usable_prequant_source ───────────────────────────────────────────────────────
+def test_usable_source_missing_path_is_none(tmp_path, monkeypatch):
+    # An allowlisted but ABSENT request-supplied path must not count as a prequant
+    # source: load_prequantized_transformer would find no file and fall back to the
+    # dense bf16 build after the resident pipeline was already evicted, so the memory
+    # planner must run the dense fit checks up front instead.
+    import os
+
+    monkeypatch.setattr(pq, "_allowed_prequant_roots", lambda: [os.path.realpath(str(tmp_path))])
+    fam = _fam(prequant_repos = (("fp8", "org/hosted-fp8"),))
+    missing = str(tmp_path / "missing.pt")
+    assert pq.usable_prequant_source(fam, "fp8", path_override = missing) is None
+
+
+def test_usable_source_disallowed_path_is_none(tmp_path, monkeypatch):
+    # A path OUTSIDE the UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH allowlist (including the
+    # default empty allowlist) is refused by the loader, so it must resolve to None
+    # here even when the file exists.
+    ckpt = tmp_path / "model.pt"
+    ckpt.write_bytes(b"x")
+    monkeypatch.setattr(pq, "_allowed_prequant_roots", lambda: [])
+    fam = _fam(prequant_repos = (("fp8", "org/hosted-fp8"),))
+    assert pq.usable_prequant_source(fam, "fp8", path_override = str(ckpt)) is None
+
+
+def test_usable_source_allowed_present_path_wins(tmp_path, monkeypatch):
+    # Allowlisted AND present: the override is usable and takes priority over the
+    # hosted repo, exactly like resolve_prequant_source.
+    import os
+
+    ckpt = tmp_path / "model.pt"
+    ckpt.write_bytes(b"x")
+    monkeypatch.setattr(pq, "_allowed_prequant_roots", lambda: [os.path.realpath(str(tmp_path))])
+    fam = _fam(prequant_repos = (("fp8", "org/hosted-fp8"),))
+    src = pq.usable_prequant_source(fam, "fp8", path_override = str(ckpt))
+    assert src == PrequantSource(kind = "path", location = str(ckpt), filename = None)
+
+
+def test_usable_source_repo_unaffected_by_allowlist(monkeypatch):
+    # Hosted-repo sources are first-party and keep resolving with no allowlist at all.
+    monkeypatch.setattr(pq, "_allowed_prequant_roots", lambda: [])
+    fam = _fam(prequant_repos = (("fp8", "org/hosted-fp8"),))
+    src = pq.usable_prequant_source(fam, "fp8")
+    assert src is not None and src.kind == "repo" and src.location == "org/hosted-fp8"
+
+
 # ── load_prequantized_transformer ────────────────────────────────────────────────
 class _FakeTransformer:
    calls: dict = {}