From b294a7fbc4d888c34d6557bb686bf88092fdde23 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Mon, 3 Aug 2026 06:24:06 -0500 Subject: [PATCH] feat(nodes): add seeded Text LLM sampling --- .../src/content/docs/features/prompt-tools.md | 6 +- invokeai/app/api/routers/utilities.py | 2 + invokeai/app/invocations/text_llm.py | 11 ++- invokeai/backend/text_llm_pipeline.py | 36 +++++++- invokeai/frontend/web/openapi.json | 28 +++++- .../frontend/web/src/services/api/schema.ts | 12 +++ .../invocations/test_text_llm_with_preset.py | 2 + .../text_llm/test_text_llm_api_models.py | 35 ++++++- .../text_llm/test_text_llm_pipeline.py | 91 ++++++++++++++++++- 9 files changed, 214 insertions(+), 9 deletions(-) diff --git a/docs/src/content/docs/features/prompt-tools.md b/docs/src/content/docs/features/prompt-tools.md index 45bfd96170a..ff1a1840e02 100644 --- a/docs/src/content/docs/features/prompt-tools.md +++ b/docs/src/content/docs/features/prompt-tools.md @@ -2,7 +2,7 @@ title: LLM Prompt Tools sidebar: order: 3 -lastUpdated: 2026-05-23 +lastUpdated: 2026-08-03 --- InvokeAI includes two built-in tools that use local language models to help you write better prompts. Both tools appear as small buttons in the top-right corner of the positive prompt area and are only visible when you have a compatible model installed. @@ -52,4 +52,6 @@ Both tools overwrite your current prompt. You can undo this change: ## Workflow Node -A **Text LLM** node is also available in the workflow editor for use in automated pipelines. It accepts a prompt string and model selection as inputs and outputs the expanded text as a string. +The workflow editor provides **Text LLM** and **Text LLM (with System Prompt Preset)** nodes for automated pipelines. Both accept a prompt, model, maximum token count, and seed, then output generated text as a string. The preset variant reads its system prompt from the System Prompts library. + +Using the same seed reproduces sampling when the model, prompt, settings, hardware, and software versions remain the same. Results may differ across devices or software versions. The prompt area's **Expand Prompt** tool chooses a fresh seed for each request so repeated expansions can vary. diff --git a/invokeai/app/api/routers/utilities.py b/invokeai/app/api/routers/utilities.py index 023f653df6c..da4d939c57b 100644 --- a/invokeai/app/api/routers/utilities.py +++ b/invokeai/app/api/routers/utilities.py @@ -20,6 +20,7 @@ from invokeai.app.services.image_files.image_files_common import ImageFileNotFoundException from invokeai.app.services.model_records.model_records_base import UnknownModelException from invokeai.app.util.dynamicprompts import find_missing_wildcards +from invokeai.app.util.misc import get_random_seed from invokeai.backend.llava_onevision_pipeline import LlavaOnevisionPipeline from invokeai.backend.model_manager.taxonomy import ModelType from invokeai.backend.text_llm_pipeline import DEFAULT_SYSTEM_PROMPT, ProgressCallback, TextLLMPipeline @@ -166,6 +167,7 @@ def _run_expand_prompt( prompt=prompt, system_prompt=system_prompt or DEFAULT_SYSTEM_PROMPT, max_new_tokens=max_tokens, + seed=get_random_seed(), device=model_device, dtype=TorchDevice.choose_torch_dtype(), progress_callback=progress_callback, diff --git a/invokeai/app/invocations/text_llm.py b/invokeai/app/invocations/text_llm.py index 8308675cc0a..036591bdd1f 100644 --- a/invokeai/app/invocations/text_llm.py +++ b/invokeai/app/invocations/text_llm.py @@ -10,6 +10,7 @@ SystemPromptNotFoundError, SystemPromptRecordDTO, ) +from invokeai.app.util.misc import SEED_MAX from invokeai.backend.model_manager.taxonomy import ModelType from invokeai.backend.text_llm_pipeline import DEFAULT_SYSTEM_PROMPT, TextLLMPipeline from invokeai.backend.util.devices import TorchDevice @@ -21,6 +22,7 @@ def _run_text_llm( prompt: str, system_prompt: str, max_tokens: int, + seed: int, ) -> str: """Shared LLM invocation body used by every text-LLM node in this module.""" model_config = context.models.get_config(text_llm_model) @@ -35,6 +37,7 @@ def _run_text_llm( prompt=prompt, system_prompt=system_prompt, max_new_tokens=max_tokens, + seed=seed, device=model_device, dtype=TorchDevice.choose_torch_dtype(), ) @@ -45,7 +48,7 @@ def _run_text_llm( title="Text LLM", tags=["llm", "text", "prompt"], category="llm", - version="1.0.0", + version="1.1.0", classification=Classification.Beta, ) class TextLLMInvocation(BaseInvocation): @@ -72,6 +75,7 @@ class TextLLMInvocation(BaseInvocation): le=2048, description="Maximum number of tokens to generate.", ) + seed: int = InputField(default=0, ge=0, le=SEED_MAX, description=FieldDescriptions.seed) @torch.no_grad() def invoke(self, context: InvocationContext) -> StringOutput: @@ -81,6 +85,7 @@ def invoke(self, context: InvocationContext) -> StringOutput: prompt=self.prompt, system_prompt=self.system_prompt, max_tokens=self.max_tokens, + seed=self.seed, ) return StringOutput(value=output) @@ -90,7 +95,7 @@ def invoke(self, context: InvocationContext) -> StringOutput: title="Text LLM (with System Prompt Preset)", tags=["llm", "text", "prompt", "preset", "template"], category="llm", - version="1.0.0", + version="1.1.0", classification=Classification.Beta, ) class TextLLMWithPresetInvocation(BaseInvocation): @@ -125,6 +130,7 @@ class TextLLMWithPresetInvocation(BaseInvocation): le=2048, description="Maximum number of tokens to generate.", ) + seed: int = InputField(default=0, ge=0, le=SEED_MAX, description=FieldDescriptions.seed) def _resolve_system_prompt(self, context: InvocationContext) -> SystemPromptRecordDTO: """Resolve the referenced preset, enforcing the same access rules as the REST API. @@ -169,5 +175,6 @@ def invoke(self, context: InvocationContext) -> StringOutput: prompt=self.prompt, system_prompt=record.content, max_tokens=self.max_tokens, + seed=self.seed, ) return StringOutput(value=output) diff --git a/invokeai/backend/text_llm_pipeline.py b/invokeai/backend/text_llm_pipeline.py index 740313385ad..78a56976c32 100644 --- a/invokeai/backend/text_llm_pipeline.py +++ b/invokeai/backend/text_llm_pipeline.py @@ -4,6 +4,7 @@ from typing import Callable import torch +from torch.overrides import TorchFunctionMode from transformers import PreTrainedModel, PreTrainedTokenizerBase, TextIteratorStreamer DEFAULT_SYSTEM_PROMPT = ( @@ -29,6 +30,37 @@ PROGRESS_EMIT_INTERVAL = 0.1 +class _SeededMultinomialMode(TorchFunctionMode): + """Use invocation-local generators for multinomial sampling in this thread.""" + + def __init__(self, seed: int): + self._seed = seed + self._generators: dict[torch.device, torch.Generator] = {} + + def _get_generator(self, device: torch.device) -> torch.Generator: + generator_device = torch.device("cpu") if device.type == "mps" else device + if generator_device not in self._generators: + self._generators[generator_device] = torch.Generator(device=generator_device).manual_seed(self._seed) + return self._generators[generator_device] + + def __torch_function__(self, func, types, args=(), kwargs=None): + kwargs = dict(kwargs or {}) + if func is not torch.multinomial or kwargs.get("generator") is not None: + return func(*args, **kwargs) + + probabilities = args[0] if args else kwargs["input"] + generator = self._get_generator(probabilities.device) + kwargs["generator"] = generator + if probabilities.device.type != "mps": + return func(*args, **kwargs) + + if args: + args = (probabilities.cpu(), *args[1:]) + else: + kwargs["input"] = probabilities.cpu() + return func(*args, **kwargs).to(probabilities.device) + + class TextLLMPipeline: """A wrapper for a causal language model + tokenizer for text generation.""" @@ -41,6 +73,7 @@ def run( prompt: str, system_prompt: str = DEFAULT_SYSTEM_PROMPT, max_new_tokens: int = 300, + seed: int = 0, device: torch.device = torch.device("cpu"), dtype: torch.dtype = torch.float16, progress_callback: ProgressCallback | None = None, @@ -91,7 +124,8 @@ def run( def _generate() -> None: try: - self._model.generate(**generation_kwargs) + with _SeededMultinomialMode(seed): + self._model.generate(**generation_kwargs) except BaseException as e: generation_error.append(e) # transformers only calls streamer.end() on the normal exit of the diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index e19d0163e31..3a652fff328 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -81824,6 +81824,18 @@ "title": "Max Tokens", "type": "integer" }, + "seed": { + "default": 0, + "description": "Seed for random number generation", + "field_kind": "input", + "input": "any", + "maximum": 4294967295, + "minimum": 0, + "orig_default": 0, + "orig_required": false, + "title": "Seed", + "type": "integer" + }, "type": { "const": "text_llm", "default": "text_llm", @@ -81836,7 +81848,7 @@ "tags": ["llm", "text", "prompt"], "title": "Text LLM", "type": "object", - "version": "1.0.0", + "version": "1.1.0", "output": { "$ref": "#/components/schemas/StringOutput" } @@ -81927,6 +81939,18 @@ "title": "Max Tokens", "type": "integer" }, + "seed": { + "default": 0, + "description": "Seed for random number generation", + "field_kind": "input", + "input": "any", + "maximum": 4294967295, + "minimum": 0, + "orig_default": 0, + "orig_required": false, + "title": "Seed", + "type": "integer" + }, "type": { "const": "text_llm_with_preset", "default": "text_llm_with_preset", @@ -81939,7 +81963,7 @@ "tags": ["llm", "text", "prompt", "preset", "template"], "title": "Text LLM (with System Prompt Preset)", "type": "object", - "version": "1.0.0", + "version": "1.1.0", "output": { "$ref": "#/components/schemas/StringOutput" } diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 596472a0d43..c9d95f615ff 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -35633,6 +35633,12 @@ export type components = { * @default 300 */ max_tokens?: number; + /** + * Seed + * @description Seed for random number generation + * @default 0 + */ + seed?: number; /** * type * @default text_llm @@ -35694,6 +35700,12 @@ export type components = { * @default 300 */ max_tokens?: number; + /** + * Seed + * @description Seed for random number generation + * @default 0 + */ + seed?: number; /** * type * @default text_llm_with_preset diff --git a/tests/app/invocations/test_text_llm_with_preset.py b/tests/app/invocations/test_text_llm_with_preset.py index e891beee836..d47bbb710bf 100644 --- a/tests/app/invocations/test_text_llm_with_preset.py +++ b/tests/app/invocations/test_text_llm_with_preset.py @@ -26,6 +26,7 @@ def _make_invocation(prompt_id: str = "test-id") -> TextLLMWithPresetInvocation: system_prompt=SystemPromptField(system_prompt_id=prompt_id), text_llm_model=ModelIdentifierField(key="dummy", hash="x", name="dummy", base="any", type="text_llm"), max_tokens=50, + seed=123, ) @@ -70,6 +71,7 @@ def test_preset_node_loads_content_from_db_and_passes_to_llm() -> None: assert kwargs["system_prompt"] == "custom system instruction" assert kwargs["prompt"] == "a cat" assert kwargs["max_tokens"] == 50 + assert kwargs["seed"] == 123 assert result.value == "expanded" diff --git a/tests/backend/text_llm/test_text_llm_api_models.py b/tests/backend/text_llm/test_text_llm_api_models.py index 99c1884acf9..b4e17faff7a 100644 --- a/tests/backend/text_llm/test_text_llm_api_models.py +++ b/tests/backend/text_llm/test_text_llm_api_models.py @@ -1,9 +1,19 @@ """Tests for TextLLM API request/response models and validation.""" +from unittest.mock import MagicMock, patch + import pytest +import torch from pydantic import ValidationError -from invokeai.app.api.routers.utilities import ExpandPromptRequest, ExpandPromptResponse, ImageToPromptRequest +from invokeai.app.api.dependencies import ApiDependencies +from invokeai.app.api.routers.utilities import ( + ExpandPromptRequest, + ExpandPromptResponse, + ImageToPromptRequest, + _run_expand_prompt, +) +from invokeai.backend.model_manager.taxonomy import ModelType class TestExpandPromptRequest: @@ -52,3 +62,26 @@ def test_success_response(self): def test_error_response(self): resp = ExpandPromptResponse(expanded_prompt="", error="Model failed") assert resp.error == "Model failed" + + +def test_expand_prompt_uses_fresh_seed() -> None: + model_config = MagicMock(type=ModelType.TextLLM, path="model") + model = MagicMock() + model.parameters.return_value = iter([torch.nn.Parameter(torch.zeros(1))]) + loaded_model = MagicMock() + loaded_model.model_on_device.return_value.__enter__.return_value = (None, model) + services = MagicMock() + services.model_manager.store.get_model.return_value = model_config + services.model_manager.load.load_model.return_value = loaded_model + + with ( + patch.object(ApiDependencies, "invoker", MagicMock(services=services), create=True), + patch("invokeai.app.api.routers.utilities._resolve_model_path", return_value="model"), + patch("invokeai.app.api.routers.utilities.AutoTokenizer.from_pretrained"), + patch("invokeai.app.api.routers.utilities.get_random_seed", return_value=123), + patch("invokeai.app.api.routers.utilities.TextLLMPipeline") as pipeline_class, + ): + pipeline_class.return_value.run.return_value = "expanded" + assert _run_expand_prompt("cat", "model", 10, None, None, "user") == "expanded" + + assert pipeline_class.return_value.run.call_args.kwargs["seed"] == 123 diff --git a/tests/backend/text_llm/test_text_llm_pipeline.py b/tests/backend/text_llm/test_text_llm_pipeline.py index 08a3bcf57cf..54e365d2eb0 100644 --- a/tests/backend/text_llm/test_text_llm_pipeline.py +++ b/tests/backend/text_llm/test_text_llm_pipeline.py @@ -3,9 +3,11 @@ import threading from unittest.mock import MagicMock, patch +import pytest import torch -from invokeai.backend.text_llm_pipeline import DEFAULT_SYSTEM_PROMPT, TextLLMPipeline +from invokeai.backend import text_llm_pipeline +from invokeai.backend.text_llm_pipeline import DEFAULT_SYSTEM_PROMPT, TextLLMPipeline, _SeededMultinomialMode def _make_mock_tokenizer(has_chat_template: bool = True) -> MagicMock: @@ -115,6 +117,93 @@ def test_pipeline_passes_generation_params(): assert "streamer" in generate_kwargs +def test_seeded_multinomial_is_repeatable_despite_concurrent_global_rng_use(): + """Unrelated RNG use must not alter sampling for a controlled seed.""" + probabilities = torch.ones(100) + global_rng_state = torch.random.get_rng_state() + + with _SeededMultinomialMode(seed=42): + expected = torch.multinomial(probabilities, num_samples=10, replacement=True) + + assert torch.equal(torch.random.get_rng_state(), global_rng_state) + + interference_done = threading.Event() + + def _interfere() -> None: + torch.multinomial(probabilities, num_samples=10, replacement=True) + interference_done.set() + + with _SeededMultinomialMode(seed=42): + thread = threading.Thread(target=_interfere) + thread.start() + assert interference_done.wait(timeout=1) + actual = torch.multinomial(probabilities, num_samples=10, replacement=True) + thread.join() + + assert torch.equal(actual, expected) + + +def test_seeded_multinomial_contexts_are_isolated_when_interleaved_on_cpu(): + """Concurrent seeded contexts must retain independent RNG sequences.""" + probabilities = torch.arange(1, 101, dtype=torch.float32) + seeds = (42, 1234) + draw_count = 100 + + expected: dict[int, list[int]] = {} + for seed in seeds: + with _SeededMultinomialMode(seed=seed): + expected[seed] = [torch.multinomial(probabilities, num_samples=1).item() for _ in range(draw_count)] + + barrier = threading.Barrier(len(seeds)) + actual: dict[int, list[int]] = {} + + def _sample(seed: int) -> None: + samples: list[int] = [] + with _SeededMultinomialMode(seed=seed): + for _ in range(draw_count): + barrier.wait(timeout=5) + samples.append(torch.multinomial(probabilities, num_samples=1).item()) + actual[seed] = samples + + threads = [threading.Thread(target=_sample, args=(seed,), daemon=True) for seed in seeds] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=5) + + assert all(not thread.is_alive() for thread in threads) + assert actual == expected + + +def test_stalled_generation_does_not_block_later_generation(monkeypatch: pytest.MonkeyPatch): + """A timed-out worker must not own shared state needed by later runs.""" + monkeypatch.setattr(text_llm_pipeline, "STREAM_TIMEOUT", 0.05) + stalled_model = MagicMock() + release_stalled_model = threading.Event() + stalled_model.generate.side_effect = lambda **kwargs: release_stalled_model.wait() + + with pytest.raises(RuntimeError, match="Text generation stalled"): + TextLLMPipeline(stalled_model, _make_mock_tokenizer()).run("test", device=torch.device("cpu")) + + healthy_model = MagicMock() + + def _generate(**kwargs): + streamer = kwargs["streamer"] + streamer.put(torch.tensor([[1, 2, 3, 4, 5]])) + streamer.put(torch.tensor([6])) + streamer.end() + + healthy_model.generate.side_effect = _generate + healthy_tokenizer = _make_mock_tokenizer() + healthy_tokenizer.decode.return_value = "ok" + try: + TextLLMPipeline(healthy_model, healthy_tokenizer).run("test", device=torch.device("cpu")) + finally: + release_stalled_model.set() + + healthy_model.generate.assert_called_once() + + def test_pipeline_returns_joined_streamed_chunks(): """Pipeline should return the concatenated, stripped streamer output.""" tokenizer = _make_mock_tokenizer(has_chat_template=True)