From 0dd38a12393c319d3b59620598870c7a01926017 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 1 Aug 2026 03:47:32 +0200 Subject: [PATCH 1/3] feat(multi-gpu): mark the remaining text encoders as idle-GPU offloadable The `idle_gpu_offloadable` flag (#9311) lets an encoder-only node run on a borrowed idle GPU when `offload_text_encoders_to_idle_gpus` is enabled. Every text encoder that existed at the time was marked; the four added since were not, so on a multi-GPU machine they always occupy the session's own GPU. Marks wan, krea2, ideogram4 and ernie_image. All four meet the flag's stated condition: each returns its result via `detach().to("cpu")` and saves only the conditioning name, doing no work on the session device. Krea-2's optional `mask` input is passed through as a TensorField and is not processed until the denoise node. No schema change: the flag is a ClassVar set by the @invocation decorator, not a pydantic field, so no node version bumps are needed. --- invokeai/app/invocations/ernie_image_text_encoder.py | 1 + invokeai/app/invocations/ideogram4_text_encoder.py | 1 + invokeai/app/invocations/krea2_text_encoder.py | 1 + invokeai/app/invocations/wan_text_encoder.py | 1 + 4 files changed, 4 insertions(+) diff --git a/invokeai/app/invocations/ernie_image_text_encoder.py b/invokeai/app/invocations/ernie_image_text_encoder.py index c5956dc51d5..5c5baa24515 100644 --- a/invokeai/app/invocations/ernie_image_text_encoder.py +++ b/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 diff --git a/invokeai/app/invocations/ideogram4_text_encoder.py b/invokeai/app/invocations/ideogram4_text_encoder.py index ca2a809f1bb..9ca08640f16 100644 --- a/invokeai/app/invocations/ideogram4_text_encoder.py +++ b/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. diff --git a/invokeai/app/invocations/krea2_text_encoder.py b/invokeai/app/invocations/krea2_text_encoder.py index 4624edbeb9f..40c1763ac8a 100644 --- a/invokeai/app/invocations/krea2_text_encoder.py +++ b/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. diff --git a/invokeai/app/invocations/wan_text_encoder.py b/invokeai/app/invocations/wan_text_encoder.py index 5039230e1a0..0eddf40a4d9 100644 --- a/invokeai/app/invocations/wan_text_encoder.py +++ b/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. From ab47529338ded240f363a5e19df7c3707e8dbbb7 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sat, 1 Aug 2026 15:07:22 -0400 Subject: [PATCH 2/3] test(multi-gpu): guard the idle-GPU offload contract for the new encoders The four encoders this PR marks store their conditioning on the CPU, which is what makes the borrowed-GPU handoff safe -- but the flag lives on the @invocation decorator, so nothing in the node bodies hints that the contract exists. Add the regression tests that already exist for flux2_klein and flux_redux: - a registry-driven guard that every *_text_encoder node carries the flag, so the next encoder to be added is caught rather than the four already fixed; - per-node tests that each of the four detaches and moves its conditioning to the CPU, plus that krea2 forwards its regional mask TensorField untouched (resolving it here would pull a tensor onto the borrowed GPU). Also correct the borrow-cost documentation. device_pool.py claimed a lending session "waits out the (short) encoder node"; the borrow actually spans the node's model load, and caches are per-device, so the first borrow of a GPU always cold-loads the encoder there. That cost amortizes across later borrows, which hit the cache -- but work that recurs per execution does not amortize. ernie_image_text_encoder runs an autoregressive generate() for its optional prompt enhancer inside the borrow, so it stalls the lent GPU on every generation rather than once. Noted for reviewers in the flag docs and the node-authoring guide. Co-Authored-By: Claude Opus 5 (1M context) --- .../development/Guides/creating-nodes.mdx | 1 + invokeai/app/invocations/baseinvocation.py | 6 +- invokeai/backend/util/device_pool.py | 14 +- ...est_idle_offload_encoder_output_devices.py | 154 ++++++++++++++++++ .../session_processor/test_encoder_offload.py | 41 +++++ 5 files changed, 212 insertions(+), 4 deletions(-) create mode 100644 tests/app/invocations/test_idle_offload_encoder_output_devices.py diff --git a/docs/src/content/docs/development/Guides/creating-nodes.mdx b/docs/src/content/docs/development/Guides/creating-nodes.mdx index abc905f6e6a..a34e7a068f9 100644 --- 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. :::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. diff --git a/invokeai/app/invocations/baseinvocation.py b/invokeai/app/invocations/baseinvocation.py index 759e7d143db..1c1a35d700e 100644 --- a/invokeai/app/invocations/baseinvocation.py +++ b/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] diff --git a/invokeai/backend/util/device_pool.py b/invokeai/backend/util/device_pool.py index 1e6675161a6..e1a3b5bd035 100644 --- a/invokeai/backend/util/device_pool.py +++ b/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 diff --git a/tests/app/invocations/test_idle_offload_encoder_output_devices.py b/tests/app/invocations/test_idle_offload_encoder_output_devices.py new file mode 100644 index 00000000000..0ca4861b10e --- /dev/null +++ b/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 diff --git a/tests/app/services/session_processor/test_encoder_offload.py b/tests/app/services/session_processor/test_encoder_offload.py index a2254584b06..526bc87c5a7 100644 --- a/tests/app/services/session_processor/test_encoder_offload.py +++ b/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" From 52b9e4d318e2534388b9dc43182ae898d26096d8 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Mon, 3 Aug 2026 02:22:59 +0200 Subject: [PATCH 3/3] feat(multi-gpu): split the ERNIE-Image prompt enhancer into its own node MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ernie_image_text_encoder` is `idle_gpu_offloadable`, so its whole execution runs on a borrowed idle GPU whose exclusive-use lock is held until the node returns. An encoder forward is short and its model load amortizes into the borrowed device's cache; the bundled prompt enhancer's autoregressive `generate()` — up to 1024 tokens — runs afresh on every generation and never amortizes, so it re-stalled the lent GPU each time. A session dequeued onto that GPU logged `Executing queue item N` and then blocked in `acquire_session()` for the length of someone else's prompt rewrite. Carve the rewrite out into `ernie_image_prompt_enhancer`, a StringOutput node that is deliberately *not* offloadable. The encoder becomes genuinely encoder-only and keeps the flag; the enhancer stays on the session's own GPU. A per-instance opt-out was not viable: `buildErnieImageGraph` emits two encoders sharing one text encoder, and gating on the enhancer would split that pair across GPUs and cold-load the same encoder twice. The graph builder now wires the enhancer between the prompt node and the positive encoder when the toggle is on, and wires the prompt straight through otherwise. With no enhancer connected the node passes the prompt through, so pipelines that ship no PE submodel behave as before. Also make the rewrite cancelable: `generate()` gets a StoppingCriteria bound to the session's cancel event, and a cancelled run raises rather than encoding the truncated prompt. `ernie_image_text_encoder` goes to 2.0.0 — the six enhancer fields are removed, which breaks saved workflows that set them. It is a Prototype node. Tests pin the split in both directions: the enhancer must not become offloadable, and the encoder must not regain enhancer fields (either change would silently invalidate its flag with nothing else failing). Plus passthrough, the token cap, cancellation, and the new graph wiring in both toggle states. --- .../development/Guides/creating-nodes.mdx | 2 +- .../ernie_image_prompt_enhancer.py | 127 ++++++++++++++ .../invocations/ernie_image_text_encoder.py | 82 +-------- invokeai/backend/util/device_pool.py | 4 +- invokeai/frontend/web/openapi.json | 156 +++++++++++++----- .../generation/buildErnieImageGraph.test.ts | 43 +++-- .../graph/generation/buildErnieImageGraph.ts | 35 ++-- .../frontend/web/src/services/api/schema.ts | 104 ++++++++---- .../test_ernie_image_prompt_enhancer.py | 131 +++++++++++++++ ...est_idle_offload_encoder_output_devices.py | 4 +- 10 files changed, 516 insertions(+), 172 deletions(-) create mode 100644 invokeai/app/invocations/ernie_image_prompt_enhancer.py create mode 100644 tests/app/invocations/test_ernie_image_prompt_enhancer.py diff --git a/docs/src/content/docs/development/Guides/creating-nodes.mdx b/docs/src/content/docs/development/Guides/creating-nodes.mdx index a34e7a068f9..2f8456478b0 100644 --- a/docs/src/content/docs/development/Guides/creating-nodes.mdx +++ b/docs/src/content/docs/development/Guides/creating-nodes.mdx @@ -49,7 +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. +- **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. diff --git a/invokeai/app/invocations/ernie_image_prompt_enhancer.py b/invokeai/app/invocations/ernie_image_prompt_enhancer.py new file mode 100644 index 00000000000..f5314b2fef1 --- /dev/null +++ b/invokeai/app/invocations/ernie_image_prompt_enhancer.py @@ -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() diff --git a/invokeai/app/invocations/ernie_image_text_encoder.py b/invokeai/app/invocations/ernie_image_text_encoder.py index 5c5baa24515..476294b0631 100644 --- a/invokeai/app/invocations/ernie_image_text_encoder.py +++ b/invokeai/app/invocations/ernie_image_text_encoder.py @@ -1,6 +1,4 @@ -import json from contextlib import ExitStack -from typing import Optional import torch from transformers import PreTrainedModel, PreTrainedTokenizerBase @@ -12,7 +10,7 @@ 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 ( @@ -20,25 +18,23 @@ 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) @@ -49,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)] @@ -84,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) diff --git a/invokeai/backend/util/device_pool.py b/invokeai/backend/util/device_pool.py index e1a3b5bd035..c88310e9b01 100644 --- a/invokeai/backend/util/device_pool.py +++ b/invokeai/backend/util/device_pool.py @@ -28,7 +28,9 @@ 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. +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 diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index add72e7ac55..a509f79080c 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -25720,11 +25720,11 @@ "title": "ErnieImageModelLoaderOutput", "type": "object" }, - "ErnieImageTextEncoderInvocation": { + "ErnieImagePromptEnhancerInvocation": { "category": "conditioning", "class": "invocation", "classification": "prototype", - "description": "Encodes a prompt for ERNIE-Image generation, optionally rewriting it via the\nbundled prompt-enhancer (Ministral3ForCausalLM) before tokenization.", + "description": "Rewrites a prompt for ERNIE-Image generation using the pipeline's bundled prompt-enhancer\n(Ministral3ForCausalLM), sized for the intended output dimensions.\n\nIf no prompt-enhancer is connected \u2014 the pipeline may not ship one \u2014 the prompt is passed\nthrough unchanged.", "node_pack": "invokeai", "properties": { "id": { @@ -25761,29 +25761,13 @@ } ], "default": null, - "description": "Text prompt to encode.", + "description": "Text prompt to rewrite.", "field_kind": "input", "input": "any", "orig_required": true, "title": "Prompt", "ui_component": "textarea" }, - "text_encoder": { - "anyOf": [ - { - "$ref": "#/components/schemas/Mistral3EncoderField" - }, - { - "type": "null" - } - ], - "default": null, - "description": "Mistral3 text encoder + tokenizer", - "field_kind": "input", - "input": "connection", - "orig_required": true, - "title": "Text Encoder" - }, "prompt_enhancer": { "anyOf": [ { @@ -25794,44 +25778,34 @@ } ], "default": null, - "description": "If connected and `use_prompt_enhancer` is true, the PE model rewrites the prompt before encoding.", + "description": "The prompt-enhancer model. If not connected, the prompt is passed through unchanged.", "field_kind": "input", "input": "connection", "orig_default": null, "orig_required": false, "title": "Prompt Enhancer" }, - "use_prompt_enhancer": { - "default": true, - "description": "Whether to run the prompt-enhancer (no-op if no PE field is connected).", - "field_kind": "input", - "input": "any", - "orig_default": true, - "orig_required": false, - "title": "Use Prompt Enhancer", - "type": "boolean" - }, - "pe_width": { + "width": { "default": 1024, - "description": "Target width passed to the prompt enhancer.", + "description": "Target width the prompt is rewritten for.", "field_kind": "input", "input": "any", "orig_default": 1024, "orig_required": false, - "title": "Pe Width", + "title": "Width", "type": "integer" }, - "pe_height": { + "height": { "default": 1024, - "description": "Target height passed to the prompt enhancer.", + "description": "Target height the prompt is rewritten for.", "field_kind": "input", "input": "any", "orig_default": 1024, "orig_required": false, - "title": "Pe Height", + "title": "Height", "type": "integer" }, - "pe_temperature": { + "temperature": { "default": 0.6, "field_kind": "input", "input": "any", @@ -25839,10 +25813,10 @@ "minimum": 0.0, "orig_default": 0.6, "orig_required": false, - "title": "Pe Temperature", + "title": "Temperature", "type": "number" }, - "pe_top_p": { + "top_p": { "default": 0.95, "field_kind": "input", "input": "any", @@ -25850,9 +25824,90 @@ "minimum": 0.0, "orig_default": 0.95, "orig_required": false, - "title": "Pe Top P", + "title": "Top P", "type": "number" }, + "type": { + "const": "ernie_image_prompt_enhancer", + "default": "ernie_image_prompt_enhancer", + "field_kind": "node_attribute", + "title": "type", + "type": "string" + } + }, + "required": ["type", "id"], + "tags": ["prompt", "ernie-image"], + "title": "Prompt Enhancer - ERNIE-Image", + "type": "object", + "version": "1.0.0", + "output": { + "$ref": "#/components/schemas/StringOutput" + } + }, + "ErnieImageTextEncoderInvocation": { + "category": "conditioning", + "class": "invocation", + "classification": "prototype", + "description": "Encodes a prompt for ERNIE-Image generation.\n\nRewriting a prompt with the pipeline's bundled prompt-enhancer is the separate\n`ernie_image_prompt_enhancer` node; connect its output to `prompt` to enhance. Keeping the two\napart is what lets this node stay `idle_gpu_offloadable` \u2014 see that node's decorator for why the\nenhancer must not be.", + "node_pack": "invokeai", + "properties": { + "id": { + "description": "The id of this instance of an invocation. Must be unique among all instances of invocations.", + "field_kind": "node_attribute", + "title": "Id", + "type": "string" + }, + "is_intermediate": { + "default": false, + "description": "Whether or not this is an intermediate invocation.", + "field_kind": "node_attribute", + "input": "direct", + "orig_required": true, + "title": "Is Intermediate", + "type": "boolean", + "ui_hidden": false, + "ui_type": "IsIntermediate" + }, + "use_cache": { + "default": true, + "description": "Whether or not to use the cache", + "field_kind": "node_attribute", + "title": "Use Cache", + "type": "boolean" + }, + "prompt": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Text prompt to encode.", + "field_kind": "input", + "input": "any", + "orig_required": true, + "title": "Prompt", + "ui_component": "textarea" + }, + "text_encoder": { + "anyOf": [ + { + "$ref": "#/components/schemas/Mistral3EncoderField" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Mistral3 text encoder + tokenizer", + "field_kind": "input", + "input": "connection", + "orig_required": true, + "title": "Text Encoder" + }, "type": { "const": "ernie_image_text_encoder", "default": "ernie_image_text_encoder", @@ -25865,7 +25920,7 @@ "tags": ["prompt", "conditioning", "ernie-image"], "title": "Prompt - ERNIE-Image", "type": "object", - "version": "1.0.0", + "version": "2.0.0", "output": { "$ref": "#/components/schemas/ErnieImageConditioningOutput" } @@ -33867,6 +33922,9 @@ { "$ref": "#/components/schemas/ErnieImageModelLoaderInvocation" }, + { + "$ref": "#/components/schemas/ErnieImagePromptEnhancerInvocation" + }, { "$ref": "#/components/schemas/ErnieImageTextEncoderInvocation" }, @@ -42381,6 +42439,9 @@ { "$ref": "#/components/schemas/ErnieImageModelLoaderInvocation" }, + { + "$ref": "#/components/schemas/ErnieImagePromptEnhancerInvocation" + }, { "$ref": "#/components/schemas/ErnieImageTextEncoderInvocation" }, @@ -43708,6 +43769,9 @@ { "$ref": "#/components/schemas/ErnieImageModelLoaderInvocation" }, + { + "$ref": "#/components/schemas/ErnieImagePromptEnhancerInvocation" + }, { "$ref": "#/components/schemas/ErnieImageTextEncoderInvocation" }, @@ -44647,6 +44711,9 @@ "ernie_image_model_loader": { "$ref": "#/components/schemas/ErnieImageModelLoaderOutput" }, + "ernie_image_prompt_enhancer": { + "$ref": "#/components/schemas/StringOutput" + }, "ernie_image_text_encoder": { "$ref": "#/components/schemas/ErnieImageConditioningOutput" }, @@ -45445,6 +45512,7 @@ "dynamic_prompt", "ernie_image_denoise", "ernie_image_model_loader", + "ernie_image_prompt_enhancer", "ernie_image_text_encoder", "ernie_image_vae_decode", "esrgan", @@ -45934,6 +46002,9 @@ { "$ref": "#/components/schemas/ErnieImageModelLoaderInvocation" }, + { + "$ref": "#/components/schemas/ErnieImagePromptEnhancerInvocation" + }, { "$ref": "#/components/schemas/ErnieImageTextEncoderInvocation" }, @@ -46976,6 +47047,9 @@ { "$ref": "#/components/schemas/ErnieImageModelLoaderInvocation" }, + { + "$ref": "#/components/schemas/ErnieImagePromptEnhancerInvocation" + }, { "$ref": "#/components/schemas/ErnieImageTextEncoderInvocation" }, diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildErnieImageGraph.test.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildErnieImageGraph.test.ts index 464871ae46e..e84b430dee5 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildErnieImageGraph.test.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildErnieImageGraph.test.ts @@ -103,6 +103,9 @@ const build = (generationMode: 'txt2img' | 'img2img' | 'inpaint' | 'outpaint' = const nodesOf = (g: BuiltGraph) => Object.values(g.getGraph().nodes) as Record[]; const positiveEncoder = (g: BuiltGraph) => nodesOf(g).find((n) => n.type === 'ernie_image_text_encoder' && String(n.id).startsWith('pos_prompt')); +const enhancer = (g: BuiltGraph) => nodesOf(g).find((n) => n.type === 'ernie_image_prompt_enhancer'); +const edgesInto = (g: BuiltGraph, node: Record | undefined, field: string) => + g.getGraph().edges.filter((e) => e.destination.node_id === node?.id && e.destination.field === field); describe('buildErnieImageGraph', () => { beforeEach(() => { @@ -130,19 +133,19 @@ describe('buildErnieImageGraph', () => { describe('prompt enhancer', () => { // The enhancer is handed the target size and rewrites the prompt to suit that aspect ratio. - // Its node defaults are 1024x1024, so leaving `pe_width`/`pe_height` unwired means a portrait + // Its node defaults are 1024x1024, so leaving `width`/`height` unwired means a portrait // generation is enhanced for a square image. it('passes the real generation dimensions to the enhancer', async () => { const { g } = await build(); - expect(positiveEncoder(g)).toMatchObject({ pe_width: 832, pe_height: 1216 }); + expect(enhancer(g)).toMatchObject({ width: 832, height: 1216 }); }); it('tracks the generation dimensions rather than hardcoding them', async () => { originalSize = { width: 1536, height: 640 }; const { g } = await build(); - expect(positiveEncoder(g)).toMatchObject({ pe_width: 1536, pe_height: 640 }); + expect(enhancer(g)).toMatchObject({ width: 1536, height: 640 }); }); it('reports the original size, not the intermediate scaled render size', async () => { @@ -150,18 +153,36 @@ describe('buildErnieImageGraph', () => { // user ends up with has the original dimensions -- that is the aspect ratio to enhance for. const { g } = await build(); - expect(positiveEncoder(g)).toMatchObject({ pe_width: originalSize.width }); - expect(positiveEncoder(g)).not.toMatchObject({ pe_width: scaledSize.width }); + expect(enhancer(g)).toMatchObject({ width: originalSize.width }); + expect(enhancer(g)).not.toMatchObject({ width: scaledSize.width }); }); - it('wires the enhancer edge only when the toggle is on', async () => { + // The enhancer is a node of its own rather than a mode of the encoder so that it keeps running + // on the session's GPU: `ernie_image_text_encoder` is `idle_gpu_offloadable`, and an + // autoregressive rewrite is far too long to hold a borrowed idle GPU's lock for. + it('runs the enhancer as its own node, feeding the positive encoder', async () => { const { g } = await build(); - expect(g.getGraph().edges.some((e) => e.destination.field === 'prompt_enhancer')).toBe(true); + const pe = enhancer(g); + + expect(pe).toBeDefined(); + expect(edgesInto(g, pe, 'prompt_enhancer')).toHaveLength(1); + // The user's prompt reaches the encoder *through* the enhancer, not directly. + const intoEncoderPrompt = edgesInto(g, positiveEncoder(g), 'prompt'); + expect(intoEncoderPrompt).toHaveLength(1); + expect(intoEncoderPrompt[0]?.source.node_id).toBe(pe?.id); + }); + it('omits the enhancer node entirely when the toggle is off', async () => { nextId = 0; usePromptEnhancer = false; - const { g: gOff } = await build(); - expect(gOff.getGraph().edges.some((e) => e.destination.field === 'prompt_enhancer')).toBe(false); + const { g } = await build(); + + expect(enhancer(g)).toBeUndefined(); + expect(g.getGraph().edges.some((e) => e.destination.field === 'prompt_enhancer')).toBe(false); + // ...and the prompt node then wires straight into the encoder, so nothing is left unwired. + const intoEncoderPrompt = edgesInto(g, positiveEncoder(g), 'prompt'); + expect(intoEncoderPrompt).toHaveLength(1); + expect(String(intoEncoderPrompt[0]?.source.node_id)).toMatch(/^positive_prompt/); }); it('never enhances the negative prompt', async () => { @@ -170,7 +191,9 @@ describe('buildErnieImageGraph', () => { (n) => n.type === 'ernie_image_text_encoder' && String(n.id).startsWith('neg_prompt') ); - expect(negEncoder).toMatchObject({ use_prompt_enhancer: false }); + // The negative encoder takes its prompt as a literal field, with no edge from the enhancer. + expect(negEncoder).toMatchObject({ prompt: 'a negative prompt' }); + expect(edgesInto(g, negEncoder, 'prompt')).toHaveLength(0); }); }); diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildErnieImageGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildErnieImageGraph.ts index 84c9ecdd205..98307ebab3c 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildErnieImageGraph.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildErnieImageGraph.ts @@ -59,7 +59,6 @@ export const buildErnieImageGraph = async (arg: GraphBuilderArg): Promise | null = null; @@ -68,8 +67,18 @@ export const buildErnieImageGraph = async (arg: GraphBuilderArg): Promise | null = null; + if (usePromptEnhancer) { + promptEnhancer = g.addNode({ + type: 'ernie_image_prompt_enhancer', + id: getPrefixedId('prompt_enhancer'), }); } @@ -96,12 +105,16 @@ export const buildErnieImageGraph = async (arg: GraphBuilderArg): Promise