Merge branch 'invoke-ai:main' into main

Author: joshistoast

USER_ISOLATION_IMPLEMENTATION.md

@@ -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)

invokeai/app/api/routers/utilities.py

@@ -1,8 +1,8 @@
 from contextlib import ExitStack
-from typing import Iterator, Literal, Optional, Tuple, Union
+from typing import Iterator, Literal, Optional, Tuple
 
 import torch
-from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer, T5TokenizerFast
+from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
 from invokeai.app.invocations.fields import (
@@ -86,7 +86,7 @@ def _t5_encode(self, context: InvocationContext) -> torch.Tensor:
             ExitStack() as exit_stack,
         ):
             assert isinstance(t5_text_encoder, T5EncoderModel)
-            assert isinstance(t5_tokenizer, (T5Tokenizer, T5TokenizerFast))
+            assert isinstance(t5_tokenizer, T5Tokenizer)
 
             # Determine if the model is quantized.
             # If the model is quantized, then we need to apply the LoRA weights as sidecar layers. This results in
@@ -186,7 +186,7 @@ def _t5_lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[ModelP
     def _log_t5_tokenization(
         self,
         context: InvocationContext,
-        tokenizer: Union[T5Tokenizer, T5TokenizerFast],
+        tokenizer: T5Tokenizer,
     ) -> None:
         """Logs the tokenization of a prompt for a T5-based model like FLUX."""
 

invokeai/app/invocations/flux_text_encoder.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:

invokeai/app/invocations/prompt.py

@@ -8,7 +8,6 @@
     CLIPTokenizer,
     T5EncoderModel,
     T5Tokenizer,
-    T5TokenizerFast,
 )
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
@@ -103,7 +102,7 @@ def _t5_encode(self, context: InvocationContext, max_seq_len: int) -> torch.Tens
         ):
             context.util.signal_progress("Running T5 encoder")
             assert isinstance(t5_text_encoder, T5EncoderModel)
-            assert isinstance(t5_tokenizer, (T5Tokenizer, T5TokenizerFast))
+            assert isinstance(t5_tokenizer, T5Tokenizer)
             t5_device = get_effective_device(t5_text_encoder)
 
             text_inputs = t5_tokenizer(

invokeai/app/invocations/sd3_text_encoder.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

invokeai/app/util/dynamicprompts.py

@@ -9,7 +9,7 @@
 import numpy as np
 from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
 from PIL import Image, ImageFilter
-from transformers import AutoFeatureExtractor
+from transformers import AutoImageProcessor
 
 import invokeai.backend.util.logging as logger
 from invokeai.app.services.config.config_default import get_config
@@ -36,14 +36,14 @@ def _load_safety_checker(cls):
         try:
             model_path = get_config().models_path / CHECKER_PATH
             if model_path.exists():
-                cls.feature_extractor = AutoFeatureExtractor.from_pretrained(model_path)
+                cls.feature_extractor = AutoImageProcessor.from_pretrained(model_path)
                 cls.safety_checker = StableDiffusionSafetyChecker.from_pretrained(model_path)
             else:
                 model_path.mkdir(parents=True, exist_ok=True)
-                cls.feature_extractor = AutoFeatureExtractor.from_pretrained(repo_id)
-                cls.feature_extractor.save_pretrained(model_path, safe_serialization=True)
+                cls.feature_extractor = AutoImageProcessor.from_pretrained(repo_id)
+                cls.feature_extractor.save_pretrained(model_path)
                 cls.safety_checker = StableDiffusionSafetyChecker.from_pretrained(repo_id)
-                cls.safety_checker.save_pretrained(model_path, safe_serialization=True)
+                cls.safety_checker.save_pretrained(model_path)
         except Exception as e:
             logger.warning(f"Could not load NSFW checker: {str(e)}")