invoke-ai · lstein · Aug 3, 2026 · Aug 1, 2026 · Aug 1, 2026 · Aug 1, 2026
diff --git a/docs/src/content/docs/development/Guides/creating-nodes.mdx b/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. If your node has both kinds of work, split them: `ernie_image_prompt_enhancer` was carved out of `ernie_image_text_encoder` for exactly this reason, leaving the encoder offloadable and keeping the enhancer's `generate()` on the session's own GPU.
 
 :::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.

@@ -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]
 

@@ -0,0 +1,127 @@
+import json
+from contextlib import ExitStack
+from typing import Optional
+
+import torch
+from transformers import PreTrainedModel, PreTrainedTokenizerBase, StoppingCriteria, StoppingCriteriaList
+
+from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
+from invokeai.app.invocations.fields import (
+    Input,
+    InputField,
+    UIComponent,
+)
+from invokeai.app.invocations.model import PromptEnhancerField
+from invokeai.app.invocations.primitives import StringOutput
+from invokeai.app.services.session_processor.session_processor_common import CanceledException
+from invokeai.app.services.shared.invocation_context import InvocationContext
+
+# Hard ceiling on the prompt-enhancer's generation length. Upstream drives `max_new_tokens` off the
+# PE tokenizer's `model_max_length`, but that is unreliable as a bound: if the tokenizer config omits
+# it, transformers substitutes a sentinel (int(1e30)), and a rewrite that never emits EOS would hang
+# the graph. A rewritten image prompt is a few hundred tokens at most, so cap it.
+PE_MAX_NEW_TOKENS = 1024
+
+
+class _CancelStoppingCriteria(StoppingCriteria):
+    """Halts `generate()` when the session's cancel event fires.
+
+    `generate()` is a single opaque call from the graph's point of view, so without this a cancel
+    only takes effect once the whole rewrite has been sampled. The caller re-checks the cancel flag
+    afterwards and raises, so the truncated sequence this returns is never used.
+    """
+
+    def __init__(self, context: InvocationContext) -> None:
+        super().__init__()
+        self._context = context
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs: object) -> bool:
+        return self._context.util.is_canceled()
+
+
+@invocation(
+    "ernie_image_prompt_enhancer",
+    title="Prompt Enhancer - ERNIE-Image",
+    tags=["prompt", "ernie-image"],
+    category="conditioning",
+    version="1.0.0",
+    classification=Classification.Prototype,
+    # Deliberately NOT idle_gpu_offloadable, unlike `ernie_image_text_encoder` which this node was
+    # split out of. Offloading holds the lent GPU's exclusive-use lock for the whole node, and this
+    # one runs an autoregressive `generate()` of up to PE_MAX_NEW_TOKENS on *every* execution — a
+    # cost that never amortizes into the borrowed device's model cache the way an encoder forward
+    # does. Borrowing here would re-stall a GPU that another session may be waiting to start on.
+)
+class ErnieImagePromptEnhancerInvocation(BaseInvocation):
+    """Rewrites a prompt for ERNIE-Image generation using the pipeline's bundled prompt-enhancer
+    (Ministral3ForCausalLM), sized for the intended output dimensions.
+
+    If no prompt-enhancer is connected — the pipeline may not ship one — the prompt is passed
+    through unchanged.
+    """
+
+    prompt: str = InputField(description="Text prompt to rewrite.", ui_component=UIComponent.Textarea)
+
+    prompt_enhancer: Optional[PromptEnhancerField] = InputField(
+        default=None,
+        title="Prompt Enhancer",
+        description="The prompt-enhancer model. If not connected, the prompt is passed through unchanged.",
+        input=Input.Connection,
+    )
+
+    width: int = InputField(default=1024, description="Target width the prompt is rewritten for.")
+    height: int = InputField(default=1024, description="Target height the prompt is rewritten for.")
+    temperature: float = InputField(default=0.6, ge=0.0, le=2.0)
+    top_p: float = InputField(default=0.95, ge=0.0, le=1.0)
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> StringOutput:
+        if self.prompt_enhancer is None:
+            return StringOutput(value=self.prompt)
+
+        enhanced = self._enhance_prompt(context, self.prompt)
+        context.logger.info(f"ERNIE-Image PE rewrote prompt -> {enhanced!r}")
+        return StringOutput(value=enhanced)
+
+    def _enhance_prompt(self, context: InvocationContext, prompt: str) -> str:
+        assert self.prompt_enhancer is not None  # checked by caller
+
+        tokenizer_info = context.models.load(self.prompt_enhancer.tokenizer)
+        lm_info = context.models.load(self.prompt_enhancer.text_encoder)
+
+        with ExitStack() as exit_stack:
+            (_, tokenizer) = exit_stack.enter_context(tokenizer_info.model_on_device())
+            (_, lm) = exit_stack.enter_context(lm_info.model_on_device())
+
+            if not isinstance(tokenizer, PreTrainedTokenizerBase):
+                raise TypeError(f"Expected tokenizer, got {type(tokenizer).__name__}")
+            if not isinstance(lm, PreTrainedModel):
+                raise TypeError(f"Expected PreTrainedModel for PE, got {type(lm).__name__}")
+
+            user_content = json.dumps(
+                {"prompt": prompt, "width": self.width, "height": self.height},
+                ensure_ascii=False,
+            )
+            input_text = tokenizer.apply_chat_template(
+                [{"role": "user", "content": user_content}],
+                tokenize=False,
+                add_generation_prompt=False,
+            )
+            inputs = tokenizer(input_text, return_tensors="pt").to(lm.device)
+            output_ids = lm.generate(
+                **inputs,
+                max_new_tokens=min(tokenizer.model_max_length, PE_MAX_NEW_TOKENS),
+                do_sample=self.temperature != 1.0 or self.top_p != 1.0,
+                temperature=self.temperature,
+                top_p=self.top_p,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                stopping_criteria=StoppingCriteriaList([_CancelStoppingCriteria(context)]),
+            )
+            # The stopping criterion above cuts generation short on cancel, leaving a partial
+            # rewrite; discard it rather than encoding a truncated prompt.
+            if context.util.is_canceled():
+                raise CanceledException
+
+            generated_ids = output_ids[0][inputs["input_ids"].shape[1] :]
+            return tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
@@ -1,6 +1,4 @@
-import json
 from contextlib import ExitStack
-from typing import Optional
 
 import torch
 from transformers import PreTrainedModel, PreTrainedTokenizerBase
@@ -12,32 +10,31 @@
     InputField,
     UIComponent,
 )
-from invokeai.app.invocations.model import Mistral3EncoderField, PromptEnhancerField
+from invokeai.app.invocations.model import Mistral3EncoderField
 from invokeai.app.invocations.primitives import ErnieImageConditioningOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
     ConditioningFieldData,
     ErnieImageConditioningInfo,
 )
 
-# Hard ceiling on the prompt-enhancer's generation length. Upstream drives `max_new_tokens` off the
-# PE tokenizer's `model_max_length`, but that is unreliable as a bound: if the tokenizer config omits
-# it, transformers substitutes a sentinel (int(1e30)), and a rewrite that never emits EOS would hang
-# the graph. A rewritten image prompt is a few hundred tokens at most, so cap it.
-PE_MAX_NEW_TOKENS = 1024
-
 
 @invocation(
     "ernie_image_text_encoder",
     title="Prompt - ERNIE-Image",
     tags=["prompt", "conditioning", "ernie-image"],
     category="conditioning",
-    version="1.0.0",
+    version="2.0.0",
     classification=Classification.Prototype,
+    idle_gpu_offloadable=True,
 )
 class ErnieImageTextEncoderInvocation(BaseInvocation):
-    """Encodes a prompt for ERNIE-Image generation, optionally rewriting it via the
-    bundled prompt-enhancer (Ministral3ForCausalLM) before tokenization.
+    """Encodes a prompt for ERNIE-Image generation.
+
+    Rewriting a prompt with the pipeline's bundled prompt-enhancer is the separate
+    `ernie_image_prompt_enhancer` node; connect its output to `prompt` to enhance. Keeping the two
+    apart is what lets this node stay `idle_gpu_offloadable` — see that node's decorator for why the
+    enhancer must not be.
     """
 
     prompt: str = InputField(description="Text prompt to encode.", ui_component=UIComponent.Textarea)
@@ -48,32 +45,9 @@ class ErnieImageTextEncoderInvocation(BaseInvocation):
         input=Input.Connection,
     )
 
-    prompt_enhancer: Optional[PromptEnhancerField] = InputField(
-        default=None,
-        title="Prompt Enhancer",
-        description="If connected and `use_prompt_enhancer` is true, the PE model rewrites the prompt before encoding.",
-        input=Input.Connection,
-    )
-
-    use_prompt_enhancer: bool = InputField(
-        default=True,
-        description="Whether to run the prompt-enhancer (no-op if no PE field is connected).",
-        title="Use Prompt Enhancer",
-    )
-
-    pe_width: int = InputField(default=1024, description="Target width passed to the prompt enhancer.")
-    pe_height: int = InputField(default=1024, description="Target height passed to the prompt enhancer.")
-    pe_temperature: float = InputField(default=0.6, ge=0.0, le=2.0)
-    pe_top_p: float = InputField(default=0.95, ge=0.0, le=1.0)
-
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ErnieImageConditioningOutput:
-        prompt = self.prompt
-        if self.use_prompt_enhancer and self.prompt_enhancer is not None:
-            prompt = self._enhance_prompt(context, prompt)
-            context.logger.info(f"ERNIE-Image PE rewrote prompt -> {prompt!r}")
-
-        prompt_embeds = self._encode_prompt(context, prompt)
+        prompt_embeds = self._encode_prompt(context, self.prompt)
         prompt_embeds = prompt_embeds.detach().to("cpu")
         conditioning_data = ConditioningFieldData(
             conditionings=[ErnieImageConditioningInfo(prompt_embeds=prompt_embeds)]
@@ -83,43 +57,6 @@ def invoke(self, context: InvocationContext) -> ErnieImageConditioningOutput:
             conditioning=ErnieImageConditioningField(conditioning_name=conditioning_name)
         )
 
-    def _enhance_prompt(self, context: InvocationContext, prompt: str) -> str:
-        assert self.prompt_enhancer is not None  # checked by caller
-
-        tokenizer_info = context.models.load(self.prompt_enhancer.tokenizer)
-        lm_info = context.models.load(self.prompt_enhancer.text_encoder)
-
-        with ExitStack() as exit_stack:
-            (_, tokenizer) = exit_stack.enter_context(tokenizer_info.model_on_device())
-            (_, lm) = exit_stack.enter_context(lm_info.model_on_device())
-
-            if not isinstance(tokenizer, PreTrainedTokenizerBase):
-                raise TypeError(f"Expected tokenizer, got {type(tokenizer).__name__}")
-            if not isinstance(lm, PreTrainedModel):
-                raise TypeError(f"Expected PreTrainedModel for PE, got {type(lm).__name__}")
-
-            user_content = json.dumps(
-                {"prompt": prompt, "width": self.pe_width, "height": self.pe_height},
-                ensure_ascii=False,
-            )
-            input_text = tokenizer.apply_chat_template(
-                [{"role": "user", "content": user_content}],
-                tokenize=False,
-                add_generation_prompt=False,
-            )
-            inputs = tokenizer(input_text, return_tensors="pt").to(lm.device)
-            output_ids = lm.generate(
-                **inputs,
-                max_new_tokens=min(tokenizer.model_max_length, PE_MAX_NEW_TOKENS),
-                do_sample=self.pe_temperature != 1.0 or self.pe_top_p != 1.0,
-                temperature=self.pe_temperature,
-                top_p=self.pe_top_p,
-                pad_token_id=tokenizer.pad_token_id,
-                eos_token_id=tokenizer.eos_token_id,
-            )
-            generated_ids = output_ids[0][inputs["input_ids"].shape[1] :]
-            return tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
-
     def _encode_prompt(self, context: InvocationContext, prompt: str) -> torch.Tensor:
         text_encoder_info = context.models.load(self.text_encoder.text_encoder)
         tokenizer_info = context.models.load(self.text_encoder.tokenizer)

@@ -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.

@@ -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.

@@ -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.

@@ -18,9 +18,19 @@
   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. When a node has both kinds of work, splitting it is the way out:
+``ernie_image_prompt_enhancer`` is a separate, deliberately un-offloadable node for this reason, so
+that ``ernie_image_text_encoder`` can stay offloadable.
 """
 
 import threading