From c00eb2095893d95f9d7fa4e7a90b710f53d193d1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 08:52:16 +0000 Subject: [PATCH] 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. --- .../core/inference/diffusion_attention.py | 8 ++++ .../core/inference/diffusion_auto_policy.py | 13 ++++-- .../backend/core/inference/diffusion_cache.py | 41 +++++++++++++++---- .../core/inference/diffusion_controlnet.py | 5 +++ .../core/inference/diffusion_prequant.py | 13 ++++++ studio/backend/core/inference/video.py | 14 +++++++ .../backend/tests/test_diffusion_attention.py | 15 +++++++ .../tests/test_diffusion_auto_policy.py | 3 +- studio/backend/tests/test_diffusion_cache.py | 37 ++++++++++++++++- .../tests/test_diffusion_controlnet.py | 10 +++++ .../backend/tests/test_diffusion_prequant.py | 16 ++++++++ studio/backend/tests/test_video_backend.py | 31 ++++++++++++++ 12 files changed, 193 insertions(+), 13 deletions(-) diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index 8e60e1a5f0..4f228a42ef 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -129,6 +129,14 @@ def select_attention_backend( backend = _ALIASES[alias] if backend == "native": return None + # AITER is the AMD ROCm kernel, not an NVIDIA one: honor it on a ROCm (AMD) CUDA + # target and drop it everywhere else (diffusers' own set-time check rejects it off + # ROCm anyway). Without this special-case the NVIDIA-only guard below would silently + # drop the one explicit backend that only ever works on ROCm. + if backend == "aiter": + if getattr(target, "device", None) == "cuda" and not _is_cuda_nvidia(target): + return backend + return None # Every explicit kernel here (cuDNN / flash* / sage) is CUDA+NVIDIA-only; on # ROCm / MPS / CPU diffusers accepts the name at set time and the first # generation crashes, so drop to the native default up front. diff --git a/studio/backend/core/inference/diffusion_auto_policy.py b/studio/backend/core/inference/diffusion_auto_policy.py index 85653364dd..77df14464c 100644 --- a/studio/backend/core/inference/diffusion_auto_policy.py +++ b/studio/backend/core/inference/diffusion_auto_policy.py @@ -188,10 +188,15 @@ def resolve_dense_quant_candidate( return None prequant_available = False try: - from .diffusion_prequant import resolve_prequant_source - prequant_available = ( - resolve_prequant_source(fam, scheme, path_override = prequant_path) is not None - ) + from .diffusion_prequant import local_prequant_path_ready, resolve_prequant_source + src = resolve_prequant_source(fam, scheme, path_override = prequant_path) + # A request-supplied local path override is only usable if the loader will accept it + # (allowlisted AND present); otherwise load_prequantized_transformer refuses it and + # rebuilds dense after the resident pipe is unloaded -- the evict-then-OOM this + # small-plan prefetch exists to avoid. Hosted-repo sources keep the existing signal. + if src is not None and src.kind == "path" and not local_prequant_path_ready(src.location): + src = None + prequant_available = src is not None except Exception: # noqa: BLE001 -- prequant probing must never sink the candidate prequant_available = False estimate = estimate_dense_quant( diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py index 758719da95..05c781a30f 100644 --- a/studio/backend/core/inference/diffusion_cache.py +++ b/studio/backend/core/inference/diffusion_cache.py @@ -64,6 +64,27 @@ def normalize_transformer_cache(value: Optional[str]) -> Optional[str]: return normalized +def _pipeline_opens_cache_context(pipe: Any) -> bool: + """Whether the pipeline enters ``transformer.cache_context(...)`` in its denoise loop. + The First-Block-Cache hook requires it at run time, and a CacheMixin transformer alone + does NOT guarantee it: Flux Kontext / img2img / inpaint / controlnet reuse the CacheMixin + FluxTransformer2DModel but never open a cache_context. Read from the pipeline ``__call__`` + source, resolved off the instance so a per-expert proxy view (``_SecondDiTView``) + delegates to the real pipe; if it cannot be read, report False so the cache stays off.""" + import inspect + + call = getattr(pipe, "__call__", None) + if call is None: + return False + try: + src = inspect.getsource(call) + except (OSError, TypeError): + return False + # Match the actual call `cache_context(` -- a bare mention in a comment/docstring lacks + # the paren, so this does not false-positive on prose. + return "cache_context(" in src + + def apply_step_cache( pipe: Any, *, @@ -89,17 +110,23 @@ 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. + # Engage only via the transformer's native enable_cache (the diffusers CacheMixin path): + # the lower-level apply_first_block_cache hook would install on a non-CacheMixin + # transformer too (e.g. Z-Image), whose pipeline opens no cache_context and would crash + # the first generation -- 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 + # A CacheMixin transformer is necessary but NOT sufficient: the First-Block-Cache hook + # raises "No context is set" on the first forward unless the PIPELINE wraps its denoise + # loop in transformer.cache_context(...). Flux Kontext / img2img / inpaint / controlnet + # reuse the CacheMixin FluxTransformer2DModel yet their __call__ opens no cache_context, + # so engaging FBCache there would crash every default generation -- run uncached instead. + if not _pipeline_opens_cache_context(pipe): + _warn(logger, mode, RuntimeError("pipeline __call__ opens no cache_context; running uncached")) + return None try: try: from diffusers import FirstBlockCacheConfig diff --git a/studio/backend/core/inference/diffusion_controlnet.py b/studio/backend/core/inference/diffusion_controlnet.py index eab7e8b92f..5f5a4b4b26 100644 --- a/studio/backend/core/inference/diffusion_controlnet.py +++ b/studio/backend/core/inference/diffusion_controlnet.py @@ -171,6 +171,11 @@ def resolve_controlnet(spec_id: str, *, family: Optional[str] = None) -> Resolve with a clear error rather than being loaded through the wrong pipeline class later. """ entry = _catalog_by_id().get(spec_id) + if entry is None: + # A curated entry addressed by its full repo id (owner/name) rather than its catalog + # id must still hit the family gate below, not slip through to the bare-repo fallback + # and get downloaded + loaded through the wrong family's ControlNet class. + entry = next((e for e in _CURATED if e.repo_id and e.repo_id == spec_id), None) if entry is not None: # A curated/local entry may declare the families it is built for. A client that # bypasses the UI filter (direct API call) could send an entry for another family; diff --git a/studio/backend/core/inference/diffusion_prequant.py b/studio/backend/core/inference/diffusion_prequant.py index 903e323d0f..2bfc72dfbb 100644 --- a/studio/backend/core/inference/diffusion_prequant.py +++ b/studio/backend/core/inference/diffusion_prequant.py @@ -86,6 +86,19 @@ def _local_prequant_path_allowed(path: str) -> bool: return any(real == r or real.startswith(r + os.sep) for r in roots) +def local_prequant_path_ready(path: str) -> bool: + """True only when a request-supplied local pre-quant path would actually load: it is + inside an allowlisted root AND the checkpoint file is present. The auto-policy planner + uses this before budgeting the small prequant plan, so it never skips the dense shards + for a path ``load_prequantized_transformer`` will refuse -- which would otherwise evict + the resident pipeline and then rebuild dense under an undersized plan (OOM).""" + import os + + if not _local_prequant_path_allowed(path): + return False + return os.path.isfile(os.path.expanduser(path)) + + @dataclass(frozen = True) class PrequantSource: """Where a pre-quantized transformer checkpoint lives. ``kind`` is "path" (a local diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index dc536bc4d6..53e9fe70a6 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -671,6 +671,15 @@ class VideoBackend: if self._load_token != token: return logger.error("video.load_failed: %s", exc) + # Free the debris of a failed construction (mirrors diffusion.py's _run_load): + # no _VideoLoadState was committed, so no later unload releases the VRAM a + # partially built pipeline (OOM in from_pretrained / quant / placement) left + # reserved in the caching allocator -- which would OOM the next load. Guarded so + # a sticky CUDA error cannot skip stamping the real error below. + try: + clear_gpu_cache() + except Exception: # noqa: BLE001 -- cleanup is best-effort + pass from utils.native_path_leases import redact_native_paths with self._lock: @@ -1644,6 +1653,11 @@ class VideoBackend: mp4_bytes = self._encode_mp4( video_frames, out_fps, audio_track, pipe if fam.has_audio else None ) + # A cancel that landed during the (blocking, uncancellable) export/mux must + # still discard the clip: cancel_generate() already reported success for it, + # so re-check here before it is returned and persisted to the gallery. + if cancel.is_set(): + raise RuntimeError(VIDEO_CANCELLED_MSG) duration_s = len(video_frames) / float(out_fps) if out_fps else 0.0 self._gen = {"active": False} return { diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py index da200f2c48..1f0fff7ccc 100644 --- a/studio/backend/tests/test_diffusion_attention.py +++ b/studio/backend/tests/test_diffusion_attention.py @@ -92,6 +92,21 @@ def test_explicit_backend_dropped_off_nvidia_cuda(monkeypatch): assert select_attention_backend(_target(device = "mps"), alias, speed_active = True) is None +def test_aiter_honored_on_rocm(monkeypatch): + # AITER is the AMD ROCm kernel; on a ROCm CUDA target it must be honored, not dropped by + # the NVIDIA-only guard -- it is the one explicit backend that only ever works on ROCm. + monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: False) # hip build + assert select_attention_backend(_target(), "aiter", speed_active = False) == "aiter" + + +def test_aiter_dropped_off_rocm(monkeypatch): + # aiter on NVIDIA CUDA (or MPS / CPU) is not usable, so it drops to the native default. + monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: True) # NVIDIA + assert select_attention_backend(_target(), "aiter", speed_active = False) is None + monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: False) + assert select_attention_backend(_target(device = "mps"), "aiter", speed_active = False) is None + + def test_explicit_native_returns_none(): # native is the default -> nothing to set. assert select_attention_backend(_target(), "native", speed_active = True) is None diff --git a/studio/backend/tests/test_diffusion_auto_policy.py b/studio/backend/tests/test_diffusion_auto_policy.py index 53c902164c..239360a82d 100644 --- a/studio/backend/tests/test_diffusion_auto_policy.py +++ b/studio/backend/tests/test_diffusion_auto_policy.py @@ -169,7 +169,8 @@ def test_candidate_none_for_an_unlisted_family(monkeypatch): def test_candidate_uses_prequant_transient_when_available(monkeypatch): - _patch_selector(monkeypatch, prequant = object()) + # A hosted-repo prequant source (kind="repo") is available without a local-path check. + _patch_selector(monkeypatch, prequant = SimpleNamespace(kind = "repo", location = "org/int8")) est = resolve_dense_quant_candidate(fam = _fam("z-image"), target = object(), requested = "int8") assert est is not None and est.prequant is True assert est.transient_transformer_mib == est.steady_transformer_mib diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py index 06838edd2c..1794482239 100644 --- a/studio/backend/tests/test_diffusion_cache.py +++ b/studio/backend/tests/test_diffusion_cache.py @@ -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) diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py index b8ac050db6..184c279a1c 100644 --- a/studio/backend/tests/test_diffusion_controlnet.py +++ b/studio/backend/tests/test_diffusion_controlnet.py @@ -58,6 +58,16 @@ def test_resolve_controlnet_enforces_family_match(): assert dc.resolve_controlnet("qwen-union").path +def test_resolve_controlnet_repo_id_still_family_gated(): + # A curated ControlNet addressed by its full repo id (not its short catalog id) must still + # hit the family gate, not slip through the bare-repo fallback and load through the wrong + # family's ControlNet class. + with pytest.raises(ValueError, match = "is for"): + dc.resolve_controlnet("InstantX/Qwen-Image-ControlNet-Union", family = "flux.1") + r = dc.resolve_controlnet("InstantX/Qwen-Image-ControlNet-Union", family = "qwen-image") + assert r.path == "InstantX/Qwen-Image-ControlNet-Union" and not r.is_local + + def test_union_control_mode_maps_only_union_entries(): # Union entries map a known control type to its integer mode; a union model always # needs a concrete mode, so an unmapped type (passthrough) defaults to 0. A non-union diff --git a/studio/backend/tests/test_diffusion_prequant.py b/studio/backend/tests/test_diffusion_prequant.py index 8b11cbd98f..b646933e06 100644 --- a/studio/backend/tests/test_diffusion_prequant.py +++ b/studio/backend/tests/test_diffusion_prequant.py @@ -59,6 +59,22 @@ def test_resolve_nothing_configured_is_none(): assert resolve_prequant_source(_fam(), "fp8", path_override = "") is None +def test_local_prequant_path_ready(tmp_path, monkeypatch): + # The auto-policy planner budgets the small prequant plan only when a request-supplied + # path would actually load: present AND inside an allowlisted root. Missing or not + # allowlisted -> not ready, else the loader refuses it and rebuilds dense after evict. + import os + + ckpt = tmp_path / "model.pt" + ckpt.write_bytes(b"x") + root = os.path.realpath(str(tmp_path)) + monkeypatch.setattr(pq, "_allowed_prequant_roots", lambda: [root]) + assert pq.local_prequant_path_ready(str(ckpt)) is True + assert pq.local_prequant_path_ready(str(tmp_path / "missing.pt")) is False + monkeypatch.setattr(pq, "_allowed_prequant_roots", lambda: []) + assert pq.local_prequant_path_ready(str(ckpt)) is False + + # ── load_prequantized_transformer ──────────────────────────────────────────────── class _FakeTransformer: calls: dict = {} diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index c6928ba936..8d8e9824fe 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -151,6 +151,12 @@ class _FakeWanDiT: def set_attention_backend(self, backend) -> None: self.attention = backend + @contextlib.contextmanager + def cache_context(self, name): + # Real Wan / HV15 / LTX pipelines open a cache_context around the denoise loop; the + # First-Block-Cache hook needs it, so the fake transformer provides it too. + yield + class _FakeWanVae: def __init__(self) -> None: @@ -229,6 +235,8 @@ class _FakeWanPipeSingle(_FakeWanPipeBase): "num_frames": num_frames, **kwargs, } + with self.transformer.cache_context("cond"): # real Wan pipeline wraps the denoise loop + pass return self._finish(num_inference_steps, num_frames, callback_on_step_end) @@ -264,6 +272,8 @@ class _FakeWanPipeMoE(_FakeWanPipeBase): "num_frames": num_frames, **kwargs, } + with self.transformer.cache_context("cond"): # real Wan pipeline wraps the denoise loop + pass return self._finish(num_inference_steps, num_frames, callback_on_step_end) @@ -344,6 +354,8 @@ class _FakeHV15Pipe: "num_frames": num_frames, **kwargs, } + with self.transformer.cache_context("cond"): # real HV15 pipeline wraps the denoise loop + pass for _ in range(int(num_inference_steps or 1)): self.scheduler.step() frames = [[object() for _ in range(int(num_frames or 1))]] @@ -1062,6 +1074,25 @@ def test_hv15_cancel_unwinds_scheduler_loop(fake_runtime): assert pipe.hooks_freed == 1 +def test_cancel_during_export_discards_clip(fake_runtime, monkeypatch): + # A cancel that lands during the (blocking, uncancellable) export/mux must still discard + # the clip: cancel_generate() already reported success for it, so generate() must raise + # the cancelled sentinel rather than return the clip to be persisted to the gallery. + backend = VideoBackend() + backend.load_pipeline( + "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v", + model_kind = "pipeline", + ) + + def _encode_and_cancel(frames, fps, audio, pipe): + backend.cancel_generate() # cancel arrives mid-mux, after the last denoise-step check + return b"MP4" + + monkeypatch.setattr(VideoBackend, "_encode_mp4", staticmethod(_encode_and_cancel)) + with pytest.raises(RuntimeError, match = VIDEO_CANCELLED_MSG): + backend.generate(prompt = "a fox", steps = 4) + + def test_singleton(): assert get_video_backend() is get_video_backend()