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:
parent
3b659ba075
commit
c00eb20958
12 changed files with 193 additions and 13 deletions
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue