{"record":{"id":"f28f47551d9549cb","repo":"Comfy-Org/ComfyUI","slug":"got-a-mask-of-shape-list-mask-shape-expected","errorCode":null,"errorMessage":"Got a mask of shape {list(mask.shape)}, expected [b, q, k] or [q, k]","messagePattern":"Got a mask of shape (.+?), expected \\[b, q, k\\] or \\[q, k\\]","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"comfy/utils.py","lineNumber":1342,"sourceCode":"            input_mask = input_mask.reshape((1, 1, -1, input_mask.shape[-2], input_mask.shape[-1]))\n        scale_mode = \"trilinear\"\n\n    mask = torch.nn.functional.interpolate(input_mask, size=output_shape[2:], mode=scale_mode)\n    if mask.shape[1] < output_shape[1]:\n        mask = mask.repeat((1, output_shape[1]) + (1,) * dims)[:,:output_shape[1]]\n    mask = repeat_to_batch_size(mask, output_shape[0])\n    return mask\n\ndef upscale_dit_mask(mask: torch.Tensor, img_size_in, img_size_out):\n        hi, wi = img_size_in\n        ho, wo = img_size_out\n        # if it's already the correct size, no need to do anything\n        if (hi, wi) == (ho, wo):\n            return mask\n        if mask.ndim == 2:\n            mask = mask.unsqueeze(0)\n        if mask.ndim != 3:\n            raise ValueError(f\"Got a mask of shape {list(mask.shape)}, expected [b, q, k] or [q, k]\")\n        txt_tokens = mask.shape[1] - (hi * wi)\n        # quadrants of the mask\n        txt_to_txt = mask[:, :txt_tokens, :txt_tokens]\n        txt_to_img = mask[:, :txt_tokens, txt_tokens:]\n        img_to_img = mask[:, txt_tokens:, txt_tokens:]\n        img_to_txt = mask[:, txt_tokens:, :txt_tokens]\n\n        # convert to 1d x 2d, interpolate, then back to 1d x 1d\n        txt_to_img = rearrange  (txt_to_img, \"b t (h w) -> b t h w\", h=hi, w=wi)\n        txt_to_img = interpolate(txt_to_img, size=img_size_out, mode=\"bilinear\")\n        txt_to_img = rearrange  (txt_to_img, \"b t h w -> b t (h w)\")\n        # this one is hard because we have to do it twice\n        # convert to 1d x 2d, interpolate, then to 2d x 1d, interpolate, then 1d x 1d\n        img_to_img = rearrange  (img_to_img, \"b hw (h w) -> b hw h w\", h=hi, w=wi)\n        img_to_img = interpolate(img_to_img, size=img_size_out, mode=\"bilinear\")\n        img_to_img = rearrange  (img_to_img, \"b (hk wk) hq wq -> b (hq wq) hk wk\", hk=hi, wk=wi)\n        img_to_img = interpolate(img_to_img, size=img_size_out, mode=\"bilinear\")\n        img_to_img = rearrange  (img_to_img, \"b (hq wq) hk wk -> b (hk wk) (hq wq)\", hq=ho, wq=wo)","sourceCodeStart":1324,"sourceCodeEnd":1360,"githubUrl":"https://github.com/Comfy-Org/ComfyUI/blob/1c6d8d45b3693bfbb32385b410d813a7fd6be216/comfy/utils.py#L1324-L1360","documentation":"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.","triggerScenarios":"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.","commonSituations":"Custom mask-producing nodes emitting per-head 4D masks; adapters that prepend extra dims; passing a mask built for a different attention API.","solutions":["Reduce the mask to [b,q,k] (e.g. average or select heads: mask.mean(dim=1)) before passing to upscale_dit_mask","If you have a plain [q,k] mask, pass it 2D — it is unsqueezed automatically","Check the producing node's output shape contract and fix it there rather than squeezing downstream"],"exampleFix":"// before\nmask = upscale_dit_mask(attn_mask_4d, (32, 32), (64, 64))  # [b,heads,q,k]\n\n# after\nmask = upscale_dit_mask(attn_mask_4d.mean(dim=1), (32, 32), (64, 64))  # [b,q,k]","handlingStrategy":"type-guard","validationCode":"if mask.ndim == 4:\n    mask = mask.mean(dim=1)  # collapse heads -> [b, q, k]\nassert mask.ndim in (2, 3), f'bad mask rank {mask.ndim}'","typeGuard":"def is_dit_mask(m: torch.Tensor) -> bool:\n    return m.ndim in (2, 3)","tryCatchPattern":null,"preventionTips":["Standardize mask-producing nodes on [b, q, k]","Collapse head dims at the source, not downstream"],"tags":["attention-mask","diit","tensor-shape","mask"],"backgroundTag":null,"analyzedSha":"1c6d8d45b3693bfbb32385b410d813a7fd6be216","analyzedAt":"2026-08-14T19:37:18.893Z","schemaVersion":2},"datasetVersion":"2026-08-15T22:17:37.221Z"}