diffusion: address review round (FBCache context guard, aiter/ROCm, video cleanup, prequant + ControlNet gating)

- diffusion_cache: do not engage FBCache when the selected pipeline opens no cache_context.
  A CacheMixin transformer is necessary but not sufficient -- Flux Kontext / img2img /
  inpaint / controlnet reuse the CacheMixin FluxTransformer2DModel yet their __call__ never
  opens a cache_context, so the First-Block-Cache hook raised 'No context is set' on the
  first forward, crashing every default FLUX.1-Kontext edit (28 steps, above the FBCache
  threshold). Detect it from the pipeline __call__ source, resolved off the instance so the
  per-expert proxy view delegates to the real pipe.
- diffusion_attention: honor an explicit aiter backend on ROCm/AMD targets instead of
  dropping it via the NVIDIA-only guard (aiter is the AMD ROCm kernel; it only works there).
- video: clear the CUDA cache on a failed load so a partially built pipeline's reserved VRAM
  does not OOM the next load (mirrors the image backend), and re-check cancellation after the
  export/mux so a clip cancelled during the blocking encode is discarded, not persisted.
- diffusion_auto_policy / diffusion_prequant: validate a request-supplied prequant path
  override (present AND allowlisted) before budgeting the small prequant plan, so the loader
  does not skip the dense shards and then rebuild dense after evicting the resident pipeline.
- diffusion_controlnet: family-gate a curated ControlNet addressed by its full repo id, not
  only its short catalog id, so a cross-family repo id 400s up front instead of downloading
  and loading through the wrong ControlNet class.
This commit is contained in:
Daniel Han 2026-07-09 08:52:16 +00:00
commit c00eb20958
12 changed files with 193 additions and 13 deletions

View file

@ -68,8 +68,32 @@ class _NonCacheMixinTransformer:
the load runs uncached instead (e.g. Z-Image)."""
class _CtxPipe:
"""A pipeline whose denoise loop opens ``transformer.cache_context(...)`` (like FluxPipeline)
-- the First-Block-Cache hook needs it, so FBCache may engage here."""
def __init__(self, transformer):
self.transformer = transformer
def __call__(self, *args, **kwargs):
with self.transformer.cache_context("cond"):
return None
class _NoCtxPipe:
"""A pipeline that never enters a caching context (like FluxKontextPipeline / img2img /
inpaint / controlnet, which reuse the CacheMixin FluxTransformer2DModel): FBCache must NOT
engage or the hook raises "No context is set" on the first forward."""
def __init__(self, transformer):
self.transformer = transformer
def __call__(self, *args, **kwargs):
return None
def _pipe(transformer):
return types.SimpleNamespace(transformer = transformer)
return _CtxPipe(transformer)
def _stub_diffusers(monkeypatch, *, hook_recorder = None):
@ -129,6 +153,17 @@ def test_non_cachemixin_runs_uncached(monkeypatch):
assert rec == {} # the standalone hook was never called
def test_pipeline_without_cache_context_runs_uncached(monkeypatch):
# A CacheMixin transformer whose PIPELINE never opens a cache_context (Flux Kontext /
# img2img / inpaint / controlnet reuse the CacheMixin FluxTransformer2DModel) must run
# uncached -- otherwise the First-Block-Cache hook raises "No context is set" on the
# first forward, crashing every default generation.
_stub_diffusers(monkeypatch)
t = _MixinTransformer()
assert apply_step_cache(_NoCtxPipe(t), mode = "fbcache") is None
assert t.enabled_with is None # enable_cache was never called
def test_incompatible_model_runs_uncached(monkeypatch):
# enable_cache raising (e.g. unrecognised block signature) must not fail the load.
_stub_diffusers(monkeypatch)