Comfy-Org/ComfyUI · error · NotImplementedError

Must flatten output.

Error message

Must flatten output.

What it means

The Mochi patch embedder only implements the flatten=True path: after patch projection it rearranges tokens to (B, T*H*W, C). The flatten=False branch deliberately raises NotImplementedError, so any construction with flatten=False cannot run a forward pass.

Source

Thrown at comfy/ldm/genmo/joint_model/layers.py:149

        )

    def forward(self, x):
        B, _C, T, H, W = x.shape
        if not self.dynamic_img_pad:
            assert H % self.patch_size[0] == 0, f"Input height ({H}) should be divisible by patch size ({self.patch_size[0]})."
            assert W % self.patch_size[1] == 0, f"Input width ({W}) should be divisible by patch size ({self.patch_size[1]})."
        else:
            pad_h = (self.patch_size[0] - H % self.patch_size[0]) % self.patch_size[0]
            pad_w = (self.patch_size[1] - W % self.patch_size[1]) % self.patch_size[1]
            x = F.pad(x, (0, pad_w, 0, pad_h))

        x = rearrange(x, "B C T H W -> (B T) C H W", B=B, T=T)
        x = comfy.ldm.common_dit.pad_to_patch_size(x, self.patch_size, padding_mode='circular')
        x = self.proj(x)

        # Flatten temporal and spatial dimensions.
        if not self.flatten:
            raise NotImplementedError("Must flatten output.")
        x = rearrange(x, "(B T) C H W -> B (T H W) C", B=B, T=T)

        x = self.norm(x)
        return x

View on GitHub (pinned to 1c6d8d45b3)

Solutions

1. Construct the patch embedder with flatten=True (the default and only supported path).
2. If you need spatial tokens, take the flattened (B, THC, C) output and rearrange it yourself afterwards.
3. Patch the module locally to implement the unflattened return if your research code truly needs it.

Example fix

# before
PatchEmbed(..., flatten=False)

# after
PatchEmbed(..., flatten=True)  # then rearrange downstream if spatial layout is needed

assert flatten is True, "Mochi PatchEmbed only supports flatten=True"

When it happens

Trigger: Building PatchEmbed/joint-model patchify layer with flatten=False (e.g. copying upstream genmo code that exposes the option) and then calling forward(). The upstream repo defaults to flatten=True, so this arises almost only via manual construction.

Common situations: Porting Mochi tooling that passes an explicit flatten flag; experiment code that wants unflattened spatial tokens (B,C,T,H,W) out of the patch embed.

Related errors

• Temperature {temperature} is not supported.
• dim_x={dim_x} should be divisible by num_heads={num_heads}
• ObjectPatchHook is not supported yet in ComfyUI.
• InjectionsHook is not supported yet in ComfyUI.
• Scalar feature is not implemented yet.

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