diff --git a/docs/src/content/docs/development/Guides/creating-nodes.mdx b/docs/src/content/docs/development/Guides/creating-nodes.mdx index abc905f6e6a..2f8456478b0 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. 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/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/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 c5956dc51d5..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,24 +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) @@ -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) 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. diff --git a/invokeai/backend/util/device_pool.py b/invokeai/backend/util/device_pool.py index 1e6675161a6..c88310e9b01 100644 --- a/invokeai/backend/util/device_pool.py +++ b/invokeai/backend/util/device_pool.py @@ -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 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 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()) + + 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"