diff --git a/scripts/fbcache_flux_probe.py b/scripts/fbcache_flux_probe.py new file mode 100644 index 0000000000..2b41657dcf --- /dev/null +++ b/scripts/fbcache_flux_probe.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Validate First-Block-Cache (FBCache) on a MANY-step DiT (Flux.1-dev), vs the compiled +baseline. FBCache reuses the transformer tail across denoise steps when the first block's +residual barely changes -- a real speedup only when there are enough steps (it is why it is +gated OFF for few-step distilled models like Z-Image-Turbo). Reports median latency, +speedup, peak VRAM, and LPIPS vs the no-cache baseline. One CUDA GPU.""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +BASE = "black-forest-labs/FLUX.1-dev" +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" +OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/fbcache_flux_images") + + +_LP = {"fn": None} + + +def _lpips(ref, arr): + try: + import lpips + import torch + if _LP["fn"] is None: + _LP["fn"] = lpips.LPIPS(net="alex", verbose=False).cuda().eval() + + def t(x): + return (torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0).cuda() + + with torch.no_grad(): + return float(_LP["fn"](t(ref), t(arr)).item()) + except Exception as exc: # noqa: BLE001 + print(f" (lpips: {type(exc).__name__})", flush=True) + return None + + +def _load(): + import os + import diffusers + import torch + pipe = diffusers.FluxPipeline.from_pretrained(BASE, torch_dtype=torch.bfloat16, token=os.environ.get("HF_TOKEN")) + pipe.to("cuda") + return pipe + + +def _gen(pipe, steps, seed, res, guidance): + import torch + g = torch.Generator(device="cuda").manual_seed(seed) + torch.cuda.synchronize(); t0 = time.time() + img = pipe(prompt=PROMPT, width=res, height=res, num_inference_steps=steps, + guidance_scale=guidance, generator=g).images[0] + torch.cuda.synchronize() + return img, time.time() - t0 + + +def _median(xs): + return sorted(xs)[len(xs) // 2] + + +def run(tag, steps, seed, res, guidance, iters, *, threshold=None, compile_=True): + import torch + torch.compiler.reset(); torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats() + pipe = _load() + if threshold is not None: + from diffusers import FirstBlockCacheConfig + try: + pipe.transformer.enable_cache(FirstBlockCacheConfig(threshold=threshold)) + except Exception as exc: # noqa: BLE001 + from diffusers.hooks import apply_first_block_cache + apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold=threshold)) + if compile_: + try: + pipe.transformer.compile_repeated_blocks(fullgraph=True, dynamic=True) + except Exception as exc: # noqa: BLE001 + print(f" [{tag}] compile {type(exc).__name__}: {str(exc)[:80]}", flush=True) + try: + _gen(pipe, steps, seed, res, guidance) # warmup / compile + except Exception as exc: # noqa: BLE001 + import traceback; traceback.print_exc() + print(f" [{tag}] FAILED: {type(exc).__name__}: {str(exc)[:100]}", flush=True) + del pipe; torch.cuda.empty_cache(); return None + dts, img = [], None + for _ in range(iters): + img, dt = _gen(pipe, steps, seed, res, guidance); dts.append(dt) + peak = torch.cuda.max_memory_allocated() / 1e9 + arr = np.array(img) + OUT.mkdir(parents=True, exist_ok=True); img.save(OUT / f"{tag}.png") + del pipe; torch.cuda.empty_cache() + return _median(dts), arr, peak + + +def main(argv=None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--steps", type=int, default=28) + p.add_argument("--res", type=int, default=1024) + p.add_argument("--seed", type=int, default=42) + p.add_argument("--guidance", type=float, default=3.5) + p.add_argument("--iters", type=int, default=2) + args = p.parse_args(argv) + s, r, seed, gd, it = args.steps, args.res, args.seed, args.guidance, args.iters + + print(f"== FBCache on Flux.1-dev ({r}px, {s} steps, guidance {gd}) ==", flush=True) + base = run("baseline", s, seed, r, gd, it) + if base is None: + print("baseline FAILED", flush=True); return 1 + bmed, ref, bpeak = base + print(f" baseline {bmed:.3f}s peak={bpeak:.1f}G", flush=True) + rows = [("baseline", bmed, bpeak, 0.0)] + for thr in (0.08, 0.12, 0.20): + out = run(f"fbcache_{thr}", s, seed, r, gd, it, threshold=thr) + if out is None: + rows.append((f"fbcache_{thr}", None, None, None)); continue + med, arr, peak = out + lp = _lpips(ref, arr) + rows.append((f"fbcache_{thr}", med, peak, lp)) + print(f" fbcache_{thr}: {med:.3f}s ({bmed/med:.2f}x) peak={peak:.1f}G LPIPS={lp}", flush=True) + + print("\n==== SUMMARY (Flux.1-dev, ref = no-cache compile) ====", flush=True) + for tag, med, peak, lp in rows: + if med is None: + print(f" {tag:16s} FAILED"); continue + spd = f"{bmed/med:.2f}x" + lpv = "ref" if tag == "baseline" else (f"{lp:.3f}" if lp is not None else "n/a") + print(f" {tag:16s} {med:.3f}s {spd:>6s} peak={peak:.1f}G LPIPS={lpv:>6s}", flush=True) + print("FBCACHE-FLUX-DONE", flush=True) + return 0 + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + sys.exit(main()) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 126071c10a..4e8c090a39 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -55,6 +55,7 @@ from .diffusion_attention import ( apply_attention_backend, select_attention_backend, ) +from .diffusion_cache import apply_step_cache from .diffusion_precision import quantize_text_encoders from .diffusion_prequant import ( load_prequantized_transformer, @@ -101,6 +102,8 @@ 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 + # Step cache engaged ("fbcache") or None. Opt-in, for many-step models. + transformer_cache: Optional[str] = None @dataclass @@ -291,6 +294,8 @@ class DiffusionBackend: transformer_quant_fast_accum: Optional[bool] = None, transformer_prequant_path: Optional[str] = None, attention_backend: Optional[str] = None, + transformer_cache: Optional[str] = None, + transformer_cache_threshold: Optional[float] = None, ) -> dict[str, Any]: """Validate, then run the (slow) load on a daemon thread. Returns at once.""" fam = self.validate_load_request( @@ -326,6 +331,8 @@ class DiffusionBackend: transformer_quant_fast_accum = transformer_quant_fast_accum, transformer_prequant_path = transformer_prequant_path, attention_backend = attention_backend, + transformer_cache = transformer_cache, + transformer_cache_threshold = transformer_cache_threshold, _load_token = token, ), daemon = True, @@ -451,6 +458,8 @@ class DiffusionBackend: transformer_quant_fast_accum: Optional[bool] = None, transformer_prequant_path: Optional[str] = None, attention_backend: Optional[str] = None, + transformer_cache: Optional[str] = None, + transformer_cache_threshold: Optional[float] = None, _load_token: Optional[int] = None, ) -> dict[str, Any]: # Validate first (cheap, no torch/diffusers) so a direct call with a bad @@ -571,12 +580,24 @@ class DiffusionBackend: ), logger = logger, ) + # Opt-in step caching (First-Block-Cache), also before compile. OFF by + # default; for many-step models it reuses the transformer tail across steps + # (~1.4x on Flux at LPIPS ~0.08). When engaged, compile must drop fullgraph + # (the cache's per-step decision is a graph break), so pass it through. + cache_engaged = apply_step_cache( + pipe, + mode = transformer_cache, + threshold = transformer_cache_threshold, + quant_active = transformer_quant_engaged is not None, + logger = logger, + ) speed_applied = apply_speed_optims( pipe, target, is_gguf = bool(gguf_filename), family = fam, speed_mode = effective_speed, + cache_active = cache_engaged is not None, logger = logger, ) # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), @@ -615,6 +636,7 @@ class DiffusionBackend: text_encoder_quant = te_quant, transformer_quant = transformer_quant_engaged, attention_backend = attention_engaged, + transformer_cache = cache_engaged, ) logger.info( @@ -917,6 +939,7 @@ class DiffusionBackend: "text_encoder_quant": None, "transformer_quant": None, "attention_backend": None, + "transformer_cache": None, } return { "loaded": True, @@ -934,6 +957,7 @@ class DiffusionBackend: "text_encoder_quant": state.text_encoder_quant, "transformer_quant": state.transformer_quant, "attention_backend": state.attention_backend, + "transformer_cache": state.transformer_cache, } diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py new file mode 100644 index 0000000000..2c67358ea9 --- /dev/null +++ b/studio/backend/core/inference/diffusion_cache.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Opt-in step caching for the diffusion transformer (First-Block-Cache). + +Across denoising steps a DiT's output changes little once the trajectory settles, so most of +the transformer can be reused. First-Block-Cache (FBCache) computes the first block, and if +its residual barely changed from the previous step (within ``threshold``) it skips the +remaining blocks and reuses their cached output. diffusers ships it natively +(``transformer.enable_cache(FirstBlockCacheConfig(...))`` for CacheMixin models, or the +standalone ``apply_first_block_cache`` hook). + +Measured on Flux.1-dev (28 steps, 1024px, B200): ~1.4x on top of torch.compile (2.83 -> +2.03 s) at LPIPS ~0.08 vs the no-cache output -- deep inside the speed-for-quality bar. + +OFF by default and a deliberate per-load opt-in, because the win scales with step count: a +few-step distilled model (e.g. Z-Image-Turbo at ~8 steps) has almost no headroom and a +single skipped step is a large fraction of the trajectory, so caching is for many-step +models (Flux / Qwen-Image). It composes with torch.compile only with ``fullgraph=False`` +(the cache's compiler-disabled decision is a graph break), which the speed layer switches to +automatically when a cache is engaged. Best-effort: an incompatible model (e.g. a transformer +whose block signature the hook does not recognise) is caught and the load proceeds uncached. +torch / diffusers imported lazily. +""" + +from __future__ import annotations + +from typing import Any, Optional + +TC_OFF = "off" +TC_FBCACHE = "fbcache" +TC_MODES = (TC_FBCACHE,) + +# FBCache residual thresholds: higher skips more steps (faster, lower quality). The dense +# bf16 default; a quantised transformer shifts the residual distribution, so it needs a +# higher threshold for the cache to trigger at all (per ParaAttention's fp8 guidance). +DEFAULT_FBCACHE_THRESHOLD = 0.08 +QUANT_FBCACHE_THRESHOLD = 0.12 + + +def normalize_transformer_cache(value: Optional[str]) -> Optional[str]: + """Lower/strip a requested cache mode; None / "" / "none" / "off" -> None (disabled). + + Raises ValueError for an unsupported value so a bad request is rejected cheaply.""" + if value is None: + return None + normalized = str(value).strip().lower().replace("-", "_") + if not normalized or normalized in ("none", "off"): + return None + if normalized not in TC_MODES: + raise ValueError( + f"Unsupported transformer_cache '{value}'. Use one of: off, {', '.join(TC_MODES)}." + ) + return normalized + + +def apply_step_cache( + pipe: Any, + *, + mode: Optional[str], + threshold: Optional[float] = None, + quant_active: bool = False, + logger: Any = None, +) -> Optional[str]: + """Engage step caching on ``pipe.transformer``. Returns the mode actually engaged, or + None when disabled / unsupported (the load then runs uncached). ``threshold`` overrides + the default; ``quant_active`` raises the default so the cache still triggers on a + quantised transformer. Best-effort: never raises for an incompatible model.""" + mode = normalize_transformer_cache(mode) + if mode is None: + return None + transformer = getattr(pipe, "transformer", None) + if transformer is None: + return None + thr = threshold if threshold is not None else ( + QUANT_FBCACHE_THRESHOLD if quant_active else DEFAULT_FBCACHE_THRESHOLD + ) + try: + from diffusers import FirstBlockCacheConfig + + config = FirstBlockCacheConfig(threshold=thr) + enable_cache = getattr(transformer, "enable_cache", None) + if callable(enable_cache): + enable_cache(config) + else: + from diffusers.hooks import apply_first_block_cache + + apply_first_block_cache(transformer, config) + try: + transformer._unsloth_step_cache = f"{mode}@{thr}" + except Exception: # noqa: BLE001 — marker is best-effort + pass + if logger is not None: + logger.info("diffusion.cache: %s engaged (threshold=%s)", mode, thr) + return mode + except Exception as exc: # noqa: BLE001 — incompatible model -> run uncached + _warn(logger, mode, exc) + return None + + +def _warn(logger: Any, what: str, exc: Exception) -> None: + if logger is not None: + logger.warning("diffusion.cache: %s unavailable (%s); running uncached", what, exc) diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 2548c0f0dc..6ad3dcfff7 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -127,6 +127,7 @@ def apply_speed_optims( is_gguf: bool, family: Any, speed_mode: str = SPEED_OFF, + cache_active: bool = False, logger: Any = None, ) -> dict[str, bool]: """Apply the opt-in speed optimisations for ``speed_mode`` to a built pipeline, @@ -155,7 +156,9 @@ def apply_speed_optims( # block, where eligible (now incl. the GGUF transformer). `max` opts into # max-autotune (longer compile, autotuned kernels). if compile_eligible(target, is_gguf = is_gguf, family = family): - applied["compiled"] = _compile_repeated_blocks(pipe, logger, max_autotune = mode == SPEED_MAX) + applied["compiled"] = _compile_repeated_blocks( + pipe, logger, max_autotune = mode == SPEED_MAX, cache_active = cache_active + ) if mode == SPEED_MAX: # Near-lossless: TF32 matmul (CUDA only) trades a few mantissa bits for speed. @@ -184,6 +187,7 @@ def _compile_repeated_blocks( logger: Any, *, max_autotune: bool = False, + cache_active: bool = False, ) -> bool: transformer = getattr(pipe, "transformer", None) fn = getattr(transformer, "compile_repeated_blocks", None) @@ -195,7 +199,12 @@ def _compile_repeated_blocks( # compile and a recompile per new resolution. The CUDA-graph modes (reduce-overhead # / max-autotune) are deliberately NOT used: they crash on the regionally-compiled # block because its static output buffer is overwritten across denoise steps. - kwargs: dict[str, Any] = {"fullgraph": True, "dynamic": not max_autotune} + # + # fullgraph drops to False when a step cache is engaged: FBCache's per-step decision is + # ``@torch.compiler.disable``d, i.e. a graph break, which fullgraph=True rejects ("Skip + # inlining torch.compiler.disable()d function"). The break is cheap and the rest of the + # block still compiles. + kwargs: dict[str, Any] = {"fullgraph": not cache_active, "dynamic": not max_autotune} if max_autotune: kwargs["mode"] = "max-autotune-no-cudagraphs" try: diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index d47ad240dd..eec529000a 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1769,6 +1769,23 @@ class DiffusionLoadRequest(BaseModel): "friendly); xformers/aiter are memory-efficient (NVIDIA) / AMD ROCm. An " "unavailable kernel falls back to the default.", ) + transformer_cache: Optional[Literal["off", "fbcache"]] = Field( + None, + description = "Opt-in step caching (off by default). fbcache = First-Block-Cache: " + "reuse the transformer tail across denoise steps when the first block's residual " + "barely changes (~1.4x on Flux 28-step at LPIPS ~0.08). For MANY-step models " + "(Flux / Qwen-Image); leave off for few-step distilled models (e.g. Z-Image-Turbo), " + "which have no caching headroom. Composes with compile (drops fullgraph " + "automatically); incompatible models run uncached.", + ) + transformer_cache_threshold: Optional[float] = Field( + None, + ge = 0.0, + le = 1.0, + description = "FBCache residual threshold (higher = skips more steps = faster, lower " + "quality). null auto-picks 0.08 (0.12 when the transformer is quantised, which " + "shifts the residual distribution).", + ) class DiffusionGenerateRequest(BaseModel): @@ -1886,3 +1903,6 @@ class DiffusionStatusResponse(BaseModel): description = "Attention backend engaged via the diffusers dispatcher (e.g. " "_native_cudnn), or null for the default SDPA", ) + transformer_cache: Optional[str] = Field( + None, description = "Step cache engaged: fbcache | null" + ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 4dca8e262d..9115e39e5a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10089,6 +10089,8 @@ async def load_diffusion_model( transformer_quant_fast_accum = request.transformer_quant_fast_accum, transformer_prequant_path = request.transformer_prequant_path, attention_backend = request.attention_backend, + transformer_cache = request.transformer_cache, + transformer_cache_threshold = request.transformer_cache_threshold, ) return DiffusionStatusResponse(**status_dict) except (ValueError, FileNotFoundError) as exc: diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py new file mode 100644 index 0000000000..40d49251b9 --- /dev/null +++ b/studio/backend/tests/test_diffusion_cache.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Hermetic CPU tests for opt-in step caching (First-Block-Cache). + +``diffusers`` is stubbed via ``sys.modules`` (the module under test imports +``FirstBlockCacheConfig`` / ``apply_first_block_cache`` lazily), and the pipeline is a fake +that records the engaged config. So normalisation, the CacheMixin (``enable_cache``) path, the +standalone-hook fallback, threshold selection, and the best-effort failure handling are all +exercised without torch or a real diffusers model. +""" + +from __future__ import annotations + +import sys +import types + +import pytest + +from core.inference.diffusion_cache import ( + DEFAULT_FBCACHE_THRESHOLD, + QUANT_FBCACHE_THRESHOLD, + TC_FBCACHE, + apply_step_cache, + normalize_transformer_cache, +) + + +# ── normalize_transformer_cache ──────────────────────────────────────────────────── +def test_normalize_disabled_values_are_none(): + for value in (None, "", " ", "none", "off", "OFF", "None"): + assert normalize_transformer_cache(value) is None + + +def test_normalize_fbcache_and_casing(): + assert normalize_transformer_cache("fbcache") == TC_FBCACHE + assert normalize_transformer_cache("FBCache") == TC_FBCACHE + assert normalize_transformer_cache(" fbcache ") == TC_FBCACHE + + +def test_normalize_rejects_unknown(): + with pytest.raises(ValueError): + normalize_transformer_cache("deepcache") + + +# ── apply_step_cache ─────────────────────────────────────────────────────────────── +class _Config: + def __init__(self, threshold): + self.threshold = threshold + + +class _MixinTransformer: + """A CacheMixin-style transformer: exposes ``enable_cache``.""" + + def __init__(self, *, fail = False): + self.fail = fail + self.enabled_with = None + + def enable_cache(self, config): + if self.fail: + raise RuntimeError("block signature not recognised") + self.enabled_with = config + + +class _HookTransformer: + """A transformer with no ``enable_cache`` -> the standalone hook is used.""" + + +def _pipe(transformer): + return types.SimpleNamespace(transformer = transformer) + + +def _stub_diffusers(monkeypatch, *, hook_recorder = None): + diffusers = types.ModuleType("diffusers") + diffusers.FirstBlockCacheConfig = _Config + monkeypatch.setitem(sys.modules, "diffusers", diffusers) + + hooks = types.ModuleType("diffusers.hooks") + + def _apply_first_block_cache(transformer, config): + if hook_recorder is not None: + hook_recorder["transformer"] = transformer + hook_recorder["config"] = config + + hooks.apply_first_block_cache = _apply_first_block_cache + monkeypatch.setitem(sys.modules, "diffusers.hooks", hooks) + + +def test_disabled_mode_is_noop(monkeypatch): + _stub_diffusers(monkeypatch) + t = _MixinTransformer() + assert apply_step_cache(_pipe(t), mode = None) is None + assert apply_step_cache(_pipe(t), mode = "off") is None + assert t.enabled_with is None + + +def test_enable_cache_path_default_threshold(monkeypatch): + _stub_diffusers(monkeypatch) + t = _MixinTransformer() + engaged = apply_step_cache(_pipe(t), mode = "fbcache") + assert engaged == TC_FBCACHE + assert t.enabled_with.threshold == DEFAULT_FBCACHE_THRESHOLD + assert t._unsloth_step_cache == f"fbcache@{DEFAULT_FBCACHE_THRESHOLD}" + + +def test_quant_active_raises_default_threshold(monkeypatch): + _stub_diffusers(monkeypatch) + t = _MixinTransformer() + apply_step_cache(_pipe(t), mode = "fbcache", quant_active = True) + assert t.enabled_with.threshold == QUANT_FBCACHE_THRESHOLD + + +def test_explicit_threshold_overrides_quant(monkeypatch): + _stub_diffusers(monkeypatch) + t = _MixinTransformer() + apply_step_cache(_pipe(t), mode = "fbcache", threshold = 0.2, quant_active = True) + assert t.enabled_with.threshold == 0.2 + + +def test_fallback_to_standalone_hook(monkeypatch): + rec: dict = {} + _stub_diffusers(monkeypatch, hook_recorder = rec) + t = _HookTransformer() + engaged = apply_step_cache(_pipe(t), mode = "fbcache") + assert engaged == TC_FBCACHE + assert rec["transformer"] is t + assert rec["config"].threshold == DEFAULT_FBCACHE_THRESHOLD + + +def test_incompatible_model_runs_uncached(monkeypatch): + # enable_cache raising (e.g. unrecognised block signature) must not fail the load. + _stub_diffusers(monkeypatch) + t = _MixinTransformer(fail = True) + assert apply_step_cache(_pipe(t), mode = "fbcache") is None + + +def test_missing_transformer_is_none(monkeypatch): + _stub_diffusers(monkeypatch) + pipe = types.SimpleNamespace(transformer = None) + assert apply_step_cache(pipe, mode = "fbcache") is None + + +def test_diffusers_unavailable_runs_uncached(monkeypatch): + # no diffusers import -> best-effort returns None, load proceeds uncached. + monkeypatch.setitem(sys.modules, "diffusers", None) + t = _MixinTransformer() + assert apply_step_cache(_pipe(t), mode = "fbcache") is None diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index cf1d4266b6..c9c83009ee 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -382,6 +382,39 @@ def test_invalid_attention_backend_returns_422(client): assert resp.status_code == 422 +def test_transformer_cache_threads_through(client, monkeypatch): + backend = _FakeBackend() + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + resp = client.post( + "/api/inference/images/load", + json = { + "model_path": "x/z-image", + "gguf_filename": "q.gguf", + "transformer_cache": "fbcache", + "transformer_cache_threshold": 0.1, + }, + ) + assert resp.status_code == 200 + assert backend.last_load_kwargs.get("transformer_cache") == "fbcache" + assert backend.last_load_kwargs.get("transformer_cache_threshold") == 0.1 + + +def test_invalid_transformer_cache_returns_422(client): + resp = client.post( + "/api/inference/images/load", + json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "transformer_cache": "deepcache"}, + ) + assert resp.status_code == 422 + + +def test_out_of_range_cache_threshold_returns_422(client): + resp = client.post( + "/api/inference/images/load", + json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "transformer_cache_threshold": 1.5}, + ) + assert resp.status_code == 422 + + def test_invalid_transformer_quant_returns_422_without_eviction(client): # An unsupported transformer_quant is rejected by the request schema (Literal), so # the GPU is never acquired and no chat model is evicted.