diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 8058e959b1..50bc621417 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -2126,14 +2126,18 @@ class DiffusionBackend: diffusers keys FBCache residuals by cache context ("cond"/"uncond") on the long-lived transformer, and neither the pipeline nor the context exit resets - them (``StateManager`` only clears via ``reset_stateful_hooks``, which no - pipeline calls). This backend reuses one resident pipe across generations, so - without a reset the next generation's first step compares its first-block - residual against the PREVIOUS request's -- a tensor-shape mismatch when the - resolution/batch changed, or a stale-cache reuse otherwise. Best-effort: a - transformer without the hook (uncached load) is a silent no-op.""" + them. The transformer-level reset entry point is ``_reset_stateful_cache`` in + diffusers 0.39 (``reset_stateful_hooks`` lives only on the HookRegistry, so a + getattr for it on the transformer is a silent no-op), and no pipeline calls it. + This backend reuses one resident pipe across generations, so without a reset the + next generation's first step compares its first-block residual against the + PREVIOUS request's -- a tensor-shape mismatch when the resolution/batch changed, + or a stale-cache reuse otherwise. Best-effort: a transformer without the hook + (uncached load) is a silent no-op.""" transformer = getattr(pipe, "transformer", None) - reset = getattr(transformer, "reset_stateful_hooks", None) + reset = getattr(transformer, "_reset_stateful_cache", None) or getattr( + transformer, "reset_stateful_hooks", None + ) if callable(reset): try: reset() diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index b26cf530ad..2a8b5922ca 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -2581,14 +2581,34 @@ def test_plan_memory_dense_replan_does_not_double_count_prefetched_transformer(m def test_reset_step_cache_helper_is_best_effort(): - # Calls the transformer's reset hook when present. + # Prefers the real diffusers CacheMixin hook (_reset_stateful_cache): on a genuine Flux / + # QwenImage transformer that is the reset entry point, and reset_stateful_hooks lives only on + # the HookRegistry (getattr for it on the transformer returns None), so the old lookup was a + # silent no-op that left stale FBCache residuals for the next generation. calls = [] pipe = types.SimpleNamespace( - transformer = types.SimpleNamespace(reset_stateful_hooks = lambda: calls.append(True)) + transformer = types.SimpleNamespace(_reset_stateful_cache = lambda: calls.append("real")) ) DiffusionBackend._reset_step_cache(pipe) - assert calls == [True] - # No transformer, or a transformer without the hook -> silent no-op (never raises). + assert calls == ["real"] + # _reset_stateful_cache wins when both are present. + calls.clear() + pipe = types.SimpleNamespace( + transformer = types.SimpleNamespace( + _reset_stateful_cache = lambda: calls.append("real"), + reset_stateful_hooks = lambda: calls.append("fallback"), + ) + ) + DiffusionBackend._reset_step_cache(pipe) + assert calls == ["real"] + # Falls back to reset_stateful_hooks for a transformer that exposes only that. + calls.clear() + pipe = types.SimpleNamespace( + transformer = types.SimpleNamespace(reset_stateful_hooks = lambda: calls.append("fallback")) + ) + DiffusionBackend._reset_step_cache(pipe) + assert calls == ["fallback"] + # No transformer, or a transformer without either hook -> silent no-op (never raises). DiffusionBackend._reset_step_cache(types.SimpleNamespace()) DiffusionBackend._reset_step_cache(types.SimpleNamespace(transformer = object())) @@ -2605,8 +2625,10 @@ def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path): family_override = "z-image", ) resets = [] + # Use the real diffusers CacheMixin entry point (_reset_stateful_cache); a genuine + # Flux/QwenImage transformer exposes this, not reset_stateful_hooks. backend._state.pipe.transformer = types.SimpleNamespace( - reset_stateful_hooks = lambda: resets.append(True) + _reset_stateful_cache = lambda: resets.append(True) ) # No cache engaged (transformer_cache is None) -> reset must NOT run. backend.generate(prompt = "a sloth")