Comfy-Org/ComfyUI · error · ValueError

attention_mask must be {expected}, got {tuple(attention_mask

Error message

attention_mask must be {expected}, got {tuple(attention_mask.shape)}

What it means

The Lens joint img/txt attention block requires an additive attention mask of exactly (bsz, 1, 1, seq_img + seq_txt) so it can be passed straight to optimized_attention with skip_reshape=True. Any other shape is rejected rather than silently broadcasting to wrong tokens.

Source

Thrown at comfy/ldm/lens/model.py:165

        txt_qkv = self.txt_qkv(encoder_hidden_states).view(bsz, seq_txt, 3, self.heads, self.dim_head)
        txt_q, txt_k, txt_v = txt_qkv.unbind(dim=2)
        txt_q = self.norm_added_q(txt_q)
        txt_k = self.norm_added_k(txt_k)

        # [B, S, H, D] → [B, H, S, D] for attention, dels to avoid VRAM peaks
        q = torch.cat([img_q, txt_q], dim=1).transpose(1, 2)
        del img_q, txt_q
        k = torch.cat([img_k, txt_k], dim=1).transpose(1, 2)
        del img_k, txt_k
        v = torch.cat([img_v, txt_v], dim=1).transpose(1, 2)
        del img_v, txt_v

        q, k = apply_rope(q, k, freqs_cis)

        if attention_mask is not None:
            expected = (bsz, 1, 1, seq_img + seq_txt)
            if attention_mask.shape != expected:
                raise ValueError(
                    f"attention_mask must be {expected}, got {tuple(attention_mask.shape)}"
                )
            attention_mask = attention_mask.to(q.dtype)

        out = optimized_attention(
            q, k, v, self.heads, mask=attention_mask, skip_reshape=True,
            transformer_options=transformer_options,
        )

        img_out = self.to_out[1](self.to_out[0](out[:, :seq_img, :]))
        txt_out = self.to_add_out(out[:, seq_img:, :])
        return img_out, txt_out


class LensTransformerBlock(nn.Module):
    def __init__(
        self,
        dim: int,

View on GitHub (pinned to 1c6d8d45b3)

Solutions

1. Expand your (B, S) mask: mask[:, None, None, :] to match (bsz, 1, 1, seq_img+seq_txt)
2. Recompute the mask after concatenation order img+txt with the current batch size
3. Pass attention_mask=None if full attention is intended

Example fix

# before
out = block(img, txt, freqs, attention_mask=mask)  # mask is (B, S)
# after
mask4 = mask[:, None, None, :].expand(bsz, 1, 1, seq_img + seq_txt)
out = block(img, txt, freqs, attention_mask=mask4)

S = seq_img + seq_txt
assert mask is None or tuple(mask.shape) == (bsz, 1, 1, S), mask.shape

When it happens

Trigger: Passing a (B, S) or (B, 1, S) mask instead of the 4D expanded form, or a mask built against text-only/img-only length; also batch-size mismatches when cond/uncond are combined.

Common situations: Custom node authors adapting masks from other architectures (Flux/Wan style masks are (B,S)), or callers reusing a mask computed before sequence lengths were final.

Related errors

• Image-token count {image_idx.shape[0]} != ViT output count {
• {name} inner dimension {inner_dim} is not divisible by head
• PixDiT_T2I requires context (text embeddings) of shape [B, L
• SeedVR2 patch input temporal size must satisfy T % {t} == 1,
• SeedVR2 expected {name} to be 5-D native latent, got shape {

AI-assisted analysis of Comfy-Org/ComfyUI@1c6d8d45b3 (2026-08-14). Data as JSON: /api/errors/6f2fad4342a29aa0. Report an issue: GitHub.