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

[Pfannkuchensack]

Summary

The idle_gpu_offloadable flag introduced in #9311 lets an encoder-only node run its whole execution on a borrowed idle GPU when offload_text_encoders_to_idle_gpus is enabled. Every text encoder that existed at the time was marked, but four have been added since and were never given the flag — so on a multi-GPU machine they always run on the session's own GPU, holding VRAM that the generation needs.

Marks the four stragglers: wan_text_encoder, krea2_text_encoder, ideogram4_text_encoder, ernie_image_text_encoder.

All four satisfy the condition the flag documents ("encoder-only nodes that store their result on the CPU and do no work on the session's own GPU"): each returns via detach().to("cpu") and persists only the conditioning name. Krea-2's optional mask input is a TensorField that is passed through untouched and first processed in krea2_denoise.

After this change 14 text-encoder invocations carry the flag (15 including flux_redux, which is an image-prompt encoder rather than a text one).

Tests

The flag lives on the @invocation decorator, so nothing inside a node body hints that the CPU-store contract exists — a future edit dropping a .to("cpu") would break multi-GPU silently. Added the regression tests that flux2_klein and flux_redux already have:

• a registry-driven guard that every *_text_encoder node carries the flag (enumerating rather than listing the four, so the next encoder added is what it catches), with a documented _NOT_OFFLOADABLE escape hatch;
• per-node tests that each of the four detaches and moves its conditioning to the CPU, plus one asserting Krea-2 forwards its regional-mask TensorField untouched — resolving it in the encoder would pull a tensor onto the borrowed GPU.

Borrow-cost documentation

device_pool.py claimed a lending session "waits out the (short) encoder node". That was already inaccurate: the borrow 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 — later borrows are sticky and hit the cache — but work that recurs per execution does not.

ernie_image_text_encoder is the one node here where that matters: its optional prompt enhancer runs an autoregressive generate() of up to 1024 tokens inside the borrow, so it re-stalls the lent GPU on every generation instead of once. Corrected the claim in device_pool.py, the idle_gpu_offloadable docstring, and the node-authoring guide so the tradeoff is visible.

Reviewers may want to weigh in on ERNIE specifically. Marking it is still a net win in the common case and it is kept as-is here, but the options are not obvious — see the review comment below for the analysis, including why a per-instance opt-out (offload only when the enhancer is off) is worse than either alternative given how buildErnieImageGraph emits its positive/negative encoder pair.

Related Issues / Discussions

Follows up #9311 (which introduced the flag) and #9263 (multi-GPU parallel session execution).

QA Instructions

Requires two or more CUDA GPUs and offload_text_encoders_to_idle_gpus: true in invokeai.yaml.

1. Run a generation with Wan, Krea-2, Ideogram4 or ERNIE-Image while a second GPU is idle.
2. With InvokeAI at DEBUG level, the session processor logs Running <node type> on idle device cuda:N (session device cuda:M) for the text-encoder node. Before this PR that line never appeared for these four.
3. Confirm the conditioning still produces identical images — the node runs on a different device but its output is moved to the CPU either way. (Not applicable to ERNIE-Image with the prompt enhancer on: the rewrite is sampled, so it is not reproducible run to run, with or without this PR.)
4. Single-GPU installs are unaffected: with no idle device to borrow, the node runs exactly as before.

Merge Plan

Nothing special. No node versions are bumped because no fields changed:
idle_gpu_offloadable is a ClassVar set by the @invocation decorator, not a
pydantic field, so it does not appear in the node schema and schema.ts is
untouched. Bumping the version here would be actively harmful — it signals a
template change to the frontend and causes saved workflows to re-instantiate the
node, for no benefit.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration — n/a, backend only
• Documentation added / updated (if applicable) — borrow-cost criterion added to the node-authoring guide
• Updated What's New copy (if doing a release after this PR)

[lstein]

Reviewed this and pushed one commit to the branch (ab47529338) — tests plus a documentation correction. No behavioural change: all four markings are as you had them. Details below, including one design question I'd like your read on.

The markings check out

I traced each of the four against the contract. All four store their conditioning via detach().to("cpu") and persist only the conditioning name; Krea-2's regional mask is a TensorField name reference that is forwarded untouched and first resolved in krea2_denoise. Attacks that failed, for the record:

• Device-dependent output dtype. krea2_text_encoder.py:157 stamps the conditioning dtype from choose_bfloat16_safe_dtype(borrowed_device), which would be wrong on a heterogeneous pool — but krea2_denoise.py:159 re-casts to the session device's inference_dtype, so it cannot mismatch.
• Loaders resolving the wrong GPU. All four loaders derive target_device from TorchDevice.choose_torch_device(), which returns the thread-local pin first, so models land on the borrowed GPU as intended.
• Borrow lifecycle. The finally in _maybe_offload_to_idle_gpu restores the pin, the cache-stats swap and the lock on both the exception and CanceledException paths; the borrow wraps only invoke_internal, so output serialization is outside it.
• Missed candidates. Every *_text_encoder node is now marked. (wan_ref_image_encoder is a VAE image encoder that does image I/O and calls TorchDevice.empty_cache() — correctly out of scope here, but perhaps worth its own look.)

Two small notes: the body said "all 12 text-encoder invocations" — it's 14 after this PR (15 marked including flux_redux, which isn't a text encoder); and QA step 3 isn't checkable for ERNIE with the enhancer on, since the rewrite is sampled. Both corrected in the description.

What I added

The flag lives on the decorator, so nothing in a node body signals that the CPU-store contract exists — dropping a .to("cpu") later would break multi-GPU silently with nothing failing. flux2_klein and flux_redux each already have a regression test for exactly this; the four new ones now do too. The marker test enumerates the registry rather than listing the four, so what it actually guards is the next encoder somebody adds — which is the failure mode this PR is fixing.

The one thing I'd like your read on: ERNIE

device_pool.py said a lending session "waits out the (short) encoder node". That was already inaccurate before this PR — the borrow 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 amortizes (borrow selection is sticky, later borrows hit the cache). What doesn't amortize is work that recurs per execution, and ernie_image_text_encoder runs the prompt enhancer's autoregressive generate() — up to 1024 tokens, not interruptible by the cancel event — inside the borrow. So it re-stalls the lent GPU on every generation: a session dequeued onto that GPU logs Executing queue item N and then blocks in acquire_session() for the length of someone else's prompt rewrite.

I tried fixing that with a per-instance opt-out (offload only when the enhancer won't run) and backed it out — it's worse than either alternative. buildErnieImageGraph.ts:54-73,99-102 emits two ernie_image_text_encoder nodes sharing one text_encoder: pos_prompt with the enhancer on, and neg_prompt with it always off (guidance_scale > 1, the default). A per-instance rule splits that pair across GPUs: the positive node runs on the session GPU, loading the encoder and the PE there — exactly what the feature exists to avoid — and the negative node then borrows the idle GPU and cold-loads the same encoder a second time, for no benefit, since it's already resident. Strictly worse than either marking ERNIE or not marking it.

So the real options are:

1. Keep it marked (what's on the branch). Both encoder nodes borrow the same idle GPU, one encoder copy, session GPU free — the full benefit. Cost: the idle GPU's lock is held through the PE on every generation. Only bites when a second session is dequeued during that window, so on a single-user multi-GPU box it costs nothing.
2. Don't mark ERNIE. No stall, no benefit; the encoder and PE stay on the session GPU contending with the ERNIE transformer.
3. Split the prompt enhancer into its own node, leaving a genuinely encoder-only ernie_image_text_encoder. That gets both properties, but it's a node-schema change and a graph-builder change — a separate PR.

I left (1) in place since it's what you wrote and it's defensible, and documented the tradeoff in device_pool.py, the flag docstring, and the node-authoring guide so the next person marking a node weighs runtime rather than just "is it an encoder". Happy to go to (2) if you'd rather be conservative, and (3) seems like the right eventual answer.

[lstein]

LGTM and am happy to merge when you give the word. However, please take a look at the comment regarding the handling of ERNIE Image and let me know if you want to make the suggested architectural change of splitting the PE into its own node.