feat(multi-gpu): mark the remaining text encoders as idle-GPU offloadable

docs/src/content/docs/development/Guides/creating-nodes.mdx

@@ -49,6 +49,7 @@ Because the whole node is moved to another device, only mark a node `idle_gpu_of
 - **It is encoder-only.** Its sole GPU work is loading one or more encoder models and running their forward pass. It must not load or run the denoise/transformer or VAE, or do any other work tied to the session's own GPU.
 - **It stores its result on the CPU before returning.** Move output tensors to the CPU (`tensor.detach().to("cpu")`) and save them as conditioning/tensors. The denoiser picks them up and moves them onto its own GPU later — this is what makes the cross-GPU handoff safe and device-agnostic.
 - **It places inputs on the loaded model's device, not a fixed device.** Resolve the device from the model you just loaded (e.g. `get_effective_device(model)` from `invokeai.backend.model_manager.load.model_cache.utils`, or `TorchDevice.choose_torch_device()`), rather than hard-coding `cuda:0`. The built-in `flux_text_encoder` and `compel` nodes are good references.
+- **Its runtime is dominated by that forward pass.** The borrow holds the lent GPU's lock for the whole node, and a session dequeued onto that GPU blocks until it is released. Model caches are per-device, so the first borrow of a GPU cold-loads the encoder there — that cost is paid once and then amortizes across later borrows, which hit the cache. Work that recurs on *every* execution does not amortize, so a node that runs something open-ended per call (an autoregressive `generate()` loop, say) will stall the lent GPU again on every generation.
 
 :::caution[Only mark encoder-only nodes]
 If a node that also runs the denoiser, VAE, or other session-GPU work is marked `idle_gpu_offloadable=True`, that work will be re-pinned to the wrong GPU and can misplace tensors or raise device-mismatch errors. When in doubt, leave it unset (the default is `False`) — the node will still work correctly, just without the offload optimization.

invokeai/app/invocations/baseinvocation.py

@@ -279,7 +279,11 @@ def invoke_internal(self, context: InvocationContext, services: "InvocationServi
     """Whether this node's entire execution may be temporarily re-pinned to an idle GPU when
     `offload_text_encoders_to_idle_gpus` is enabled in multi-GPU mode. Only set this to True on nodes
     that exclusively load encoder model(s), run a forward pass, and store their result on the CPU —
-    i.e. nodes that do no work tied to the session's own GPU. Set via the `@invocation` decorator."""
+    i.e. nodes that do no work tied to the session's own GPU. Set via the `@invocation` decorator.
+
+    Weigh the node's runtime before setting this: the borrow holds the lent GPU's exclusive-use lock
+    for the *whole* node — model load included — and a session dequeued onto that GPU blocks until it
+    is released. See `invokeai/backend/util/device_pool.py`."""
 
     UIConfig: ClassVar[UIConfigBase]
 

invokeai/app/invocations/ernie_image_text_encoder.py

@@ -34,6 +34,7 @@
     category="conditioning",
     version="1.0.0",
     classification=Classification.Prototype,
+    idle_gpu_offloadable=True,
 )
 class ErnieImageTextEncoderInvocation(BaseInvocation):
     """Encodes a prompt for ERNIE-Image generation, optionally rewriting it via the

invokeai/app/invocations/ideogram4_text_encoder.py

@@ -21,6 +21,7 @@
     category="conditioning",
     version="1.0.0",
     classification=Classification.Prototype,
+    idle_gpu_offloadable=True,
 )
 class Ideogram4TextEncoderInvocation(BaseInvocation):
     """Encodes a prompt for Ideogram 4 using the Qwen3-VL encoder.

invokeai/app/invocations/krea2_text_encoder.py

@@ -47,6 +47,7 @@
     category="conditioning",
     version="1.1.0",
     classification=Classification.Prototype,
+    idle_gpu_offloadable=True,
 )
 class Krea2TextEncoderInvocation(BaseInvocation):
     """Encodes a text prompt for Krea-2 using the Qwen3-VL text encoder.

invokeai/app/invocations/wan_text_encoder.py

@@ -31,6 +31,7 @@
     category="conditioning",
     version="1.0.0",
     classification=Classification.Prototype,
+    idle_gpu_offloadable=True,
 )
 class WanTextEncoderInvocation(BaseInvocation):
     """Encodes a text prompt for Wan 2.2 using the UMT5-XXL encoder.

invokeai/backend/util/device_pool.py

@@ -18,9 +18,17 @@
   encoder runs on the worker's own GPU instead.
 
 Because borrows are non-blocking try-acquires and a session only ever blocking-acquires its *own*
-device lock, there is no lock-ordering cycle — the design is deadlock-free. The only cost is that,
-in the startup race where a borrow wins the lock a moment before the lent GPU's own session starts,
-that session waits out the (short) encoder node before beginning.
+device lock, there is no lock-ordering cycle — the design is deadlock-free. The cost is that, in the
+startup race where a borrow wins the lock a moment before the lent GPU's own session starts, that
+session waits out the whole borrowed node before beginning.
+
+Note that "the whole borrowed node" includes the encoder's *model load*. Caches are per-device, so
+the first borrow of a given GPU always cold-loads the encoder into that GPU's cache — seconds, not
+milliseconds. Subsequent borrows of the same GPU hit that cache (borrow selection is sticky for this
+reason), so the stall amortizes. Keep that in mind before marking a node ``idle_gpu_offloadable``:
+the cost is bounded by how long the node runs, and a node that does substantial work *per execution*
+— an autoregressive ``generate()`` loop, say — makes the stall recur on every generation instead of
+amortizing away.
 """
 
 import threading

tests/app/invocations/test_idle_offload_encoder_output_devices.py

@@ -0,0 +1,154 @@
+"""The Wan, Krea-2, Ideogram 4 and ERNIE-Image text encoders are idle_gpu_offloadable: each may run
+on a borrowed idle GPU whose device-pool lock is released the moment the node returns. Like the
+other offloadable encoders (see test_flux2_klein_output_device.py, test_flux_redux_output_device.py),
+their saved conditioning must be detached and moved to the CPU — otherwise the embeddings stay
+resident on the borrowed device, pinning VRAM on a GPU another session may immediately start using,
+and the denoise loop that picks them up on the session's own GPU sees mixed devices.
+
+The flag is declared on the @invocation decorator, so nothing inside these node bodies hints that
+the contract exists; these tests are the guard against a future edit dropping the .to("cpu").
+"""
+
+from unittest.mock import MagicMock
+
+import torch
+
+
+def _gpu_tensor_yielding(cpu_tensor: MagicMock) -> MagicMock:
+    """A stand-in for a tensor on a (borrowed) GPU: .detach().to("cpu") yields the CPU copy."""
+    gpu_tensor = MagicMock(spec=torch.Tensor)
+    gpu_tensor.detach.return_value.to.return_value = cpu_tensor
+    return gpu_tensor
+
+
+def test_wan_conditioning_is_saved_on_cpu(monkeypatch):
+    from invokeai.app.invocations.wan_text_encoder import WanTextEncoderInvocation
+
+    invocation = WanTextEncoderInvocation.model_construct(prompt="a prompt", wan_t5_encoder=MagicMock())
+
+    cpu_embeds = MagicMock(spec=torch.Tensor)
+    cpu_mask = MagicMock(spec=torch.Tensor)
+    gpu_embeds = _gpu_tensor_yielding(cpu_embeds)
+    gpu_mask = _gpu_tensor_yielding(cpu_mask)
+    monkeypatch.setattr(invocation, "_encode", lambda context: (gpu_embeds, gpu_mask))
+
+    context = MagicMock()
+    context.conditioning.save.return_value = "cond-name"
+
+    invocation.invoke(context)
+
+    gpu_embeds.detach.return_value.to.assert_called_once_with("cpu")
+    gpu_mask.detach.return_value.to.assert_called_once_with("cpu")
+    info = context.conditioning.save.call_args.args[0].conditionings[0]
+    assert info.prompt_embeds is cpu_embeds
+    assert info.prompt_attention_mask is cpu_mask
+
+
+def test_wan_conditioning_tolerates_a_full_attention_mask(monkeypatch):
+    """_encode returns None for the mask when every token is valid; that must not blow up the
+    CPU move."""
+    from invokeai.app.invocations.wan_text_encoder import WanTextEncoderInvocation
+
+    invocation = WanTextEncoderInvocation.model_construct(prompt="a prompt", wan_t5_encoder=MagicMock())
+
+    cpu_embeds = MagicMock(spec=torch.Tensor)
+    monkeypatch.setattr(invocation, "_encode", lambda context: (_gpu_tensor_yielding(cpu_embeds), None))
+
+    context = MagicMock()
+    context.conditioning.save.return_value = "cond-name"
+
+    invocation.invoke(context)
+
+    info = context.conditioning.save.call_args.args[0].conditionings[0]
+    assert info.prompt_embeds is cpu_embeds
+    assert info.prompt_attention_mask is None
+
+
+def test_krea2_conditioning_is_saved_on_cpu(monkeypatch):
+    from invokeai.app.invocations.krea2_text_encoder import Krea2TextEncoderInvocation
+
+    invocation = Krea2TextEncoderInvocation.model_construct(prompt="a prompt", mask=None, qwen3_vl_encoder=MagicMock())
+
+    cpu_embeds = MagicMock(spec=torch.Tensor)
+    cpu_mask = MagicMock(spec=torch.Tensor)
+    gpu_embeds = _gpu_tensor_yielding(cpu_embeds)
+    gpu_mask = _gpu_tensor_yielding(cpu_mask)
+    monkeypatch.setattr(invocation, "_encode", lambda context: (gpu_embeds, gpu_mask))
+
+    context = MagicMock()
+    context.conditioning.save.return_value = "cond-name"
+
+    invocation.invoke(context)
+
+    gpu_embeds.detach.return_value.to.assert_called_once_with("cpu")
+    gpu_mask.detach.return_value.to.assert_called_once_with("cpu")
+    info = context.conditioning.save.call_args.args[0].conditionings[0]
+    assert info.prompt_embeds is cpu_embeds
+    assert info.prompt_embeds_mask is cpu_mask
+
+
+def test_krea2_regional_mask_is_passed_through_untouched(monkeypatch):
+    """The optional regional `mask` is a TensorField (a name reference), not a tensor: the encoder
+    must forward it as-is and leave the load/device placement to krea2_denoise. Touching it here
+    would resolve a tensor onto the borrowed GPU."""
+    from invokeai.app.invocations.fields import TensorField
+    from invokeai.app.invocations.krea2_text_encoder import Krea2TextEncoderInvocation
+
+    mask_field = TensorField(tensor_name="regional-mask")
+    invocation = Krea2TextEncoderInvocation.model_construct(
+        prompt="a prompt", mask=mask_field, qwen3_vl_encoder=MagicMock()
+    )
+    monkeypatch.setattr(
+        invocation, "_encode", lambda context: (_gpu_tensor_yielding(MagicMock(spec=torch.Tensor)), None)
+    )
+
+    context = MagicMock()
+    context.conditioning.save.return_value = "cond-name"
+
+    output = invocation.invoke(context)
+
+    assert output.conditioning.mask is mask_field
+    context.tensors.load.assert_not_called()
+
+
+def test_ideogram4_conditioning_is_saved_on_cpu(monkeypatch):
+    import invokeai.app.invocations.ideogram4_text_encoder as ideogram4_module
+    from invokeai.app.invocations.ideogram4_text_encoder import Ideogram4TextEncoderInvocation
+
+    invocation = Ideogram4TextEncoderInvocation.model_construct(prompt="a prompt", qwen3_encoder=MagicMock())
+
+    cpu_embeds = MagicMock(spec=torch.Tensor)
+    gpu_embeds = _gpu_tensor_yielding(cpu_embeds)
+    monkeypatch.setattr(ideogram4_module, "encode_qwen3vl_prompt", lambda prompt, tokenizer, text_encoder: gpu_embeds)
+
+    context = MagicMock()
+    # `with model_on_device() as (_, model)` — the mock's __enter__ must unpack into two values.
+    context.models.load.return_value.model_on_device.return_value.__enter__.return_value = (None, MagicMock())
+    context.conditioning.save.return_value = "cond-name"
+
+    invocation.invoke(context)
+
+    gpu_embeds.detach.return_value.to.assert_called_once_with("cpu")
+    info = context.conditioning.save.call_args.args[0].conditionings[0]
+    assert info.prompt_embeds is cpu_embeds
+
+
+def test_ernie_image_conditioning_is_saved_on_cpu(monkeypatch):
+    from invokeai.app.invocations.ernie_image_text_encoder import ErnieImageTextEncoderInvocation
+
+    invocation = ErnieImageTextEncoderInvocation.model_construct(
+        prompt="a prompt", text_encoder=MagicMock(), prompt_enhancer=None, use_prompt_enhancer=False
+    )
+
+    cpu_embeds = MagicMock(spec=torch.Tensor)
+    gpu_embeds = _gpu_tensor_yielding(cpu_embeds)
+    monkeypatch.setattr(invocation, "_encode_prompt", lambda context, prompt: gpu_embeds)
+
+    context = MagicMock()
+    context.conditioning.save.return_value = "cond-name"
+
+    invocation.invoke(context)
+
+    gpu_embeds.detach.return_value.to.assert_called_once_with("cpu")
+    info = context.conditioning.save.call_args.args[0].conditionings[0]
+    assert info.prompt_embeds is cpu_embeds

tests/app/services/session_processor/test_encoder_offload.py

@@ -239,3 +239,44 @@ def test_real_nodes_declare_the_marker_correctly():
     assert CompelInvocation.idle_gpu_offloadable is True
     # A non-encoder node defaults to False (never re-pinned to a borrowed GPU).
     assert IntegerInvocation.idle_gpu_offloadable is False
+
+
+def test_every_text_encoder_node_declares_the_marker():
+    """Every `*_text_encoder` node must carry the flag.
+
+    Four encoders (wan, krea2, ideogram4, ernie_image) were added after the flag landed and went
+    unmarked for several releases — on a multi-GPU box they silently ran on the session's own GPU,
+    holding VRAM the denoise model needed. Enumerating the registry rather than listing the known
+    encoders is deliberate: the next encoder to be added is the one this guards.
+
+    A text encoder that genuinely should not be offloaded (it does session-GPU work, or runs long
+    enough that holding the lent GPU's lock would stall a session dequeued onto it) belongs in
+    `_NOT_OFFLOADABLE` with a comment saying why.
+    """
+    import importlib
+
+    import invokeai.app.invocations as invocations_package
+    from invokeai.app.invocations.baseinvocation import InvocationRegistry
+
+    # Encoders that must NOT be offloadable. Empty today; add with a justification.
+    _NOT_OFFLOADABLE: set[str] = set()
+
+    for module_name in invocations_package.__all__:
+        importlib.import_module(f"invokeai.app.invocations.{module_name}")
+
+    encoders = {
+        cls.get_type(): cls.idle_gpu_offloadable
+        for cls in InvocationRegistry.get_invocation_classes()
+        if cls.get_type().endswith("_text_encoder")
+    }
+    # Sanity-check the enumeration itself: if this drops to nothing, the assertion below is vacuous.
+    assert len(encoders) >= 11, f"expected the known text encoders to be registered, got {sorted(encoders)}"
+
+    unmarked = sorted(t for t, flag in encoders.items() if not flag and t not in _NOT_OFFLOADABLE)
+    assert not unmarked, (
+        f"text-encoder nodes missing idle_gpu_offloadable=True: {unmarked}. "
+        "Add the flag to the @invocation decorator, or add the type to _NOT_OFFLOADABLE with a reason."
+    )
+    # The four that prompted this guard, named explicitly so a rename cannot silently drop them.
+    for node_type in ("wan_text_encoder", "krea2_text_encoder", "ideogram4_text_encoder", "ernie_image_text_encoder"):
+        assert encoders.get(node_type) is True, f"{node_type} is not registered as an offloadable text encoder"