diff --git a/modules/attention/context.py b/modules/attention/context.py index 71442e678..6134dbc9d 100644 --- a/modules/attention/context.py +++ b/modules/attention/context.py @@ -27,22 +27,27 @@ def denoiser_name(pipe) -> str | None: return None +# every slot a pipeline can enter once per denoising step, from sd_offload_state.group_offload_main. The aux +# components on that list (decoder, controlnet, prior) are left alone: they pack no attention sequence, and a +# publication from one would clear the layout the denoiser just set +DENOISER_SLOTS = ('transformer', 'unet', 'transformer_2', 'transformer_ref', 'unconditional_transformer') + + def install_layout_hook(pipe) -> None: """Let a classic pipeline's denoiser publish its own packing: the modular path has its own hook, this is the rest.""" from modules import shared if pipe is None or not getattr(shared.opts, 'sparse_attention_enabled', False): return - module = getattr(pipe, 'transformer', None) - if module is None: - module = getattr(pipe, 'unet', None) - if module is None or getattr(module, 'sdnext_layout_hook', None) is not None or getattr(module, 'sdnext_state_hook', None) is not None: - return from modules.attention.sparse import layout as sparse_layout def publish(denoiser, args, kwargs): # pylint: disable=unused-argument set_layout(sparse_layout.layout_from_kwargs(kwargs, denoiser.__class__.__name__)) - module.sdnext_layout_hook = module.register_forward_pre_hook(publish, with_kwargs=True) + for name in DENOISER_SLOTS: + module = getattr(pipe, name, None) + if module is None or getattr(module, 'sdnext_layout_hook', None) is not None or getattr(module, 'sdnext_state_hook', None) is not None: + continue + module.sdnext_layout_hook = module.register_forward_pre_hook(publish, with_kwargs=True) def begin(pipe, steps: int = 0) -> None: diff --git a/test/test-attention-sparse.py b/test/test-attention-sparse.py index 59e571e18..f3fcf7b45 100644 --- a/test/test-attention-sparse.py +++ b/test/test-attention-sparse.py @@ -283,11 +283,13 @@ def test_layout_hook_publishes_from_the_denoiser_kwargs(): return hidden_states class Pipe: - def __init__(self, transformer): + def __init__(self, transformer, second): self.transformer = transformer + self.unconditional_transformer = second # ideogram runs a second denoiser, wan a14b a transformer_2 denoiser = FluxTransformer2DModel() - pipe = Pipe(denoiser) + second = FluxTransformer2DModel() + pipe = Pipe(denoiser, second) previous = getattr(shared.opts, 'sparse_attention_enabled', False) try: shared.opts.data['sparse_attention_enabled'] = False @@ -297,6 +299,7 @@ def test_layout_hook_publishes_from_the_denoiser_kwargs(): ctx.install_layout_hook(pipe) ctx.install_layout_hook(pipe) assert getattr(denoiser, 'sdnext_layout_hook', None) is not None, 'the denoiser is hooked once' + assert getattr(second, 'sdnext_layout_hook', None) is not None, 'every denoiser slot is hooked, not just the first' ctx.set_layout(None) denoiser(hidden_states=torch.zeros(1, 4096, 4, device=device), txt_ids=torch.zeros(512, 3, device=device), img_ids=torch.zeros(4096, 3, device=device)) published = ctx.current.layout