{"record":{"id":"8494261b79e2a836","repo":"invoke-ai/InvokeAI","slug":"wan-memory-optimization-context-cannot-be-nested","errorCode":null,"errorMessage":"Wan memory optimization context cannot be nested.","messagePattern":"Wan memory optimization context cannot be nested\\.","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"invokeai/backend/wan/memory_optimization.py","lineNumber":217,"sourceCode":"    *,\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)\n        for block in blocks:\n            original_forward = block.forward\n            had_instance_forward = \"forward\" in block.__dict__\n            block._invokeai_original_forward = original_forward\n            block._invokeai_activation_chunk_size = activation_chunk_size","sourceCodeStart":199,"sourceCodeEnd":235,"githubUrl":"https://github.com/invoke-ai/InvokeAI/blob/0b6a024f2ff6a86bfb953dcdb9cc504ef7397a06/invokeai/backend/wan/memory_optimization.py#L199-L235","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Remove the nested context so only one wan_memory_optimization block wraps the transformer at a time","Refactor so the inner code reuses the caller's active context instead of opening its own","Ensure the outer context exits (restoring originals) before opening a new one on the same transformer"],"exampleFix":"// before\nwith wan_memory_optimization(transformer, ...):\n    with wan_memory_optimization(transformer, ...):  # RuntimeError\n        run_diffusion()\n// after\nwith wan_memory_optimization(transformer, ...):\n    run_diffusion()","handlingStrategy":"validation","validationCode":"def can_enter_memory_optimization(transformer) -> bool:\n    blocks = getattr(transformer, \"blocks\", None)\n    if blocks is None:\n        return False\n    return not (hasattr(transformer, \"_invokeai_original_forward\")\n                or any(hasattr(b, \"_invokeai_original_forward\") for b in blocks))\n\nassert can_enter_memory_optimization(transformer), \"already inside wan_memory_optimization\"","typeGuard":"def is_wan_memory_optimized(obj: Any) -> bool:\n    return hasattr(obj, \"_invokeai_original_forward\")","tryCatchPattern":"try:\n    with wan_memory_optimization(transformer, chunk_size):\n        run_diffusion(transformer)\nexcept RuntimeError as e:\n    if \"cannot be nested\" in str(e):\n        run_diffusion(transformer)  # context already active; reuse it\n    else:\n        raise","preventionTips":["Never wrap wan_memory_optimization inside another one on the same transformer","Have helper functions accept an \"already optimized\" flag rather than opening their own context","Keep patch/unpatch in a single owned code path and document it","Add an assert/guard for _invokeai_original_forward before entering the context in shared code"],"tags":["runtime-error","nested-context","state-management","wan"],"backgroundTag":"nested-context-not-allowed","analyzedSha":"0b6a024f2ff6a86bfb953dcdb9cc504ef7397a06","analyzedAt":"2026-08-29T04:46:49.967Z","schemaVersion":2},"datasetVersion":"2026-08-29T07:17:48.351Z"}