invoke-ai/InvokeAI · error · ValueError

Krea-2 regional attention mask shape {tuple(regional_attenti

Error message

Krea-2 regional attention mask shape {tuple(regional_attention_mask.shape)} does not match the transformer sequence length {hidden_states.shape[1]}.

What it means

Krea-2 regional prompting builds a custom attention mask covering the full transformer sequence (text embeddings + image latents). Before applying it, each attention block validates that the mask is square and matches hidden_states.shape[1] (the sequence length). A mismatch means the mask was computed for a different sequence length — usually because latent resolution, number of text conditionings, or token counts changed after the mask was built.

Source

Thrown at invokeai/backend/krea2/attention.py:55


class Krea2MemoryEfficientAttnProcessor:
    """Drop-in replacement for ``Krea2AttnProcessor`` that avoids the ``enable_gqa`` math fallback."""

    def __init__(self, regional_prompting_state: Krea2RegionalPromptingState | None = None) -> None:
        self.regional_prompting_state = regional_prompting_state

    def __call__(
        self,
        attn,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
    ) -> torch.Tensor:
        if self.regional_prompting_state is not None and self.regional_prompting_state.attention_mask is not None:
            regional_attention_mask = self.regional_prompting_state.attention_mask
            if regional_attention_mask.shape != (hidden_states.shape[1], hidden_states.shape[1]):
                raise ValueError(
                    f"Krea-2 regional attention mask shape {tuple(regional_attention_mask.shape)} does not match "
                    f"the transformer sequence length {hidden_states.shape[1]}."
                )
            attention_mask = (
                regional_attention_mask if attention_mask is None else attention_mask & regional_attention_mask
            )

        query = attn.to_q(hidden_states).unflatten(-1, (attn.num_heads, attn.head_dim))
        key = attn.to_k(hidden_states).unflatten(-1, (attn.num_kv_heads, attn.head_dim))
        value = attn.to_v(hidden_states).unflatten(-1, (attn.num_kv_heads, attn.head_dim))
        gate = attn.to_gate(hidden_states)

        query = attn.norm_q(query)
        key = attn.norm_k(key)

        if image_rotary_emb is not None:
            query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
            key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)

View on GitHub (pinned to 0b6a024f2f)

Solutions

  1. Rebuild the regional prompting extension via Krea2RegionalPromptingExtension.from_text_conditionings(...) with the same image_seq_len used by the current transformer pass.
  2. Ensure image dimensions (latent size) are fixed for the lifetime of the regional attention mask.
  3. Pass attention_mask=None or a fresh regional state if you intentionally changed resolution or prompts.
  4. Log hidden_states.shape[1] and the mask shape to find where they diverge.

Example fix

// before
ext = build_regional_extension(text_conds, image_seq_len=4096)
# ... later run at 2x resolution -> seq_len=16384, mask is 4096x4096
// after
ext = Krea2RegionalPromptingExtension.from_text_conditionings(text_conds, image_seq_len=transformer_seq_len)
Defensive patterns

Strategy: try-catch

Validate before calling

if regional_state and regional_state.attention_mask is not None:
    assert regional_state.attention_mask.shape == (seq_len, seq_len), \
        f'mask {regional_state.attention_mask.shape} != ({seq_len}, {seq_len})'

Try / catch

try:
    output = transformer(...)
except ValueError as e:
    if 'regional attention mask shape' in str(e):
        regional_state = Krea2RegionalPromptingExtension.from_text_conditionings(
            text_conditionings, image_seq_len=hidden_states_len)
        output = transformer(...)  # rebuild and retry once
    else:
        raise

Prevention

When it happens

Trigger: Running the Krea-2 transformer __call__ with regional_prompting_state.attention_mask whose shape differs from (seq_len, seq_len): changed image height/width after the extension was constructed, different text conditionings, or a mask reused across denoising stages with differing sequence lengths.

Common situations: Caching a Krea2RegionalPromptingExtension built for one resolution and reusing it at another; mixing regional masks between batches; changing prompt tokenization after mask construction.

Related errors


AI-assisted analysis of invoke-ai/InvokeAI@0b6a024f2f (2026-08-29). Data as JSON: /api/errors/8d5e962ef22056b7. Report an issue: GitHub.