diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 01bf37a418..c0473e774b 100644 --- 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), ) diff --git a/studio/backend/core/inference/diffusion_auto_policy.py b/studio/backend/core/inference/diffusion_auto_policy.py index 254ccbafc5..d1158b881c 100644 --- 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 diff --git a/studio/backend/core/inference/diffusion_prequant.py b/studio/backend/core/inference/diffusion_prequant.py index 2bfc72dfbb..111017ecae 100644 --- 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, diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 2d18ba0772..35cbaab0e1 100644 --- 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 ): diff --git a/studio/backend/tests/test_diffusion_prequant.py b/studio/backend/tests/test_diffusion_prequant.py index b646933e06..7935f2862b 100644 --- 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 = {}