feat(attention): publish the token layout for classic pipelines

Only the modular path had a per-forward kwargs hook, so a classic pipeline had
no layout and the selector sparsified its whole sequence, prompt conditioning
included. A denoiser pre-hook installed per generation now reads whichever
convention the model uses: the *_indices tensors, or the txt_ids and img_ids a
joint transformer is given.

Packing order is not derivable from the call and differs between architectures,
so only classes verified to pack text first publish; anything else falls back
as before rather than pinning the wrong half of the sequence dense.

modules/attention/context.py

@@ -27,12 +27,31 @@ def denoiser_name(pipe) -> str | None:
     return None
 
 
+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)
+
+
 def begin(pipe, steps: int = 0) -> None:
     from modules import devices
     current.active = True
     current.role = 'transformer'
     current.layout = None
     current.model_key = (pipe.__class__.__name__, denoiser_name(pipe)) if pipe is not None else None
+    install_layout_hook(pipe)
     device = devices.device if devices.device is not None else torch.device('cpu')
     if current.step_buffer is None or current.step_buffer.device != device:
         current.step_buffer = torch.zeros((), dtype=torch.int64, device=device)

modules/attention/sparse/__init__.py

@@ -1,8 +1,8 @@
 """Block-sparse attention: the selector, the token layout it respects, and the consumers that apply it."""
 from modules.attention.sparse.selector import BlockSelection, SparseSpec, block_count, radial_blocks, schedule, select_blocks
-from modules.attention.sparse.layout import Span, TokenLayout, block_pins, layout_from_index_kwargs, layout_from_prefix, layout_from_segments, publish_segments, segments_from_live
+from modules.attention.sparse.layout import Span, TokenLayout, block_pins, layout_from_index_kwargs, layout_from_kwargs, layout_from_prefix, layout_from_segments, publish_segments, segments_from_live
 
 __all__ = [
     'BlockSelection', 'SparseSpec', 'block_count', 'radial_blocks', 'schedule', 'select_blocks',
-    'Span', 'TokenLayout', 'block_pins', 'layout_from_index_kwargs', 'layout_from_prefix', 'layout_from_segments', 'publish_segments', 'segments_from_live',
+    'Span', 'TokenLayout', 'block_pins', 'layout_from_index_kwargs', 'layout_from_kwargs', 'layout_from_prefix', 'layout_from_segments', 'publish_segments', 'segments_from_live',
 ]

modules/attention/sparse/layout.py

@@ -71,6 +71,30 @@ def layout_from_index_kwargs(kwargs: dict, length: int | None = None) -> TokenLa
     return TokenLayout(spans=tuple(spans), length=length if length is not None else spans[-1].end, source='indices')
 
 
+# how an architecture orders its joint sequence, which the call itself does not reveal. Verified against the
+# diffusers transformers that take txt_ids and img_ids; HiDream packs image first and is deliberately absent, so
+# it falls back rather than being pinned backwards. An unlisted class publishes nothing.
+JOINT_TEXT_FIRST = frozenset({
+    'FluxTransformer2DModel', 'Flux2Transformer2DModel', 'ChromaTransformer2DModel', 'BriaTransformer2DModel',
+    'BriaFiboTransformer2DModel', 'LongCatImageTransformer2DModel', 'OvisImageTransformer2DModel',
+})
+
+
+def layout_from_stream_ids(kwargs: dict, cls_name: str | None) -> TokenLayout | None:
+    """Read the stream lengths off the rotary id tensors a joint transformer is given by name."""
+    if cls_name not in JOINT_TEXT_FIRST:
+        return None
+    text, image = kwargs.get('txt_ids'), kwargs.get('img_ids')
+    if not torch.is_tensor(text) or not torch.is_tensor(image) or text.dim() < 2 or image.dim() < 2:
+        return None
+    return layout_from_segments((('text', text.shape[-2]), ('image', image.shape[-2])), source='stream-ids')
+
+
+def layout_from_kwargs(kwargs: dict, cls_name: str | None = None) -> TokenLayout | None:
+    """Whatever the denoiser says about its own packing, by whichever convention it uses."""
+    return layout_from_index_kwargs(kwargs or {}) or layout_from_stream_ids(kwargs or {}, cls_name)
+
+
 def layout_from_segments(segments, length: int | None = None, source: str = 'segments') -> TokenLayout:
     """Build a layout from ordered (kind, count) pairs, the form a transformer knows at its packing site."""
     spans: list[Span] = []