fix(attention): hook every denoiser slot for layout publication

install_layout_hook took transformer or unet and stopped, so a pipeline
with a second denoiser left that one unhooked and its forwards ran against
whatever the first had published. The slot list follows the denoiser names
in sd_offload_state; the aux components on that list pack no attention
sequence, and a publication from one would clear the layout.

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:

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