diff --git a/invokeai/app/api/routers/utilities.py b/invokeai/app/api/routers/utilities.py index 568546603ab..7356e688512 100644 --- a/invokeai/app/api/routers/utilities.py +++ b/invokeai/app/api/routers/utilities.py @@ -17,6 +17,7 @@ from invokeai.app.api.routers._access import assert_image_read_access 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.backend.llava_onevision_pipeline import LlavaOnevisionPipeline from invokeai.backend.model_manager.taxonomy import ModelType from invokeai.backend.text_llm_pipeline import DEFAULT_SYSTEM_PROMPT, TextLLMPipeline @@ -52,8 +53,19 @@ async def parse_dynamicprompts( """Creates a batch process""" max_prompts = min(max_prompts, 10000) generator: Union[RandomPromptGenerator, CombinatorialPromptGenerator] + error: Optional[str] = None + + # An unknown wildcard used as a variant value sends the combinatorial generator into an infinite + # loop, so bail out early with a clear message instead of hanging the request (and with it the UI + # preview). The random generator handles unknown wildcards gracefully, so only the combinatorial + # path is guarded. + if combinatorial: + missing_wildcards = find_missing_wildcards(prompt) + if missing_wildcards: + wildcards = ", ".join(missing_wildcards) + return DynamicPromptsResponse(prompts=[prompt], error=f"No values found for wildcard(s): {wildcards}") + try: - error: Optional[str] = None if combinatorial: generator = CombinatorialPromptGenerator() prompts = generator.generate(prompt, max_prompts=max_prompts) diff --git a/invokeai/app/invocations/prompt.py b/invokeai/app/invocations/prompt.py index 48eec0ac0ef..6658b0da7d0 100644 --- a/invokeai/app/invocations/prompt.py +++ b/invokeai/app/invocations/prompt.py @@ -9,6 +9,7 @@ from invokeai.app.invocations.fields import InputField, UIComponent from invokeai.app.invocations.primitives import StringCollectionOutput from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.app.util.dynamicprompts import find_missing_wildcards @invocation( @@ -31,6 +32,13 @@ class DynamicPromptInvocation(BaseInvocation): def invoke(self, context: InvocationContext) -> StringCollectionOutput: if self.combinatorial: + # An unknown wildcard used as a variant value sends the combinatorial generator into an + # infinite loop, so fail fast with a clear message instead of hanging the invocation. The + # random generator handles unknown wildcards gracefully and needs no guard. + missing_wildcards = find_missing_wildcards(self.prompt) + if missing_wildcards: + raise ValueError(f"No values found for wildcard(s): {', '.join(missing_wildcards)}") + generator = CombinatorialPromptGenerator() prompts = generator.generate(self.prompt, max_prompts=self.max_prompts) else: diff --git a/invokeai/app/util/dynamicprompts.py b/invokeai/app/util/dynamicprompts.py new file mode 100644 index 00000000000..5c922318218 --- /dev/null +++ b/invokeai/app/util/dynamicprompts.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +from collections.abc import Iterator + +from dynamicprompts.commands import ( + Command, + SequenceCommand, + VariantCommand, + WildcardCommand, + WrapCommand, +) +from dynamicprompts.parser.parse import parse +from dynamicprompts.wildcards import WildcardManager +from pyparsing import ParseException + + +def _iter_wildcard_names(command: Command, in_variant: bool = False) -> Iterator[str]: + """Recursively yield the statically-known wildcard names that appear as (part of) a variant value. + + Only wildcards reachable from a `VariantCommand` value are yielded: those are the references that + hang the combinatorial generator (see `find_missing_wildcards`). The same wildcard used as plain + literal text (e.g. `a __nope__ b`) generates fine, so it is intentionally not reported. + """ + if isinstance(command, WildcardCommand): + # The wildcard name may itself be a dynamic Command (e.g. `__${var}__`). Only plain string + # names can be validated ahead of time, so the dynamic case is intentionally skipped. + if in_variant and isinstance(command.wildcard, str): + yield command.wildcard + elif isinstance(command, SequenceCommand): + for token in command.tokens: + yield from _iter_wildcard_names(token, in_variant) + elif isinstance(command, VariantCommand): + # Everything below a variant value is a variant-nested reference, even across sequences. + for value in command.values: + yield from _iter_wildcard_names(value, in_variant=True) + elif isinstance(command, WrapCommand): + yield from _iter_wildcard_names(command.wrapper, in_variant) + yield from _iter_wildcard_names(command.inner, in_variant) + # LiteralCommand and variable commands reference no wildcards we can resolve statically. + + +def find_missing_wildcards(prompt: str, wildcard_manager: WildcardManager | None = None) -> list[str]: + """Return the unique unknown wildcard names in `prompt` that hang the combinatorial generator. + + Referencing an unknown wildcard *as a variant value* (e.g. `{__nope__|x}`) makes dynamicprompts' + combinatorial generator loop forever: its not-found fallback (`get_wildcard_not_found_fallback`) + yields the wrapped wildcard infinitely, and the combinatorial variant logic dedupes those + duplicates away without ever advancing. Detecting these names up front lets callers report a clear + error instead of hanging. Only the combinatorial generator is affected, and only for wildcards + nested in a variant — a bare `a __nope__ b` generates fine and is not reported. + + Without a configured `wildcard_manager`, an empty one is used so that every referenced wildcard is + treated as missing (wildcards are not resolved against any files here). + """ + if wildcard_manager is None: + wildcard_manager = WildcardManager() + + try: + tree = parse(prompt) + except ParseException: + # Malformed prompts are surfaced separately by the generators; nothing to validate here. + return [] + + missing: list[str] = [] + for name in _iter_wildcard_names(tree): + if name not in missing and not wildcard_manager.get_values(name): + missing.append(name) + return missing diff --git a/tests/app/routers/test_utilities.py b/tests/app/routers/test_utilities.py index ce91f2efd24..84a69696881 100644 --- a/tests/app/routers/test_utilities.py +++ b/tests/app/routers/test_utilities.py @@ -67,6 +67,50 @@ def test_dynamicprompts_works_for_user(client: TestClient, user1_token: str): assert "prompts" in body +def test_dynamicprompts_unknown_wildcard_returns_error_without_hanging(client: TestClient, user1_token: str): + """An unknown wildcard used as a variant value would otherwise loop forever in the combinatorial generator. + + The endpoint must instead return promptly with a clear error and the original prompt echoed back. + """ + r = client.post( + "/api/v1/utilities/dynamicprompts", + json={"prompt": "{__random__8chan|fenster|stuff}"}, + headers={"Authorization": f"Bearer {user1_token}"}, + ) + assert r.status_code == status.HTTP_200_OK + body = r.json() + assert body["error"] is not None + assert "random" in body["error"] + assert body["prompts"] == ["{__random__8chan|fenster|stuff}"] + + +def test_dynamicprompts_bare_unknown_wildcard_still_generates(client: TestClient, user1_token: str): + """A wildcard used as plain literal text (not a variant value) does not hang and must not error.""" + r = client.post( + "/api/v1/utilities/dynamicprompts", + json={"prompt": "a photo, __my_style__"}, + headers={"Authorization": f"Bearer {user1_token}"}, + ) + assert r.status_code == status.HTTP_200_OK + body = r.json() + assert body["error"] is None + assert body["prompts"] # non-empty + assert all(p == "a photo, __my_style__" for p in body["prompts"]) + + +def test_dynamicprompts_random_generator_ignores_unknown_wildcard(client: TestClient, user1_token: str): + """The random generator handles unknown wildcards gracefully, so the guard must not fire for it.""" + r = client.post( + "/api/v1/utilities/dynamicprompts", + json={"prompt": "{__random__8chan|fenster|stuff}", "combinatorial": False, "seed": 0}, + headers={"Authorization": f"Bearer {user1_token}"}, + ) + assert r.status_code == status.HTTP_200_OK + body = r.json() + assert body["error"] is None + assert body["prompts"] # non-empty; the random generator expanded the variant + + # ----------------------------- image_to_prompt: ownership / read-access ----------------------------- diff --git a/tests/app/util/test_dynamicprompts.py b/tests/app/util/test_dynamicprompts.py new file mode 100644 index 00000000000..83ff02f3c4c --- /dev/null +++ b/tests/app/util/test_dynamicprompts.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +import pytest + +from invokeai.app.util.dynamicprompts import find_missing_wildcards + + +def test_find_missing_wildcards_detects_unknown_wildcard_in_variant() -> None: + # Regression: `__random__` inside a variant is parsed as a wildcard reference. Left unchecked it + # sends the combinatorial generator into an infinite loop, so it must be reported up front. + assert find_missing_wildcards("{__random__8chan|fenster|stuff}") == ["random"] + + +def test_find_missing_wildcards_detects_unknown_wildcard_nested_in_sequence_in_variant() -> None: + # The wildcard hangs the generator even when wrapped in other text inside the variant value. + assert find_missing_wildcards("{a __nope__|b}") == ["nope"] + + +@pytest.mark.parametrize("prompt", ["a __nope__ b", "__nope__", "a photo, __my_style__"]) +def test_find_missing_wildcards_ignores_wildcards_outside_variants(prompt: str) -> None: + # A wildcard used as plain literal text generates fine (no hang), so it must not be reported. + assert find_missing_wildcards(prompt) == [] + + +@pytest.mark.parametrize("prompt", ["plain text", "{a|b|c}", "a {2$$x|y|z}"]) +def test_find_missing_wildcards_ignores_prompts_without_wildcards(prompt: str) -> None: + assert find_missing_wildcards(prompt) == [] + + +def test_find_missing_wildcards_dedupes_repeated_unknown_wildcards() -> None: + assert find_missing_wildcards("{__nope__|a} {__nope__|b} {__other__|c}") == ["nope", "other"]