diff --git a/scripts/fbcache_flux_probe.py b/scripts/fbcache_flux_probe.py index 31d96a1c40..c0462f1947 100644 --- a/scripts/fbcache_flux_probe.py +++ b/scripts/fbcache_flux_probe.py @@ -18,7 +18,7 @@ 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") +OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "fbcache_flux_images" _LP = {"fn": None} @@ -101,8 +101,13 @@ def run( from diffusers.hooks import apply_first_block_cache apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold = threshold)) if compile_: + # FBCache's per-step decision is a graph break, so a cached run must compile with + # fullgraph=False (mirroring the production path); fullgraph=True would fail the + # warmup compile and the row would silently fall back to an eager cached run, + # producing misleading speedup numbers. + fullgraph = threshold is None try: - pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True) + pipe.transformer.compile_repeated_blocks(fullgraph = fullgraph, dynamic = True) except Exception as exc: # noqa: BLE001 print(f" [{tag}] compile {type(exc).__name__}: {str(exc)[:80]}", flush = True) try: diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 4e8c090a39..bd45439dc1 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -588,7 +588,10 @@ class DiffusionBackend: pipe, mode = transformer_cache, threshold = transformer_cache_threshold, - quant_active = transformer_quant_engaged is not None, + # GGUF transformers are quantized too (the default Studio path), so the + # cache needs the higher quantized threshold to still trigger -- not just + # the dense-quant fast path. + quant_active = transformer_quant_engaged is not None or bool(gguf_filename), logger = logger, ) speed_applied = apply_speed_optims( diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py index 4fd5815e16..b7a2b7ac45 100644 --- a/studio/backend/core/inference/diffusion_cache.py +++ b/studio/backend/core/inference/diffusion_cache.py @@ -77,16 +77,22 @@ def apply_step_cache( if threshold is not None else (QUANT_FBCACHE_THRESHOLD if quant_active else DEFAULT_FBCACHE_THRESHOLD) ) + # Only engage via the transformer's native enable_cache (the diffusers CacheMixin path). + # That mixin is present exactly when the pipeline wraps the transformer call in a + # cache_context, which the First-Block-Cache hook requires at run time. The lower-level + # apply_first_block_cache hook would install on a non-CacheMixin transformer too (e.g. + # Z-Image), but its pipeline opens no cache_context, so the first generation would crash + # inside the hook -- so a model without enable_cache runs uncached per the best-effort + # contract instead of being reported as cached and then failing. + enable_cache = getattr(transformer, "enable_cache", None) + if not callable(enable_cache): + _warn(logger, mode, RuntimeError("transformer has no cache_context (not a CacheMixin)")) + return None 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) + enable_cache(config) try: transformer._unsloth_step_cache = f"{mode}@{thr}" except Exception: # noqa: BLE001 — marker is best-effort diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py index 40d49251b9..62071d9aa6 100644 --- a/studio/backend/tests/test_diffusion_cache.py +++ b/studio/backend/tests/test_diffusion_cache.py @@ -4,10 +4,9 @@ """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. +``FirstBlockCacheConfig`` lazily), and the pipeline is a fake that records the engaged config. +So normalisation, the CacheMixin (``enable_cache``) gating, threshold selection, and the +best-effort failure handling are all exercised without torch or a real diffusers model. """ from __future__ import annotations @@ -62,8 +61,11 @@ class _MixinTransformer: self.enabled_with = config -class _HookTransformer: - """A transformer with no ``enable_cache`` -> the standalone hook is used.""" +class _NonCacheMixinTransformer: + """A transformer with no ``enable_cache`` (not a CacheMixin) -> must run uncached. + + Its pipeline opens no ``cache_context``, so installing FBCache would crash at generation; + the load runs uncached instead (e.g. Z-Image).""" def _pipe(transformer): @@ -117,14 +119,14 @@ def test_explicit_threshold_overrides_quant(monkeypatch): assert t.enabled_with.threshold == 0.2 -def test_fallback_to_standalone_hook(monkeypatch): +def test_non_cachemixin_runs_uncached(monkeypatch): + # A transformer without enable_cache (e.g. Z-Image) must NOT install the standalone hook + # -- its pipeline opens no cache_context, so it runs uncached instead of crashing at gen. 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 + t = _NonCacheMixinTransformer() + assert apply_step_cache(_pipe(t), mode = "fbcache") is None + assert rec == {} # the standalone hook was never called def test_incompatible_model_runs_uncached(monkeypatch):