invoke-ai/InvokeAI · error · RuntimeError

Wan memory optimization context cannot be nested.

Error message

Wan memory optimization context cannot be nested.

What it means

The Wan memory-optimization context manager patches the transformer and its blocks by stashing `_invokeai_original_forward` attributes; a nested `with wan_memory_optimization(...)` would overwrite that state and corrupt the outer context, so the library refuses to nest. It detects nesting by checking whether the transformer or any block already carries the original-forward marker.

Source

Thrown at invokeai/backend/wan/memory_optimization.py:217

    *,
    enabled: bool,
    activation_chunk_size: int = WAN_ACTIVATION_CHUNK_SIZE,
) -> Iterator[None]:
    """Temporarily chunk Wan transformer pointwise activations during inference."""
    if not enabled:
        yield
        return
    if activation_chunk_size <= 0:
        raise ValueError("activation_chunk_size must be positive")

    blocks: Any = getattr(transformer, "blocks", None)
    if blocks is None:
        raise TypeError(f"Expected a Wan transformer with blocks, got {type(transformer).__name__}.")
    blocks = list(blocks)
    if hasattr(transformer, "_invokeai_original_forward") or any(
        hasattr(block, "_invokeai_original_forward") for block in blocks
    ):
        raise RuntimeError("Wan memory optimization context cannot be nested.")

    patched_blocks: list[tuple[torch.nn.Module, Any, bool]] = []
    original_transformer_forward = transformer.forward
    transformer_had_instance_forward = "forward" in transformer.__dict__
    patch_transformer_forward = all(
        hasattr(transformer, name)
        for name in ("condition_embedder", "patch_embedding", "proj_out", "rope", "scale_shift_table")
    )
    try:
        if patch_transformer_forward:
            transformer._invokeai_original_forward = original_transformer_forward
            transformer._invokeai_activation_chunk_size = activation_chunk_size
            transformer.forward = MethodType(_optimized_wan_transformer_forward, transformer)
        for block in blocks:
            original_forward = block.forward
            had_instance_forward = "forward" in block.__dict__
            block._invokeai_original_forward = original_forward
            block._invokeai_activation_chunk_size = activation_chunk_size

View on GitHub (pinned to 0b6a024f2f)

Solutions

  1. Remove the nested context so only one wan_memory_optimization block wraps the transformer at a time
  2. Refactor so the inner code reuses the caller's active context instead of opening its own
  3. Ensure the outer context exits (restoring originals) before opening a new one on the same transformer

Example fix

// before
with wan_memory_optimization(transformer, ...):
    with wan_memory_optimization(transformer, ...):  # RuntimeError
        run_diffusion()
// after
with wan_memory_optimization(transformer, ...):
    run_diffusion()
Defensive patterns

Strategy: validation

Validate before calling

def can_enter_memory_optimization(transformer) -> bool:
    blocks = getattr(transformer, "blocks", None)
    if blocks is None:
        return False
    return not (hasattr(transformer, "_invokeai_original_forward")
                or any(hasattr(b, "_invokeai_original_forward") for b in blocks))

assert can_enter_memory_optimization(transformer), "already inside wan_memory_optimization"

Type guard

def is_wan_memory_optimized(obj: Any) -> bool:
    return hasattr(obj, "_invokeai_original_forward")

Try / catch

try:
    with wan_memory_optimization(transformer, chunk_size):
        run_diffusion(transformer)
except RuntimeError as e:
    if "cannot be nested" in str(e):
        run_diffusion(transformer)  # context already active; reuse it
    else:
        raise

Prevention

When it happens

Trigger: Calling `with wan_memory_optimization(transformer, ...)` while the same transformer (or one of its blocks) is already patched inside an active outer memory-optimization context, e.g. two nested with-blocks or calling a helper that itself opens the context while the caller also opens it.

Common situations: Composing two pipelines/features that each wrap diffusion in their own wan_memory_optimization context; accidentally wrapping the same context twice in shared diffusion code; a wrapper function calling wan_memory_optimization when the caller already did.

Related errors


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