{"record":{"id":"9856a06f27966432","repo":"invoke-ai/InvokeAI","slug":"expected-a-wan-transformer-with-blocks-got-type","errorCode":null,"errorMessage":"Expected a Wan transformer with blocks, got {type(transformer).__name__}.","messagePattern":"Expected a Wan transformer with blocks, got (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"invokeai/backend/wan/memory_optimization.py","lineNumber":212,"sourceCode":"\n\n@contextmanager\ndef wan_memory_optimization(\n    transformer: torch.nn.Module,\n    *,\n    enabled: bool,\n    activation_chunk_size: int = WAN_ACTIVATION_CHUNK_SIZE,\n) -> Iterator[None]:\n    \"\"\"Temporarily chunk Wan transformer pointwise activations during inference.\"\"\"\n    if not enabled:\n        yield\n        return\n    if activation_chunk_size <= 0:\n        raise ValueError(\"activation_chunk_size must be positive\")\n\n    blocks: Any = getattr(transformer, \"blocks\", None)\n    if blocks is None:\n        raise TypeError(f\"Expected a Wan transformer with blocks, got {type(transformer).__name__}.\")\n    blocks = list(blocks)\n    if hasattr(transformer, \"_invokeai_original_forward\") or any(\n        hasattr(block, \"_invokeai_original_forward\") for block in blocks\n    ):\n        raise RuntimeError(\"Wan memory optimization context cannot be nested.\")\n\n    patched_blocks: list[tuple[torch.nn.Module, Any, bool]] = []\n    original_transformer_forward = transformer.forward\n    transformer_had_instance_forward = \"forward\" in transformer.__dict__\n    patch_transformer_forward = all(\n        hasattr(transformer, name)\n        for name in (\"condition_embedder\", \"patch_embedding\", \"proj_out\", \"rope\", \"scale_shift_table\")\n    )\n    try:\n        if patch_transformer_forward:\n            transformer._invokeai_original_forward = original_transformer_forward\n            transformer._invokeai_activation_chunk_size = activation_chunk_size\n            transformer.forward = MethodType(_optimized_wan_transformer_forward, transformer)","sourceCodeStart":194,"sourceCodeEnd":230,"githubUrl":"https://github.com/invoke-ai/InvokeAI/blob/0b6a024f2ff6a86bfb953dcdb9cc504ef7397a06/invokeai/backend/wan/memory_optimization.py#L194-L230","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","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."],"exampleFix":"// before\nwith wan_memory_optimization(t, True, 8):\n    with wan_memory_optimization(t, True, 16):  # TypeError\n        run()\n// after\nwith wan_memory_optimization(t, True, 16):\n    run()","handlingStrategy":"try-catch","validationCode":"def is_patched(transformer) -> bool:\n    return hasattr(transformer, \"_invokeai_original_forward\") or any(\n        hasattr(b, \"_invokeai_original_forward\") for b in getattr(transformer, \"blocks\", [])\n    )\n\nif not is_patched(transformer):\n    with wan_memory_optimization(transformer, True, 16):\n        run_diffusion()","typeGuard":null,"tryCatchPattern":"try:\n    with wan_memory_optimization(transformer, True, 16):\n        run_diffusion()\nexcept TypeError as e:\n    if \"Expected a Wan transformer with blocks\" in str(e):\n        log.warning(\"optimization already active; running without nesting\")\n        run_diffusion()","preventionTips":["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."],"tags":["python","context-manager","monkey-patching","memory-optimization","wan"],"backgroundTag":"nested-context-not-allowed","analyzedSha":"0b6a024f2ff6a86bfb953dcdb9cc504ef7397a06","analyzedAt":"2026-08-29T04:46:49.967Z","schemaVersion":2},"datasetVersion":"2026-08-29T07:17:48.351Z"}