invoke-ai/InvokeAI · error · TypeError

Expected a Wan transformer with blocks, got {type(transforme

Error message

Expected a Wan transformer with blocks, got {type(transformer).__name__}.

What it means

wan_memory_optimization monkey-patches each Wan transformer block's forward method and stores originals in a _invokeai_original_forward attribute. If the transformer or its blocks already carry that attribute, a previous optimization context is still active (or was not cleanly restored), and nesting would double-wrap forwards and corrupt restore logic, so a TypeError is raised.

Source

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


@contextmanager
def wan_memory_optimization(
    transformer: torch.nn.Module,
    *,
    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)

View on GitHub (pinned to 0b6a024f2f)

Solutions

Remove the nested inner wan_memory_optimization call; reuse the outer context.
Ensure the outer with-block fully exits (no exception swallowed mid-patch) before enabling again.
If state is stale after a crash, reload/reinstantiate the transformer or manually restore _invokeai_original_forward on the transformer and each block.

Example fix

// before
with wan_memory_optimization(t, True, 8):
    with wan_memory_optimization(t, True, 16):  # TypeError
        run()
// after
with wan_memory_optimization(t, True, 16):
    run()

Defensive patterns

Strategy: try-catch

Validate before calling

def is_patched(transformer) -> bool:
    return hasattr(transformer, "_invokeai_original_forward") or any(
        hasattr(b, "_invokeai_original_forward") for b in getattr(transformer, "blocks", [])
    )

if not is_patched(transformer):
    with wan_memory_optimization(transformer, True, 16):
        run_diffusion()

Try / catch

try:
    with wan_memory_optimization(transformer, True, 16):
        run_diffusion()
except TypeError as e:
    if "Expected a Wan transformer with blocks" in str(e):
        log.warning("optimization already active; running without nesting")
        run_diffusion()

Prevention

Never enter wan_memory_optimization while another instance is active on the same transformer.
Structure code so the context is entered at exactly one level (e.g. only in _run_diffusion).
After any crash inside the context, reload the transformer to clear stale _invokeai_original_forward attributes.
Avoid concurrent inference on a shared transformer with the optimization enabled.

When it happens

Trigger: Entering a second (nested) wan_memory_optimization context on the same transformer while one is already active; reusing a transformer whose patching failed mid-way and left _invokeai_original_forward behind (e.g. an exception inside a previous with-block before cleanup).

Common situations: Calling _run_diffusion inside an outer context that also enables the optimization, running inference concurrently on the same transformer from two threads, a prior crash leaving stale patched state on a long-lived model object.

Related errors

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