invoke-ai/InvokeAI · error · RuntimeError

Expected 4-D latent (B, C, H, W) after extraction, got shape

Error message

Expected 4-D latent (B, C, H, W) after extraction, got shape {latent.shape}

What it means

extract_latent unpacks a pipeline's raw latent output (optionally squeezing a temporal dim for video models) and requires the final tensor to be 4-D in (B, C, H, W) layout. Any other rank means the pipeline returned an unexpected structure, so a RuntimeError is raised rather than decoding garbage.

Source

Thrown at invokeai/backend/pid/_src/inference/pipeline_registry.py:316

        # _unpack_latents_with_ids returns a list/stacked tensor (B, C, H, W)
        latent = result if isinstance(result, torch.Tensor) else torch.stack(result, dim=0)
    elif cfg.name == "qwenimage":
        # QwenImage: packed (B, seq_len, C) → (B, C, 1, H, W) with temporal dim
        from diffusers.pipelines.qwenimage.pipeline_qwenimage import QwenImagePipeline

        latent = QwenImagePipeline._unpack_latents(
            latent,
            height=height,
            width=width,
            vae_scale_factor=pipeline.vae_scale_factor,
        )
        # Squeeze temporal dim: (B, C, 1, H, W) → (B, C, H, W)
        latent = latent.squeeze(2)

    # ZImage: already (B, C, H, W), no unpacking needed.

    if latent.ndim != 4:
        raise RuntimeError(f"Expected 4-D latent (B, C, H, W) after extraction, got shape {latent.shape}")
    return latent


def decode_with_pipeline_vae(pipeline, latent: torch.Tensor, cfg: DiffusionPipelineConfig) -> torch.Tensor:
    """Standard VAE decode using the pipeline's own VAE.

    Takes the *normalized* latent (as returned by output_type="latent"),
    denormalizes it, and decodes to pixel space.

    Returns: (B, 3, H, W) float tensor in [0, 1].
    """
    raw_latent = denormalize_latent(pipeline, latent, cfg)

    if cfg.uses_bn_normalization:
        # Flux2 VAE: unpatch before decoding.
        # raw_latent is (B, C_packed, pH, pW) — C_packed = latent_channels * patch_h * patch_w.
        # Must undo patchification to get (B, latent_channels, H/8, W/8) before vae.decode().
        from diffusers.pipelines.flux2.pipeline_flux2 import Flux2Pipeline

View on GitHub (pinned to 0b6a024f2f)

Solutions

  1. Index/select the latent you want first — for video latents pass a single frame (squeeze/slice the temporal dim to size 1).
  2. Confirm the pipeline is one of the supported backbones in PIPELINE_REGISTRY; wrap custom pipelines to emit (B, C, H, W).
  3. If you get a 3-D tensor, unsqueeze a batch dim: latent.unsqueeze(0).

Example fix

// before: 5-D video latent
latent = pipe_output.frames_latent  # (B, C, T, H, W), T > 1
extract_latent(latent)
// after
latent = pipe_output.frames_latent[:, :, 0]  # (B, C, H, W)
extract_latent(latent)
Defensive patterns

Strategy: type-guard

Validate before calling

def ensure_bchw(latent: 'torch.Tensor') -> 'torch.Tensor':
    if latent.ndim == 5 and latent.shape[2] == 1:
        latent = latent.squeeze(2)
    if latent.ndim == 3:
        latent = latent.unsqueeze(0)
    if latent.ndim != 4:
        raise ValueError(f'cannot normalize latent to (B,C,H,W): {latent.shape}')
    return latent

Type guard

def is_bchw(latent: 'torch.Tensor') -> bool:
    return latent.ndim == 4  # (B, C, H, W)

Try / catch

try:
    latent = extract_latent(raw)
except RuntimeError as e:
    if 'Expected 4-D latent' in str(e):
        logger.error('Unexpected latent shape: %s', e)
        latent = extract_latent(normalize_to_bchw(raw))
    else:
        raise

Prevention

When it happens

Trigger: Calling extract_latent with a pipeline whose latent output is not (B, C, H, W) and not a (B, C, 1, H, W) video latent — e.g. a 3-D latent (C, H, W) from a batch-size-1 pipeline that skips the batch dim, a 5-D latent with temporal dim != 1, or a custom pipeline returning tuples/dicts the extractor mis-unpacked.

Common situations: Using a custom or third-party pipeline whose output layout differs from the supported backbones; batch-of-videos latents (B, C, T, H, W) with T > 1; forgetting to index into a pipeline output before passing it to extract_latent.

Related errors


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