diffusion: fix step-cache Off, kernel-install dep clobber, retry-under-lock, effective steps
- Step cache 'Off' is preserved: the frontend defaulted to 'off' and mapped it to an omitted transformer_cache, which the backend now reads as 'auto', so leaving the control at Off silently enabled FBCache on 20+ step families. Default the control to Auto, add an explicit Auto option, and send auto -> omitted so Off maps to an explicit cache-off. - Kernel auto-install adds --no-deps: 'pip install --only-binary :all: xformers' resolves xformers' pinned torch and replaces the running torch/triton. --no-deps installs only the best-effort kernel wheel; an ABI mismatch just fails to import and falls back to native, never clobbering core deps. - Do not retry a failed kernel install under the load lock: the pre-install runs outside the locks, then the in-lock resolve re-attempts pip (up to 600s) while holding _generate_lock/_lock and blocking unload/cancel/new loads. Record the attempt in a process-level set so the in-lock call short-circuits to native. - Cache auto-toggle keys on effective denoise steps: an image-conditioned run with strength < 1 (upscale default 0.35) denoises a fraction of the requested steps, so a 28-step request runs ~10 steps. Compute the effective count the way diffusers get_timesteps does and gate FBCache on it, only when strength is actually applied.
This commit is contained in:
Daniel Han
parent
4373ebc13b
commit
7c8f4919d6
6 changed files with 146 additions and 6 deletions
|
|
@ -78,6 +78,7 @@ from .diffusion_cache import (
|
|||
TC_AUTO,
|
||||
TC_FBCACHE,
|
||||
apply_step_cache,
|
||||
effective_denoise_steps,
|
||||
maybe_toggle_step_cache,
|
||||
normalize_transformer_cache,
|
||||
)
|
||||
|
|
@ -2273,9 +2274,26 @@ class DiffusionBackend:
|
|||
# schedule kept it off, and a few-step turbo request drops it (skipping
|
||||
# a step there is a large quality hit). Explicit choices never toggle.
|
||||
if state.cache_auto:
|
||||
# Key the policy on the EFFECTIVE denoise steps: an img2img/upscale/
|
||||
# inpaint request at strength < 1 only denoises a fraction of `steps`
|
||||
# (e.g. a 28-step upscale at strength 0.35 runs ~10 steps), so passing
|
||||
# the raw request would wrongly engage FBCache on exactly the short
|
||||
# trajectory the policy keeps uncached. Only fold in `strength` when it
|
||||
# is ACTUALLY applied to the pipe (same gate as the kwarg below), so a
|
||||
# stray strength on a txt2img request never shortens the count.
|
||||
strength_applied = (
|
||||
strength
|
||||
if (
|
||||
strength is not None
|
||||
and init_pil is not None
|
||||
and "strength" in call_params
|
||||
)
|
||||
else None
|
||||
)
|
||||
denoise_steps = effective_denoise_steps(steps, strength_applied)
|
||||
toggled = maybe_toggle_step_cache(
|
||||
state.pipe,
|
||||
steps = steps,
|
||||
steps = denoise_steps,
|
||||
quant_active = state.cache_quant_active,
|
||||
threshold = state.cache_threshold,
|
||||
logger = logger,
|
||||
|
|
@ -2289,7 +2307,7 @@ class DiffusionBackend:
|
|||
if isinstance(entry, dict):
|
||||
entry["value"] = toggled or "off"
|
||||
entry["reason"] = (
|
||||
f"auto: {steps}-step generation "
|
||||
f"auto: {denoise_steps}-step generation "
|
||||
+ ("reaches" if toggled else "is below")
|
||||
+ f" {FBCACHE_MIN_STEPS}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -178,6 +178,15 @@ _INSTALLABLE_BACKENDS: dict[str, tuple[str, str]] = {
|
|||
# 0 - never install; a missing kernel falls back to native
|
||||
_ATTENTION_INSTALL_ENV = "UNSLOTH_DIFFUSION_ATTENTION_INSTALL"
|
||||
|
||||
# Packages a pip install has already been attempted for in THIS process (success or
|
||||
# failure). The loader pre-installs the kernel OUTSIDE its locks and then re-resolves the
|
||||
# same backend under _generate_lock, where apply_attention_backend would otherwise call
|
||||
# pip a SECOND time -- for a package with no matching wheel / an offline host that repeat
|
||||
# runs the full (up to 600s) install while holding the load lock, blocking unload/cancel/
|
||||
# new loads for exactly the failure the pre-install was added to keep off the lock. Record
|
||||
# each attempt so a retry is a no-op and set_attention_backend falls back to native at once.
|
||||
_INSTALL_ATTEMPTED: set[str] = set()
|
||||
|
||||
|
||||
def _ensure_attention_backend_installed(backend: str, logger: Any = None) -> None:
|
||||
"""Best-effort wheel-only install of the package ``backend`` needs, when allowed.
|
||||
|
|
@ -202,6 +211,14 @@ def _ensure_attention_backend_installed(backend: str, logger: Any = None) -> Non
|
|||
return
|
||||
except Exception: # noqa: BLE001 — a broken install probes as missing; try the install
|
||||
pass
|
||||
# Only ever attempt each package's install once per process. The loader pre-installs
|
||||
# this backend outside its locks; if that failed (no wheel / offline) the module is
|
||||
# still missing here, so without this guard the in-lock apply path would re-run the
|
||||
# whole install under _generate_lock and block unload/cancel. A recorded attempt makes
|
||||
# the retry a no-op -> set_attention_backend raises on the missing package -> native.
|
||||
if package in _INSTALL_ATTEMPTED:
|
||||
return
|
||||
_INSTALL_ATTEMPTED.add(package)
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
|
|
@ -211,7 +228,16 @@ def _ensure_attention_backend_installed(backend: str, logger: Any = None) -> Non
|
|||
)
|
||||
try:
|
||||
subprocess.run(
|
||||
[sys.executable, "-m", "pip", "install", "--only-binary", ":all:", package],
|
||||
# --no-deps: install ONLY this best-effort kernel wheel, never its declared
|
||||
# dependencies. xformers/flash-attn pin an exact torch (e.g. torch==2.x), so
|
||||
# normal resolution would upgrade/replace the running torch/triton and leave
|
||||
# later loads on a different, possibly CUDA-mismatched dependency stack. Without
|
||||
# its deps an ABI-incompatible kernel simply fails to import -> native fallback,
|
||||
# which is the same best-effort outcome as an uninstallable wheel.
|
||||
[
|
||||
sys.executable, "-m", "pip", "install",
|
||||
"--only-binary", ":all:", "--no-deps", package,
|
||||
],
|
||||
capture_output = True,
|
||||
timeout = 600,
|
||||
check = True,
|
||||
|
|
|
|||
|
|
@ -126,6 +126,24 @@ def apply_step_cache(
|
|||
return None
|
||||
|
||||
|
||||
def effective_denoise_steps(steps: int, strength: Optional[float]) -> int:
|
||||
"""The number of steps diffusers ACTUALLY denoises for a request.
|
||||
|
||||
An image-conditioned workflow with ``strength`` < 1 (img2img / upscale / inpaint) runs
|
||||
only a fraction of ``num_inference_steps``: diffusers' ``get_timesteps`` computes
|
||||
``init_timestep = min(num_inference_steps * strength, num_inference_steps)`` and denoises
|
||||
``num_inference_steps - int(num_inference_steps - init_timestep)`` steps. The auto
|
||||
step-cache policy must key on THIS count -- e.g. a 28-step upscale at strength 0.35 runs
|
||||
~10 real steps, exactly the short trajectory FBCache should stay off (each skipped step
|
||||
is a large quality hit). ``strength`` None (txt2img / reference) or >= 1 -> the full count.
|
||||
"""
|
||||
s = int(steps)
|
||||
if strength is None or float(strength) >= 1.0:
|
||||
return s
|
||||
init = min(s * float(strength), s)
|
||||
return max(1, s - int(max(s - init, 0)))
|
||||
|
||||
|
||||
def maybe_toggle_step_cache(
|
||||
pipe: Any,
|
||||
*,
|
||||
|
|
|
|||
|
|
@ -241,6 +241,10 @@ def _no_real_installs(monkeypatch):
|
|||
# backends (sage/flash*), so hard-disable the gate; install tests re-enable it
|
||||
# with a stubbed subprocess.
|
||||
monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "0")
|
||||
# The install once-per-process memo is module state; clear it so each test starts
|
||||
# with a fresh "not yet attempted" set (otherwise an earlier test's attempt would
|
||||
# make a later install a no-op).
|
||||
att._INSTALL_ATTEMPTED.clear()
|
||||
|
||||
|
||||
class _Recorder:
|
||||
|
|
@ -290,6 +294,45 @@ def test_install_runs_wheel_only_for_missing_kernel(monkeypatch):
|
|||
assert "--only-binary" in cmd and ":all:" in cmd and "sageattention" in cmd
|
||||
|
||||
|
||||
def test_install_uses_no_deps_to_protect_core_deps(monkeypatch):
|
||||
# A kernel add-on (xformers/flash-attn) pins an exact torch, so a normal install would
|
||||
# upgrade/replace the running torch/triton. --no-deps installs only the kernel wheel;
|
||||
# an ABI-incompatible one fails to import and falls back to native rather than clobbering
|
||||
# the environment's core deps.
|
||||
monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "auto")
|
||||
import importlib.util
|
||||
|
||||
monkeypatch.setattr(importlib.util, "find_spec", lambda name: None)
|
||||
run = _Recorder()
|
||||
_stub_subprocess(monkeypatch, run)
|
||||
att._ensure_attention_backend_installed("xformers")
|
||||
assert len(run.calls) == 1
|
||||
assert "--no-deps" in run.calls[0]
|
||||
|
||||
|
||||
def test_failed_install_not_retried_in_same_process(monkeypatch):
|
||||
# The loader pre-installs the kernel OUTSIDE its locks and then re-resolves the same
|
||||
# backend under _generate_lock; if the pre-install failed (no wheel / offline) the
|
||||
# in-lock apply path must NOT re-run pip (a second up-to-600s install holding the load
|
||||
# lock blocks unload/cancel). The once-per-process memo makes the retry a no-op.
|
||||
monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "auto")
|
||||
import importlib.util
|
||||
import subprocess as sp
|
||||
|
||||
monkeypatch.setattr(importlib.util, "find_spec", lambda name: None) # stays missing
|
||||
|
||||
calls: list[list[str]] = []
|
||||
|
||||
def _boom(cmd, **kwargs):
|
||||
calls.append(list(cmd))
|
||||
raise sp.CalledProcessError(returncode = 1, cmd = cmd)
|
||||
|
||||
_stub_subprocess(monkeypatch, _boom)
|
||||
att._ensure_attention_backend_installed("sage") # pre-install attempt (outside lock)
|
||||
att._ensure_attention_backend_installed("sage") # in-lock retry -> must be skipped
|
||||
assert len(calls) == 1
|
||||
|
||||
|
||||
def test_install_invalidates_import_caches_on_success(monkeypatch):
|
||||
# A wheel written to site-packages after the finder cached that directory can be
|
||||
# missed by the very next import, so a successful install must invalidate the caches
|
||||
|
|
|
|||
|
|
@ -179,10 +179,44 @@ def test_diffusers_unavailable_runs_uncached(monkeypatch):
|
|||
from core.inference.diffusion_cache import ( # noqa: E402
|
||||
FBCACHE_MIN_STEPS,
|
||||
TC_AUTO,
|
||||
effective_denoise_steps,
|
||||
maybe_toggle_step_cache,
|
||||
)
|
||||
|
||||
|
||||
# ── effective_denoise_steps (strength-aware step count for the auto policy) ─────────
|
||||
def test_effective_steps_txt2img_is_full_count():
|
||||
# No strength (txt2img / reference) -> the full requested step count.
|
||||
assert effective_denoise_steps(28, None) == 28
|
||||
assert effective_denoise_steps(28, 1.0) == 28 # full redraw denoises every step
|
||||
|
||||
|
||||
def test_effective_steps_low_strength_shrinks_below_the_bar():
|
||||
# A 28-step upscale at strength 0.35 denoises ~10 steps (diffusers get_timesteps),
|
||||
# which is below FBCACHE_MIN_STEPS -> the auto policy must NOT engage FBCache there.
|
||||
eff = effective_denoise_steps(28, 0.35)
|
||||
assert eff == 10
|
||||
assert eff < FBCACHE_MIN_STEPS
|
||||
|
||||
|
||||
def test_effective_steps_matches_diffusers_get_timesteps():
|
||||
# Mirror diffusers exactly: num_inference_steps - int(num_inference_steps -
|
||||
# min(num_inference_steps * strength, num_inference_steps)).
|
||||
for steps, strength in [(28, 0.35), (28, 0.8), (50, 0.5), (20, 0.99), (30, 0.1)]:
|
||||
init = min(steps * strength, steps)
|
||||
expected = max(1, steps - int(max(steps - init, 0)))
|
||||
assert effective_denoise_steps(steps, strength) == expected
|
||||
|
||||
|
||||
def test_toggle_stays_off_for_low_strength_workflow(monkeypatch):
|
||||
# End to end: a 28-step request would engage FBCache, but at strength 0.35 the
|
||||
# effective ~10 steps keep it uncached.
|
||||
_stub_diffusers(monkeypatch)
|
||||
t = _ToggleTransformer()
|
||||
mode = maybe_toggle_step_cache(_pipe(t), steps = effective_denoise_steps(28, 0.35))
|
||||
assert mode is None and t.enables == 0
|
||||
|
||||
|
||||
class _ToggleTransformer(_MixinTransformer):
|
||||
"""CacheMixin-style fake with the disable side too, counting transitions."""
|
||||
|
||||
|
|
|
|||
|
|
@ -992,7 +992,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
|
|||
"auto",
|
||||
);
|
||||
const [memoryMode, setMemoryMode] = useState<"auto" | "fast" | "balanced" | "low_vram">("auto");
|
||||
const [transformerCache, setTransformerCache] = useState<"off" | "fbcache">("off");
|
||||
const [transformerCache, setTransformerCache] = useState<"auto" | "off" | "fbcache">("auto");
|
||||
const [cpuOffload, setCpuOffload] = useState(false);
|
||||
// The last load descriptor, so "Reapply" can reload the same model with new advanced
|
||||
// options without the user re-picking it from the dropdown.
|
||||
|
|
@ -1475,7 +1475,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
|
|||
transformer_quant: transformerQuant === "auto" ? undefined : transformerQuant,
|
||||
attention_backend: attentionBackend === "auto" ? undefined : attentionBackend,
|
||||
memory_mode: memoryMode === "auto" ? undefined : memoryMode,
|
||||
transformer_cache: transformerCache === "off" ? undefined : transformerCache,
|
||||
transformer_cache: transformerCache === "auto" ? undefined : transformerCache,
|
||||
});
|
||||
} catch (err) {
|
||||
dismissLoadToast();
|
||||
|
|
@ -1925,10 +1925,11 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
|
|||
/>
|
||||
<AdvancedSelect
|
||||
label="Step cache"
|
||||
hint="First-Block-Cache reuses the transformer tail across steps for many-step models (~1.4x). Leave off for few-step distilled models."
|
||||
hint="First-Block-Cache reuses the transformer tail across steps for many-step models (~1.4x). Auto enables it for many-step schedules and skips it for few-step distilled models; Off disables it entirely."
|
||||
value={transformerCache}
|
||||
onValueChange={(v) => setTransformerCache(v as typeof transformerCache)}
|
||||
options={[
|
||||
["auto", "Auto"],
|
||||
["off", "Off"],
|
||||
["fbcache", "First-Block-Cache"],
|
||||
]}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue