Comfy-Org/ComfyUI · error · ValueError

Got a mask of shape {list(mask.shape)}, expected [b, q, k] o

Error message

Got a mask of shape {list(mask.shape)}, expected [b, q, k] or [q, k]

What it means

upscale_dit_mask spatially upscales the image-token quadrant of a DiT attention mask when the latent resolution changes between passes. It accepts 2D [q,k] (auto-unsqueezed to batch 1) or 3D [b,q,k] masks; anything with ndim other than 2/3 (e.g. 4D) cannot be interpreted as token-to-token attention and is rejected.

Source

Thrown at comfy/utils.py:1342

            input_mask = input_mask.reshape((1, 1, -1, input_mask.shape[-2], input_mask.shape[-1]))
        scale_mode = "trilinear"

    mask = torch.nn.functional.interpolate(input_mask, size=output_shape[2:], mode=scale_mode)
    if mask.shape[1] < output_shape[1]:
        mask = mask.repeat((1, output_shape[1]) + (1,) * dims)[:,:output_shape[1]]
    mask = repeat_to_batch_size(mask, output_shape[0])
    return mask

def upscale_dit_mask(mask: torch.Tensor, img_size_in, img_size_out):
        hi, wi = img_size_in
        ho, wo = img_size_out
        # if it's already the correct size, no need to do anything
        if (hi, wi) == (ho, wo):
            return mask
        if mask.ndim == 2:
            mask = mask.unsqueeze(0)
        if mask.ndim != 3:
            raise ValueError(f"Got a mask of shape {list(mask.shape)}, expected [b, q, k] or [q, k]")
        txt_tokens = mask.shape[1] - (hi * wi)
        # quadrants of the mask
        txt_to_txt = mask[:, :txt_tokens, :txt_tokens]
        txt_to_img = mask[:, :txt_tokens, txt_tokens:]
        img_to_img = mask[:, txt_tokens:, txt_tokens:]
        img_to_txt = mask[:, txt_tokens:, :txt_tokens]

        # convert to 1d x 2d, interpolate, then back to 1d x 1d
        txt_to_img = rearrange  (txt_to_img, "b t (h w) -> b t h w", h=hi, w=wi)
        txt_to_img = interpolate(txt_to_img, size=img_size_out, mode="bilinear")
        txt_to_img = rearrange  (txt_to_img, "b t h w -> b t (h w)")
        # this one is hard because we have to do it twice
        # convert to 1d x 2d, interpolate, then to 2d x 1d, interpolate, then 1d x 1d
        img_to_img = rearrange  (img_to_img, "b hw (h w) -> b hw h w", h=hi, w=wi)
        img_to_img = interpolate(img_to_img, size=img_size_out, mode="bilinear")
        img_to_img = rearrange  (img_to_img, "b (hk wk) hq wq -> b (hq wq) hk wk", hk=hi, wk=wi)
        img_to_img = interpolate(img_to_img, size=img_size_out, mode="bilinear")
        img_to_img = rearrange  (img_to_img, "b (hq wq) hk wk -> b (hk wk) (hq wq)", hq=ho, wq=wo)

View on GitHub (pinned to 1c6d8d45b3)

Solutions

1. Reduce the mask to [b,q,k] (e.g. average or select heads: mask.mean(dim=1)) before passing to upscale_dit_mask
2. If you have a plain [q,k] mask, pass it 2D — it is unsqueezed automatically
3. Check the producing node's output shape contract and fix it there rather than squeezing downstream

Example fix

// before
mask = upscale_dit_mask(attn_mask_4d, (32, 32), (64, 64))  # [b,heads,q,k]

# after
mask = upscale_dit_mask(attn_mask_4d.mean(dim=1), (32, 32), (64, 64))  # [b,q,k]

if mask.ndim == 4:
    mask = mask.mean(dim=1)  # collapse heads -> [b, q, k]
assert mask.ndim in (2, 3), f'bad mask rank {mask.ndim}'

def is_dit_mask(m: torch.Tensor) -> bool:
    return m.ndim in (2, 3)

When it happens

Trigger: Calling utils.upscale_dit_mask with a mask whose ndim is not 2 or 3 — for example a 4D [b,heads,q,k] mask passed through from a model patch, or a 1D/0D tensor from a malformed mask construction. Called from model_base.py when attention_mask is present and ref size differs from target token size.

Common situations: Custom mask-producing nodes emitting per-head 4D masks; adapters that prepend extra dims; passing a mask built for a different attention API.

Related errors

• Input img and txt tensors must have 3 dimensions.
• Input img tensor must be in [B, C, H, W] format.
• Input txt tensors must have 3 dimensions.
• Input img and txt tensors must have 3 dimensions.
• Input img and txt tensors must have 3 dimensions.

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