Studio diffusion (Phase 12): only engage FBCache on context-aware transformers; quantized threshold for GGUF

- apply_step_cache now engages only via the transformer's native enable_cache (the diffusers
  CacheMixin path), which exists exactly when the pipeline wraps the transformer call in a
  cache_context. The standalone apply_first_block_cache fallback installed on non-CacheMixin
  transformers too (e.g. Z-Image), whose pipeline opens no cache_context, so the load reported
  transformer_cache=fbcache and then the first generation crashed inside the hook. Such a model
  now runs uncached per the best-effort contract.
- GGUF transformers are quantized (the default Studio load path), so they now use the higher
  quantized FBCache threshold when the caller leaves it unset, instead of the dense default
  that could keep the cache from triggering.
- fbcache_flux_probe.py: compile cached runs with fullgraph=False (FBCache is a graph break, so
  fullgraph=True failed warmup and silently measured an eager cached run); output dir is now
  relative to the script, not a hardcoded path.
This commit is contained in:
Daniel Han 2026-06-28 06:02:12 +00:00
commit 1d289aab61
4 changed files with 37 additions and 21 deletions

View file

@ -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: