diff --git a/invokeai/app/api/routers/utilities.py b/invokeai/app/api/routers/utilities.py index f77f77a8534..43c1eda6c93 100644 --- a/invokeai/app/api/routers/utilities.py +++ b/invokeai/app/api/routers/utilities.py @@ -6,15 +6,26 @@ import torch from dynamicprompts.generators import CombinatorialPromptGenerator, RandomPromptGenerator -from fastapi import Body, HTTPException +from dynamicprompts.wildcards import WildcardManager +from fastapi import Body, HTTPException, Query from fastapi.routing import APIRouter from pydantic import BaseModel, Field from pyparsing import ParseException from transformers import AutoProcessor, AutoTokenizer, LlavaOnevisionForConditionalGeneration, LlavaOnevisionProcessor +from invokeai.app.api.auth_dependencies import CurrentUserOrDefault from invokeai.app.api.dependencies import ApiDependencies 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.wildcards import ( + WildcardsResponse, + WildcardValuesResponse, + clean_dynamic_prompt_outputs, + find_missing_wildcard_references, + get_wildcard_values, + get_wildcards_path, + index_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 @@ -31,6 +42,41 @@ class DynamicPromptsResponse(BaseModel): prompts: list[str] error: Optional[str] = None + warnings: list[str] = Field(default_factory=list) + missing_wildcards: list[str] = Field(default_factory=list) + + +@utilities_router.get( + "/wildcards", + operation_id="list_wildcards", + responses={ + 200: {"model": WildcardsResponse}, + }, +) +async def list_wildcards(_: CurrentUserOrDefault) -> WildcardsResponse: + """List local dynamic prompt wildcards from INVOKEAI_ROOT/wildcards.""" + wildcards_path = get_wildcards_path(ApiDependencies.invoker.services.configuration.root_path) + return index_wildcards(wildcards_path) + + +@utilities_router.get( + "/wildcards/values", + operation_id="get_wildcard_values", + responses={ + 200: {"model": WildcardValuesResponse}, + }, +) +async def list_wildcard_values( + _: CurrentUserOrDefault, + path: str = Query(description="The relative wildcard path to read values for"), + limit: int = Query(default=200, ge=1, le=1000, description="The max number of wildcard values to return"), +) -> WildcardValuesResponse: + """List values for a single local dynamic prompt wildcard.""" + wildcards_path = get_wildcards_path(ApiDependencies.invoker.services.configuration.root_path) + values = get_wildcard_values(wildcards_path, path, limit) + if values is None: + raise HTTPException(status_code=404, detail=f"Wildcard '{path}' not found") + return values @utilities_router.post( @@ -41,6 +87,7 @@ class DynamicPromptsResponse(BaseModel): }, ) async def parse_dynamicprompts( + _: CurrentUserOrDefault, prompt: str = Body(description="The prompt to parse with dynamicprompts"), max_prompts: int = Body(ge=1, le=10000, default=1000, description="The max number of prompts to generate"), combinatorial: bool = Body(default=True, description="Whether to use the combinatorial generator"), @@ -49,18 +96,30 @@ async def parse_dynamicprompts( """Creates a batch process""" max_prompts = min(max_prompts, 10000) generator: Union[RandomPromptGenerator, CombinatorialPromptGenerator] + warnings: list[str] = [] + wildcards_path = get_wildcards_path(ApiDependencies.invoker.services.configuration.root_path) + wildcard_index = index_wildcards(wildcards_path) + missing_wildcards = find_missing_wildcard_references(prompt, wildcard_index.wildcards) + if wildcard_index.errors: + warnings.append("Some wildcard files could not be indexed.") try: error: Optional[str] = None + wildcard_manager = WildcardManager(wildcards_path) if wildcards_path.is_dir() else None if combinatorial: - generator = CombinatorialPromptGenerator() + generator = CombinatorialPromptGenerator(wildcard_manager=wildcard_manager) prompts = generator.generate(prompt, max_prompts=max_prompts) else: - generator = RandomPromptGenerator(seed=seed) + generator = RandomPromptGenerator(wildcard_manager=wildcard_manager, seed=seed) prompts = generator.generate(prompt, num_images=max_prompts) except ParseException as e: prompts = [prompt] error = str(e) - return DynamicPromptsResponse(prompts=prompts if prompts else [""], error=error) + return DynamicPromptsResponse( + prompts=clean_dynamic_prompt_outputs(prompts) if prompts else [""], + error=error, + warnings=warnings, + missing_wildcards=missing_wildcards, + ) # --- Expand Prompt --- diff --git a/invokeai/app/invocations/prompt.py b/invokeai/app/invocations/prompt.py index 48eec0ac0ef..5e3112b8fba 100644 --- a/invokeai/app/invocations/prompt.py +++ b/invokeai/app/invocations/prompt.py @@ -3,12 +3,14 @@ import numpy as np from dynamicprompts.generators import CombinatorialPromptGenerator, RandomPromptGenerator +from dynamicprompts.wildcards import WildcardManager from pydantic import field_validator from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation 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.wildcards import clean_dynamic_prompt_outputs, get_wildcards_path @invocation( @@ -30,14 +32,16 @@ class DynamicPromptInvocation(BaseInvocation): combinatorial: bool = InputField(default=False, description="Whether to use the combinatorial generator") def invoke(self, context: InvocationContext) -> StringCollectionOutput: + wildcards_path = get_wildcards_path(context.config.get().root_path) + wildcard_manager = WildcardManager(wildcards_path) if wildcards_path.is_dir() else None if self.combinatorial: - generator = CombinatorialPromptGenerator() + generator = CombinatorialPromptGenerator(wildcard_manager=wildcard_manager) prompts = generator.generate(self.prompt, max_prompts=self.max_prompts) else: - generator = RandomPromptGenerator() + generator = RandomPromptGenerator(wildcard_manager=wildcard_manager) prompts = generator.generate(self.prompt, num_images=self.max_prompts) - return StringCollectionOutput(collection=prompts) + return StringCollectionOutput(collection=clean_dynamic_prompt_outputs(prompts)) @invocation( diff --git a/invokeai/app/util/wildcards.py b/invokeai/app/util/wildcards.py new file mode 100644 index 00000000000..5548ea57701 --- /dev/null +++ b/invokeai/app/util/wildcards.py @@ -0,0 +1,360 @@ +import fnmatch +import json +import re +from pathlib import Path +from typing import Any, Literal + +import yaml +from pydantic import BaseModel, Field + +WILDCARDS_DIR_NAME = "wildcards" +SUPPORTED_WILDCARD_EXTENSIONS = { + ".txt": "txt", + ".json": "json", + ".yaml": "yaml", + ".yml": "yaml", +} +WILDCARD_REFERENCE_PATTERN = re.compile(r"__([^\r\n]+?)__") +VISIBLE_BOM_PREFIXES = ("\ufeff", "\u00ef\u00bb\u00bf", "\u00c3\u00af\u00c2\u00bb\u00c2\u00bf") + + +class WildcardIndexItem(BaseModel): + token: str + path: str + label: str + file_type: Literal["txt", "json", "yaml"] + value_count: int = Field(ge=0) + samples: list[str] + + +class WildcardIndexError(BaseModel): + path: str + message: str + + +class WildcardsResponse(BaseModel): + wildcard_dir: str = WILDCARDS_DIR_NAME + wildcards: list[WildcardIndexItem] + errors: list[WildcardIndexError] + + +class WildcardValuesResponse(BaseModel): + token: str + path: str + label: str + file_type: Literal["txt", "json", "yaml"] + value_count: int = Field(ge=0) + values: list[str] + truncated: bool + + +def get_wildcards_path(root_path: Path) -> Path: + return root_path / WILDCARDS_DIR_NAME + + +def index_wildcards(wildcards_path: Path, max_samples: int = 5) -> WildcardsResponse: + wildcards: list[WildcardIndexItem] = [] + errors: list[WildcardIndexError] = [] + seen_paths: dict[str, str] = {} + + if not wildcards_path.exists(): + return WildcardsResponse(wildcards=wildcards, errors=errors) + if not wildcards_path.is_dir(): + return WildcardsResponse( + wildcards=wildcards, + errors=[ + WildcardIndexError(path=WILDCARDS_DIR_NAME, message="Wildcards path exists but is not a folder"), + ], + ) + + root = wildcards_path.resolve() + for path in sorted(root.rglob("*")): + if path.is_dir(): + continue + + file_type = SUPPORTED_WILDCARD_EXTENSIONS.get(path.suffix.lower()) + if file_type is None: + continue + + rel_path = _get_safe_relative_path(path, root) + if rel_path is None: + errors.append(WildcardIndexError(path=path.name, message="Wildcard file is outside the wildcards folder")) + continue + + if path.is_symlink(): + errors.append(WildcardIndexError(path=_as_posix(rel_path), message="Symlinks are not supported")) + continue + + try: + if file_type == "txt": + items = _index_txt_wildcard(path, rel_path, max_samples) + else: + items = _index_structured_wildcard(path, rel_path, file_type, max_samples) + for item in items: + existing_path = seen_paths.get(item.path) + if existing_path is not None: + errors.append( + WildcardIndexError( + path=_as_posix(rel_path), + message=f"Duplicate wildcard path '{item.path}' already indexed from '{existing_path}'", + ) + ) + continue + seen_paths[item.path] = _as_posix(rel_path) + wildcards.append(item) + except Exception as e: + errors.append(WildcardIndexError(path=_as_posix(rel_path), message=str(e))) + + wildcards.sort(key=lambda item: item.path) + return WildcardsResponse(wildcards=wildcards, errors=errors) + + +def get_wildcard_references(prompt: str) -> list[str]: + refs: list[str] = [] + seen: set[str] = set() + for match in WILDCARD_REFERENCE_PATTERN.finditer(prompt): + path = normalize_wildcard_reference(match.group(1)) + if path and path not in seen: + seen.add(path) + refs.append(path) + return refs + + +def find_missing_wildcard_references(prompt: str, wildcards: list[WildcardIndexItem]) -> list[str]: + available_paths = {wildcard.path for wildcard in wildcards} + missing: list[str] = [] + for ref in get_wildcard_references(prompt): + if _reference_exists(ref, available_paths): + continue + missing.append(ref) + return missing + + +def clean_dynamic_prompt_outputs(prompts: list[str]) -> list[str]: + return [_clean_wildcard_value(prompt) for prompt in prompts] + + +def get_wildcard_values(wildcards_path: Path, reference: str, limit: int = 200) -> WildcardValuesResponse | None: + if _has_absolute_wildcard_reference(reference): + return None + wildcard_path = normalize_wildcard_reference(reference) + if not _is_safe_wildcard_path(wildcard_path) or not wildcards_path.exists() or not wildcards_path.is_dir(): + return None + + root = wildcards_path.resolve() + for path in sorted(root.rglob("*")): + if path.is_dir(): + continue + + file_type = SUPPORTED_WILDCARD_EXTENSIONS.get(path.suffix.lower()) + if file_type is None: + continue + + rel_path = _get_safe_relative_path(path, root) + if rel_path is None or path.is_symlink(): + continue + + try: + if file_type == "txt": + item = _get_txt_wildcard_values(path, rel_path, wildcard_path, limit) + else: + item = _get_structured_wildcard_values(path, rel_path, file_type, wildcard_path, limit) + except Exception: + continue + + if item is not None: + return item + + return None + + +def normalize_wildcard_reference(reference: str) -> str: + path = reference.strip() + if path.startswith("~") or path.startswith("@"): + path = path[1:] + if "(" in path: + path = path.split("(", 1)[0] + path = path.replace("\\", "/").strip("/") + return path + + +def _is_safe_wildcard_path(path: str) -> bool: + if not path or "*" in path: + return False + candidate = Path(path) + return not candidate.is_absolute() and ".." not in candidate.parts + + +def _has_absolute_wildcard_reference(reference: str) -> bool: + normalized = reference.strip().replace("\\", "/") + if normalized.startswith("~") or normalized.startswith("@"): + normalized = normalized[1:] + return normalized.startswith("/") + + +def _reference_exists(reference: str, available_paths: set[str]) -> bool: + if "*" in reference: + return any(fnmatch.fnmatch(path, reference) for path in available_paths) + return reference in available_paths + + +def _get_safe_relative_path(path: Path, root: Path) -> Path | None: + resolved = path.resolve() + if not resolved.is_relative_to(root): + return None + return resolved.relative_to(root) + + +def _index_txt_wildcard(path: Path, rel_path: Path, max_samples: int) -> list[WildcardIndexItem]: + values = _read_txt_values(path) + wildcard_path = _as_posix(rel_path.with_suffix("")) + return [ + WildcardIndexItem( + token=f"__{wildcard_path}__", + path=wildcard_path, + label=Path(wildcard_path).name, + file_type="txt", + value_count=len(values), + samples=values[:max_samples], + ) + ] + + +def _get_txt_wildcard_values( + path: Path, rel_path: Path, wildcard_path: str, limit: int +) -> WildcardValuesResponse | None: + indexed_path = _as_posix(rel_path.with_suffix("")) + if indexed_path != wildcard_path: + return None + values = _read_txt_values(path) + return _make_values_response(indexed_path, "txt", values, limit) + + +def _read_txt_values(path: Path) -> list[str]: + values: list[str] = [] + for line in path.read_text(encoding="utf-8").splitlines(): + value = _clean_wildcard_value(line) + if value and not value.startswith("#"): + values.append(value) + return values + + +def _index_structured_wildcard( + path: Path, rel_path: Path, file_type: Literal["json", "yaml"], max_samples: int +) -> list[WildcardIndexItem]: + data = _read_structured_data(path, file_type) + + folder_prefix = _as_posix(rel_path.parent) if rel_path.parent != Path(".") else "" + file_stem_path = _as_posix(rel_path.with_suffix("")) + leaves = _collect_wildcard_leaves(data) + items: list[WildcardIndexItem] = [] + + if isinstance(data, list): + values = _stringify_values(data) + items.append(_make_index_item(file_stem_path, file_type, values, max_samples)) + return items + + for key_path, values in leaves: + wildcard_path = "/".join(part for part in [folder_prefix, key_path] if part) + items.append(_make_index_item(wildcard_path, file_type, values, max_samples)) + + return items + + +def _get_structured_wildcard_values( + path: Path, + rel_path: Path, + file_type: Literal["json", "yaml"], + wildcard_path: str, + limit: int, +) -> WildcardValuesResponse | None: + data = _read_structured_data(path, file_type) + folder_prefix = _as_posix(rel_path.parent) if rel_path.parent != Path(".") else "" + file_stem_path = _as_posix(rel_path.with_suffix("")) + + if isinstance(data, list): + if file_stem_path != wildcard_path: + return None + return _make_values_response(file_stem_path, file_type, _stringify_values(data), limit) + + for key_path, values in _collect_wildcard_leaves(data): + indexed_path = "/".join(part for part in [folder_prefix, key_path] if part) + if indexed_path == wildcard_path: + return _make_values_response(indexed_path, file_type, values, limit) + + return None + + +def _read_structured_data(path: Path, file_type: Literal["json", "yaml"]) -> Any: + text = path.read_text(encoding="utf-8") + if file_type == "json": + return json.loads(text) + return yaml.safe_load(text) + + +def _make_index_item( + wildcard_path: str, file_type: Literal["txt", "json", "yaml"], values: list[str], max_samples: int +) -> WildcardIndexItem: + return WildcardIndexItem( + token=f"__{wildcard_path}__", + path=wildcard_path, + label=Path(wildcard_path).name, + file_type=file_type, + value_count=len(values), + samples=values[:max_samples], + ) + + +def _make_values_response( + wildcard_path: str, file_type: Literal["txt", "json", "yaml"], values: list[str], limit: int +) -> WildcardValuesResponse: + return WildcardValuesResponse( + token=f"__{wildcard_path}__", + path=wildcard_path, + label=Path(wildcard_path).name, + file_type=file_type, + value_count=len(values), + values=values[:limit], + truncated=len(values) > limit, + ) + + +def _collect_wildcard_leaves(data: Any, prefix: str = "") -> list[tuple[str, list[str]]]: + if isinstance(data, list): + return [(prefix, _stringify_values(data))] if prefix else [] + + if not isinstance(data, dict): + return [] + + leaves: list[tuple[str, list[str]]] = [] + for key, value in data.items(): + if not isinstance(key, str): + continue + key_path = f"{prefix}/{key}" if prefix else key + leaves.extend(_collect_wildcard_leaves(value, key_path)) + return leaves + + +def _stringify_values(values: list[Any]) -> list[str]: + result: list[str] = [] + for value in values: + if isinstance(value, str): + result.append(_clean_wildcard_value(value)) + elif isinstance(value, (int, float, bool)): + result.append(str(value)) + elif isinstance(value, dict): + result.extend(_clean_wildcard_value(key) for key in value.keys() if isinstance(key, str)) + if result and result[-1] == "": + result.pop() + return result + + +def _clean_wildcard_value(value: str) -> str: + cleaned = value.strip() + for prefix in VISIBLE_BOM_PREFIXES: + cleaned = cleaned.replace(prefix, "") + return cleaned.strip() + + +def _as_posix(path: Path) -> str: + return path.as_posix() diff --git a/invokeai/frontend/web/public/locales/en.json b/invokeai/frontend/web/public/locales/en.json index d99bb04a631..58d51776421 100644 --- a/invokeai/frontend/web/public/locales/en.json +++ b/invokeai/frontend/web/public/locales/en.json @@ -1753,6 +1753,29 @@ "dynamicPrompts": "Dynamic Prompts", "maxPrompts": "Max Prompts", "promptsPreview": "Prompts Preview", + "promptExpansionsWithCount_one": "Prompt Expansion ({{count}})", + "promptExpansionsWithCount_other": "Prompt Expansions ({{count}})", + "maxCombinations": "Max Combinations", + "randomSamples": "Random Samples", + "preview": "Preview", + "reshuffleNow": "Reshuffle Now", + "resolveFailed": "Failed to resolve dynamic prompts", + "mode": { + "label": "Mode", + "randomLabel": "Random Sample", + "randomDesc": "Sample prompts. Randomness applies to random wildcards.", + "combinatorialLabel": "All Combinations", + "combinatorialDesc": "Preview and queue every combination up to the limit." + }, + "randomness": { + "label": "Randomness", + "perImageLabel": "Per Image", + "perImageDesc": "Roll a new random sample for each generated image.", + "perInvokeLabel": "Per Invoke", + "perInvokeDesc": "Random wildcards roll once per Invoke; cyclic wildcards still advance per queued output.", + "manualLabel": "Locked Preview", + "manualDesc": "Keep the preview fixed until Reshuffle is used." + }, "seedBehaviour": { "label": "Seed Behaviour", "perIterationLabel": "Seed per Iteration", @@ -1764,6 +1787,110 @@ "problemGeneratingPrompts": "Problem generating prompts", "promptsToGenerate": "Prompts to Generate" }, + "promptWorkbench": { + "panel": { + "title": "Prompt intent" + }, + "autocomplete": { + "loading": "Loading local wildcards...", + "unavailable": "Wildcard index unavailable. Restart the backend or check the wildcard endpoint.", + "empty": "No local wildcards found.", + "noMatches": "No local wildcards match \"{{query}}\".", + "noWildcardMatches": "No local wildcard matches.", + "matching": "Matching \"{{query}}\"", + "localWildcards": "Local wildcards", + "insertWithBehaviorAria": "Insert {{path}} with wildcard behavior" + }, + "actions": { + "wildcardBehavior": "Wildcard behavior", + "wildcardActions": "Wildcard actions", + "weightActions": "Weight actions", + "randomWildcard": "Random wildcard", + "randomWildcardTooltip": "Use this wildcard as a random token. Random cadence is controlled by the prompt.", + "cycle": "Cycle through values", + "cycleTooltip": "Cycles through values across generated outputs.", + "pickFixed": "Pick fixed value", + "remove": "Remove", + "openWildcardActionsAria": "Open {{path}} wildcard actions", + "openWeightActionsAria": "Open {{text}} weight actions" + }, + "behavior": { + "randomImageLabel": "Random per Image", + "randomInvokeLabel": "Random per Invoke", + "randomPreviewLabel": "Random preview", + "randomImageShort": "Random/image", + "randomInvokeShort": "Random/invoke", + "previewShort": "Preview", + "cycleLabel": "Cycle", + "cycleShort": "Cycle", + "allCombinationsLabel": "All combinations", + "allShort": "All", + "missingLabel": "Missing", + "unavailableLabel": "Unavailable" + }, + "diagnostics": { + "weightsLiteralLabel": "Weights literal?", + "weightsLiteralDesc": "This prompt uses weight syntax, but the selected model may encode it as literal text.", + "wildcardErrorLabel": "Wildcard error", + "wildcardErrorDesc": "The local wildcard index could not be loaded. Restart the backend so /api/v1/utilities/wildcards is available.", + "missingLabel_one": "Missing {{count}}", + "missingLabel_other": "Missing {{count}}", + "missingDesc_one": "Missing wildcard: {{wildcards}}", + "missingDesc_other": "Missing wildcards: {{wildcards}}", + "wildcardsFoundLabel_one": "Wildcards {{count}}", + "wildcardsFoundLabel_other": "Wildcards {{count}}", + "wildcardsFoundDesc": "All referenced wildcards are available locally.", + "indexErrorsLabel_one": "Index error {{count}}", + "indexErrorsLabel_other": "Index errors {{count}}", + "indexErrorsDesc": "Some local wildcard files could not be read.", + "dynamicErrorLabel": "Dynamic error", + "dynamicErrorDesc": "Dynamic prompt parser error: {{error}}", + "dynamicAllLabel_one": "All {{count}}", + "dynamicAllLabel_other": "All {{count}}", + "dynamicMixedLabel": "Mixed dynamic", + "dynamicCycleLabel_one": "Cycle {{count}}", + "dynamicCycleLabel_other": "Cycle {{count}}", + "dynamicAllDesc": "All-combinations prompt expansion is active for this prompt.", + "dynamicMixedDesc": "Cyclic wildcards advance per output; random wildcards follow the selected randomness mode.", + "dynamicCycleDesc": "Cyclic wildcard expansion is deterministic and cycles through values across generated outputs.", + "dynamicRandomImageDesc": "Random prompt sampling will roll fresh values for each generated image.", + "dynamicRandomInvokeDesc": "Random prompt sampling will roll once when generation is queued.", + "dynamicRandomPreviewDesc": "Random prompt sampling is fixed to the current preview until Reshuffle is used.", + "changeRandomBehaviorTooltip": "{{description}} Change random wildcard behavior.", + "openDynamicPreviewTooltip": "{{description}} Open dynamic prompt preview." + }, + "rows": { + "wildcardIntentAria": "{{path}} wildcard intent", + "weightIntentAria": "{{text}} prompt weight", + "missingWildcard": "Missing wildcard", + "wildcardIndexUnavailable": "Wildcard index unavailable", + "randomWildcard": "Random wildcard", + "randomWildcardWithCount_one": "Random wildcard · {{count}} value", + "randomWildcardWithCount_other": "Random wildcard · {{count}} values", + "cycle": "Cycle", + "cycleWithCount_one": "Cycle · {{count}} value", + "cycleWithCount_other": "Cycle · {{count}} values", + "allCombinations": "All combinations", + "allCombinationsWithCount_one": "All combinations · {{count}} value", + "allCombinationsWithCount_other": "All combinations · {{count}} values" + }, + "values": { + "loading": "Loading values..." + }, + "weight": { + "promptWeight": "prompt weight", + "supportedLabel": "Weight supported", + "literalLabel": "Weight may be literal", + "literalShort": "Literal?", + "valueLabel": "{{value}}", + "supportedDescription": "Prompt weights are supported for this model.", + "literalDescription": "Prompt weight syntax may be treated as literal text by this model." + }, + "header": { + "randomImageTooltip": "Random wildcards roll once per generated image.", + "randomInvokeTooltip": "Random wildcards roll once per Invoke; cyclic wildcards still advance per generated output." + } + }, "sdxl": { "cfgScale": "CFG Scale", "concatPromptStyle": "Linking Prompt & Style", diff --git a/invokeai/frontend/web/src/common/util/promptAST.test.ts b/invokeai/frontend/web/src/common/util/promptAST.test.ts index 32ad9dc09fb..82091bb86ad 100644 --- a/invokeai/frontend/web/src/common/util/promptAST.test.ts +++ b/invokeai/frontend/web/src/common/util/promptAST.test.ts @@ -70,6 +70,18 @@ describe('promptAST', () => { ]); }); + it('should not tokenize hyphenated words as weights', () => { + const tokens = tokenize('dance-more rose-'); + expect(tokens).toEqual([ + { type: 'word', value: 'dance', start: 0, end: 5 }, + { type: 'word', value: '-', start: 5, end: 6 }, + { type: 'word', value: 'more', start: 6, end: 10 }, + { type: 'whitespace', value: ' ', start: 10, end: 11 }, + { type: 'word', value: 'rose', start: 11, end: 15 }, + { type: 'weight', value: '-', start: 15, end: 16 }, + ]); + }); + it('should tokenize embeddings', () => { const tokens = tokenize(''); expect(tokens).toEqual([ @@ -200,6 +212,17 @@ describe('promptAST', () => { expect(ast).toEqual([{ type: 'word', text: 'cat', attention: '+', range: { start: 0, end: 4 } }]); }); + it('should parse hyphenated words without attention', () => { + const ast = parseTokens(tokenize('dance-more rose-')); + expect(ast).toEqual([ + { type: 'word', text: 'dance', range: { start: 0, end: 5 }, attention: undefined }, + { type: 'word', text: '-', range: { start: 5, end: 6 }, attention: undefined }, + { type: 'word', text: 'more', range: { start: 6, end: 10 }, attention: undefined }, + { type: 'whitespace', value: ' ', range: { start: 10, end: 11 } }, + { type: 'word', text: 'rose', attention: '-', range: { start: 11, end: 16 } }, + ]); + }); + it('should parse embeddings', () => { const tokens = tokenize(''); const ast = parseTokens(tokens); diff --git a/invokeai/frontend/web/src/common/util/promptAST.ts b/invokeai/frontend/web/src/common/util/promptAST.ts index 0a1af621224..bdcc5747d9c 100644 --- a/invokeai/frontend/web/src/common/util/promptAST.ts +++ b/invokeai/frontend/web/src/common/util/promptAST.ts @@ -222,7 +222,7 @@ function tokenizeWord(prompt: string, i: number): TokenizeResult { // Check for weight immediately after word (e.g., "Lorem+", "consectetur-") const weightMatch = prompt.slice(j).match(WEIGHT_PATTERN); - if (weightMatch && weightMatch[0]) { + if (weightMatch && weightMatch[0] && getIsWordWeightSuffix(prompt, j, weightMatch[0])) { const weightEnd = j + weightMatch[0].length; return { token: { type: 'word', value: word, start: i, end: j }, @@ -239,6 +239,14 @@ function tokenizeWord(prompt: string, i: number): TokenizeResult { return null; } +function getIsWordWeightSuffix(prompt: string, start: number, weight: string): boolean { + const nextChar = prompt[start + weight.length]; + if (/^[+-]+$/.test(weight) && nextChar && WORD_CHAR_PATTERN.test(nextChar)) { + return false; + } + return true; +} + function tokenizeEmbedding(char: string, i: number): TokenizeResult { if (char === '<') { return { diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/components/DynamicPromptsPreviewModal.tsx b/invokeai/frontend/web/src/features/dynamicPrompts/components/DynamicPromptsPreviewModal.tsx index 2e62742b7e3..d26d7d3365a 100644 --- a/invokeai/frontend/web/src/features/dynamicPrompts/components/DynamicPromptsPreviewModal.tsx +++ b/invokeai/frontend/web/src/features/dynamicPrompts/components/DynamicPromptsPreviewModal.tsx @@ -13,7 +13,10 @@ import { memo } from 'react'; import { useTranslation } from 'react-i18next'; import ParamDynamicPromptsMaxPrompts from './ParamDynamicPromptsMaxPrompts'; +import ParamDynamicPromptsMode from './ParamDynamicPromptsMode'; import ParamDynamicPromptsPreview from './ParamDynamicPromptsPreview'; +import ParamDynamicPromptsRandomRefreshMode from './ParamDynamicPromptsRandomRefreshMode'; +import ParamDynamicPromptsReshuffle from './ParamDynamicPromptsReshuffle'; import ParamDynamicPromptsSeedBehaviour from './ParamDynamicPromptsSeedBehaviour'; export const DynamicPromptsModal = memo(() => { @@ -29,8 +32,11 @@ export const DynamicPromptsModal = memo(() => { + + + diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsMaxPrompts.tsx b/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsMaxPrompts.tsx index e17743207fa..ab53d16f97b 100644 --- a/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsMaxPrompts.tsx +++ b/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsMaxPrompts.tsx @@ -2,11 +2,13 @@ import { CompositeNumberInput, CompositeSlider, FormControl, FormLabel } from '@ import { useAppDispatch, useAppSelector } from 'app/store/storeHooks'; import { InformationalPopover } from 'common/components/InformationalPopover/InformationalPopover'; import { - maxPromptsChanged, - selectDynamicPromptsCombinatorial, - selectDynamicPromptsMaxPrompts, + maxCombinationsChanged, + randomSamplesChanged, + selectDynamicPromptsMaxCombinations, + selectDynamicPromptsMode, + selectDynamicPromptsRandomSamples, } from 'features/dynamicPrompts/store/dynamicPromptsSlice'; -import { memo, useCallback } from 'react'; +import { memo, useCallback, useMemo } from 'react'; import { useTranslation } from 'react-i18next'; const CONSTRAINTS = { @@ -20,36 +22,46 @@ const CONSTRAINTS = { }; const ParamDynamicPromptsMaxPrompts = () => { - const maxPrompts = useAppSelector(selectDynamicPromptsMaxPrompts); - const combinatorial = useAppSelector(selectDynamicPromptsCombinatorial); - const dispatch = useAppDispatch(); const { t } = useTranslation(); + const mode = useAppSelector(selectDynamicPromptsMode); + const randomSamples = useAppSelector(selectDynamicPromptsRandomSamples); + const maxCombinations = useAppSelector(selectDynamicPromptsMaxCombinations); + const dispatch = useAppDispatch(); + const value = mode === 'combinatorial' ? maxCombinations : randomSamples; + const label = useMemo( + () => (mode === 'combinatorial' ? t('dynamicPrompts.maxCombinations') : t('dynamicPrompts.randomSamples')), + [mode, t] + ); const handleChange = useCallback( (v: number) => { - dispatch(maxPromptsChanged(v)); + if (mode === 'combinatorial') { + dispatch(maxCombinationsChanged(v)); + } else { + dispatch(randomSamplesChanged(v)); + } }, - [dispatch] + [dispatch, mode] ); return ( - + - {t('dynamicPrompts.maxPrompts')} + {label} diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsMode.tsx b/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsMode.tsx new file mode 100644 index 00000000000..b7422c23528 --- /dev/null +++ b/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsMode.tsx @@ -0,0 +1,53 @@ +import type { ComboboxOnChange, ComboboxOption } from '@invoke-ai/ui-library'; +import { Combobox, FormControl, FormLabel } from '@invoke-ai/ui-library'; +import { useAppDispatch, useAppSelector } from 'app/store/storeHooks'; +import { + isDynamicPromptMode, + modeChanged, + selectDynamicPromptsMode, +} from 'features/dynamicPrompts/store/dynamicPromptsSlice'; +import { memo, useCallback, useMemo } from 'react'; +import { useTranslation } from 'react-i18next'; + +const ParamDynamicPromptsMode = () => { + const { t } = useTranslation(); + const dispatch = useAppDispatch(); + const mode = useAppSelector(selectDynamicPromptsMode); + + const options = useMemo( + () => [ + { + value: 'random', + label: t('dynamicPrompts.mode.randomLabel'), + description: t('dynamicPrompts.mode.randomDesc'), + }, + { + value: 'combinatorial', + label: t('dynamicPrompts.mode.combinatorialLabel'), + description: t('dynamicPrompts.mode.combinatorialDesc'), + }, + ], + [t] + ); + + const handleChange = useCallback( + (v) => { + if (!isDynamicPromptMode(v?.value)) { + return; + } + dispatch(modeChanged(v.value)); + }, + [dispatch] + ); + + const value = useMemo(() => options.find((o) => o.value === mode), [mode, options]); + + return ( + + {t('dynamicPrompts.mode.label')} + + + ); +}; + +export default memo(ParamDynamicPromptsMode); diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsPreview.tsx b/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsPreview.tsx index d635a27a3d9..39943d81e7f 100644 --- a/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsPreview.tsx +++ b/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsPreview.tsx @@ -26,7 +26,7 @@ const ParamDynamicPromptsPreview = () => { const prompts = useAppSelector(selectDynamicPromptsPrompts); const label = useMemo(() => { - let _label = `${t('dynamicPrompts.promptsPreview')} (${prompts.length})`; + let _label = t('dynamicPrompts.promptExpansionsWithCount', { count: prompts.length }); if (parsingError) { _label += ` - ${parsingError}`; } diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsRandomRefreshMode.tsx b/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsRandomRefreshMode.tsx new file mode 100644 index 00000000000..f68851abcc6 --- /dev/null +++ b/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsRandomRefreshMode.tsx @@ -0,0 +1,64 @@ +import type { ComboboxOnChange, ComboboxOption } from '@invoke-ai/ui-library'; +import { Combobox, FormControl, FormLabel } from '@invoke-ai/ui-library'; +import { useAppDispatch, useAppSelector } from 'app/store/storeHooks'; +import { + isDynamicPromptRandomRefreshMode, + randomRefreshModeChanged, + selectDynamicPromptsMode, + selectDynamicPromptsRandomRefreshMode, +} from 'features/dynamicPrompts/store/dynamicPromptsSlice'; +import { memo, useCallback, useMemo } from 'react'; +import { useTranslation } from 'react-i18next'; + +const ParamDynamicPromptsRandomRefreshMode = () => { + const { t } = useTranslation(); + const dispatch = useAppDispatch(); + const mode = useAppSelector(selectDynamicPromptsMode); + const randomRefreshMode = useAppSelector(selectDynamicPromptsRandomRefreshMode); + + const options = useMemo( + () => [ + { + value: 'per_image', + label: t('dynamicPrompts.randomness.perImageLabel'), + description: t('dynamicPrompts.randomness.perImageDesc'), + }, + { + value: 'per_enqueue', + label: t('dynamicPrompts.randomness.perInvokeLabel'), + description: t('dynamicPrompts.randomness.perInvokeDesc'), + }, + { + value: 'manual', + label: t('dynamicPrompts.randomness.manualLabel'), + description: t('dynamicPrompts.randomness.manualDesc'), + }, + ], + [t] + ); + + const handleChange = useCallback( + (v) => { + if (!isDynamicPromptRandomRefreshMode(v?.value)) { + return; + } + dispatch(randomRefreshModeChanged(v.value)); + }, + [dispatch] + ); + + const value = useMemo(() => options.find((o) => o.value === randomRefreshMode), [options, randomRefreshMode]); + + if (mode !== 'random') { + return null; + } + + return ( + + {t('dynamicPrompts.randomness.label')} + + + ); +}; + +export default memo(ParamDynamicPromptsRandomRefreshMode); diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsReshuffle.tsx b/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsReshuffle.tsx new file mode 100644 index 00000000000..78d947f5651 --- /dev/null +++ b/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsReshuffle.tsx @@ -0,0 +1,30 @@ +import { Button, FormControl, FormLabel } from '@invoke-ai/ui-library'; +import { useAppDispatch, useAppSelector } from 'app/store/storeHooks'; +import { randomSeedChanged, selectDynamicPromptsMode } from 'features/dynamicPrompts/store/dynamicPromptsSlice'; +import { memo, useCallback } from 'react'; +import { useTranslation } from 'react-i18next'; + +const ParamDynamicPromptsReshuffle = () => { + const { t } = useTranslation(); + const dispatch = useAppDispatch(); + const mode = useAppSelector(selectDynamicPromptsMode); + + const reshuffle = useCallback(() => { + dispatch(randomSeedChanged(Date.now())); + }, [dispatch]); + + if (mode !== 'random') { + return null; + } + + return ( + + {t('dynamicPrompts.preview')} + + + ); +}; + +export default memo(ParamDynamicPromptsReshuffle); diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsSeedBehaviour.tsx b/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsSeedBehaviour.tsx index 252f3aa05fc..bc9791b6f38 100644 --- a/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsSeedBehaviour.tsx +++ b/invokeai/frontend/web/src/features/dynamicPrompts/components/ParamDynamicPromptsSeedBehaviour.tsx @@ -5,6 +5,8 @@ import { InformationalPopover } from 'common/components/InformationalPopover/Inf import { isSeedBehaviour, seedBehaviourChanged, + selectDynamicPromptsMode, + selectDynamicPromptsRandomRefreshMode, selectDynamicPromptsSeedBehaviour, } from 'features/dynamicPrompts/store/dynamicPromptsSlice'; import { memo, useCallback, useMemo } from 'react'; @@ -13,6 +15,8 @@ import { useTranslation } from 'react-i18next'; const ParamDynamicPromptsSeedBehaviour = () => { const dispatch = useAppDispatch(); const { t } = useTranslation(); + const mode = useAppSelector(selectDynamicPromptsMode); + const randomRefreshMode = useAppSelector(selectDynamicPromptsRandomRefreshMode); const seedBehaviour = useAppSelector(selectDynamicPromptsSeedBehaviour); const options = useMemo(() => { @@ -42,6 +46,10 @@ const ParamDynamicPromptsSeedBehaviour = () => { const value = useMemo(() => options.find((o) => o.value === seedBehaviour), [options, seedBehaviour]); + if (mode === 'random' && randomRefreshMode === 'per_image') { + return null; + } + return ( diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/components/ShowDynamicPromptsPreviewButton.tsx b/invokeai/frontend/web/src/features/dynamicPrompts/components/ShowDynamicPromptsPreviewButton.tsx index 1a23f0d375c..b954cbf16ab 100644 --- a/invokeai/frontend/web/src/features/dynamicPrompts/components/ShowDynamicPromptsPreviewButton.tsx +++ b/invokeai/frontend/web/src/features/dynamicPrompts/components/ShowDynamicPromptsPreviewButton.tsx @@ -6,7 +6,7 @@ import { selectDynamicPromptsIsError, selectDynamicPromptsIsLoading, } from 'features/dynamicPrompts/store/dynamicPromptsSlice'; -import { memo } from 'react'; +import { memo, useCallback } from 'react'; import { useTranslation } from 'react-i18next'; import { PiBracketsCurlyBold } from 'react-icons/pi'; @@ -20,6 +20,10 @@ export const ShowDynamicPromptsPreviewButton = memo(() => { const isError = useAppSelector(selectDynamicPromptsIsError); const { isOpen, onOpen } = useDynamicPromptsModal(); + const onClick = useCallback(() => { + onOpen(); + }, [onOpen]); + return ( { isDisabled={isOpen} aria-label={t('dynamicPrompts.showDynamicPrompts')} icon={} - onClick={onOpen} + onClick={onClick} sx={isLoading ? loadingStyles : undefined} colorScheme={isError && !isLoading ? 'error' : 'base'} /> diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/hooks/useDynamicPromptsWatcher.tsx b/invokeai/frontend/web/src/features/dynamicPrompts/hooks/useDynamicPromptsWatcher.tsx index 2fd6b18b06a..16f6c7ac6bc 100644 --- a/invokeai/frontend/web/src/features/dynamicPrompts/hooks/useDynamicPromptsWatcher.tsx +++ b/invokeai/frontend/web/src/features/dynamicPrompts/hooks/useDynamicPromptsWatcher.tsx @@ -1,16 +1,22 @@ import { useAppSelector, useAppStore } from 'app/store/storeHooks'; import { debounce } from 'es-toolkit/compat'; +import { selectIterations } from 'features/controlLayers/store/paramsSlice'; import { + type DynamicPromptMode, isErrorChanged, isLoadingChanged, parsingErrorChanged, promptsChanged, - selectDynamicPromptsMaxPrompts, + selectDynamicPromptsMaxCombinations, + selectDynamicPromptsMode, + selectDynamicPromptsRandomRefreshMode, + selectDynamicPromptsRandomSamples, + selectDynamicPromptsRandomSeed, } from 'features/dynamicPrompts/store/dynamicPromptsSlice'; import { getShouldProcessPrompt } from 'features/dynamicPrompts/util/getShouldProcessPrompt'; +import { resolveDynamicPrompts } from 'features/dynamicPrompts/util/resolveDynamicPrompts'; import { selectPresetModifiedPrompts } from 'features/nodes/util/graph/graphBuilderUtils'; import { useEffect, useMemo } from 'react'; -import { utilitiesApi } from 'services/api/endpoints/utilities'; const DYNAMIC_PROMPTS_DEBOUNCE_MS = 1000; @@ -21,52 +27,54 @@ export const useDynamicPromptsWatcher = () => { const { getState, dispatch } = useAppStore(); // The prompt to process is derived from the preset-modified prompts const presetModifiedPrompts = useAppSelector(selectPresetModifiedPrompts); - const maxPrompts = useAppSelector(selectDynamicPromptsMaxPrompts); + const mode = useAppSelector(selectDynamicPromptsMode); + const randomSamples = useAppSelector(selectDynamicPromptsRandomSamples); + const randomRefreshMode = useAppSelector(selectDynamicPromptsRandomRefreshMode); + const maxCombinations = useAppSelector(selectDynamicPromptsMaxCombinations); + const randomSeed = useAppSelector(selectDynamicPromptsRandomSeed); + const iterations = useAppSelector(selectIterations); const debouncedUpdateDynamicPrompts = useMemo( () => - debounce(async (positivePrompt: string, maxPrompts: number) => { - // Try to fetch the dynamic prompts and store in state - try { - const req = dispatch( - utilitiesApi.endpoints.dynamicPrompts.initiate( - { - prompt: positivePrompt, - max_prompts: maxPrompts, - }, - { subscribe: false } - ) - ); + debounce( + async ( + positivePrompt: string, + mode: DynamicPromptMode, + randomSamples: number, + maxCombinations: number, + randomSeed: number, + randomRefreshMode: 'manual' | 'per_enqueue' | 'per_image', + iterations: number + ) => { + // Try to fetch the dynamic prompts and store in state + try { + const res = await resolveDynamicPrompts({ + dispatch, + prompt: positivePrompt, + mode, + randomSamples, + maxCombinations, + randomSeed, + randomRefreshMode, + iterations, + }); - const res = await req.unwrap(); - - dispatch(promptsChanged(res.prompts)); - dispatch(parsingErrorChanged(res.error)); - dispatch(isErrorChanged(false)); - } catch { - dispatch(isErrorChanged(true)); - dispatch(isLoadingChanged(false)); - } - }, DYNAMIC_PROMPTS_DEBOUNCE_MS), + dispatch(promptsChanged(res.prompts)); + dispatch(parsingErrorChanged(res.error)); + dispatch(isErrorChanged(false)); + } catch { + dispatch(isErrorChanged(true)); + dispatch(isLoadingChanged(false)); + } + }, + DYNAMIC_PROMPTS_DEBOUNCE_MS + ), [dispatch] ); useEffect(() => { - // Before we execute, imperatively check the dynamic prompts query cache to see if we have already fetched this prompt const state = getState(); - const cachedPrompts = utilitiesApi.endpoints.dynamicPrompts.select({ - prompt: presetModifiedPrompts.positive, - max_prompts: maxPrompts, - })(state).data; - - if (cachedPrompts) { - // Yep we already did this prompt, use the cached result - dispatch(promptsChanged(cachedPrompts.prompts)); - dispatch(parsingErrorChanged(cachedPrompts.error)); - return; - } - // If the prompt is not in the cache, check if we should process it - this is just looking for dynamic prompts syntax if (!getShouldProcessPrompt(presetModifiedPrompts.positive)) { dispatch(promptsChanged([presetModifiedPrompts.positive])); @@ -80,6 +88,25 @@ export const useDynamicPromptsWatcher = () => { dispatch(isLoadingChanged(true)); } - debouncedUpdateDynamicPrompts(presetModifiedPrompts.positive, maxPrompts); - }, [debouncedUpdateDynamicPrompts, dispatch, getState, maxPrompts, presetModifiedPrompts]); + debouncedUpdateDynamicPrompts( + presetModifiedPrompts.positive, + mode, + randomSamples, + maxCombinations, + randomSeed, + randomRefreshMode, + iterations + ); + }, [ + debouncedUpdateDynamicPrompts, + dispatch, + getState, + iterations, + maxCombinations, + mode, + presetModifiedPrompts, + randomSamples, + randomRefreshMode, + randomSeed, + ]); }; diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/store/dynamicPromptsSlice.test.ts b/invokeai/frontend/web/src/features/dynamicPrompts/store/dynamicPromptsSlice.test.ts new file mode 100644 index 00000000000..121be735e1b --- /dev/null +++ b/invokeai/frontend/web/src/features/dynamicPrompts/store/dynamicPromptsSlice.test.ts @@ -0,0 +1,124 @@ +import type { Equals } from 'tsafe'; +import { assert } from 'tsafe'; +import { describe, expect, it } from 'vitest'; +import type { z } from 'zod'; + +import { dynamicPromptsSliceConfig, modeChanged, randomRefreshModeChanged } from './dynamicPromptsSlice'; + +describe('dynamicPromptsSlice', () => { + type InitialState = ReturnType; + type SchemaState = z.infer; + + const { reducer } = dynamicPromptsSliceConfig.slice; + const migrate = dynamicPromptsSliceConfig.persistConfig?.migrate; + + it('keeps the initial state aligned with the persisted schema', () => { + assert>(); + }); + + it('defaults fresh users to random mode', () => { + expect(dynamicPromptsSliceConfig.getInitialState()).toMatchObject({ + _version: 4, + mode: 'random', + randomSamples: 1, + maxCombinations: 100, + randomSeed: 0, + randomRefreshMode: 'per_image', + }); + }); + + it('can switch to all-combinations mode', () => { + const state = reducer(dynamicPromptsSliceConfig.getInitialState(), modeChanged('combinatorial')); + + expect(state.mode).toBe('combinatorial'); + }); + + it('can switch random refresh behavior', () => { + const state = reducer(dynamicPromptsSliceConfig.getInitialState(), randomRefreshModeChanged('per_enqueue')); + + expect(state.randomRefreshMode).toBe('per_enqueue'); + }); + + it('migrates existing combinatorial users to all-combinations mode', () => { + expect(migrate).toBeDefined(); + + const state = migrate?.({ + _version: 1, + maxPrompts: 250, + combinatorial: true, + seedBehaviour: 'PER_PROMPT', + }); + + expect(state).toMatchObject({ + _version: 4, + mode: 'combinatorial', + maxCombinations: 250, + randomSamples: 1, + randomRefreshMode: 'manual', + seedBehaviour: 'PER_PROMPT', + }); + }); + + it('migrates existing random users to random mode', () => { + expect(migrate).toBeDefined(); + + const state = migrate?.({ + _version: 1, + maxPrompts: 12, + combinatorial: false, + }); + + expect(state).toMatchObject({ + _version: 4, + mode: 'random', + maxCombinations: 12, + randomSamples: 1, + randomRefreshMode: 'per_image', + }); + }); + + it('migrates version 2 random users to per-image random', () => { + expect(migrate).toBeDefined(); + + const state = migrate?.({ + _version: 2, + mode: 'random', + randomSamples: 1, + maxCombinations: 100, + randomSeed: 123, + prompts: ['test'], + isError: false, + isLoading: false, + seedBehaviour: 'PER_ITERATION', + }); + + expect(state).toMatchObject({ + _version: 4, + mode: 'random', + randomRefreshMode: 'per_image', + }); + }); + + it('preserves version 3 locked-preview random users', () => { + expect(migrate).toBeDefined(); + + const state = migrate?.({ + _version: 3, + mode: 'random', + randomSamples: 1, + maxCombinations: 100, + randomSeed: 123, + randomRefreshMode: 'manual', + prompts: ['test'], + isError: false, + isLoading: false, + seedBehaviour: 'PER_ITERATION', + }); + + expect(state).toMatchObject({ + _version: 4, + mode: 'random', + randomRefreshMode: 'manual', + }); + }); +}); diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/store/dynamicPromptsSlice.ts b/invokeai/frontend/web/src/features/dynamicPrompts/store/dynamicPromptsSlice.ts index 8e071261224..8238f90e89a 100644 --- a/invokeai/frontend/web/src/features/dynamicPrompts/store/dynamicPromptsSlice.ts +++ b/invokeai/frontend/web/src/features/dynamicPrompts/store/dynamicPromptsSlice.ts @@ -11,10 +11,21 @@ const zSeedBehaviour = z.enum(['PER_ITERATION', 'PER_PROMPT']); export const isSeedBehaviour = buildZodTypeGuard(zSeedBehaviour); export type SeedBehaviour = z.infer; +const zDynamicPromptMode = z.enum(['random', 'combinatorial']); +export const isDynamicPromptMode = buildZodTypeGuard(zDynamicPromptMode); +export type DynamicPromptMode = z.infer; + +const zDynamicPromptRandomRefreshMode = z.enum(['manual', 'per_enqueue', 'per_image']); +export const isDynamicPromptRandomRefreshMode = buildZodTypeGuard(zDynamicPromptRandomRefreshMode); +export type DynamicPromptRandomRefreshMode = z.infer; + const zDynamicPromptsState = z.object({ - _version: z.literal(1), - maxPrompts: z.number().int().min(1).max(1000), - combinatorial: z.boolean(), + _version: z.literal(4), + mode: zDynamicPromptMode, + randomSamples: z.number().int().min(1).max(1000), + maxCombinations: z.number().int().min(1).max(10000), + randomSeed: z.number().int().min(0), + randomRefreshMode: zDynamicPromptRandomRefreshMode, prompts: z.array(z.string()), parsingError: z.string().nullish(), isError: z.boolean(), @@ -24,9 +35,12 @@ const zDynamicPromptsState = z.object({ export type DynamicPromptsState = z.infer; const getInitialState = (): DynamicPromptsState => ({ - _version: 1, - maxPrompts: 100, - combinatorial: true, + _version: 4, + mode: 'random', + randomSamples: 1, + maxCombinations: 100, + randomSeed: 0, + randomRefreshMode: 'per_image', prompts: [], parsingError: undefined, isError: false, @@ -38,8 +52,20 @@ const slice = createSlice({ name: 'dynamicPrompts', initialState: getInitialState(), reducers: { - maxPromptsChanged: (state, action: PayloadAction) => { - state.maxPrompts = action.payload; + modeChanged: (state, action: PayloadAction) => { + state.mode = action.payload; + }, + randomSamplesChanged: (state, action: PayloadAction) => { + state.randomSamples = action.payload; + }, + maxCombinationsChanged: (state, action: PayloadAction) => { + state.maxCombinations = action.payload; + }, + randomSeedChanged: (state, action: PayloadAction) => { + state.randomSeed = action.payload; + }, + randomRefreshModeChanged: (state, action: PayloadAction) => { + state.randomRefreshMode = action.payload; }, promptsChanged: (state, action: PayloadAction) => { state.prompts = action.payload; @@ -61,7 +87,11 @@ const slice = createSlice({ }); export const { - maxPromptsChanged, + modeChanged, + randomSamplesChanged, + maxCombinationsChanged, + randomSeedChanged, + randomRefreshModeChanged, promptsChanged, parsingErrorChanged, isErrorChanged, @@ -76,24 +106,72 @@ export const dynamicPromptsSliceConfig: SliceConfig = { persistConfig: { migrate: (state) => { assert(isPlainObject(state)); - if (!('_version' in state)) { - state._version = 1; + const initialState = getInitialState(); + if (state._version === 4) { + return zDynamicPromptsState.parse({ ...initialState, ...state }); + } + if (state._version === 2 || state._version === 3) { + const mode = isDynamicPromptMode(state.mode) ? state.mode : initialState.mode; + return zDynamicPromptsState.parse({ + ...initialState, + ...state, + _version: 4, + mode, + randomRefreshMode: getMigratedRandomRefreshMode(mode, state.randomRefreshMode), + seedBehaviour: isSeedBehaviour(state.seedBehaviour) ? state.seedBehaviour : initialState.seedBehaviour, + }); } - return zDynamicPromptsState.parse(state); + const legacyCombinatorial = state.combinatorial === true; + const legacyMode = legacyCombinatorial ? 'combinatorial' : 'random'; + const legacyMaxPrompts = typeof state.maxPrompts === 'number' ? state.maxPrompts : initialState.maxCombinations; + return zDynamicPromptsState.parse({ + ...initialState, + mode: legacyMode, + maxCombinations: legacyMaxPrompts, + randomRefreshMode: getMigratedRandomRefreshMode(legacyMode, state.randomRefreshMode), + seedBehaviour: isSeedBehaviour(state.seedBehaviour) ? state.seedBehaviour : initialState.seedBehaviour, + }); }, persistDenylist: ['prompts', 'parsingError', 'isError', 'isLoading'], }, }; +const getMigratedRandomRefreshMode = ( + mode: DynamicPromptMode, + value: unknown +): DynamicPromptRandomRefreshMode => { + if (mode !== 'random') { + return 'manual'; + } + if (value === 'manual') { + return 'manual'; + } + return 'per_image'; +}; + export const selectDynamicPromptsSlice = (state: RootState) => state.dynamicPrompts; const createDynamicPromptsSelector = (selector: Selector) => createSelector(selectDynamicPromptsSlice, selector); +export const selectDynamicPromptsMode = createDynamicPromptsSelector((dynamicPrompts) => dynamicPrompts.mode); +export const selectDynamicPromptsRandomSamples = createDynamicPromptsSelector( + (dynamicPrompts) => dynamicPrompts.randomSamples +); +export const selectDynamicPromptsMaxCombinations = createDynamicPromptsSelector( + (dynamicPrompts) => dynamicPrompts.maxCombinations +); +export const selectDynamicPromptsRandomSeed = createDynamicPromptsSelector( + (dynamicPrompts) => dynamicPrompts.randomSeed +); +export const selectDynamicPromptsRandomRefreshMode = createDynamicPromptsSelector( + (dynamicPrompts) => dynamicPrompts.randomRefreshMode +); export const selectDynamicPromptsMaxPrompts = createDynamicPromptsSelector( - (dynamicPrompts) => dynamicPrompts.maxPrompts + (dynamicPrompts) => + dynamicPrompts.mode === 'combinatorial' ? dynamicPrompts.maxCombinations : dynamicPrompts.randomSamples ); export const selectDynamicPromptsCombinatorial = createDynamicPromptsSelector( - (dynamicPrompts) => dynamicPrompts.combinatorial + (dynamicPrompts) => dynamicPrompts.mode === 'combinatorial' ); export const selectDynamicPromptsPrompts = createDynamicPromptsSelector((dynamicPrompts) => dynamicPrompts.prompts); export const selectDynamicPromptsParsingError = createDynamicPromptsSelector( diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/util/getDynamicPromptsQueryArg.test.ts b/invokeai/frontend/web/src/features/dynamicPrompts/util/getDynamicPromptsQueryArg.test.ts new file mode 100644 index 00000000000..ba6c191b73a --- /dev/null +++ b/invokeai/frontend/web/src/features/dynamicPrompts/util/getDynamicPromptsQueryArg.test.ts @@ -0,0 +1,38 @@ +import { describe, expect, it } from 'vitest'; + +import { getDynamicPromptsQueryArg } from './getDynamicPromptsQueryArg'; + +describe('getDynamicPromptsQueryArg', () => { + it('passes random mode with an explicit seed', () => { + expect( + getDynamicPromptsQueryArg({ + prompt: '__camera/lens__', + mode: 'random', + randomSamples: 1, + maxCombinations: 100, + randomSeed: 42, + }) + ).toEqual({ + prompt: '__camera/lens__', + max_prompts: 1, + combinatorial: false, + seed: 42, + }); + }); + + it('passes all-combinations mode explicitly', () => { + expect( + getDynamicPromptsQueryArg({ + prompt: '__camera/lens__', + mode: 'combinatorial', + randomSamples: 1, + maxCombinations: 100, + randomSeed: 42, + }) + ).toEqual({ + prompt: '__camera/lens__', + max_prompts: 100, + combinatorial: true, + }); + }); +}); diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/util/getDynamicPromptsQueryArg.ts b/invokeai/frontend/web/src/features/dynamicPrompts/util/getDynamicPromptsQueryArg.ts new file mode 100644 index 00000000000..6cdb0f47a63 --- /dev/null +++ b/invokeai/frontend/web/src/features/dynamicPrompts/util/getDynamicPromptsQueryArg.ts @@ -0,0 +1,32 @@ +import type { DynamicPromptMode } from 'features/dynamicPrompts/store/dynamicPromptsSlice'; + +type GetDynamicPromptsQueryArgArg = { + prompt: string; + mode: DynamicPromptMode; + randomSamples: number; + maxCombinations: number; + randomSeed: number; +}; + +export const getDynamicPromptsQueryArg = ({ + prompt, + mode, + randomSamples, + maxCombinations, + randomSeed, +}: GetDynamicPromptsQueryArgArg) => { + if (mode === 'combinatorial') { + return { + prompt, + max_prompts: maxCombinations, + combinatorial: true, + }; + } + + return { + prompt, + max_prompts: randomSamples, + combinatorial: false, + seed: randomSeed, + }; +}; diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/util/getShouldProcessPrompt.test.ts b/invokeai/frontend/web/src/features/dynamicPrompts/util/getShouldProcessPrompt.test.ts new file mode 100644 index 00000000000..6bc69d41c32 --- /dev/null +++ b/invokeai/frontend/web/src/features/dynamicPrompts/util/getShouldProcessPrompt.test.ts @@ -0,0 +1,17 @@ +import { describe, expect, it } from 'vitest'; + +import { getShouldProcessPrompt } from './getShouldProcessPrompt'; + +describe('getShouldProcessPrompt', () => { + it('detects variant syntax', () => { + expect(getShouldProcessPrompt('a {red|blue} car')).toBe(true); + }); + + it('detects wildcard syntax', () => { + expect(getShouldProcessPrompt('a __color__ car')).toBe(true); + }); + + it('ignores plain prompts', () => { + expect(getShouldProcessPrompt('a red car')).toBe(false); + }); +}); diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/util/getShouldProcessPrompt.ts b/invokeai/frontend/web/src/features/dynamicPrompts/util/getShouldProcessPrompt.ts index 416363fec1f..a55bf623bf2 100644 --- a/invokeai/frontend/web/src/features/dynamicPrompts/util/getShouldProcessPrompt.ts +++ b/invokeai/frontend/web/src/features/dynamicPrompts/util/getShouldProcessPrompt.ts @@ -1,2 +1,2 @@ -const hasOpenCloseCurlyBracesRegex = /.*\{[\s\S]*\}.*/; -export const getShouldProcessPrompt = (prompt: string): boolean => hasOpenCloseCurlyBracesRegex.test(prompt); +const hasDynamicPromptSyntaxRegex = /\{[\s\S]*\}|__[^\r\n]+?__|\$\{[\s\S]*\}/; +export const getShouldProcessPrompt = (prompt: string): boolean => hasDynamicPromptSyntaxRegex.test(prompt); diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/util/promptIntent.test.ts b/invokeai/frontend/web/src/features/dynamicPrompts/util/promptIntent.test.ts new file mode 100644 index 00000000000..d471cb66083 --- /dev/null +++ b/invokeai/frontend/web/src/features/dynamicPrompts/util/promptIntent.test.ts @@ -0,0 +1,26 @@ +import { describe, expect, it } from 'vitest'; + +import { + getHasMixedCyclicAndNonCyclicDynamicPromptSyntax, + getHasNonCyclicDynamicPromptSyntax, + getIsCycleOnlyDynamicPrompt, +} from './promptIntent'; + +describe('promptIntent', () => { + it('detects cycle-only dynamic prompts', () => { + expect(getIsCycleOnlyDynamicPrompt('portrait __@lighting/studio__')).toBe(true); + expect(getIsCycleOnlyDynamicPrompt('portrait (__@lighting/studio__)++')).toBe(true); + }); + + it('detects mixed cyclic and non-cyclic dynamic prompts', () => { + const prompt = '__@lighting/studio__, __camera/lens__'; + + expect(getIsCycleOnlyDynamicPrompt(prompt)).toBe(false); + expect(getHasNonCyclicDynamicPromptSyntax(prompt)).toBe(true); + expect(getHasMixedCyclicAndNonCyclicDynamicPromptSyntax(prompt)).toBe(true); + }); + + it('treats braced variants as non-cyclic dynamic syntax', () => { + expect(getHasNonCyclicDynamicPromptSyntax('__@lighting/studio__ {warm|cool}')).toBe(true); + }); +}); diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/util/promptIntent.ts b/invokeai/frontend/web/src/features/dynamicPrompts/util/promptIntent.ts new file mode 100644 index 00000000000..fb0bf7d0d92 --- /dev/null +++ b/invokeai/frontend/web/src/features/dynamicPrompts/util/promptIntent.ts @@ -0,0 +1,29 @@ +const WILDCARD_REFERENCE_PATTERN = /__([^\r\n]+?)__/g; +const CYCLIC_WILDCARD_PATTERN = /__@[^\r\n_][^\r\n]*?__/; +const BRACED_DYNAMIC_PATTERN = /\{[\s\S]*\}|\$\{[\s\S]*\}/; + +export const getHasCyclicWildcardSyntax = (prompt: string): boolean => CYCLIC_WILDCARD_PATTERN.test(prompt); + +export const getHasNonCyclicDynamicPromptSyntax = (prompt: string): boolean => { + if (BRACED_DYNAMIC_PATTERN.test(prompt)) { + return true; + } + + for (const match of prompt.matchAll(WILDCARD_REFERENCE_PATTERN)) { + const reference = match[1]?.trim() ?? ''; + if (!reference.startsWith('@')) { + return true; + } + } + + return false; +}; + +export const getHasDynamicPromptSyntax = (prompt: string): boolean => + getHasCyclicWildcardSyntax(prompt) || getHasNonCyclicDynamicPromptSyntax(prompt); + +export const getHasMixedCyclicAndNonCyclicDynamicPromptSyntax = (prompt: string): boolean => + getHasCyclicWildcardSyntax(prompt) && getHasNonCyclicDynamicPromptSyntax(prompt); + +export const getIsCycleOnlyDynamicPrompt = (prompt: string): boolean => + getHasCyclicWildcardSyntax(prompt) && !getHasNonCyclicDynamicPromptSyntax(prompt); diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/util/refreshDynamicPromptsForEnqueue.test.ts b/invokeai/frontend/web/src/features/dynamicPrompts/util/refreshDynamicPromptsForEnqueue.test.ts new file mode 100644 index 00000000000..388471ab43d --- /dev/null +++ b/invokeai/frontend/web/src/features/dynamicPrompts/util/refreshDynamicPromptsForEnqueue.test.ts @@ -0,0 +1,194 @@ +import type { RootState } from 'app/store/store'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +const mocks = vi.hoisted(() => ({ + dynamicPromptsInitiate: vi.fn(), + getShouldProcessPrompt: vi.fn((prompt: string) => prompt.includes('__')), + selectPresetModifiedPrompts: vi.fn((state: RootState) => ({ positive: state.params.positivePrompt })), +})); + +vi.mock('features/dynamicPrompts/util/getShouldProcessPrompt', () => ({ + getShouldProcessPrompt: mocks.getShouldProcessPrompt, +})); + +vi.mock('features/nodes/util/graph/graphBuilderUtils', () => ({ + selectPresetModifiedPrompts: mocks.selectPresetModifiedPrompts, +})); + +vi.mock('services/api/endpoints/utilities', () => ({ + utilitiesApi: { + endpoints: { + dynamicPrompts: { + initiate: mocks.dynamicPromptsInitiate, + }, + }, + }, +})); + +import { getDynamicPromptsOutputCount } from './resolveDynamicPrompts'; +import { + getShouldRefreshDynamicPromptsForEnqueue, + refreshDynamicPromptsForEnqueue, +} from './refreshDynamicPromptsForEnqueue'; + +const buildState = (overrides: { + mode?: 'random' | 'combinatorial'; + randomRefreshMode?: 'manual' | 'per_enqueue' | 'per_image'; + randomSamples?: number; + iterations?: number; + positivePrompt?: string; +}): RootState => + ({ + params: { + iterations: overrides.iterations ?? 3, + positivePrompt: overrides.positivePrompt ?? '__camera/lens__', + }, + dynamicPrompts: { + mode: overrides.mode ?? 'random', + randomRefreshMode: overrides.randomRefreshMode ?? 'per_image', + randomSamples: overrides.randomSamples ?? 1, + maxCombinations: 100, + }, + }) as RootState; + +describe('refreshDynamicPromptsForEnqueue', () => { + beforeEach(() => { + vi.restoreAllMocks(); + mocks.dynamicPromptsInitiate.mockReset(); + mocks.getShouldProcessPrompt.mockClear(); + mocks.selectPresetModifiedPrompts.mockClear(); + }); + + it('refreshes random prompts for per-image and per-invoke modes only', () => { + expect(getShouldRefreshDynamicPromptsForEnqueue(buildState({ randomRefreshMode: 'per_image' }))).toBe(true); + expect(getShouldRefreshDynamicPromptsForEnqueue(buildState({ randomRefreshMode: 'per_enqueue' }))).toBe(true); + expect(getShouldRefreshDynamicPromptsForEnqueue(buildState({ randomRefreshMode: 'manual' }))).toBe(false); + expect(getShouldRefreshDynamicPromptsForEnqueue(buildState({ mode: 'combinatorial' }))).toBe(false); + }); + + it('requests randomSamples times iterations for per-image enqueue refresh', () => { + expect( + getDynamicPromptsOutputCount({ + prompt: '__camera/lens__', + randomRefreshMode: 'per_image', + randomSamples: 2, + iterations: 3, + }) + ).toBe(6); + }); + + it('requests only randomSamples for per-invoke enqueue refresh', () => { + expect( + getDynamicPromptsOutputCount({ + prompt: '__camera/lens__', + randomRefreshMode: 'per_enqueue', + randomSamples: 2, + iterations: 3, + }) + ).toBe(2); + }); + + it('requests randomSamples times iterations for per-invoke cycle-only enqueue refresh', () => { + expect( + getDynamicPromptsOutputCount({ + prompt: '__@lighting/studio__', + randomRefreshMode: 'per_enqueue', + randomSamples: 2, + iterations: 3, + }) + ).toBe(6); + }); + + it('counts cycle-only locked-preview prompts per generated output', () => { + expect( + getDynamicPromptsOutputCount({ + prompt: '__@lighting/studio__', + randomRefreshMode: 'manual', + randomSamples: 2, + iterations: 3, + }) + ).toBe(6); + }); + + it('passes combinatorial false and a seed when refreshing per-image prompts', async () => { + const request = { unwrap: vi.fn().mockResolvedValue({ prompts: ['35mm', '85mm', '50mm'], error: undefined }) }; + mocks.dynamicPromptsInitiate.mockReturnValue(request); + vi.spyOn(Date, 'now').mockReturnValue(1000); + vi.spyOn(Math, 'random').mockReturnValue(0); + + const state = buildState({ randomRefreshMode: 'per_image', randomSamples: 1, iterations: 3 }); + const dispatch = vi.fn((action) => action); + const store = { + dispatch, + getState: vi.fn(() => state), + }; + + await refreshDynamicPromptsForEnqueue(store as never); + + expect(mocks.dynamicPromptsInitiate).toHaveBeenCalledWith( + { + prompt: '__camera/lens__', + max_prompts: 3, + combinatorial: false, + seed: 1000, + }, + { subscribe: false, forceRefetch: true } + ); + expect(dispatch).toHaveBeenCalledWith(expect.objectContaining({ type: 'dynamicPrompts/promptsChanged' })); + }); + + it('resolves mixed cycle and random prompts with fixed random tokens for per-invoke enqueue refresh', async () => { + const requests = [ + { unwrap: vi.fn().mockResolvedValue({ prompts: ['<>, pearl earrings'], error: undefined }) }, + { + unwrap: vi.fn().mockResolvedValue({ + prompts: ['35mm lens, pearl earrings', '50mm lens, pearl earrings', '85mm lens, pearl earrings'], + error: undefined, + }), + }, + ]; + mocks.dynamicPromptsInitiate.mockReturnValueOnce(requests[0]).mockReturnValueOnce(requests[1]); + vi.spyOn(Date, 'now').mockReturnValue(1000); + vi.spyOn(Math, 'random').mockReturnValue(0); + + const state = buildState({ + randomRefreshMode: 'per_enqueue', + iterations: 3, + positivePrompt: '__@camera/lens__, __character/accessory__', + }); + const dispatch = vi.fn((action) => action); + const store = { + dispatch, + getState: vi.fn(() => state), + }; + + await refreshDynamicPromptsForEnqueue(store as never); + + expect(mocks.dynamicPromptsInitiate).toHaveBeenNthCalledWith( + 1, + { + prompt: '<>, __character/accessory__', + max_prompts: 1, + combinatorial: false, + seed: 1000, + }, + { subscribe: false, forceRefetch: true } + ); + expect(mocks.dynamicPromptsInitiate).toHaveBeenNthCalledWith( + 2, + { + prompt: '__@camera/lens__, pearl earrings', + max_prompts: 3, + combinatorial: false, + seed: 1000, + }, + { subscribe: false, forceRefetch: true } + ); + expect(dispatch).toHaveBeenCalledWith( + expect.objectContaining({ + payload: ['35mm lens, pearl earrings', '50mm lens, pearl earrings', '85mm lens, pearl earrings'], + type: 'dynamicPrompts/promptsChanged', + }) + ); + }); +}); diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/util/refreshDynamicPromptsForEnqueue.ts b/invokeai/frontend/web/src/features/dynamicPrompts/util/refreshDynamicPromptsForEnqueue.ts new file mode 100644 index 00000000000..5736d575d62 --- /dev/null +++ b/invokeai/frontend/web/src/features/dynamicPrompts/util/refreshDynamicPromptsForEnqueue.ts @@ -0,0 +1,67 @@ +import type { AppStore, RootState } from 'app/store/store'; +import { + isErrorChanged, + isLoadingChanged, + parsingErrorChanged, + promptsChanged, + randomSeedChanged, +} from 'features/dynamicPrompts/store/dynamicPromptsSlice'; +import { getShouldProcessPrompt } from 'features/dynamicPrompts/util/getShouldProcessPrompt'; +import { resolveDynamicPrompts } from 'features/dynamicPrompts/util/resolveDynamicPrompts'; +import { selectPresetModifiedPrompts } from 'features/nodes/util/graph/graphBuilderUtils'; + +const getRandomSeed = () => Date.now() + Math.floor(Math.random() * 1_000_000); + +export const getShouldRefreshDynamicPromptsForEnqueue = (state: RootState): boolean => { + const { dynamicPrompts } = state; + + if (dynamicPrompts.mode !== 'random' || dynamicPrompts.randomRefreshMode === 'manual') { + return false; + } + + return getShouldProcessPrompt(selectPresetModifiedPrompts(state).positive); +}; + +/** + * Dynamic prompts refresh before queueing unless the preview is locked. Per-image random mode and cycle-only per-invoke + * mode ask for one resolved prompt per generated output, then the batch builder zips those prompts with seeds one-to-one. + */ +export const refreshDynamicPromptsForEnqueue = async (store: AppStore): Promise => { + const { dispatch, getState } = store; + const state = getState(); + + if (!getShouldRefreshDynamicPromptsForEnqueue(state)) { + return state; + } + + const prompt = selectPresetModifiedPrompts(state).positive; + const { maxCombinations, randomRefreshMode, randomSamples } = state.dynamicPrompts; + const randomSeed = getRandomSeed(); + + dispatch(randomSeedChanged(randomSeed)); + dispatch(isLoadingChanged(true)); + + try { + const res = await resolveDynamicPrompts({ + dispatch, + prompt, + mode: 'random', + randomSamples, + maxCombinations, + randomSeed, + randomRefreshMode, + iterations: state.params.iterations, + forceRefetch: true, + }); + + dispatch(promptsChanged(res.prompts)); + dispatch(parsingErrorChanged(res.error)); + dispatch(isErrorChanged(false)); + + return getState(); + } catch { + dispatch(isErrorChanged(true)); + dispatch(isLoadingChanged(false)); + return null; + } +}; diff --git a/invokeai/frontend/web/src/features/dynamicPrompts/util/resolveDynamicPrompts.ts b/invokeai/frontend/web/src/features/dynamicPrompts/util/resolveDynamicPrompts.ts new file mode 100644 index 00000000000..8bb70c4040e --- /dev/null +++ b/invokeai/frontend/web/src/features/dynamicPrompts/util/resolveDynamicPrompts.ts @@ -0,0 +1,177 @@ +import type { AppDispatch } from 'app/store/store'; +import type { + DynamicPromptMode, + DynamicPromptRandomRefreshMode, +} from 'features/dynamicPrompts/store/dynamicPromptsSlice'; +import { getDynamicPromptsQueryArg } from 'features/dynamicPrompts/util/getDynamicPromptsQueryArg'; +import { + getHasCyclicWildcardSyntax, + getHasMixedCyclicAndNonCyclicDynamicPromptSyntax, +} from 'features/dynamicPrompts/util/promptIntent'; +import { utilitiesApi } from 'services/api/endpoints/utilities'; +import type { paths } from 'services/api/schema'; + +type DynamicPromptsResponse = + paths['/api/v1/utilities/dynamicprompts']['post']['responses']['200']['content']['application/json']; + +type ResolveDynamicPromptsArg = { + dispatch: AppDispatch; + prompt: string; + mode: DynamicPromptMode; + randomSamples: number; + maxCombinations: number; + randomSeed: number; + randomRefreshMode: DynamicPromptRandomRefreshMode; + iterations: number; + forceRefetch?: boolean; +}; + +const CYCLIC_WILDCARD_PATTERN = /__@[^\r\n_][^\r\n]*?__/g; +const MAX_DYNAMIC_PROMPTS = 10000; + +export const getDynamicPromptsOutputCount = ({ + iterations, + prompt, + randomRefreshMode, + randomSamples, +}: { + iterations: number; + prompt: string; + randomRefreshMode: DynamicPromptRandomRefreshMode; + randomSamples: number; +}): number => { + if (randomRefreshMode === 'per_image' || getHasCyclicWildcardSyntax(prompt)) { + return randomSamples * Math.max(iterations, 1); + } + + return randomSamples; +}; + +export const resolveDynamicPrompts = async ({ + dispatch, + prompt, + mode, + randomSamples, + maxCombinations, + randomSeed, + randomRefreshMode, + iterations, + forceRefetch = false, +}: ResolveDynamicPromptsArg): Promise => { + if (mode !== 'random') { + return fetchDynamicPrompts({ + dispatch, + prompt, + mode, + randomSamples, + maxCombinations, + randomSeed, + forceRefetch, + }); + } + + const outputCount = Math.min( + getDynamicPromptsOutputCount({ iterations, prompt, randomRefreshMode, randomSamples }), + MAX_DYNAMIC_PROMPTS + ); + + if (randomRefreshMode === 'per_image' || !getHasMixedCyclicAndNonCyclicDynamicPromptSyntax(prompt)) { + return fetchDynamicPrompts({ + dispatch, + prompt, + mode, + randomSamples: outputCount, + maxCombinations, + randomSeed, + forceRefetch, + }); + } + + const protectedPrompt = protectCyclicWildcards(prompt); + const randomResponse = await fetchDynamicPrompts({ + dispatch, + prompt: protectedPrompt.prompt, + mode, + randomSamples, + maxCombinations, + randomSeed, + forceRefetch, + }); + + if (randomResponse.error || randomResponse.prompts.length === 0) { + return randomResponse; + } + + const cyclePromptCount = Math.max(1, Math.floor(outputCount / randomResponse.prompts.length)); + const cycleResponses = await Promise.all( + randomResponse.prompts.map((randomPrompt) => + fetchDynamicPrompts({ + dispatch, + prompt: restoreCyclicWildcards(randomPrompt, protectedPrompt.tokens), + mode, + randomSamples: cyclePromptCount, + maxCombinations, + randomSeed, + forceRefetch, + }) + ) + ); + + return mergeDynamicPromptResponses( + [randomResponse, ...cycleResponses], + cycleResponses.flatMap((res) => res.prompts) + ); +}; + +const fetchDynamicPrompts = ({ + dispatch, + prompt, + mode, + randomSamples, + maxCombinations, + randomSeed, + forceRefetch, +}: { + dispatch: AppDispatch; + prompt: string; + mode: DynamicPromptMode; + randomSamples: number; + maxCombinations: number; + randomSeed: number; + forceRefetch: boolean; +}): Promise => { + const queryArg = getDynamicPromptsQueryArg({ + prompt, + mode, + randomSamples, + maxCombinations, + randomSeed, + }); + const req = dispatch(utilitiesApi.endpoints.dynamicPrompts.initiate(queryArg, { subscribe: false, forceRefetch })); + return req.unwrap(); +}; + +const protectCyclicWildcards = (prompt: string): { prompt: string; tokens: string[] } => { + const tokens: string[] = []; + const protectedPrompt = prompt.replace(CYCLIC_WILDCARD_PATTERN, (token) => { + const index = tokens.push(token) - 1; + return getCyclePlaceholder(index); + }); + + return { prompt: protectedPrompt, tokens }; +}; + +const restoreCyclicWildcards = (prompt: string, tokens: string[]): string => + tokens.reduce((nextPrompt, token, index) => nextPrompt.replaceAll(getCyclePlaceholder(index), token), prompt); + +const getCyclePlaceholder = (index: number): string => `<>`; + +const mergeDynamicPromptResponses = ( + responses: DynamicPromptsResponse[], + prompts: string[] +): DynamicPromptsResponse => ({ + prompts: prompts.length > 0 ? prompts : [''], + error: responses.map((res) => res.error).find(Boolean), + warnings: [...new Set(responses.flatMap((res) => res.warnings ?? []))], + missing_wildcards: [...new Set(responses.flatMap((res) => res.missing_wildcards ?? []))], +}); diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/buildLinearBatchConfig.test.ts b/invokeai/frontend/web/src/features/nodes/util/graph/buildLinearBatchConfig.test.ts new file mode 100644 index 00000000000..b5169533c24 --- /dev/null +++ b/invokeai/frontend/web/src/features/nodes/util/graph/buildLinearBatchConfig.test.ts @@ -0,0 +1,209 @@ +import type { RootState } from 'app/store/store'; +import type { Graph } from 'features/nodes/util/graph/generation/Graph'; +import type { Invocation } from 'services/api/types'; +import { describe, expect, it, vi } from 'vitest'; + +const mocks = vi.hoisted(() => ({ + getShouldProcessPrompt: vi.fn((prompt: string) => prompt.includes('__')), + selectPresetModifiedPrompts: vi.fn((state: RootState) => ({ positive: state.params.positivePrompt })), +})); + +vi.mock('features/dynamicPrompts/util/getShouldProcessPrompt', () => ({ + getShouldProcessPrompt: mocks.getShouldProcessPrompt, +})); + +vi.mock('features/nodes/util/graph/graphBuilderUtils', () => ({ + selectPresetModifiedPrompts: mocks.selectPresetModifiedPrompts, +})); + +import { prepareLinearUIBatch } from './buildLinearBatchConfig'; + +const g = { + getGraph: () => ({ id: 'graph', nodes: {}, edges: [] }), +} as unknown as Graph; + +const positivePromptNode = { id: 'positive_prompt' } as Invocation<'string'>; +const seedNode = { id: 'seed' } as Invocation<'integer'>; + +const buildState = (overrides: { + iterations?: number; + positivePrompt?: string; + prompts?: string[]; + randomRefreshMode?: 'manual' | 'per_enqueue' | 'per_image'; + randomSamples?: number; + seedBehaviour?: 'PER_ITERATION' | 'PER_PROMPT'; +}): RootState => + ({ + params: { + iterations: overrides.iterations ?? 3, + shouldRandomizeSeed: false, + seed: 100, + positivePrompt: overrides.positivePrompt ?? '__lighting/studio__', + }, + dynamicPrompts: { + mode: 'random', + randomRefreshMode: overrides.randomRefreshMode ?? 'per_image', + randomSamples: overrides.randomSamples ?? 1, + prompts: overrides.prompts ?? ['softbox', 'window light', 'rim light'], + seedBehaviour: overrides.seedBehaviour ?? 'PER_ITERATION', + }, + }) as RootState; + +describe('prepareLinearUIBatch', () => { + it('zips per-image random prompts with one seed per output', () => { + const batch = prepareLinearUIBatch({ + state: buildState({ iterations: 3, prompts: ['softbox', 'window light', 'rim light'] }), + g, + prepend: false, + base: 'sdxl', + positivePromptNode, + seedNode, + origin: 'generate', + destination: 'generate', + }).batch; + + expect(batch.runs).toBe(1); + expect(batch.data).toEqual([ + [ + { node_path: 'seed', field_name: 'value', items: [100, 101, 102] }, + { node_path: 'positive_prompt', field_name: 'value', items: ['softbox', 'window light', 'rim light'] }, + ], + ]); + }); + + it('ignores seedBehaviour in per-output mode because prompts are already final queued outputs', () => { + const batch = prepareLinearUIBatch({ + state: buildState({ + seedBehaviour: 'PER_PROMPT', + positivePrompt: '__@lighting/studio__', + prompts: ['softbox', 'window light', 'softbox'], + iterations: 3, + }), + g, + prepend: false, + base: 'sdxl', + positivePromptNode, + seedNode, + origin: 'generate', + destination: 'generate', + }).batch; + + expect(batch.data).toEqual([ + [ + { node_path: 'seed', field_name: 'value', items: [100, 101, 102] }, + { node_path: 'positive_prompt', field_name: 'value', items: ['softbox', 'window light', 'softbox'] }, + ], + ]); + }); + + it('queues randomSamples times iterations outputs for per-image random prompts', () => { + const prompts = ['a', 'b', 'c', 'd', 'e', 'f']; + const batch = prepareLinearUIBatch({ + state: buildState({ iterations: 3, randomSamples: 2, prompts }), + g, + prepend: false, + base: 'sdxl', + positivePromptNode, + seedNode, + origin: 'generate', + destination: 'generate', + }).batch; + + expect(batch.data).toEqual([ + [ + { node_path: 'seed', field_name: 'value', items: [100, 101, 102, 103, 104, 105] }, + { node_path: 'positive_prompt', field_name: 'value', items: prompts }, + ], + ]); + }); + + it('keeps static prompts on normal iteration batching even when per-image is selected', () => { + const batch = prepareLinearUIBatch({ + state: buildState({ positivePrompt: 'portrait', prompts: ['portrait'], iterations: 3 }), + g, + prepend: false, + base: 'sdxl', + positivePromptNode, + seedNode, + origin: 'generate', + destination: 'generate', + }).batch; + + expect(batch.data).toEqual([ + [{ node_path: 'seed', field_name: 'value', items: [100, 101, 102] }], + [{ node_path: 'positive_prompt', field_name: 'value', items: ['portrait'] }], + ]); + }); + + it('keeps per-invoke random prompts on normal prompt by iteration batching', () => { + const batch = prepareLinearUIBatch({ + state: buildState({ randomRefreshMode: 'per_enqueue', prompts: ['softbox'], iterations: 3 }), + g, + prepend: false, + base: 'sdxl', + positivePromptNode, + seedNode, + origin: 'generate', + destination: 'generate', + }).batch; + + expect(batch.data).toEqual([ + [{ node_path: 'seed', field_name: 'value', items: [100, 101, 102] }], + [{ node_path: 'positive_prompt', field_name: 'value', items: ['softbox'] }], + ]); + }); + + it('zips per-invoke cycle-only prompts with one seed per output', () => { + const batch = prepareLinearUIBatch({ + state: buildState({ + randomRefreshMode: 'per_enqueue', + positivePrompt: '__@lighting/studio__', + prompts: ['softbox', 'window light', 'rim light'], + iterations: 3, + }), + g, + prepend: false, + base: 'sdxl', + positivePromptNode, + seedNode, + origin: 'generate', + destination: 'generate', + }).batch; + + expect(batch.data).toEqual([ + [ + { node_path: 'seed', field_name: 'value', items: [100, 101, 102] }, + { node_path: 'positive_prompt', field_name: 'value', items: ['softbox', 'window light', 'rim light'] }, + ], + ]); + }); + + it('zips mixed cycle and random prompts with one seed per output', () => { + const batch = prepareLinearUIBatch({ + state: buildState({ + randomRefreshMode: 'per_enqueue', + positivePrompt: '__@lighting/studio__, __camera/lens__', + prompts: ['softbox, 50mm', 'window light, 50mm', 'rim light, 50mm'], + iterations: 3, + }), + g, + prepend: false, + base: 'sdxl', + positivePromptNode, + seedNode, + origin: 'generate', + destination: 'generate', + }).batch; + + expect(batch.data).toEqual([ + [ + { node_path: 'seed', field_name: 'value', items: [100, 101, 102] }, + { + node_path: 'positive_prompt', + field_name: 'value', + items: ['softbox, 50mm', 'window light, 50mm', 'rim light, 50mm'], + }, + ], + ]); + }); +}); diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/buildLinearBatchConfig.ts b/invokeai/frontend/web/src/features/nodes/util/graph/buildLinearBatchConfig.ts index 36b6ad4b738..a59ab1765f4 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/buildLinearBatchConfig.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/buildLinearBatchConfig.ts @@ -1,25 +1,14 @@ import type { RootState } from 'app/store/store'; import { generateSeeds } from 'common/util/generateSeeds'; -import { range } from 'es-toolkit/compat'; -import type { SeedBehaviour } from 'features/dynamicPrompts/store/dynamicPromptsSlice'; import type { BaseModelType } from 'features/nodes/types/common'; +import { + getExtendedDynamicPrompts, + getShouldUsePerOutputDynamicPrompts, +} from 'features/nodes/util/graph/dynamicPromptBatching'; import type { Graph } from 'features/nodes/util/graph/generation/Graph'; import type { components } from 'services/api/schema'; import type { Batch, EnqueueBatchArg, Invocation } from 'services/api/types'; -const getExtendedPrompts = (arg: { - seedBehaviour: SeedBehaviour; - iterations: number; - prompts: string[]; - base: BaseModelType; -}): string[] => { - const { seedBehaviour, iterations, prompts } = arg; - if (seedBehaviour === 'PER_PROMPT') { - return range(iterations).flatMap(() => prompts); - } - return prompts; -}; - export const prepareLinearUIBatch = (arg: { state: RootState; g: Graph; @@ -30,7 +19,7 @@ export const prepareLinearUIBatch = (arg: { origin: string; destination: string; }): EnqueueBatchArg => { - const { state, g, base, prepend, positivePromptNode, seedNode, origin, destination } = arg; + const { state, g, prepend, positivePromptNode, seedNode, origin, destination } = arg; const { iterations, shouldRandomizeSeed, seed } = state.params; const { prompts, seedBehaviour } = state.dynamicPrompts; @@ -38,6 +27,42 @@ export const prepareLinearUIBatch = (arg: { const firstBatchDatumList: components['schemas']['BatchDatum'][] = []; const secondBatchDatumList: components['schemas']['BatchDatum'][] = []; + if (getShouldUsePerOutputDynamicPrompts(state)) { + const perImageBatchDatumList: components['schemas']['BatchDatum'][] = []; + + if (seedNode) { + // Per-output dynamic prompt mode has already resolved the final queued outputs, so seeds are zipped one-to-one + // with those outputs instead of applying seedBehaviour as a second batching dimension. + perImageBatchDatumList.push({ + node_path: seedNode.id, + field_name: 'value', + items: generateSeeds({ + count: prompts.length, + start: shouldRandomizeSeed ? undefined : seed, + }), + }); + } + + perImageBatchDatumList.push({ + node_path: positivePromptNode.id, + field_name: 'value', + items: prompts, + }); + + data.push(perImageBatchDatumList); + + return { + prepend, + batch: { + graph: g.getGraph(), + runs: 1, + data, + origin, + destination, + }, + }; + } + // add seeds first to ensure the output order groups the prompts if (seedNode && seedBehaviour === 'PER_PROMPT') { const seeds = generateSeeds({ @@ -65,7 +90,7 @@ export const prepareLinearUIBatch = (arg: { data.push(secondBatchDatumList); } - const extendedPrompts = getExtendedPrompts({ seedBehaviour, iterations, prompts, base }); + const extendedPrompts = getExtendedDynamicPrompts({ seedBehaviour, iterations, prompts }); // zipped batch of prompts firstBatchDatumList.push({ diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/dynamicPromptBatching.ts b/invokeai/frontend/web/src/features/nodes/util/graph/dynamicPromptBatching.ts new file mode 100644 index 00000000000..c082f66c5be --- /dev/null +++ b/invokeai/frontend/web/src/features/nodes/util/graph/dynamicPromptBatching.ts @@ -0,0 +1,32 @@ +import type { RootState } from 'app/store/store'; +import { range } from 'es-toolkit/compat'; +import type { SeedBehaviour } from 'features/dynamicPrompts/store/dynamicPromptsSlice'; +import { getShouldProcessPrompt } from 'features/dynamicPrompts/util/getShouldProcessPrompt'; +import { getHasCyclicWildcardSyntax } from 'features/dynamicPrompts/util/promptIntent'; +import { selectPresetModifiedPrompts } from 'features/nodes/util/graph/graphBuilderUtils'; + +export const getExtendedDynamicPrompts = ({ + seedBehaviour, + iterations, + prompts, +}: { + seedBehaviour: SeedBehaviour; + iterations: number; + prompts: string[]; +}): string[] => { + if (seedBehaviour === 'PER_PROMPT') { + return range(iterations).flatMap(() => prompts); + } + return prompts; +}; + +export const getShouldUsePerOutputDynamicPrompts = (state: RootState): boolean => { + const prompt = selectPresetModifiedPrompts(state).positive; + const { mode, randomRefreshMode } = state.dynamicPrompts; + + if (mode !== 'random' || !getShouldProcessPrompt(prompt)) { + return false; + } + + return randomRefreshMode === 'per_image' || getHasCyclicWildcardSyntax(prompt); +}; diff --git a/invokeai/frontend/web/src/features/parameters/components/Core/ParamPositivePrompt.tsx b/invokeai/frontend/web/src/features/parameters/components/Core/ParamPositivePrompt.tsx index 955eeb4abf0..8f607c99d5c 100644 --- a/invokeai/frontend/web/src/features/parameters/components/Core/ParamPositivePrompt.tsx +++ b/invokeai/frontend/web/src/features/parameters/components/Core/ParamPositivePrompt.tsx @@ -28,6 +28,7 @@ import { PromptPopover } from 'features/prompt/PromptPopover'; import { clearPromptUndo, consumePromptUndo } from 'features/prompt/promptUndo'; import { usePrompt } from 'features/prompt/usePrompt'; import { usePromptAttentionHotkeys } from 'features/prompt/usePromptAttentionHotkeys'; +import { PromptWorkbench } from 'features/promptWorkbench/PromptWorkbench'; import { selectStylePresetActivePresetId, selectStylePresetViewMode, @@ -322,6 +323,7 @@ export const ParamPositivePrompt = memo(() => { + {hasLlavaModels && } ); diff --git a/invokeai/frontend/web/src/features/promptWorkbench/PromptInspector.tsx b/invokeai/frontend/web/src/features/promptWorkbench/PromptInspector.tsx new file mode 100644 index 00000000000..c1f2cdd88b9 --- /dev/null +++ b/invokeai/frontend/web/src/features/promptWorkbench/PromptInspector.tsx @@ -0,0 +1,703 @@ +import { + Box, + Button, + Flex, + IconButton, + Menu, + MenuButton, + MenuItem, + MenuList, + Text, + Tooltip, +} from '@invoke-ai/ui-library'; +import type { DynamicPromptRandomRefreshMode } from 'features/dynamicPrompts/store/dynamicPromptsSlice'; +import type { MouseEvent, ReactElement, ReactNode } from 'react'; +import { memo, useCallback } from 'react'; +import { useTranslation } from 'react-i18next'; +import { + PiCaretDownBold, + PiDiceFiveBold, + PiDotsThreeVerticalBold, + PiPushPinSimpleBold, + PiRepeatBold, + PiScalesBold, + PiTrashSimpleBold, +} from 'react-icons/pi'; + +import { type PromptWorkbenchTranslation, tx } from './i18n'; +import { + getWeightBehaviorLabel, + getWeightShortLabel, + getWildcardBehaviorLabel, + getWildcardBehaviorShortLabel, + type PromptRange, + type PromptWeightOccurrence, + type PromptWildcardOccurrence, + type PromptWorkbenchOccurrence, + type WildcardBehaviorAction, +} from './occurrences'; +import type { PromptWorkbenchBadgeTone } from './PromptWorkbenchBadge'; + +type PromptInspectorProps = { + occurrences: PromptWorkbenchOccurrence[]; + randomRefreshMode: DynamicPromptRandomRefreshMode; + fixedWildcardOccurrenceId: string | null; + fixedWildcardValues: string[] | null; + isFetchingFixedWildcardValues: boolean; + activeFixedValueIndex: number; + onSelectRange: (range: PromptRange) => void; + onWildcardBehaviorAction: (occurrence: PromptWildcardOccurrence, action: WildcardBehaviorAction) => void; + onRemoveWeightOccurrence: (occurrence: PromptWeightOccurrence) => void; + onFixedValue: (value: string) => void; + setFixedValueElement: (index: number, element: HTMLElement | null) => void; +}; + +const ROW_GRID_TEMPLATE = '2rem minmax(0, 1fr) minmax(5.75rem, 6.25rem) 2rem 1.75rem'; +const INTENT_CARD_BG = 'linear-gradient(180deg, rgba(25, 35, 45, 0.94) 0%, rgba(17, 26, 35, 0.96) 100%)'; +const INTENT_CARD_HOVER_BG = 'linear-gradient(180deg, rgba(29, 41, 53, 0.98) 0%, rgba(20, 31, 41, 0.98) 100%)'; +const INTENT_CARD_BORDER = 'rgba(126, 143, 164, 0.3)'; +const INTENT_CARD_BORDER_HOVER = 'rgba(153, 170, 191, 0.42)'; +const INTENT_CARD_ERROR_BORDER = 'rgba(248, 113, 113, 0.5)'; +const INTENT_CARD_ERROR_BORDER_HOVER = 'rgba(252, 165, 165, 0.58)'; + +export const PromptInspector = memo( + ({ + occurrences, + randomRefreshMode, + fixedWildcardOccurrenceId, + fixedWildcardValues, + isFetchingFixedWildcardValues, + activeFixedValueIndex, + onSelectRange, + onWildcardBehaviorAction, + onRemoveWeightOccurrence, + onFixedValue, + setFixedValueElement, + }: PromptInspectorProps) => { + if (occurrences.length === 0) { + return null; + } + + return ( + + + {occurrences.map((occurrence) => + occurrence.type === 'wildcard' ? ( + + ) : ( + + ) + )} + + + ); + } +); + +PromptInspector.displayName = 'PromptInspector'; + +type WildcardInspectorRowProps = { + occurrence: PromptWildcardOccurrence; + randomRefreshMode: DynamicPromptRandomRefreshMode; + fixedValues: string[] | null; + isFetchingFixedValues: boolean; + activeFixedValueIndex: number; + onSelectRange: (range: PromptRange) => void; + onWildcardBehaviorAction: (occurrence: PromptWildcardOccurrence, action: WildcardBehaviorAction) => void; + onFixedValue: (value: string) => void; + setFixedValueElement: (index: number, element: HTMLElement | null) => void; +}; + +const WildcardInspectorRow = memo( + ({ + occurrence, + randomRefreshMode, + fixedValues, + isFetchingFixedValues, + activeFixedValueIndex, + onSelectRange, + onWildcardBehaviorAction, + onFixedValue, + setFixedValueElement, + }: WildcardInspectorRowProps) => { + const { t } = useTranslation(); + const hasWildcard = occurrence.wildcard !== null; + const isActionable = occurrence.behavior !== 'missing' && occurrence.behavior !== 'unavailable'; + const selectionRange = occurrence.weight?.range ?? occurrence.range; + + const onSelectMouseDown = useCallback( + (e: MouseEvent) => { + e.preventDefault(); + onSelectRange(selectionRange); + }, + [onSelectRange, selectionRange] + ); + + const onBehaviorAction = useCallback( + (action: WildcardBehaviorAction) => { + onWildcardBehaviorAction(occurrence, action); + }, + [occurrence, onWildcardBehaviorAction] + ); + + return ( + + + } /> + + + + + + {t(getWildcardSecondaryText(occurrence).key, getWildcardSecondaryText(occurrence).options)} + + + + + + + {occurrence.weight ? ( + + ) : ( + + - + + )} + + + + + + {(fixedValues || isFetchingFixedValues) && ( + + {isFetchingFixedValues && ( + + {t('promptWorkbench.values.loading')} + + )} + {fixedValues?.map((value, index) => ( + + ))} + + )} + + ); + } +); + +WildcardInspectorRow.displayName = 'WildcardInspectorRow'; + +const WildcardBehaviorControl = memo( + ({ + occurrence, + randomRefreshMode, + isActionable, + canPickFixedValue, + onAction, + }: { + occurrence: PromptWildcardOccurrence; + randomRefreshMode: DynamicPromptRandomRefreshMode; + isActionable: boolean; + canPickFixedValue: boolean; + onAction: (action: WildcardBehaviorAction) => void; + }) => { + const { t } = useTranslation(); + const onButtonMouseDown = useCallback((e: MouseEvent) => { + e.preventDefault(); + }, []); + + const onRandom = useCallback(() => { + onAction('random'); + }, [onAction]); + + const onCycle = useCallback(() => { + onAction('cycle'); + }, [onAction]); + + const onFixed = useCallback(() => { + onAction('fixed'); + }, [onAction]); + + if (!isActionable) { + return ( + + + {t( + getWildcardBehaviorShortLabel(occurrence, randomRefreshMode).key, + getWildcardBehaviorShortLabel(occurrence, randomRefreshMode).options + )} + + + ); + } + + return ( + + + } + onMouseDown={onButtonMouseDown} + > + + {t( + getWildcardBehaviorShortLabel(occurrence, randomRefreshMode).key, + getWildcardBehaviorShortLabel(occurrence, randomRefreshMode).options + )} + + + + + } + onClick={onRandom} + title={t('promptWorkbench.actions.randomWildcardTooltip')} + > + {t('promptWorkbench.actions.randomWildcard')} + + } onClick={onCycle} title={t('promptWorkbench.actions.cycleTooltip')}> + {t('promptWorkbench.actions.cycle')} + + } onClick={onFixed} isDisabled={!canPickFixedValue}> + {t('promptWorkbench.actions.pickFixed')} + + + + ); + } +); + +WildcardBehaviorControl.displayName = 'WildcardBehaviorControl'; + +const WildcardActionsMenu = memo( + ({ + occurrence, + onAction, + }: { + occurrence: PromptWildcardOccurrence; + onAction: (action: WildcardBehaviorAction) => void; + }) => { + const { t } = useTranslation(); + const onButtonMouseDown = useCallback((e: MouseEvent) => { + e.preventDefault(); + }, []); + + const onRemove = useCallback(() => { + onAction('remove'); + }, [onAction]); + + return ( + + } + onMouseDown={onButtonMouseDown} + /> + + } color="error.300" onClick={onRemove}> + {t('promptWorkbench.actions.remove')} + + + + ); + } +); + +WildcardActionsMenu.displayName = 'WildcardActionsMenu'; + +const FixedValueButton = memo( + ({ + value, + index, + isActive, + onFixedValue, + setFixedValueElement, + }: { + value: string; + index: number; + isActive: boolean; + onFixedValue: (value: string) => void; + setFixedValueElement: (index: number, element: HTMLElement | null) => void; + }) => { + const onMouseDown = useCallback( + (e: MouseEvent) => { + e.preventDefault(); + onFixedValue(value); + }, + [onFixedValue, value] + ); + + const setElement = useCallback( + (element: HTMLElement | null) => { + setFixedValueElement(index, element); + }, + [index, setFixedValueElement] + ); + + return ( + + ); + } +); + +FixedValueButton.displayName = 'FixedValueButton'; + +type WeightInspectorRowProps = { + occurrence: PromptWeightOccurrence; + onSelectRange: (range: PromptRange) => void; + onRemoveWeightOccurrence: (occurrence: PromptWeightOccurrence) => void; +}; + +const WeightInspectorRow = memo(({ occurrence, onSelectRange, onRemoveWeightOccurrence }: WeightInspectorRowProps) => { + const { t } = useTranslation(); + const onSelectMouseDown = useCallback( + (e: MouseEvent) => { + e.preventDefault(); + onSelectRange(occurrence.range); + }, + [occurrence.range, onSelectRange] + ); + + const onRemove = useCallback(() => { + onRemoveWeightOccurrence(occurrence); + }, [occurrence, onRemoveWeightOccurrence]); + + return ( + + + } /> + + + + + + {t('promptWorkbench.weight.promptWeight')} + + + + + + + + + + + + + + ); +}); + +WeightInspectorRow.displayName = 'WeightInspectorRow'; + +const WeightActionsMenu = memo( + ({ occurrence, onRemove }: { occurrence: PromptWeightOccurrence; onRemove: () => void }) => { + const { t } = useTranslation(); + const onButtonMouseDown = useCallback((e: MouseEvent) => { + e.preventDefault(); + }, []); + + return ( + + } + onMouseDown={onButtonMouseDown} + /> + + } color="error.300" onClick={onRemove}> + {t('promptWorkbench.actions.remove')} + + + + ); + } +); + +WeightActionsMenu.displayName = 'WeightActionsMenu'; + +const WeightValue = memo(({ occurrence }: { occurrence: PromptWeightOccurrence }) => { + const { t } = useTranslation(); + const behaviorLabel = getWeightBehaviorLabel(occurrence); + const shortLabel = getWeightShortLabel(occurrence); + + return ( + + + {t(shortLabel.key, shortLabel.options)} + + + ); +}); + +WeightValue.displayName = 'WeightValue'; + +const TypeCell = memo(({ color, icon }: { color: string; icon: ReactElement }) => ( + + {icon} + +)); + +TypeCell.displayName = 'TypeCell'; + +const LabelCell = memo(({ children }: { children: ReactNode }) => ( + + {children} + +)); + +LabelCell.displayName = 'LabelCell'; + +const BehaviorCell = memo(({ children }: { children: ReactNode }) => ( + + {children} + +)); + +BehaviorCell.displayName = 'BehaviorCell'; + +const WeightCell = memo(({ children }: { children: ReactNode }) => ( + + {children} + +)); + +WeightCell.displayName = 'WeightCell'; + +const ActionsCell = memo(({ children }: { children: ReactNode }) => ( + + {children} + +)); + +ActionsCell.displayName = 'ActionsCell'; + +const IntentRowBox = memo( + ({ + children, + markerColor, + tone = 'neutral', + }: { + children: ReactNode; + markerColor: string; + tone?: PromptWorkbenchBadgeTone; + }) => ( + + {children} + + ) +); + +IntentRowBox.displayName = 'IntentRowBox'; + +const getWildcardSecondaryText = (occurrence: PromptWildcardOccurrence): PromptWorkbenchTranslation => { + if (occurrence.behavior === 'missing') { + return tx('promptWorkbench.rows.missingWildcard'); + } + if (occurrence.behavior === 'unavailable') { + return tx('promptWorkbench.rows.wildcardIndexUnavailable'); + } + + const options = occurrence.valueCount !== null ? { count: occurrence.valueCount } : undefined; + + switch (occurrence.behavior) { + case 'random': + return tx( + occurrence.valueCount !== null + ? 'promptWorkbench.rows.randomWildcardWithCount' + : 'promptWorkbench.rows.randomWildcard', + options + ); + case 'cycle': + return tx( + occurrence.valueCount !== null ? 'promptWorkbench.rows.cycleWithCount' : 'promptWorkbench.rows.cycle', + options + ); + case 'all': + return tx( + occurrence.valueCount !== null + ? 'promptWorkbench.rows.allCombinationsWithCount' + : 'promptWorkbench.rows.allCombinations', + options + ); + } +}; + +const getWildcardIconColor = (occurrence: PromptWildcardOccurrence): string => { + switch (occurrence.behavior) { + case 'random': + case 'cycle': + case 'all': + return occurrence.weight?.isSupported === false ? 'warning.300' : 'base.300'; + case 'missing': + case 'unavailable': + return 'error.300'; + } +}; + +const getWildcardMarkerColor = (occurrence: PromptWildcardOccurrence): string => { + switch (occurrence.behavior) { + case 'random': + case 'cycle': + case 'all': + return occurrence.weight?.isSupported === false ? 'warning.500' : 'invokeBlue.400'; + case 'missing': + case 'unavailable': + return 'error.500'; + } +}; + +const getWildcardRowTone = (occurrence: PromptWildcardOccurrence): PromptWorkbenchBadgeTone => { + switch (occurrence.behavior) { + case 'missing': + case 'unavailable': + return 'error'; + case 'random': + case 'cycle': + case 'all': + return 'neutral'; + } +}; diff --git a/invokeai/frontend/web/src/features/promptWorkbench/PromptWildcardBehaviorMenu.tsx b/invokeai/frontend/web/src/features/promptWorkbench/PromptWildcardBehaviorMenu.tsx new file mode 100644 index 00000000000..9d0906709de --- /dev/null +++ b/invokeai/frontend/web/src/features/promptWorkbench/PromptWildcardBehaviorMenu.tsx @@ -0,0 +1,118 @@ +import { IconButton, Menu, MenuButton, MenuItem, MenuList } from '@invoke-ai/ui-library'; +import type { MouseEvent, ReactElement } from 'react'; +import { memo, useCallback } from 'react'; +import { useTranslation } from 'react-i18next'; +import { + PiDiceFiveBold, + PiPushPinSimpleBold, + PiRepeatBold, + PiSquaresFourBold, + PiTrashSimpleBold, + PiWarningBold, +} from 'react-icons/pi'; + +import type { WildcardBehaviorAction, WildcardBehaviorIconType } from './occurrences'; + +type PromptWildcardBehaviorMenuProps = { + ariaLabel: string; + tooltip: string; + iconType: WildcardBehaviorIconType; + isActionable: boolean; + canPickFixedValue: boolean; + includeRemove: boolean; + onAction: (action: WildcardBehaviorAction) => void; + onOpen?: () => void; + onClose?: () => void; +}; + +const BUTTON_SIZE = 7; + +export const PromptWildcardBehaviorMenu = memo( + ({ + ariaLabel, + tooltip, + iconType, + isActionable, + canPickFixedValue, + includeRemove, + onAction, + onOpen, + onClose, + }: PromptWildcardBehaviorMenuProps) => { + const { t } = useTranslation(); + const onButtonMouseDown = useCallback((e: MouseEvent) => { + e.preventDefault(); + }, []); + + const onRandom = useCallback(() => { + onAction('random'); + }, [onAction]); + + const onCycle = useCallback(() => { + onAction('cycle'); + }, [onAction]); + + const onFixed = useCallback(() => { + onAction('fixed'); + }, [onAction]); + + const onRemove = useCallback(() => { + onAction('remove'); + }, [onAction]); + + return ( + + + + {isActionable && ( + <> + } + onClick={onRandom} + title={t('promptWorkbench.actions.randomWildcardTooltip')} + > + {t('promptWorkbench.actions.randomWildcard')} + + } onClick={onCycle} title={t('promptWorkbench.actions.cycleTooltip')}> + {t('promptWorkbench.actions.cycle')} + + } onClick={onFixed} isDisabled={!canPickFixedValue}> + {t('promptWorkbench.actions.pickFixed')} + + + )} + {includeRemove && ( + } color="error.300" onClick={onRemove}> + {t('promptWorkbench.actions.remove')} + + )} + + + ); + } +); + +PromptWildcardBehaviorMenu.displayName = 'PromptWildcardBehaviorMenu'; + +const getBehaviorMenuIcon = (iconType: WildcardBehaviorIconType): ReactElement => { + switch (iconType) { + case 'random': + return ; + case 'cycle': + return ; + case 'all': + return ; + case 'warning': + return ; + } +}; diff --git a/invokeai/frontend/web/src/features/promptWorkbench/PromptWorkbench.tsx b/invokeai/frontend/web/src/features/promptWorkbench/PromptWorkbench.tsx new file mode 100644 index 00000000000..1788d251531 --- /dev/null +++ b/invokeai/frontend/web/src/features/promptWorkbench/PromptWorkbench.tsx @@ -0,0 +1,950 @@ +import { Box, Button, Flex, Menu, MenuButton, MenuItem, MenuList, Text, Tooltip } from '@invoke-ai/ui-library'; +import { useAppDispatch, useAppSelector } from 'app/store/storeHooks'; +import { adjustPromptAttention } from 'common/util/promptAttention'; +import { selectModel } from 'features/controlLayers/store/paramsSlice'; +import { useDynamicPromptsModal } from 'features/dynamicPrompts/hooks/useDynamicPromptsModal'; +import { + modeChanged, + randomRefreshModeChanged, + selectDynamicPromptsMode, + selectDynamicPromptsParsingError, + selectDynamicPromptsPrompts, + selectDynamicPromptsRandomRefreshMode, +} from 'features/dynamicPrompts/store/dynamicPromptsSlice'; +import { selectSystemPrefersNumericAttentionWeights } from 'features/system/store/systemSlice'; +import type { MouseEvent, RefObject } from 'react'; +import { memo, useCallback, useEffect, useMemo, useRef, useState } from 'react'; +import { flushSync } from 'react-dom'; +import { useTranslation } from 'react-i18next'; +import { PiCubeBold, PiDiceFiveBold } from 'react-icons/pi'; +import type { WildcardIndexItem } from 'services/api/endpoints/utilities'; +import { useLazyWildcardValuesQuery, useWildcardsQuery } from 'services/api/endpoints/utilities'; + +import { getPromptDiagnostics, type PromptDiagnosticSeverity } from './diagnostics'; +import type { PromptWorkbenchTranslation } from './i18n'; +import { clampNavigationIndex, getNextNavigationIndex, getPromptWorkbenchKeyboardIntent } from './keyboardNavigation'; +import { getPromptModelCapabilities } from './modelCapabilities'; +import { + getPromptWorkbenchOccurrences, + getWildcardBehaviorActionIntent, + type PromptRange, + type PromptWeightOccurrence, + type PromptWildcardOccurrence, + removePromptRange, + replacePromptRange, + type WildcardBehaviorAction, +} from './occurrences'; +import { PromptInspector } from './PromptInspector'; +import { PromptWildcardBehaviorMenu } from './PromptWildcardBehaviorMenu'; +import { PromptWorkbenchBadge, type PromptWorkbenchBadgeTone } from './PromptWorkbenchBadge'; +import { + applyWildcardCompletion, + filterWildcardOptions, + getWildcardAutocompleteStatusMessage, + getWildcardCompletionContext, + getWildcardDisplayPath, + type WildcardCompletionContext, +} from './wildcards'; + +type PromptWorkbenchProps = { + prompt: string; + textareaRef: RefObject; + onPromptChange: (prompt: string) => void; +}; + +type SelectionRange = { + start: number; + end: number; +}; + +const EMPTY_WILDCARDS: [] = []; +const getCompletionContextKey = (context: WildcardCompletionContext): string => + `${context.start}:${context.end}:${context.query}`; +const PROMPT_INTENT_PANEL_BG = 'linear-gradient(180deg, rgba(15, 23, 31, 0.92) 0%, rgba(11, 18, 25, 0.96) 100%)'; +const PROMPT_INTENT_PANEL_BORDER = 'rgba(126, 143, 164, 0.28)'; + +export const PromptWorkbench = memo(({ prompt, textareaRef, onPromptChange }: PromptWorkbenchProps) => { + const { t } = useTranslation(); + const dispatch = useAppDispatch(); + const { onOpen: onOpenDynamicPromptsModal } = useDynamicPromptsModal(); + const model = useAppSelector(selectModel); + const prefersNumericWeights = useAppSelector(selectSystemPrefersNumericAttentionWeights); + const dynamicPromptMode = useAppSelector(selectDynamicPromptsMode); + const dynamicPromptRandomRefreshMode = useAppSelector(selectDynamicPromptsRandomRefreshMode); + const dynamicPrompts = useAppSelector(selectDynamicPromptsPrompts); + const dynamicPromptError = useAppSelector(selectDynamicPromptsParsingError); + const { + data: wildcardsData, + isError: isWildcardIndexUnavailable, + isFetching: isFetchingWildcards, + } = useWildcardsQuery(); + const [loadWildcardValues, wildcardValuesResult] = useLazyWildcardValuesQuery(); + const wildcards = wildcardsData?.wildcards ?? EMPTY_WILDCARDS; + const wildcardIndexErrorCount = wildcardsData?.errors.length ?? 0; + const capabilities = useMemo(() => getPromptModelCapabilities(model?.base), [model?.base]); + const [caret, setCaret] = useState(0); + const [isFocused, setIsFocused] = useState(false); + const [fixedWildcardPath, setFixedWildcardPath] = useState(null); + const [fixedWildcardContext, setFixedWildcardContext] = useState(null); + const [fixedWildcardOccurrence, setFixedWildcardOccurrence] = useState(null); + const [isAutocompleteBehaviorMenuOpen, setIsAutocompleteBehaviorMenuOpen] = useState(false); + const [dismissedCompletionContextKey, setDismissedCompletionContextKey] = useState(null); + const [activeWildcardIndex, setActiveWildcardIndex] = useState(0); + const [activeFixedValueIndex, setActiveFixedValueIndex] = useState(0); + const selectionRef = useRef({ start: 0, end: 0 }); + const [selection, setSelection] = useState({ start: 0, end: 0 }); + const completionContextRef = useRef(null); + const wildcardOptionElementsRef = useRef>([]); + const fixedValueElementsRef = useRef>([]); + const translate = useCallback( + (translation: PromptWorkbenchTranslation) => t(translation.key, translation.options), + [t] + ); + + const diagnostics = useMemo( + () => + getPromptDiagnostics({ + prompt, + modelBase: model?.base, + wildcards, + wildcardIndexUnavailable: isWildcardIndexUnavailable, + wildcardIndexErrorCount, + dynamicPromptCount: dynamicPrompts.length, + dynamicPromptMode, + dynamicPromptRandomRefreshMode, + dynamicPromptError, + }), + [ + dynamicPromptError, + dynamicPromptMode, + dynamicPromptRandomRefreshMode, + dynamicPrompts.length, + isWildcardIndexUnavailable, + model?.base, + prompt, + wildcardIndexErrorCount, + wildcards, + ] + ); + + const rawCompletionContextBase = useMemo( + () => (isFocused ? getWildcardCompletionContext(prompt, caret) : null), + [caret, isFocused, prompt] + ); + + const rawCompletionContext = useMemo(() => { + if (!rawCompletionContextBase) { + return null; + } + + return getCompletionContextKey(rawCompletionContextBase) === dismissedCompletionContextKey + ? null + : rawCompletionContextBase; + }, [dismissedCompletionContextKey, rawCompletionContextBase]); + + useEffect(() => { + if (rawCompletionContext) { + completionContextRef.current = rawCompletionContext; + } + }, [rawCompletionContext]); + + const completionContext = useMemo( + () => + rawCompletionContext ?? + fixedWildcardContext ?? + (isAutocompleteBehaviorMenuOpen ? completionContextRef.current : null), + [fixedWildcardContext, isAutocompleteBehaviorMenuOpen, rawCompletionContext] + ); + + const wildcardOptions = useMemo(() => { + if (!completionContext) { + return []; + } + return filterWildcardOptions(wildcards, completionContext.query); + }, [completionContext, wildcards]); + + const promptWorkbenchOccurrences = useMemo( + () => + getPromptWorkbenchOccurrences({ + prompt, + wildcards, + wildcardIndexUnavailable: isWildcardIndexUnavailable, + dynamicPromptMode, + supportsAttentionWeights: capabilities.supportsAttentionWeights, + }), + [capabilities.supportsAttentionWeights, dynamicPromptMode, isWildcardIndexUnavailable, prompt, wildcards] + ); + + const wildcardStatusMessage = useMemo(() => { + if (!completionContext) { + return null; + } + + return getWildcardAutocompleteStatusMessage({ + isLoading: isFetchingWildcards, + isUnavailable: isWildcardIndexUnavailable, + optionCount: wildcardOptions.length, + query: completionContext.query, + wildcardCount: wildcards.length, + }); + }, [completionContext, isFetchingWildcards, isWildcardIndexUnavailable, wildcardOptions.length, wildcards.length]); + + const fixedWildcardValues = + wildcardValuesResult.currentData?.path === fixedWildcardPath ? wildcardValuesResult.currentData.values : null; + + const activeFixedWildcardOccurrenceId = fixedWildcardOccurrence?.id ?? null; + const hasFixedValuePicker = Boolean(fixedWildcardPath && (fixedWildcardValues || wildcardValuesResult.isFetching)); + const fixedValueCount = fixedWildcardValues?.length ?? 0; + + const setWildcardOptionElement = useCallback((index: number, element: HTMLElement | null) => { + wildcardOptionElementsRef.current[index] = element; + }, []); + + const setFixedValueElement = useCallback((index: number, element: HTMLElement | null) => { + fixedValueElementsRef.current[index] = element; + }, []); + + const wildcardOptionElementSetters = useMemo( + () => + wildcardOptions.map((_wildcard, index) => (element: HTMLElement | null) => { + setWildcardOptionElement(index, element); + }), + [setWildcardOptionElement, wildcardOptions] + ); + + const fixedValueElementSetters = useMemo( + () => + (fixedWildcardValues ?? []).map((_value, index) => (element: HTMLElement | null) => { + setFixedValueElement(index, element); + }), + [fixedWildcardValues, setFixedValueElement] + ); + + useEffect(() => { + setActiveWildcardIndex((currentIndex) => clampNavigationIndex(currentIndex, wildcardOptions.length)); + wildcardOptionElementsRef.current = wildcardOptionElementsRef.current.slice(0, wildcardOptions.length); + }, [wildcardOptions.length]); + + useEffect(() => { + setActiveWildcardIndex(0); + }, [completionContext?.query, completionContext?.start]); + + useEffect(() => { + setActiveFixedValueIndex((currentIndex) => clampNavigationIndex(currentIndex, fixedValueCount)); + fixedValueElementsRef.current = fixedValueElementsRef.current.slice(0, fixedValueCount); + }, [fixedValueCount]); + + useEffect(() => { + setActiveFixedValueIndex(0); + }, [fixedWildcardPath]); + + useEffect(() => { + wildcardOptionElementsRef.current[activeWildcardIndex]?.scrollIntoView({ block: 'nearest' }); + }, [activeWildcardIndex]); + + useEffect(() => { + fixedValueElementsRef.current[activeFixedValueIndex]?.scrollIntoView({ block: 'nearest' }); + }, [activeFixedValueIndex]); + + const syncSelection = useCallback(() => { + const textarea = textareaRef.current; + if (!textarea) { + return; + } + const isTextareaFocused = document.activeElement === textarea; + if (!isTextareaFocused) { + setSelection({ start: 0, end: 0 }); + setIsFocused(false); + return; + } + const start = textarea.selectionStart ?? 0; + const end = textarea.selectionEnd ?? start; + selectionRef.current = { start, end }; + setSelection({ start, end }); + setCaret(end); + setIsFocused(true); + }, [textareaRef]); + + useEffect(() => { + const textarea = textareaRef.current; + if (!textarea) { + return; + } + let blurTimeout: number | undefined; + let selectionFrame: number | undefined; + const syncSelectionSoon = () => { + selectionFrame = window.requestAnimationFrame(syncSelection); + }; + const syncBlurred = () => { + blurTimeout = window.setTimeout(() => { + setSelection({ start: 0, end: 0 }); + setIsFocused(false); + }, 150); + }; + + textarea.addEventListener('keyup', syncSelectionSoon); + textarea.addEventListener('click', syncSelectionSoon); + textarea.addEventListener('mouseup', syncSelectionSoon); + textarea.addEventListener('input', syncSelectionSoon); + textarea.addEventListener('select', syncSelectionSoon); + textarea.addEventListener('focus', syncSelectionSoon); + textarea.addEventListener('blur', syncBlurred); + document.addEventListener('selectionchange', syncSelectionSoon); + + return () => { + textarea.removeEventListener('keyup', syncSelectionSoon); + textarea.removeEventListener('click', syncSelectionSoon); + textarea.removeEventListener('mouseup', syncSelectionSoon); + textarea.removeEventListener('input', syncSelectionSoon); + textarea.removeEventListener('select', syncSelectionSoon); + textarea.removeEventListener('focus', syncSelectionSoon); + textarea.removeEventListener('blur', syncBlurred); + document.removeEventListener('selectionchange', syncSelectionSoon); + if (blurTimeout !== undefined) { + window.clearTimeout(blurTimeout); + } + if (selectionFrame !== undefined) { + window.cancelAnimationFrame(selectionFrame); + } + }; + }, [syncSelection, textareaRef]); + + useEffect(() => { + if (!completionContext && !isAutocompleteBehaviorMenuOpen && !fixedWildcardContext) { + setFixedWildcardContext(null); + if (!fixedWildcardOccurrence) { + setFixedWildcardPath(null); + } + } + }, [completionContext, fixedWildcardContext, fixedWildcardOccurrence, isAutocompleteBehaviorMenuOpen]); + + useEffect(() => { + if ( + fixedWildcardOccurrence && + !promptWorkbenchOccurrences.some((occurrence) => occurrence.id === fixedWildcardOccurrence.id) + ) { + setFixedWildcardOccurrence(null); + } + }, [fixedWildcardOccurrence, promptWorkbenchOccurrences]); + + const focusPromptRange = useCallback( + (range: PromptRange) => { + requestAnimationFrame(() => { + const textarea = textareaRef.current; + textarea?.focus(); + textarea?.setSelectionRange(range.start, range.end); + selectionRef.current = { start: range.start, end: range.end }; + setSelection({ start: range.start, end: range.end }); + setCaret(range.end); + setIsFocused(document.activeElement === textarea); + }); + }, + [textareaRef] + ); + + const applyPromptReplacement = useCallback( + (nextPrompt: string, selection: SelectionRange) => { + flushSync(() => { + onPromptChange(nextPrompt); + }); + + requestAnimationFrame(() => { + const textarea = textareaRef.current; + textarea?.focus(); + textarea?.setSelectionRange(selection.start, selection.end); + selectionRef.current = selection; + setSelection(selection); + setCaret(selection.end); + setFixedWildcardPath(null); + setFixedWildcardContext(null); + setFixedWildcardOccurrence(null); + }); + }, + [onPromptChange, textareaRef] + ); + + const replaceOccurrenceRange = useCallback( + (range: PromptRange, replacement: string) => { + const result = replacePromptRange(prompt, range, replacement); + applyPromptReplacement(result.prompt, { start: result.caret, end: result.caret }); + }, + [applyPromptReplacement, prompt] + ); + + const replaceCompletion = useCallback( + (replacement: string, context: WildcardCompletionContext) => { + const result = applyWildcardCompletion(prompt, context, replacement); + + flushSync(() => { + onPromptChange(result.prompt); + }); + + requestAnimationFrame(() => { + const textarea = textareaRef.current; + textarea?.focus(); + textarea?.setSelectionRange(result.caret, result.caret); + selectionRef.current = { start: result.caret, end: result.caret }; + setSelection({ start: result.caret, end: result.caret }); + setCaret(result.caret); + setFixedWildcardPath(null); + setFixedWildcardContext(null); + setFixedWildcardOccurrence(null); + }); + }, + [onPromptChange, prompt, textareaRef] + ); + + const onWildcardMouseDown = useCallback( + (token: string) => (e: MouseEvent) => { + e.preventDefault(); + if (!completionContext) { + return; + } + dispatch(modeChanged('random')); + dispatch(randomRefreshModeChanged('per_image')); + replaceCompletion(token, completionContext); + }, + [completionContext, dispatch, replaceCompletion] + ); + + const openFixedValuesForAutocompleteWildcard = useCallback( + (wildcard: WildcardIndexItem, context: WildcardCompletionContext) => { + setFixedWildcardPath(wildcard.path); + setFixedWildcardContext(context); + setFixedWildcardOccurrence(null); + setActiveFixedValueIndex(0); + loadWildcardValues({ path: wildcard.path, limit: 200 }); + }, + [loadWildcardValues] + ); + + const onAutocompleteWildcardBehaviorAction = useCallback( + (wildcard: WildcardIndexItem) => (action: WildcardBehaviorAction) => { + const context = completionContext ?? completionContextRef.current; + if (!context) { + return; + } + + const intent = getWildcardBehaviorActionIntent(action, wildcard.path); + if (intent.opensFixedValues) { + openFixedValuesForAutocompleteWildcard(wildcard, context); + return; + } + + if (intent.replacement) { + replaceCompletion(intent.replacement, context); + } + }, + [completionContext, openFixedValuesForAutocompleteWildcard, replaceCompletion] + ); + + const onAutocompleteBehaviorMenuOpen = useCallback(() => { + setIsAutocompleteBehaviorMenuOpen(true); + }, []); + + const onAutocompleteBehaviorMenuClose = useCallback(() => { + setIsAutocompleteBehaviorMenuOpen(false); + }, []); + + const applyFixedValue = useCallback( + (value: string) => { + if (fixedWildcardContext) { + replaceCompletion(value, fixedWildcardContext); + return; + } + if (fixedWildcardOccurrence) { + replaceOccurrenceRange(fixedWildcardOccurrence.range, value); + } + }, + [fixedWildcardContext, fixedWildcardOccurrence, replaceCompletion, replaceOccurrenceRange] + ); + + const onFixedValueMouseDown = useCallback( + (value: string) => (e: MouseEvent) => { + e.preventDefault(); + applyFixedValue(value); + }, + [applyFixedValue] + ); + + const onInspectorFixedValue = useCallback( + (value: string) => { + applyFixedValue(value); + }, + [applyFixedValue] + ); + + const dismissCompletion = useCallback(() => { + const context = fixedWildcardContext ?? rawCompletionContextBase ?? completionContextRef.current; + if (context) { + setDismissedCompletionContextKey(getCompletionContextKey(context)); + } + setFixedWildcardPath(null); + setFixedWildcardContext(null); + setFixedWildcardOccurrence(null); + setIsAutocompleteBehaviorMenuOpen(false); + + requestAnimationFrame(() => { + textareaRef.current?.focus(); + }); + }, [fixedWildcardContext, rawCompletionContextBase, textareaRef]); + + const onPromptWorkbenchKeyDown = useCallback( + (e: KeyboardEvent) => { + const keyboardTarget = hasFixedValuePicker + ? 'fixed_values' + : completionContext && wildcardOptions.length > 0 + ? 'autocomplete' + : null; + + if (!keyboardTarget) { + return; + } + + const intent = getPromptWorkbenchKeyboardIntent({ + key: e.key, + shiftKey: e.shiftKey, + target: keyboardTarget, + }); + + if (!intent) { + return; + } + + e.preventDefault(); + e.stopPropagation(); + + if (intent === 'dismiss') { + dismissCompletion(); + return; + } + + if (intent === 'next' || intent === 'previous') { + if (keyboardTarget === 'fixed_values') { + setActiveFixedValueIndex((currentIndex) => + getNextNavigationIndex({ + currentIndex, + direction: intent, + itemCount: fixedValueCount, + }) + ); + return; + } + + setActiveWildcardIndex((currentIndex) => + getNextNavigationIndex({ + currentIndex, + direction: intent, + itemCount: wildcardOptions.length, + }) + ); + return; + } + + if (intent === 'insert_fixed_value') { + const value = fixedWildcardValues?.[activeFixedValueIndex]; + if (value) { + applyFixedValue(value); + } + return; + } + + const wildcard = wildcardOptions[activeWildcardIndex]; + if (!wildcard || !completionContext) { + return; + } + + if (intent === 'open_fixed_values') { + openFixedValuesForAutocompleteWildcard(wildcard, completionContext); + return; + } + + if (intent === 'insert_wildcard') { + dispatch(modeChanged('random')); + dispatch(randomRefreshModeChanged('per_image')); + replaceCompletion(wildcard.token, completionContext); + } + }, + [ + activeFixedValueIndex, + activeWildcardIndex, + applyFixedValue, + completionContext, + dismissCompletion, + dispatch, + fixedValueCount, + fixedWildcardValues, + hasFixedValuePicker, + openFixedValuesForAutocompleteWildcard, + replaceCompletion, + wildcardOptions, + ] + ); + + useEffect(() => { + const textarea = textareaRef.current; + if (!textarea) { + return; + } + + textarea.addEventListener('keydown', onPromptWorkbenchKeyDown); + + return () => { + textarea.removeEventListener('keydown', onPromptWorkbenchKeyDown); + }; + }, [onPromptWorkbenchKeyDown, textareaRef]); + + const adjustWeight = useCallback( + (direction: 'increment' | 'decrement') => { + const textarea = textareaRef.current; + const selection = selectionRef.current; + const result = adjustPromptAttention(prompt, selection.start, selection.end, direction, prefersNumericWeights); + + flushSync(() => { + onPromptChange(result.prompt); + }); + + requestAnimationFrame(() => { + textarea?.focus(); + textarea?.setSelectionRange(result.selectionStart, result.selectionEnd); + selectionRef.current = { start: result.selectionStart, end: result.selectionEnd }; + setSelection({ start: result.selectionStart, end: result.selectionEnd }); + setCaret(result.selectionEnd); + }); + }, + [onPromptChange, prefersNumericWeights, prompt, textareaRef] + ); + + const onDecrementMouseDown = useCallback((e: MouseEvent) => { + e.preventDefault(); + }, []); + + const onIncrementMouseDown = useCallback((e: MouseEvent) => { + e.preventDefault(); + }, []); + + const onDecrementClick = useCallback(() => { + adjustWeight('decrement'); + }, [adjustWeight]); + + const onIncrementClick = useCallback(() => { + adjustWeight('increment'); + }, [adjustWeight]); + + const onOpenQueuedOutputsMouseDown = useCallback((e: MouseEvent) => { + e.preventDefault(); + }, []); + + const onOpenQueuedOutputsClick = useCallback(() => { + onOpenDynamicPromptsModal(); + }, [onOpenDynamicPromptsModal]); + + const onRandomPerImageClick = useCallback(() => { + dispatch(modeChanged('random')); + dispatch(randomRefreshModeChanged('per_image')); + }, [dispatch]); + + const onRandomPerInvokeClick = useCallback(() => { + dispatch(modeChanged('random')); + dispatch(randomRefreshModeChanged('per_enqueue')); + }, [dispatch]); + + const onRemoveWildcardOccurrence = useCallback( + (occurrence: PromptWildcardOccurrence) => { + const result = removePromptRange(prompt, occurrence.weight?.range ?? occurrence.range); + applyPromptReplacement(result.prompt, { start: result.caret, end: result.caret }); + }, + [applyPromptReplacement, prompt] + ); + + const onRemoveWeightOccurrence = useCallback( + (occurrence: PromptWeightOccurrence) => { + const result = removePromptRange(prompt, occurrence.range); + applyPromptReplacement(result.prompt, { start: result.caret, end: result.caret }); + }, + [applyPromptReplacement, prompt] + ); + + const onInspectorWildcardBehaviorAction = useCallback( + (occurrence: PromptWildcardOccurrence, action: WildcardBehaviorAction) => { + const intent = getWildcardBehaviorActionIntent(action, occurrence.path); + + if (intent.opensFixedValues) { + if (!occurrence.wildcard) { + return; + } + setFixedWildcardPath(occurrence.path); + setFixedWildcardContext(null); + setFixedWildcardOccurrence(occurrence); + setActiveFixedValueIndex(0); + loadWildcardValues({ path: occurrence.path, limit: 200 }); + return; + } + + if (intent.removesPrompt) { + onRemoveWildcardOccurrence(occurrence); + return; + } + + if (intent.replacement) { + replaceOccurrenceRange(occurrence.range, intent.replacement); + } + }, + [loadWildcardValues, onRemoveWildcardOccurrence, replaceOccurrenceRange] + ); + + const hasPromptWorkbenchEntries = promptWorkbenchOccurrences.length > 0; + const hasRandomWildcardOccurrences = promptWorkbenchOccurrences.some( + (occurrence) => occurrence.type === 'wildcard' && occurrence.behavior === 'random' + ); + const canShowWeightControls = capabilities.supportsAttentionWeights && hasPromptWorkbenchEntries; + const canAdjustSelectionWeight = canShowWeightControls && selection.start !== selection.end; + const hasPromptWorkbenchPanel = hasPromptWorkbenchEntries || diagnostics.length > 0; + + return ( + + {completionContext && (wildcardOptions.length > 0 || wildcardStatusMessage) && ( + + + {wildcardStatusMessage && ( + + + {translate(wildcardStatusMessage)} + + + )} + {wildcardOptions.map((wildcard, index) => ( + + + + + + + {wildcard.value_count} + + + + {fixedWildcardPath === wildcard.path && ( + + {wildcardValuesResult.isFetching && ( + + {t('promptWorkbench.values.loading')} + + )} + {fixedWildcardValues?.map((value, index) => ( + + ))} + + )} + + ))} + + + {completionContext.query + ? t('promptWorkbench.autocomplete.matching', { query: completionContext.query }) + : t('promptWorkbench.autocomplete.localWildcards')} + + + )} + {hasPromptWorkbenchPanel && ( + + + + + + {t('promptWorkbench.panel.title')} + + {canShowWeightControls && ( + + + + + + + + + )} + + + {diagnostics.map((diagnostic) => + diagnostic.code === 'dynamic-active' && hasRandomWildcardOccurrences ? ( + + + + + {translate(diagnostic.label)} + + + + + + {t('promptWorkbench.behavior.randomImageShort')} + + + {t('promptWorkbench.behavior.randomInvokeShort')} + + + + ) : diagnostic.code === 'dynamic-active' ? ( + + + + ) : ( + + : undefined} + > + {translate(diagnostic.label)} + + + ) + )} + + + {hasPromptWorkbenchEntries && ( + + + + )} + + )} + + ); +}); + +PromptWorkbench.displayName = 'PromptWorkbench'; + +const getDiagnosticBadgeTone = (severity: PromptDiagnosticSeverity): PromptWorkbenchBadgeTone => { + switch (severity) { + case 'ok': + return 'neutral'; + case 'warning': + return 'warning'; + case 'error': + return 'error'; + case 'info': + return 'neutral'; + } +}; diff --git a/invokeai/frontend/web/src/features/promptWorkbench/PromptWorkbenchBadge.tsx b/invokeai/frontend/web/src/features/promptWorkbench/PromptWorkbenchBadge.tsx new file mode 100644 index 00000000000..f5d961ca598 --- /dev/null +++ b/invokeai/frontend/web/src/features/promptWorkbench/PromptWorkbenchBadge.tsx @@ -0,0 +1,54 @@ +import { Badge } from '@invoke-ai/ui-library'; +import type { PropsWithChildren, ReactElement } from 'react'; +import { forwardRef, memo } from 'react'; + +export type PromptWorkbenchBadgeTone = 'neutral' | 'warning' | 'error'; + +type PromptWorkbenchBadgeProps = PropsWithChildren<{ + tone?: PromptWorkbenchBadgeTone; + icon?: ReactElement; +}>; + +export const PromptWorkbenchBadge = memo( + forwardRef(({ children, icon, tone = 'neutral' }, ref) => { + const colors = getBadgeColors(tone); + + return ( + + {icon} + {children} + + ); + }) +); + +PromptWorkbenchBadge.displayName = 'PromptWorkbenchBadge'; + +const getBadgeColors = (tone: PromptWorkbenchBadgeTone): { color: string; borderColor: string } => { + switch (tone) { + case 'warning': + return { color: 'warning.300', borderColor: 'warning.700' }; + case 'error': + return { color: 'error.300', borderColor: 'error.700' }; + case 'neutral': + return { color: 'base.300', borderColor: 'base.700' }; + } +}; diff --git a/invokeai/frontend/web/src/features/promptWorkbench/diagnostics.test.ts b/invokeai/frontend/web/src/features/promptWorkbench/diagnostics.test.ts new file mode 100644 index 00000000000..3409df18a75 --- /dev/null +++ b/invokeai/frontend/web/src/features/promptWorkbench/diagnostics.test.ts @@ -0,0 +1,250 @@ +import { describe, expect, it } from 'vitest'; + +import { getPromptDiagnostics } from './diagnostics'; + +const wildcards = [ + { + token: '__camera/lens__', + path: 'camera/lens', + label: 'lens', + file_type: 'txt' as const, + value_count: 2, + samples: ['50mm', '85mm'], + }, +]; + +describe('prompt workbench diagnostics', () => { + it('hides supported attention weight status for SDXL', () => { + const diagnostics = getPromptDiagnostics({ + prompt: '(face:1.2)', + modelBase: 'sdxl', + wildcards, + wildcardIndexErrorCount: 0, + dynamicPromptCount: 1, + dynamicPromptMode: 'random', + }); + + expect(diagnostics.find((diagnostic) => diagnostic.code.startsWith('attention'))).toBeUndefined(); + }); + + it('warns when attention syntax is used with FLUX-like models', () => { + const diagnostics = getPromptDiagnostics({ + prompt: '(face:1.2)', + modelBase: 'flux', + wildcards, + wildcardIndexErrorCount: 0, + dynamicPromptCount: 1, + dynamicPromptMode: 'random', + }); + + expect(diagnostics.find((diagnostic) => diagnostic.code === 'attention-unsupported')).toMatchObject({ + label: { key: 'promptWorkbench.diagnostics.weightsLiteralLabel' }, + severity: 'warning', + }); + }); + + it('reports missing wildcard references', () => { + const diagnostics = getPromptDiagnostics({ + prompt: 'portrait with __missing__', + modelBase: 'sdxl', + wildcards, + wildcardIndexErrorCount: 0, + dynamicPromptCount: 1, + dynamicPromptMode: 'random', + }); + + expect(diagnostics.find((diagnostic) => diagnostic.code === 'wildcards-missing')).toMatchObject({ + label: { key: 'promptWorkbench.diagnostics.missingLabel', options: { count: 1 } }, + severity: 'error', + }); + }); + + it('reports unavailable wildcard index instead of false missing references', () => { + const diagnostics = getPromptDiagnostics({ + prompt: 'portrait with __camera/lens__', + modelBase: 'sdxl', + wildcards: [], + wildcardIndexUnavailable: true, + wildcardIndexErrorCount: 0, + dynamicPromptCount: 1, + dynamicPromptMode: 'random', + }); + + expect(diagnostics.find((diagnostic) => diagnostic.code === 'wildcards-unavailable')).toMatchObject({ + label: { key: 'promptWorkbench.diagnostics.wildcardErrorLabel' }, + severity: 'error', + }); + expect(diagnostics.find((diagnostic) => diagnostic.code === 'wildcards-missing')).toBeUndefined(); + }); + + it('hides generic available wildcard status when the prompt has no references', () => { + const diagnostics = getPromptDiagnostics({ + prompt: 'portrait', + modelBase: 'sdxl', + wildcards, + wildcardIndexErrorCount: 0, + dynamicPromptCount: 1, + dynamicPromptMode: 'random', + }); + + expect(diagnostics.find((diagnostic) => diagnostic.code === 'wildcards-available')).toBeUndefined(); + expect( + diagnostics.find((diagnostic) => diagnostic.label.key === 'promptWorkbench.diagnostics.wildcardsLabel') + ).toBeUndefined(); + }); + + it('reports referenced wildcards when present', () => { + const diagnostics = getPromptDiagnostics({ + prompt: 'portrait __camera/lens__', + modelBase: 'sdxl', + wildcards, + wildcardIndexErrorCount: 0, + dynamicPromptCount: 1, + dynamicPromptMode: 'random', + }); + + expect(diagnostics.find((diagnostic) => diagnostic.code === 'wildcards-found')).toMatchObject({ + label: { key: 'promptWorkbench.diagnostics.wildcardsFoundLabel', options: { count: 1 } }, + severity: 'ok', + }); + }); + + it('keeps wildcard index error warnings visible', () => { + const diagnostics = getPromptDiagnostics({ + prompt: 'portrait', + modelBase: 'sdxl', + wildcards, + wildcardIndexErrorCount: 2, + dynamicPromptCount: 1, + dynamicPromptMode: 'random', + }); + + expect(diagnostics.find((diagnostic) => diagnostic.code === 'wildcards-index-errors')).toMatchObject({ + label: { key: 'promptWorkbench.diagnostics.indexErrorsLabel', options: { count: 2 } }, + severity: 'warning', + }); + }); + + it('reports dynamic prompt count by mode', () => { + const diagnostics = getPromptDiagnostics({ + prompt: 'portrait __camera/lens__', + modelBase: 'sdxl', + wildcards, + wildcardIndexErrorCount: 0, + dynamicPromptCount: 1, + dynamicPromptMode: 'random', + }); + + expect(diagnostics.find((diagnostic) => diagnostic.code === 'dynamic-active')).toMatchObject({ + label: { key: 'promptWorkbench.behavior.randomImageShort' }, + severity: 'ok', + }); + }); + + it('reports random refresh per image', () => { + const diagnostics = getPromptDiagnostics({ + prompt: 'portrait __camera/lens__', + modelBase: 'sdxl', + wildcards, + wildcardIndexErrorCount: 0, + dynamicPromptCount: 1, + dynamicPromptMode: 'random', + dynamicPromptRandomRefreshMode: 'per_image', + }); + + expect(diagnostics.find((diagnostic) => diagnostic.code === 'dynamic-active')).toMatchObject({ + label: { key: 'promptWorkbench.behavior.randomImageShort' }, + severity: 'ok', + }); + }); + + it('reports random refresh per invoke', () => { + const diagnostics = getPromptDiagnostics({ + prompt: 'portrait __camera/lens__', + modelBase: 'sdxl', + wildcards, + wildcardIndexErrorCount: 0, + dynamicPromptCount: 1, + dynamicPromptMode: 'random', + dynamicPromptRandomRefreshMode: 'per_enqueue', + }); + + expect(diagnostics.find((diagnostic) => diagnostic.code === 'dynamic-active')).toMatchObject({ + label: { key: 'promptWorkbench.behavior.randomInvokeShort' }, + severity: 'ok', + }); + }); + + it('reports locked preview random refresh', () => { + const diagnostics = getPromptDiagnostics({ + prompt: 'portrait __camera/lens__', + modelBase: 'sdxl', + wildcards, + wildcardIndexErrorCount: 0, + dynamicPromptCount: 1, + dynamicPromptMode: 'random', + dynamicPromptRandomRefreshMode: 'manual', + }); + + expect(diagnostics.find((diagnostic) => diagnostic.code === 'dynamic-active')).toMatchObject({ + label: { key: 'promptWorkbench.behavior.randomPreviewLabel' }, + severity: 'ok', + }); + }); + + it('reports cyclic wildcard prompts as deterministic', () => { + const diagnostics = getPromptDiagnostics({ + prompt: 'portrait __@camera/lens__', + modelBase: 'sdxl', + wildcards, + wildcardIndexErrorCount: 0, + dynamicPromptCount: 1, + dynamicPromptMode: 'random', + dynamicPromptRandomRefreshMode: 'per_enqueue', + }); + + expect(diagnostics.find((diagnostic) => diagnostic.code === 'dynamic-active')).toMatchObject({ + label: { key: 'promptWorkbench.diagnostics.dynamicCycleLabel', options: { count: 1 } }, + severity: 'ok', + }); + }); + + it('reports mixed cyclic and random dynamic syntax', () => { + const diagnostics = getPromptDiagnostics({ + prompt: 'portrait __@camera/lens__ {warm|cool}', + modelBase: 'sdxl', + wildcards, + wildcardIndexErrorCount: 0, + dynamicPromptCount: 2, + dynamicPromptMode: 'random', + dynamicPromptRandomRefreshMode: 'per_enqueue', + }); + + expect(diagnostics.find((diagnostic) => diagnostic.code === 'dynamic-active')).toMatchObject({ + label: { key: 'promptWorkbench.diagnostics.dynamicMixedLabel' }, + severity: 'ok', + description: { key: 'promptWorkbench.diagnostics.dynamicMixedDesc' }, + }); + }); + + it('surfaces dynamic prompt parser errors in the tooltip description', () => { + const diagnostics = getPromptDiagnostics({ + prompt: 'portrait {broken', + modelBase: 'sdxl', + wildcards, + wildcardIndexErrorCount: 0, + dynamicPromptCount: 1, + dynamicPromptMode: 'random', + dynamicPromptError: 'Could not parse prompt', + }); + + expect(diagnostics.find((diagnostic) => diagnostic.code === 'dynamic-error')).toMatchObject({ + label: { key: 'promptWorkbench.diagnostics.dynamicErrorLabel' }, + severity: 'error', + description: { + key: 'promptWorkbench.diagnostics.dynamicErrorDesc', + options: { error: 'Could not parse prompt' }, + }, + }); + }); +}); diff --git a/invokeai/frontend/web/src/features/promptWorkbench/diagnostics.ts b/invokeai/frontend/web/src/features/promptWorkbench/diagnostics.ts new file mode 100644 index 00000000000..bcdd39e63c4 --- /dev/null +++ b/invokeai/frontend/web/src/features/promptWorkbench/diagnostics.ts @@ -0,0 +1,200 @@ +import type { DynamicPromptRandomRefreshMode } from 'features/dynamicPrompts/store/dynamicPromptsSlice'; +import { + getHasCyclicWildcardSyntax, + getHasDynamicPromptSyntax, + getHasMixedCyclicAndNonCyclicDynamicPromptSyntax, +} from 'features/dynamicPrompts/util/promptIntent'; +import type { BaseModelType } from 'features/nodes/types/common'; +import type { WildcardIndexItem } from 'services/api/endpoints/utilities'; + +import { type PromptWorkbenchTranslation, tx } from './i18n'; +import { getPromptModelCapabilities } from './modelCapabilities'; +import { getMissingWildcardReferences, getWildcardReferences } from './wildcards'; + +const NUMERIC_ATTENTION_PATTERN = /\([^()\r\n]+:[+-]?\d+(?:\.\d+)?\)/; +const COMPEL_ATTENTION_PATTERN = /\)(?:[+-]+|[+-]?\d+(?:\.\d+)?)/; +const BRACKET_ATTENTION_PATTERN = /\[[^\]\r\n]+\]/; + +export type PromptDiagnosticSeverity = 'ok' | 'info' | 'warning' | 'error'; + +export type PromptDiagnostic = { + code: string; + label: PromptWorkbenchTranslation; + severity: PromptDiagnosticSeverity; + description: PromptWorkbenchTranslation; +}; + +type GetPromptDiagnosticsArg = { + prompt: string; + modelBase: BaseModelType | null | undefined; + wildcards: WildcardIndexItem[]; + wildcardIndexUnavailable?: boolean; + wildcardIndexErrorCount: number; + dynamicPromptCount: number; + dynamicPromptMode: 'random' | 'combinatorial'; + dynamicPromptRandomRefreshMode?: DynamicPromptRandomRefreshMode; + dynamicPromptError?: string | null; +}; + +export const getPromptDiagnostics = ({ + prompt, + modelBase, + wildcards, + wildcardIndexUnavailable = false, + wildcardIndexErrorCount, + dynamicPromptCount, + dynamicPromptMode, + dynamicPromptRandomRefreshMode = 'per_image', + dynamicPromptError, +}: GetPromptDiagnosticsArg): PromptDiagnostic[] => { + const capabilities = getPromptModelCapabilities(modelBase); + const hasAttentionSyntax = getHasAttentionSyntax(prompt); + const wildcardReferences = getWildcardReferences(prompt); + const missingWildcards = wildcardIndexUnavailable ? [] : getMissingWildcardReferences(prompt, wildcards); + const diagnostics: PromptDiagnostic[] = []; + + if (hasAttentionSyntax && !capabilities.supportsAttentionWeights) { + diagnostics.push({ + code: 'attention-unsupported', + label: tx('promptWorkbench.diagnostics.weightsLiteralLabel'), + severity: 'warning', + description: tx('promptWorkbench.diagnostics.weightsLiteralDesc'), + }); + } + + if (wildcardIndexUnavailable) { + diagnostics.push({ + code: 'wildcards-unavailable', + label: tx('promptWorkbench.diagnostics.wildcardErrorLabel'), + severity: 'error', + description: tx('promptWorkbench.diagnostics.wildcardErrorDesc'), + }); + } else if (missingWildcards.length > 0) { + diagnostics.push({ + code: 'wildcards-missing', + label: tx('promptWorkbench.diagnostics.missingLabel', { count: missingWildcards.length }), + severity: 'error', + description: tx('promptWorkbench.diagnostics.missingDesc', { + count: missingWildcards.length, + wildcards: missingWildcards.join(', '), + }), + }); + } else if (wildcardReferences.length > 0) { + diagnostics.push({ + code: 'wildcards-found', + label: tx('promptWorkbench.diagnostics.wildcardsFoundLabel', { count: wildcardReferences.length }), + severity: 'ok', + description: tx('promptWorkbench.diagnostics.wildcardsFoundDesc'), + }); + } + + if (!wildcardIndexUnavailable && wildcardIndexErrorCount > 0) { + diagnostics.push({ + code: 'wildcards-index-errors', + label: tx('promptWorkbench.diagnostics.indexErrorsLabel', { count: wildcardIndexErrorCount }), + severity: 'warning', + description: tx('promptWorkbench.diagnostics.indexErrorsDesc'), + }); + } + + if (dynamicPromptError) { + diagnostics.push({ + code: 'dynamic-error', + label: tx('promptWorkbench.diagnostics.dynamicErrorLabel'), + severity: 'error', + description: tx('promptWorkbench.diagnostics.dynamicErrorDesc', { error: dynamicPromptError }), + }); + } else if (getHasDynamicPromptSyntax(prompt)) { + const count = Math.max(dynamicPromptCount, 1); + const isCombinatorial = dynamicPromptMode === 'combinatorial'; + const hasCyclicWildcard = getHasCyclicWildcardSyntax(prompt); + const hasMixedCyclicAndNonCyclicSyntax = getHasMixedCyclicAndNonCyclicDynamicPromptSyntax(prompt); + diagnostics.push({ + code: 'dynamic-active', + label: getDynamicPromptLabel({ + count, + isCombinatorial, + hasCyclicWildcard, + hasMixedCyclicAndNonCyclicSyntax, + randomRefreshMode: dynamicPromptRandomRefreshMode, + }), + severity: 'ok', + description: getDynamicPromptDescription({ + isCombinatorial, + hasCyclicWildcard, + hasMixedCyclicAndNonCyclicSyntax, + randomRefreshMode: dynamicPromptRandomRefreshMode, + }), + }); + } + + return diagnostics; +}; + +export const getHasAttentionSyntax = (prompt: string): boolean => + NUMERIC_ATTENTION_PATTERN.test(prompt) || + COMPEL_ATTENTION_PATTERN.test(prompt) || + BRACKET_ATTENTION_PATTERN.test(prompt); + +const getDynamicPromptLabel = (arg: { + count: number; + isCombinatorial: boolean; + hasCyclicWildcard: boolean; + hasMixedCyclicAndNonCyclicSyntax: boolean; + randomRefreshMode: DynamicPromptRandomRefreshMode; +}): PromptWorkbenchTranslation => { + const { count, isCombinatorial, hasCyclicWildcard, hasMixedCyclicAndNonCyclicSyntax, randomRefreshMode } = arg; + + if (isCombinatorial) { + return tx('promptWorkbench.diagnostics.dynamicAllLabel', { count }); + } + + if (hasMixedCyclicAndNonCyclicSyntax) { + return tx('promptWorkbench.diagnostics.dynamicMixedLabel'); + } + + if (hasCyclicWildcard) { + return tx('promptWorkbench.diagnostics.dynamicCycleLabel', { count }); + } + + if (randomRefreshMode === 'per_image') { + return tx('promptWorkbench.behavior.randomImageShort'); + } + + if (randomRefreshMode === 'per_enqueue') { + return tx('promptWorkbench.behavior.randomInvokeShort'); + } + + return tx('promptWorkbench.behavior.randomPreviewLabel'); +}; + +const getDynamicPromptDescription = (arg: { + isCombinatorial: boolean; + hasCyclicWildcard: boolean; + hasMixedCyclicAndNonCyclicSyntax: boolean; + randomRefreshMode: DynamicPromptRandomRefreshMode; +}): PromptWorkbenchTranslation => { + const { isCombinatorial, hasCyclicWildcard, hasMixedCyclicAndNonCyclicSyntax, randomRefreshMode } = arg; + + if (isCombinatorial) { + return tx('promptWorkbench.diagnostics.dynamicAllDesc'); + } + + if (hasMixedCyclicAndNonCyclicSyntax) { + return tx('promptWorkbench.diagnostics.dynamicMixedDesc'); + } + + if (hasCyclicWildcard) { + return tx('promptWorkbench.diagnostics.dynamicCycleDesc'); + } + + if (randomRefreshMode === 'per_image') { + return tx('promptWorkbench.diagnostics.dynamicRandomImageDesc'); + } + + if (randomRefreshMode === 'per_enqueue') { + return tx('promptWorkbench.diagnostics.dynamicRandomInvokeDesc'); + } + + return tx('promptWorkbench.diagnostics.dynamicRandomPreviewDesc'); +}; diff --git a/invokeai/frontend/web/src/features/promptWorkbench/i18n.ts b/invokeai/frontend/web/src/features/promptWorkbench/i18n.ts new file mode 100644 index 00000000000..dd4f2db4dd0 --- /dev/null +++ b/invokeai/frontend/web/src/features/promptWorkbench/i18n.ts @@ -0,0 +1,9 @@ +export type PromptWorkbenchTranslation = { + key: string; + options?: Record; +}; + +export const tx = (key: string, options?: Record): PromptWorkbenchTranslation => ({ + key, + options, +}); diff --git a/invokeai/frontend/web/src/features/promptWorkbench/keyboardNavigation.test.ts b/invokeai/frontend/web/src/features/promptWorkbench/keyboardNavigation.test.ts new file mode 100644 index 00000000000..3a213384ee2 --- /dev/null +++ b/invokeai/frontend/web/src/features/promptWorkbench/keyboardNavigation.test.ts @@ -0,0 +1,47 @@ +import { describe, expect, it } from 'vitest'; + +import { + clampNavigationIndex, + getNextNavigationIndex, + getPromptWorkbenchKeyboardIntent, +} from './keyboardNavigation'; + +describe('prompt workbench keyboard navigation', () => { + it('wraps arrow movement through autocomplete options', () => { + expect(getNextNavigationIndex({ currentIndex: 0, direction: 'next', itemCount: 3 })).toBe(1); + expect(getNextNavigationIndex({ currentIndex: 2, direction: 'next', itemCount: 3 })).toBe(0); + expect(getNextNavigationIndex({ currentIndex: 0, direction: 'previous', itemCount: 3 })).toBe(2); + }); + + it('clamps stale active indices when option counts change', () => { + expect(clampNavigationIndex(5, 3)).toBe(2); + expect(clampNavigationIndex(-1, 3)).toBe(0); + expect(clampNavigationIndex(2, 0)).toBe(0); + }); + + it('maps autocomplete enter and tab to default wildcard insertion', () => { + expect(getPromptWorkbenchKeyboardIntent({ key: 'Enter', shiftKey: false, target: 'autocomplete' })).toBe( + 'insert_wildcard' + ); + expect(getPromptWorkbenchKeyboardIntent({ key: 'Tab', shiftKey: false, target: 'autocomplete' })).toBe( + 'insert_wildcard' + ); + }); + + it('maps autocomplete shift-enter to fixed-value expansion', () => { + expect(getPromptWorkbenchKeyboardIntent({ key: 'Enter', shiftKey: true, target: 'autocomplete' })).toBe( + 'open_fixed_values' + ); + }); + + it('maps fixed-value enter to concrete value insertion', () => { + expect(getPromptWorkbenchKeyboardIntent({ key: 'Enter', shiftKey: false, target: 'fixed_values' })).toBe( + 'insert_fixed_value' + ); + }); + + it('maps escape to dismiss for both keyboard targets', () => { + expect(getPromptWorkbenchKeyboardIntent({ key: 'Escape', shiftKey: false, target: 'autocomplete' })).toBe('dismiss'); + expect(getPromptWorkbenchKeyboardIntent({ key: 'Escape', shiftKey: false, target: 'fixed_values' })).toBe('dismiss'); + }); +}); diff --git a/invokeai/frontend/web/src/features/promptWorkbench/keyboardNavigation.ts b/invokeai/frontend/web/src/features/promptWorkbench/keyboardNavigation.ts new file mode 100644 index 00000000000..eeab89dc18b --- /dev/null +++ b/invokeai/frontend/web/src/features/promptWorkbench/keyboardNavigation.ts @@ -0,0 +1,68 @@ +export type PromptWorkbenchKeyboardTarget = 'autocomplete' | 'fixed_values'; +export type PromptWorkbenchKeyboardIntent = + | 'dismiss' + | 'insert_fixed_value' + | 'insert_wildcard' + | 'next' + | 'open_fixed_values' + | 'previous'; + +export const getNextNavigationIndex = ({ + currentIndex, + direction, + itemCount, +}: { + currentIndex: number; + direction: 'next' | 'previous'; + itemCount: number; +}): number => { + if (itemCount <= 0) { + return 0; + } + + const delta = direction === 'next' ? 1 : -1; + return (currentIndex + delta + itemCount) % itemCount; +}; + +export const clampNavigationIndex = (index: number, itemCount: number): number => { + if (itemCount <= 0) { + return 0; + } + return Math.min(Math.max(index, 0), itemCount - 1); +}; + +export const getPromptWorkbenchKeyboardIntent = ({ + key, + shiftKey, + target, +}: { + key: string; + shiftKey: boolean; + target: PromptWorkbenchKeyboardTarget; +}): PromptWorkbenchKeyboardIntent | null => { + if (key === 'Escape') { + return 'dismiss'; + } + + if (key === 'ArrowDown') { + return 'next'; + } + + if (key === 'ArrowUp') { + return 'previous'; + } + + if (target === 'fixed_values' && key === 'Enter') { + return 'insert_fixed_value'; + } + + if (target === 'autocomplete' && key === 'Enter' && shiftKey) { + return 'open_fixed_values'; + } + + if (target === 'autocomplete' && (key === 'Enter' || key === 'Tab')) { + return 'insert_wildcard'; + } + + return null; +}; diff --git a/invokeai/frontend/web/src/features/promptWorkbench/modelCapabilities.ts b/invokeai/frontend/web/src/features/promptWorkbench/modelCapabilities.ts new file mode 100644 index 00000000000..c38c5c9dccf --- /dev/null +++ b/invokeai/frontend/web/src/features/promptWorkbench/modelCapabilities.ts @@ -0,0 +1,19 @@ +import type { BaseModelType } from 'features/nodes/types/common'; + +export type PromptModelCapabilities = { + supportsAttentionWeights: boolean; + attentionWeightsLabelKey: string; +}; + +const ATTENTION_WEIGHT_BASES = new Set(['sd-1', 'sd-2', 'sdxl']); + +export const getPromptModelCapabilities = (base: BaseModelType | null | undefined): PromptModelCapabilities => { + const supportsAttentionWeights = !!base && ATTENTION_WEIGHT_BASES.has(base); + + return { + supportsAttentionWeights, + attentionWeightsLabelKey: supportsAttentionWeights + ? 'promptWorkbench.weight.supportedDescription' + : 'promptWorkbench.weight.literalDescription', + }; +}; diff --git a/invokeai/frontend/web/src/features/promptWorkbench/occurrences.test.ts b/invokeai/frontend/web/src/features/promptWorkbench/occurrences.test.ts new file mode 100644 index 00000000000..d4e1580a3bd --- /dev/null +++ b/invokeai/frontend/web/src/features/promptWorkbench/occurrences.test.ts @@ -0,0 +1,284 @@ +import { describe, expect, it } from 'vitest'; + +import { + getPromptWeightOccurrences, + getPromptWildcardOccurrences, + getPromptWorkbenchOccurrences, + getWeightShortLabel, + getWildcardBehaviorActionIntent, + getWildcardBehaviorIconType, + getWildcardBehaviorLabel, + getWildcardBehaviorShortLabel, + removePromptRange, + replacePromptRange, +} from './occurrences'; + +const wildcards = [ + { + token: '__camera/lens__', + path: 'camera/lens', + label: 'lens', + file_type: 'txt' as const, + value_count: 4, + samples: ['35mm lens', '50mm lens'], + }, + { + token: '__lighting/studio__', + path: 'lighting/studio', + label: 'studio', + file_type: 'txt' as const, + value_count: 3, + samples: ['softbox'], + }, +]; + +describe('prompt workbench occurrences', () => { + it('parses random wildcard occurrences', () => { + const occurrences = getPromptWildcardOccurrences({ + prompt: 'portrait __camera/lens__', + wildcards, + wildcardIndexUnavailable: false, + dynamicPromptMode: 'random', + }); + + expect(occurrences).toHaveLength(1); + expect(occurrences[0]).toMatchObject({ + type: 'wildcard', + path: 'camera/lens', + behavior: 'random', + valueCount: 4, + }); + expect(getWildcardBehaviorLabel(occurrences[0]!, 'per_image')).toEqual({ + key: 'promptWorkbench.behavior.randomImageLabel', + options: undefined, + }); + expect(getWildcardBehaviorLabel(occurrences[0]!, 'per_enqueue')).toEqual({ + key: 'promptWorkbench.behavior.randomInvokeLabel', + options: undefined, + }); + expect(getWildcardBehaviorShortLabel(occurrences[0]!, 'per_image')).toEqual({ + key: 'promptWorkbench.behavior.randomImageShort', + options: undefined, + }); + expect(getWildcardBehaviorShortLabel(occurrences[0]!, 'per_enqueue')).toEqual({ + key: 'promptWorkbench.behavior.randomInvokeShort', + options: undefined, + }); + expect(getWildcardBehaviorShortLabel(occurrences[0]!, 'manual')).toEqual({ + key: 'promptWorkbench.behavior.previewShort', + options: undefined, + }); + expect(getWildcardBehaviorIconType(occurrences[0]!)).toBe('random'); + }); + + it('parses cyclic wildcard occurrences', () => { + const occurrence = getPromptWildcardOccurrences({ + prompt: 'portrait __@camera/lens__', + wildcards, + wildcardIndexUnavailable: false, + dynamicPromptMode: 'random', + })[0]!; + + expect(occurrence).toMatchObject({ + path: 'camera/lens', + behavior: 'cycle', + }); + expect(getWildcardBehaviorShortLabel(occurrence, 'per_image')).toEqual({ + key: 'promptWorkbench.behavior.cycleShort', + options: undefined, + }); + expect(getWildcardBehaviorIconType(occurrence)).toBe('cycle'); + }); + + it('marks unknown wildcards as missing', () => { + const occurrence = getPromptWildcardOccurrences({ + prompt: 'portrait __missing/path__', + wildcards, + wildcardIndexUnavailable: false, + dynamicPromptMode: 'random', + })[0]!; + + expect(occurrence).toMatchObject({ + path: 'missing/path', + behavior: 'missing', + valueCount: null, + }); + expect(getWildcardBehaviorShortLabel(occurrence, 'per_image')).toEqual({ + key: 'promptWorkbench.behavior.missingLabel', + options: undefined, + }); + expect(getWildcardBehaviorIconType(occurrence)).toBe('warning'); + }); + + it('replaces a specific duplicate wildcard occurrence', () => { + const prompt = '__camera/lens__, __camera/lens__'; + const occurrences = getPromptWildcardOccurrences({ + prompt, + wildcards, + wildcardIndexUnavailable: false, + dynamicPromptMode: 'random', + }); + + expect(replacePromptRange(prompt, occurrences[1]!.range, '__@camera/lens__').prompt).toBe( + '__camera/lens__, __@camera/lens__' + ); + }); + + it('removes wildcard tokens cleanly from comma-separated prompts', () => { + const prompt = 'portrait, __camera/lens__, studio'; + const occurrence = getPromptWildcardOccurrences({ + prompt, + wildcards, + wildcardIndexUnavailable: false, + dynamicPromptMode: 'random', + })[0]!; + + expect(removePromptRange(prompt, occurrence.range).prompt).toBe('portrait, studio'); + }); + + it('removes wildcard tokens cleanly when the comma follows the token', () => { + const prompt = 'portrait __camera/lens__, studio'; + const occurrence = getPromptWildcardOccurrences({ + prompt, + wildcards, + wildcardIndexUnavailable: false, + dynamicPromptMode: 'random', + })[0]!; + + expect(removePromptRange(prompt, occurrence.range).prompt).toBe('portrait, studio'); + }); + + it('detects weighted spans and unsupported model warnings', () => { + expect(getPromptWeightOccurrences({ prompt: '(face)++ cat+', supportsAttentionWeights: false })).toMatchObject([ + { type: 'weight', text: '(face)++', attention: '++', isSupported: false }, + { type: 'weight', text: 'cat+', attention: '+', isSupported: false }, + ]); + }); + + it('does not report hyphenated words as prompt weights', () => { + expect(getPromptWeightOccurrences({ prompt: 'dance-more, rose-', supportsAttentionWeights: true })).toMatchObject([ + { type: 'weight', text: 'rose-', attention: '-' }, + ]); + }); + + it('returns mixed prompt workbench occurrences in prompt order', () => { + expect( + getPromptWorkbenchOccurrences({ + prompt: '(face)++ __camera/lens__', + wildcards, + wildcardIndexUnavailable: false, + dynamicPromptMode: 'combinatorial', + supportsAttentionWeights: true, + }).map((occurrence) => occurrence.type) + ).toEqual(['weight', 'wildcard']); + }); + + it('returns compact wildcard labels for all-combinations prompts', () => { + const occurrence = getPromptWildcardOccurrences({ + prompt: 'portrait __camera/lens__', + wildcards, + wildcardIndexUnavailable: false, + dynamicPromptMode: 'combinatorial', + })[0]!; + + expect(occurrence.behavior).toBe('all'); + expect(getWildcardBehaviorShortLabel(occurrence, 'per_image')).toEqual({ + key: 'promptWorkbench.behavior.allShort', + options: undefined, + }); + expect(getWildcardBehaviorIconType(occurrence)).toBe('all'); + }); + + it('returns compact weight labels for supported and unsupported weights', () => { + const supported = getPromptWeightOccurrences({ prompt: '(face)++', supportsAttentionWeights: true }); + const unsupported = getPromptWeightOccurrences({ prompt: '(face)+', supportsAttentionWeights: false }); + + expect(getWeightShortLabel(supported[0]!)).toEqual({ + key: 'promptWorkbench.weight.valueLabel', + options: { value: '++' }, + }); + expect(getWeightShortLabel({ ...supported[0]!, attention: 1.2 })).toEqual({ + key: 'promptWorkbench.weight.valueLabel', + options: { value: '1.2' }, + }); + expect(getWeightShortLabel(unsupported[0]!)).toEqual({ + key: 'promptWorkbench.weight.literalShort', + options: undefined, + }); + }); + + it('maps wildcard behavior menu actions to prompt intent', () => { + expect(getWildcardBehaviorActionIntent('random', 'camera/lens')).toEqual({ + replacement: '__camera/lens__', + }); + expect(getWildcardBehaviorActionIntent('cycle', 'camera/lens')).toEqual({ + replacement: '__@camera/lens__', + }); + expect(getWildcardBehaviorActionIntent('fixed', 'camera/lens')).toEqual({ opensFixedValues: true }); + expect(getWildcardBehaviorActionIntent('remove', 'camera/lens')).toEqual({ removesPrompt: true }); + }); + + it('merges wrapper-weighted wildcards into one wildcard occurrence', () => { + const occurrences = getPromptWorkbenchOccurrences({ + prompt: '(__lighting/studio__)++, (face)+', + wildcards, + wildcardIndexUnavailable: false, + dynamicPromptMode: 'random', + supportsAttentionWeights: true, + }); + + expect(occurrences).toHaveLength(2); + expect(occurrences[0]).toMatchObject({ + type: 'wildcard', + path: 'lighting/studio', + weight: { + type: 'weight', + text: '(__lighting/studio__)++', + attention: '++', + isSupported: true, + }, + }); + expect(occurrences[1]).toMatchObject({ + type: 'weight', + text: '(face)+', + }); + }); + + it('marks wrapper-weighted wildcards as unsupported when the model does not support weights', () => { + const occurrences = getPromptWorkbenchOccurrences({ + prompt: '(__lighting/studio__)++', + wildcards, + wildcardIndexUnavailable: false, + dynamicPromptMode: 'random', + supportsAttentionWeights: false, + }); + + expect(occurrences[0]).toMatchObject({ + type: 'wildcard', + weight: { + isSupported: false, + }, + }); + }); + + it('keeps token and wrapper ranges separate for weighted wildcard edits', () => { + const prompt = '(__lighting/studio__)++, (face)+'; + const occurrence = getPromptWorkbenchOccurrences({ + prompt, + wildcards, + wildcardIndexUnavailable: false, + dynamicPromptMode: 'random', + supportsAttentionWeights: true, + })[0]; + + expect(occurrence).toMatchObject({ type: 'wildcard' }); + if (occurrence?.type !== 'wildcard') { + throw new Error('Expected wildcard occurrence'); + } + + expect(replacePromptRange(prompt, occurrence.range, '__@lighting/studio__').prompt).toBe( + '(__@lighting/studio__)++, (face)+' + ); + expect(removePromptRange(prompt, occurrence.weight!.range).prompt).toBe('(face)+'); + }); +}); diff --git a/invokeai/frontend/web/src/features/promptWorkbench/occurrences.ts b/invokeai/frontend/web/src/features/promptWorkbench/occurrences.ts new file mode 100644 index 00000000000..51f67dcbf14 --- /dev/null +++ b/invokeai/frontend/web/src/features/promptWorkbench/occurrences.ts @@ -0,0 +1,380 @@ +import { type ASTNode, type Attention, parseTokens, tokenize } from 'common/util/promptAST'; +import type { DynamicPromptRandomRefreshMode } from 'features/dynamicPrompts/store/dynamicPromptsSlice'; +import type { WildcardIndexItem } from 'services/api/endpoints/utilities'; + +import { type PromptWorkbenchTranslation, tx } from './i18n'; +import { getCyclicWildcardToken, normalizeWildcardReference } from './wildcards'; + +const WILDCARD_OCCURRENCE_REGEX = /__([^\r\n]+?)__/g; +const WRAPPED_WILDCARD_WEIGHT_SUFFIX_REGEX = /^[ \t]*\)(?:[+-]+|[+-]?\d+(?:\.\d+)?)$/; + +export type PromptRange = { + start: number; + end: number; +}; + +export type PromptWildcardBehavior = 'random' | 'cycle' | 'all' | 'missing' | 'unavailable'; +export type WildcardBehaviorIconType = 'random' | 'cycle' | 'all' | 'warning'; +export type WildcardBehaviorAction = 'random' | 'cycle' | 'fixed' | 'remove'; + +export type PromptWildcardOccurrence = { + id: string; + type: 'wildcard'; + token: string; + path: string; + rawReference: string; + range: PromptRange; + behavior: PromptWildcardBehavior; + wildcard: WildcardIndexItem | null; + valueCount: number | null; + weight: PromptWeightOccurrence | null; +}; + +export type PromptWeightOccurrence = { + id: string; + type: 'weight'; + text: string; + attention: Attention; + range: PromptRange; + isSupported: boolean; +}; + +export type PromptWorkbenchOccurrence = PromptWildcardOccurrence | PromptWeightOccurrence; + +type GetPromptWorkbenchOccurrencesArg = { + prompt: string; + wildcards: WildcardIndexItem[]; + wildcardIndexUnavailable: boolean; + dynamicPromptMode: 'random' | 'combinatorial'; + supportsAttentionWeights: boolean; +}; + +type PromptReplacementResult = { + prompt: string; + caret: number; +}; + +type WildcardBehaviorActionIntent = { + replacement?: string; + opensFixedValues?: boolean; + removesPrompt?: boolean; +}; + +export const getPromptWorkbenchOccurrences = ({ + prompt, + wildcards, + wildcardIndexUnavailable, + dynamicPromptMode, + supportsAttentionWeights, +}: GetPromptWorkbenchOccurrencesArg): PromptWorkbenchOccurrence[] => { + const wildcardOccurrences = getPromptWildcardOccurrences({ + prompt, + wildcards, + wildcardIndexUnavailable, + dynamicPromptMode, + }); + const weightOccurrences = getPromptWeightOccurrences({ prompt, supportsAttentionWeights }); + const attachedWeightIds = new Set(); + const weightedWildcardOccurrences = wildcardOccurrences.map((occurrence) => { + const weight = getWrappingWeightOccurrence(prompt, occurrence, weightOccurrences); + if (weight) { + attachedWeightIds.add(weight.id); + } + return { ...occurrence, weight }; + }); + + return [ + ...weightedWildcardOccurrences, + ...weightOccurrences.filter((occurrence) => !attachedWeightIds.has(occurrence.id)), + ].sort((a, b) => getOccurrenceSortRange(a).start - getOccurrenceSortRange(b).start || a.range.end - b.range.end); +}; + +export const getPromptWildcardOccurrences = ({ + prompt, + wildcards, + wildcardIndexUnavailable, + dynamicPromptMode, +}: Omit): PromptWildcardOccurrence[] => { + const occurrences: PromptWildcardOccurrence[] = []; + + for (const match of prompt.matchAll(WILDCARD_OCCURRENCE_REGEX)) { + const token = match[0]; + const rawReference = match[1] ?? ''; + const start = match.index ?? 0; + const end = start + token.length; + const path = normalizeWildcardReference(rawReference); + const matchingWildcards = findMatchingWildcards(path, wildcards); + const exactWildcard = wildcards.find((wildcard) => wildcard.path === path) ?? null; + const isCycle = rawReference.trim().startsWith('@'); + const isMissing = !wildcardIndexUnavailable && matchingWildcards.length === 0; + + occurrences.push({ + id: `wildcard:${start}:${end}`, + type: 'wildcard', + token, + path, + rawReference, + range: { start, end }, + behavior: getWildcardBehavior({ + dynamicPromptMode, + isCycle, + isMissing, + wildcardIndexUnavailable, + }), + wildcard: exactWildcard, + weight: null, + valueCount: + matchingWildcards.length > 0 + ? matchingWildcards.reduce((total, wildcard) => total + wildcard.value_count, 0) + : null, + }); + } + + return occurrences; +}; + +export const getWildcardBehaviorLabel = ( + occurrence: PromptWildcardOccurrence, + randomRefreshMode: DynamicPromptRandomRefreshMode +): PromptWorkbenchTranslation => { + switch (occurrence.behavior) { + case 'random': + if (randomRefreshMode === 'per_image') { + return tx('promptWorkbench.behavior.randomImageLabel'); + } + if (randomRefreshMode === 'per_enqueue') { + return tx('promptWorkbench.behavior.randomInvokeLabel'); + } + return tx('promptWorkbench.behavior.randomPreviewLabel'); + case 'cycle': + return tx('promptWorkbench.behavior.cycleLabel'); + case 'all': + return tx('promptWorkbench.behavior.allCombinationsLabel'); + case 'missing': + return tx('promptWorkbench.behavior.missingLabel'); + case 'unavailable': + return tx('promptWorkbench.behavior.unavailableLabel'); + } +}; + +export const getWildcardBehaviorShortLabel = ( + occurrence: PromptWildcardOccurrence, + randomRefreshMode: DynamicPromptRandomRefreshMode +): PromptWorkbenchTranslation => { + switch (occurrence.behavior) { + case 'random': + if (randomRefreshMode === 'per_image') { + return tx('promptWorkbench.behavior.randomImageShort'); + } + if (randomRefreshMode === 'per_enqueue') { + return tx('promptWorkbench.behavior.randomInvokeShort'); + } + return tx('promptWorkbench.behavior.previewShort'); + case 'cycle': + return tx('promptWorkbench.behavior.cycleShort'); + case 'all': + return tx('promptWorkbench.behavior.allShort'); + case 'missing': + return tx('promptWorkbench.behavior.missingLabel'); + case 'unavailable': + return tx('promptWorkbench.behavior.unavailableLabel'); + } +}; + +export const getWildcardBehaviorIconType = (occurrence: PromptWildcardOccurrence): WildcardBehaviorIconType => { + switch (occurrence.behavior) { + case 'random': + return 'random'; + case 'cycle': + return 'cycle'; + case 'all': + return 'all'; + case 'missing': + case 'unavailable': + return 'warning'; + } +}; + +export const getPromptWeightOccurrences = ({ + prompt, + supportsAttentionWeights, +}: { + prompt: string; + supportsAttentionWeights: boolean; +}): PromptWeightOccurrence[] => { + try { + const ast = parseTokens(tokenize(prompt)); + const occurrences: PromptWeightOccurrence[] = []; + collectPromptWeightOccurrences(ast, prompt, supportsAttentionWeights, occurrences); + return occurrences; + } catch { + return []; + } +}; + +export const getWeightBehaviorLabel = (occurrence: PromptWeightOccurrence): PromptWorkbenchTranslation => + occurrence.isSupported ? tx('promptWorkbench.weight.supportedLabel') : tx('promptWorkbench.weight.literalLabel'); + +export const getWeightShortLabel = (occurrence: PromptWeightOccurrence): PromptWorkbenchTranslation => + occurrence.isSupported + ? tx('promptWorkbench.weight.valueLabel', { value: String(occurrence.attention) }) + : tx('promptWorkbench.weight.literalShort'); + +export const getWildcardBehaviorActionIntent = ( + action: WildcardBehaviorAction, + wildcardPath: string +): WildcardBehaviorActionIntent => { + switch (action) { + case 'random': + return { replacement: `__${wildcardPath}__` }; + case 'cycle': + return { replacement: getCyclicWildcardToken(wildcardPath) }; + case 'fixed': + return { opensFixedValues: true }; + case 'remove': + return { removesPrompt: true }; + } +}; + +export const replacePromptRange = ( + prompt: string, + range: PromptRange, + replacement: string +): PromptReplacementResult => { + return { + prompt: `${prompt.slice(0, range.start)}${replacement}${prompt.slice(range.end)}`, + caret: range.start + replacement.length, + }; +}; + +export const removePromptRange = (prompt: string, range: PromptRange): PromptReplacementResult => { + const nextPrompt = cleanPromptAfterRemoval(`${prompt.slice(0, range.start)}${prompt.slice(range.end)}`); + return { + prompt: nextPrompt, + caret: Math.min(range.start, nextPrompt.length), + }; +}; + +const getWildcardBehavior = ({ + dynamicPromptMode, + isCycle, + isMissing, + wildcardIndexUnavailable, +}: { + dynamicPromptMode: 'random' | 'combinatorial'; + isCycle: boolean; + isMissing: boolean; + wildcardIndexUnavailable: boolean; +}): PromptWildcardBehavior => { + if (wildcardIndexUnavailable) { + return 'unavailable'; + } + + if (isMissing) { + return 'missing'; + } + + if (isCycle) { + return 'cycle'; + } + + if (dynamicPromptMode === 'combinatorial') { + return 'all'; + } + + return 'random'; +}; + +const findMatchingWildcards = (reference: string, wildcards: WildcardIndexItem[]): WildcardIndexItem[] => { + if (!reference.includes('*')) { + return wildcards.filter((wildcard) => wildcard.path === reference); + } + + const regex = new RegExp(`^${reference.split('*').map(escapeRegExp).join('.*')}$`); + return wildcards.filter((wildcard) => regex.test(wildcard.path)); +}; + +const getWrappingWeightOccurrence = ( + prompt: string, + wildcardOccurrence: PromptWildcardOccurrence, + weightOccurrences: PromptWeightOccurrence[] +): PromptWeightOccurrence | null => { + const candidates = weightOccurrences + .filter((weightOccurrence) => getDoesWeightCleanlyWrapWildcard(prompt, weightOccurrence, wildcardOccurrence)) + .sort((a, b) => a.range.end - a.range.start - (b.range.end - b.range.start)); + + return candidates[0] ?? null; +}; + +const getDoesWeightCleanlyWrapWildcard = ( + prompt: string, + weightOccurrence: PromptWeightOccurrence, + wildcardOccurrence: PromptWildcardOccurrence +): boolean => { + if ( + weightOccurrence.range.start >= wildcardOccurrence.range.start || + weightOccurrence.range.end <= wildcardOccurrence.range.end + ) { + return false; + } + + return ( + /^[ \t]*\([ \t]*$/.test(prompt.slice(weightOccurrence.range.start, wildcardOccurrence.range.start)) && + WRAPPED_WILDCARD_WEIGHT_SUFFIX_REGEX.test(prompt.slice(wildcardOccurrence.range.end, weightOccurrence.range.end)) + ); +}; + +const getOccurrenceSortRange = (occurrence: PromptWorkbenchOccurrence): PromptRange => + occurrence.type === 'wildcard' && occurrence.weight ? occurrence.weight.range : occurrence.range; + +const collectPromptWeightOccurrences = ( + nodes: ASTNode[], + prompt: string, + supportsAttentionWeights: boolean, + occurrences: PromptWeightOccurrence[] +) => { + for (const node of nodes) { + if (node.type === 'word' && node.attention !== undefined) { + occurrences.push({ + id: `weight:${node.range.start}:${node.range.end}`, + type: 'weight', + text: prompt.slice(node.range.start, node.range.end), + attention: node.attention, + range: node.range, + isSupported: supportsAttentionWeights, + }); + } + + if (node.type === 'group') { + if (node.attention !== undefined) { + occurrences.push({ + id: `weight:${node.range.start}:${node.range.end}`, + type: 'weight', + text: prompt.slice(node.range.start, node.range.end), + attention: node.attention, + range: node.range, + isSupported: supportsAttentionWeights, + }); + } + collectPromptWeightOccurrences(node.children, prompt, supportsAttentionWeights, occurrences); + } + + if (node.type === 'prompt_function') { + for (const arg of node.promptArgs) { + collectPromptWeightOccurrences(arg.nodes, prompt, supportsAttentionWeights, occurrences); + } + } + } +}; + +const cleanPromptAfterRemoval = (prompt: string): string => { + return prompt + .replace(/[ \t]*,[ \t]*,[ \t]*/g, ', ') + .replace(/[ \t]+,[ \t]*/g, ', ') + .replace(/^[ \t]*,[ \t]*/g, '') + .replace(/[ \t]*,[ \t]*$/g, '') + .replace(/[ \t]{2,}/g, ' ') + .trim(); +}; + +const escapeRegExp = (value: string): string => value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); diff --git a/invokeai/frontend/web/src/features/promptWorkbench/wildcards.test.ts b/invokeai/frontend/web/src/features/promptWorkbench/wildcards.test.ts new file mode 100644 index 00000000000..890a7244f21 --- /dev/null +++ b/invokeai/frontend/web/src/features/promptWorkbench/wildcards.test.ts @@ -0,0 +1,101 @@ +import { describe, expect, it } from 'vitest'; + +import { + applyWildcardCompletion, + filterWildcardOptions, + getCyclicWildcardToken, + getMissingWildcardReferences, + getWildcardAutocompleteStatusMessage, + getWildcardCompletionContext, + getWildcardDisplayPath, +} from './wildcards'; + +const wildcards = [ + { + token: '__camera/lens__', + path: 'camera/lens', + label: 'lens', + file_type: 'txt' as const, + value_count: 2, + samples: ['50mm', '85mm'], + }, + { + token: '__lighting/studio__', + path: 'lighting/studio', + label: 'studio', + file_type: 'yaml' as const, + value_count: 1, + samples: ['softbox'], + }, +]; + +describe('prompt workbench wildcards', () => { + it('detects wildcard autocomplete context after double underscores', () => { + expect(getWildcardCompletionContext('portrait __camera', 17)).toEqual({ + start: 9, + end: 17, + query: 'camera', + }); + }); + + it('does not detect autocomplete context after a closed wildcard', () => { + expect(getWildcardCompletionContext('portrait __camera/lens__', 24)).toBeNull(); + }); + + it('filters wildcard options by path and samples', () => { + expect(filterWildcardOptions(wildcards, 'soft').map((wildcard) => wildcard.path)).toEqual(['lighting/studio']); + expect(filterWildcardOptions(wildcards, 'lens').map((wildcard) => wildcard.path)).toEqual(['camera/lens']); + }); + + it('displays wildcard paths without syntax delimiters', () => { + expect(getWildcardDisplayPath(wildcards[0]!)).toBe('camera/lens'); + }); + + it('explains why autocomplete has no visible options', () => { + expect( + getWildcardAutocompleteStatusMessage({ + isLoading: false, + isUnavailable: true, + optionCount: 0, + query: 'camera', + wildcardCount: 0, + }) + ).toEqual({ + key: 'promptWorkbench.autocomplete.unavailable', + options: undefined, + }); + + expect( + getWildcardAutocompleteStatusMessage({ + isLoading: false, + isUnavailable: false, + optionCount: 0, + query: 'missing', + wildcardCount: wildcards.length, + }) + ).toEqual({ + key: 'promptWorkbench.autocomplete.noMatches', + options: { query: 'missing' }, + }); + }); + + it('reports missing wildcard references while allowing glob references', () => { + expect(getMissingWildcardReferences('__camera/*__ __missing__', wildcards)).toEqual(['missing']); + }); + + it('applies wildcard, cyclic wildcard, and fixed-value completions', () => { + const context = getWildcardCompletionContext('portrait __camera', 17); + + expect(context).not.toBeNull(); + expect(applyWildcardCompletion('portrait __camera', context!, '__camera/lens__')).toEqual({ + prompt: 'portrait __camera/lens__', + caret: 24, + }); + expect(applyWildcardCompletion('portrait __camera', context!, getCyclicWildcardToken('camera/lens')).prompt).toBe( + 'portrait __@camera/lens__' + ); + expect(applyWildcardCompletion('portrait __camera', context!, '85mm portrait lens').prompt).toBe( + 'portrait 85mm portrait lens' + ); + }); +}); diff --git a/invokeai/frontend/web/src/features/promptWorkbench/wildcards.ts b/invokeai/frontend/web/src/features/promptWorkbench/wildcards.ts new file mode 100644 index 00000000000..c894e28ed44 --- /dev/null +++ b/invokeai/frontend/web/src/features/promptWorkbench/wildcards.ts @@ -0,0 +1,148 @@ +import type { WildcardIndexItem } from 'services/api/endpoints/utilities'; + +import { type PromptWorkbenchTranslation, tx } from './i18n'; + +const WILDCARD_REFERENCE_REGEX = /__([^\r\n]+?)__/g; +const WILDCARD_COMPLETION_STOP_CHARS = new Set([',', ';', '{', '}', '[', ']', '(', ')']); + +export type WildcardCompletionContext = { + start: number; + end: number; + query: string; +}; + +export type WildcardCompletionResult = { + prompt: string; + caret: number; +}; + +type GetWildcardAutocompleteStatusMessageArg = { + isLoading: boolean; + isUnavailable: boolean; + optionCount: number; + query: string; + wildcardCount: number; +}; + +export const getWildcardCompletionContext = (prompt: string, caret: number): WildcardCompletionContext | null => { + const beforeCaret = prompt.slice(0, caret); + const delimiterCount = beforeCaret.match(/__/g)?.length ?? 0; + if (delimiterCount % 2 === 0) { + return null; + } + + const start = beforeCaret.lastIndexOf('__'); + + if (start < 0) { + return null; + } + + const query = beforeCaret.slice(start + 2); + if (query.includes('__') || [...query].some((char) => WILDCARD_COMPLETION_STOP_CHARS.has(char) || /\s/.test(char))) { + return null; + } + + return { start, end: caret, query }; +}; + +export const filterWildcardOptions = (wildcards: WildcardIndexItem[], query: string): WildcardIndexItem[] => { + const normalizedQuery = query.toLowerCase(); + return wildcards + .filter((wildcard) => { + if (!normalizedQuery) { + return true; + } + return ( + wildcard.path.toLowerCase().includes(normalizedQuery) || + wildcard.label.toLowerCase().includes(normalizedQuery) || + wildcard.samples.some((sample) => sample.toLowerCase().includes(normalizedQuery)) + ); + }) + .slice(0, 30); +}; + +export const applyWildcardCompletion = ( + prompt: string, + context: WildcardCompletionContext, + replacement: string +): WildcardCompletionResult => { + return { + prompt: `${prompt.slice(0, context.start)}${replacement}${prompt.slice(context.end)}`, + caret: context.start + replacement.length, + }; +}; + +export const getCyclicWildcardToken = (wildcardPath: string): string => `__@${wildcardPath}__`; + +export const getWildcardDisplayPath = (wildcard: Pick): string => wildcard.path; + +export const getWildcardAutocompleteStatusMessage = ({ + isLoading, + isUnavailable, + optionCount, + query, + wildcardCount, +}: GetWildcardAutocompleteStatusMessageArg): PromptWorkbenchTranslation | null => { + if (optionCount > 0) { + return null; + } + + if (isLoading) { + return tx('promptWorkbench.autocomplete.loading'); + } + + if (isUnavailable) { + return tx('promptWorkbench.autocomplete.unavailable'); + } + + if (wildcardCount === 0) { + return tx('promptWorkbench.autocomplete.empty'); + } + + if (query) { + return tx('promptWorkbench.autocomplete.noMatches', { query }); + } + + return tx('promptWorkbench.autocomplete.noWildcardMatches'); +}; + +export const getWildcardReferences = (prompt: string): string[] => { + const refs: string[] = []; + const seen = new Set(); + for (const match of prompt.matchAll(WILDCARD_REFERENCE_REGEX)) { + const ref = normalizeWildcardReference(match[1] ?? ''); + if (!ref || seen.has(ref)) { + continue; + } + seen.add(ref); + refs.push(ref); + } + return refs; +}; + +export const getMissingWildcardReferences = (prompt: string, wildcards: WildcardIndexItem[]): string[] => { + const available = new Set(wildcards.map((wildcard) => wildcard.path)); + return getWildcardReferences(prompt).filter((ref) => !wildcardReferenceExists(ref, available)); +}; + +export const normalizeWildcardReference = (reference: string): string => { + let path = reference.trim(); + if (path.startsWith('~') || path.startsWith('@')) { + path = path.slice(1); + } + if (path.includes('(')) { + path = path.split('(', 1)[0] ?? ''; + } + return path.replaceAll('\\', '/').replace(/^\/+|\/+$/g, ''); +}; + +const wildcardReferenceExists = (reference: string, available: Set): boolean => { + if (!reference.includes('*')) { + return available.has(reference); + } + + const regex = new RegExp(`^${reference.split('*').map(escapeRegExp).join('.*')}$`); + return [...available].some((path) => regex.test(path)); +}; + +const escapeRegExp = (value: string): string => value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); diff --git a/invokeai/frontend/web/src/features/queue/hooks/useEnqueueCanvas.ts b/invokeai/frontend/web/src/features/queue/hooks/useEnqueueCanvas.ts index 68e1e9a382e..0fb6997c144 100644 --- a/invokeai/frontend/web/src/features/queue/hooks/useEnqueueCanvas.ts +++ b/invokeai/frontend/web/src/features/queue/hooks/useEnqueueCanvas.ts @@ -7,6 +7,7 @@ import { withResult, withResultAsync } from 'common/util/result'; import { useCanvasManagerSafe } from 'features/controlLayers/contexts/CanvasManagerProviderGate'; import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager'; import { positivePromptAddedToHistory, selectPositivePrompt } from 'features/controlLayers/store/paramsSlice'; +import { refreshDynamicPromptsForEnqueue } from 'features/dynamicPrompts/util/refreshDynamicPromptsForEnqueue'; import type { BaseModelType } from 'features/nodes/types/common'; import { prepareLinearUIBatch } from 'features/nodes/util/graph/buildLinearBatchConfig'; import { buildAnimaGraph } from 'features/nodes/util/graph/generation/buildAnimaGraph'; @@ -23,16 +24,29 @@ import type { GraphBuilderArg } from 'features/nodes/util/graph/types'; import { UnsupportedGenerationModeError } from 'features/nodes/util/graph/types'; import { toast } from 'features/toast/toast'; import { useCallback } from 'react'; +import { useTranslation } from 'react-i18next'; import { serializeError } from 'serialize-error'; import { enqueueMutationFixedCacheKeyOptions, queueApi } from 'services/api/endpoints/queue'; import { assert, AssertionError } from 'tsafe'; const log = logger('generation'); -const enqueueCanvas = async (store: AppStore, canvasManager: CanvasManager, prepend: boolean) => { - const { dispatch, getState } = store; +const enqueueCanvas = async ( + store: AppStore, + canvasManager: CanvasManager, + prepend: boolean, + dynamicPromptsErrorTitle: string +) => { + const { dispatch } = store; - const state = getState(); + const state = await refreshDynamicPromptsForEnqueue(store); + if (!state) { + toast({ + status: 'error', + title: dynamicPromptsErrorTitle, + }); + return; + } const destination = selectCanvasDestination(state); @@ -133,6 +147,7 @@ const enqueueCanvas = async (store: AppStore, canvasManager: CanvasManager, prep }; export const useEnqueueCanvas = () => { + const { t } = useTranslation(); const store = useAppStore(); const canvasManager = useCanvasManagerSafe(); const enqueue = useCallback( @@ -141,9 +156,9 @@ export const useEnqueueCanvas = () => { log.error('Canvas manager is not available'); return; } - return enqueueCanvas(store, canvasManager, prepend); + return enqueueCanvas(store, canvasManager, prepend, t('dynamicPrompts.resolveFailed')); }, - [canvasManager, store] + [canvasManager, store, t] ); return enqueue; }; diff --git a/invokeai/frontend/web/src/features/queue/hooks/useEnqueueGenerate.ts b/invokeai/frontend/web/src/features/queue/hooks/useEnqueueGenerate.ts index 54b37e1b95e..41c912bdf7f 100644 --- a/invokeai/frontend/web/src/features/queue/hooks/useEnqueueGenerate.ts +++ b/invokeai/frontend/web/src/features/queue/hooks/useEnqueueGenerate.ts @@ -5,6 +5,7 @@ import { useAppStore } from 'app/store/storeHooks'; import { extractMessageFromAssertionError } from 'common/util/extractMessageFromAssertionError'; import { withResult, withResultAsync } from 'common/util/result'; import { positivePromptAddedToHistory, selectPositivePrompt } from 'features/controlLayers/store/paramsSlice'; +import { refreshDynamicPromptsForEnqueue } from 'features/dynamicPrompts/util/refreshDynamicPromptsForEnqueue'; import type { BaseModelType } from 'features/nodes/types/common'; import { prepareLinearUIBatch } from 'features/nodes/util/graph/buildLinearBatchConfig'; import { buildAnimaGraph } from 'features/nodes/util/graph/generation/buildAnimaGraph'; @@ -20,16 +21,24 @@ import type { GraphBuilderArg } from 'features/nodes/util/graph/types'; import { UnsupportedGenerationModeError } from 'features/nodes/util/graph/types'; import { toast } from 'features/toast/toast'; import { useCallback } from 'react'; +import { useTranslation } from 'react-i18next'; import { serializeError } from 'serialize-error'; import { enqueueMutationFixedCacheKeyOptions, queueApi } from 'services/api/endpoints/queue'; import { assert, AssertionError } from 'tsafe'; const log = logger('generation'); -const enqueueGenerate = async (store: AppStore, prepend: boolean) => { - const { dispatch, getState } = store; +const enqueueGenerate = async (store: AppStore, prepend: boolean, dynamicPromptsErrorTitle: string) => { + const { dispatch } = store; - const state = getState(); + const state = await refreshDynamicPromptsForEnqueue(store); + if (!state) { + toast({ + status: 'error', + title: dynamicPromptsErrorTitle, + }); + return; + } const model = state.params.model; if (!model) { @@ -126,12 +135,13 @@ const enqueueGenerate = async (store: AppStore, prepend: boolean) => { }; export const useEnqueueGenerate = () => { + const { t } = useTranslation(); const store = useAppStore(); const enqueue = useCallback( (prepend: boolean) => { - return enqueueGenerate(store, prepend); + return enqueueGenerate(store, prepend, t('dynamicPrompts.resolveFailed')); }, - [store] + [store, t] ); return enqueue; }; diff --git a/invokeai/frontend/web/src/features/queue/hooks/useEnqueueUpscaling.ts b/invokeai/frontend/web/src/features/queue/hooks/useEnqueueUpscaling.ts index 318c1fd5ff0..bfe7226d758 100644 --- a/invokeai/frontend/web/src/features/queue/hooks/useEnqueueUpscaling.ts +++ b/invokeai/frontend/web/src/features/queue/hooks/useEnqueueUpscaling.ts @@ -2,18 +2,28 @@ import { logger } from 'app/logging/logger'; import type { AppStore } from 'app/store/store'; import { useAppStore } from 'app/store/storeHooks'; import { positivePromptAddedToHistory, selectPositivePrompt } from 'features/controlLayers/store/paramsSlice'; +import { refreshDynamicPromptsForEnqueue } from 'features/dynamicPrompts/util/refreshDynamicPromptsForEnqueue'; import type { BaseModelType } from 'features/nodes/types/common'; import { prepareLinearUIBatch } from 'features/nodes/util/graph/buildLinearBatchConfig'; import { buildMultidiffusionUpscaleGraph } from 'features/nodes/util/graph/buildMultidiffusionUpscaleGraph'; +import { toast } from 'features/toast/toast'; import { useCallback } from 'react'; +import { useTranslation } from 'react-i18next'; import { enqueueMutationFixedCacheKeyOptions, queueApi } from 'services/api/endpoints/queue'; const log = logger('generation'); -const enqueueUpscaling = async (store: AppStore, prepend: boolean) => { - const { dispatch, getState } = store; +const enqueueUpscaling = async (store: AppStore, prepend: boolean, dynamicPromptsErrorTitle: string) => { + const { dispatch } = store; - const state = getState(); + const state = await refreshDynamicPromptsForEnqueue(store); + if (!state) { + toast({ + status: 'error', + title: dynamicPromptsErrorTitle, + }); + return; + } const model = state.params.model; if (!model) { @@ -47,12 +57,13 @@ const enqueueUpscaling = async (store: AppStore, prepend: boolean) => { }; export const useEnqueueUpscaling = () => { + const { t } = useTranslation(); const store = useAppStore(); const enqueue = useCallback( (prepend: boolean) => { - return enqueueUpscaling(store, prepend); + return enqueueUpscaling(store, prepend, t('dynamicPrompts.resolveFailed')); }, - [store] + [store, t] ); return enqueue; }; diff --git a/invokeai/frontend/web/src/features/queue/store/readiness.test.ts b/invokeai/frontend/web/src/features/queue/store/readiness.test.ts index 632006050e6..134456fa0ec 100644 --- a/invokeai/frontend/web/src/features/queue/store/readiness.test.ts +++ b/invokeai/frontend/web/src/features/queue/store/readiness.test.ts @@ -52,9 +52,12 @@ const kleinVaeModel = { key: 'vae', name: 'VAE', base: 'flux2', type: 'vae' }; const kleinQwen3Model = { key: 'qwen3', name: 'Qwen3', base: 'flux2', type: 'qwen3_encoder' }; const baseDynamicPrompts: DynamicPromptsState = { - _version: 1, - maxPrompts: 100, - combinatorial: false, + _version: 4, + mode: 'random', + randomSamples: 1, + maxCombinations: 100, + randomSeed: 0, + randomRefreshMode: 'per_image', prompts: ['test prompt'], parsingError: undefined, isError: false, diff --git a/invokeai/frontend/web/src/services/api/endpoints/utilities.ts b/invokeai/frontend/web/src/services/api/endpoints/utilities.ts index 44e1bedcc26..62c334901cd 100644 --- a/invokeai/frontend/web/src/services/api/endpoints/utilities.ts +++ b/invokeai/frontend/web/src/services/api/endpoints/utilities.ts @@ -35,8 +35,29 @@ type ImageToPromptResponse = { error?: string | null; }; +export type WildcardsResponse = + paths['/api/v1/utilities/wildcards']['get']['responses']['200']['content']['application/json']; +export type WildcardIndexItem = WildcardsResponse['wildcards'][number]; +export type WildcardValuesResponse = + paths['/api/v1/utilities/wildcards/values']['get']['responses']['200']['content']['application/json']; + export const utilitiesApi = api.injectEndpoints({ endpoints: (build) => ({ + wildcards: build.query({ + query: () => ({ + url: buildUtilitiesUrl('wildcards'), + method: 'GET', + }), + providesTags: ['FetchOnReconnect'], + }), + wildcardValues: build.query({ + query: ({ path, limit = 200 }) => ({ + url: buildUtilitiesUrl('wildcards/values'), + method: 'GET', + params: { path, limit }, + }), + providesTags: ['FetchOnReconnect'], + }), dynamicPrompts: build.query< paths['/api/v1/utilities/dynamicprompts']['post']['responses']['200']['content']['application/json'], paths['/api/v1/utilities/dynamicprompts']['post']['requestBody']['content']['application/json'] @@ -67,4 +88,5 @@ export const utilitiesApi = api.injectEndpoints({ }), }); -export const { useExpandPromptMutation, useImageToPromptMutation } = utilitiesApi; +export const { useExpandPromptMutation, useImageToPromptMutation, useLazyWildcardValuesQuery, useWildcardsQuery } = + utilitiesApi; diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 68f24a26ec1..0019a95f975 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -285,6 +285,46 @@ export type paths = { patch: operations["update_user_api_v1_auth_users__user_id__patch"]; trace?: never; }; + "/api/v1/utilities/wildcards": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * List Wildcards + * @description List local dynamic prompt wildcards from INVOKEAI_ROOT/wildcards. + */ + get: operations["list_wildcards"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/api/v1/utilities/wildcards/values": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * List Wildcard Values + * @description List values for a single local dynamic prompt wildcard. + */ + get: operations["get_wildcard_values"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/api/v1/utilities/dynamicprompts": { parameters: { query?: never; @@ -8916,6 +8956,10 @@ export type components = { prompts: string[]; /** Error */ error?: string | null; + /** Warnings */ + warnings?: string[]; + /** Missing Wildcards */ + missing_wildcards?: string[]; }; /** * Upscale (RealESRGAN) @@ -16289,14 +16333,14 @@ export type components = { * Convert Cache Dir * Format: path * @description Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions). - * @default models/.convert_cache + * @default models\.convert_cache */ convert_cache_dir?: string; /** * Download Cache Dir * Format: path * @description Path to the directory that contains dynamically downloaded models. - * @default models/.download_cache + * @default models\.download_cache */ download_cache_dir?: string; /** @@ -31694,6 +31738,63 @@ export type components = { */ cover_image_name?: string | null; }; + /** WildcardIndexError */ + WildcardIndexError: { + /** Path */ + path: string; + /** Message */ + message: string; + }; + /** WildcardIndexItem */ + WildcardIndexItem: { + /** Token */ + token: string; + /** Path */ + path: string; + /** Label */ + label: string; + /** + * File Type + * @enum {string} + */ + file_type: "txt" | "json" | "yaml"; + /** Value Count */ + value_count: number; + /** Samples */ + samples: string[]; + }; + /** WildcardValuesResponse */ + WildcardValuesResponse: { + /** Token */ + token: string; + /** Path */ + path: string; + /** Label */ + label: string; + /** + * File Type + * @enum {string} + */ + file_type: "txt" | "json" | "yaml"; + /** Value Count */ + value_count: number; + /** Values */ + values: string[]; + /** Truncated */ + truncated: boolean; + }; + /** WildcardsResponse */ + WildcardsResponse: { + /** + * Wildcard Dir + * @default wildcards + */ + wildcard_dir?: string; + /** Wildcards */ + wildcards: components["schemas"]["WildcardIndexItem"][]; + /** Errors */ + errors: components["schemas"]["WildcardIndexError"][]; + }; /** Workflow */ Workflow: { /** @@ -33167,6 +33268,60 @@ export interface operations { }; }; }; + list_wildcards: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["WildcardsResponse"]; + }; + }; + }; + }; + get_wildcard_values: { + parameters: { + query: { + /** @description The relative wildcard path to read values for */ + path: string; + /** @description The max number of wildcard values to return */ + limit?: number; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["WildcardValuesResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; parse_dynamicprompts: { parameters: { query?: never; diff --git a/tests/app/routers/test_multiuser_authorization.py b/tests/app/routers/test_multiuser_authorization.py index 3461f37e7e9..9a588c952ed 100644 --- a/tests/app/routers/test_multiuser_authorization.py +++ b/tests/app/routers/test_multiuser_authorization.py @@ -179,6 +179,7 @@ def enable_multiuser(monkeypatch: Any, mock_invoker: Invoker): monkeypatch.setattr("invokeai.app.api.routers.recall_parameters.ApiDependencies", mock_deps) monkeypatch.setattr("invokeai.app.api.routers.model_manager.ApiDependencies", mock_deps) monkeypatch.setattr("invokeai.app.api.routers.custom_nodes.ApiDependencies", mock_deps) + monkeypatch.setattr("invokeai.app.api.routers.utilities.ApiDependencies", mock_deps) yield @@ -1821,6 +1822,33 @@ def test_queue_cleared_still_broadcast(self, socketio: Any) -> None: assert "default" in rooms_emitted_to +class TestUtilitiesAuthorization: + """Tests that wildcard and dynamic prompt utility endpoints require auth in multiuser mode.""" + + def test_list_wildcards_requires_auth(self, enable_multiuser: Any, client: TestClient) -> None: + r = client.get("/api/v1/utilities/wildcards") + + assert r.status_code == status.HTTP_401_UNAUTHORIZED + + def test_list_wildcard_values_requires_auth(self, enable_multiuser: Any, client: TestClient) -> None: + r = client.get("/api/v1/utilities/wildcards/values", params={"path": "camera/lens"}) + + assert r.status_code == status.HTTP_401_UNAUTHORIZED + + def test_parse_dynamicprompts_requires_auth(self, enable_multiuser: Any, client: TestClient) -> None: + r = client.post( + "/api/v1/utilities/dynamicprompts", + json={"prompt": "__camera/lens__", "max_prompts": 1, "combinatorial": False}, + ) + + assert r.status_code == status.HTTP_401_UNAUTHORIZED + + def test_list_wildcards_allows_authenticated_user(self, client: TestClient, user1_token: str) -> None: + r = client.get("/api/v1/utilities/wildcards", headers=_auth(user1_token)) + + assert r.status_code == status.HTTP_200_OK + + class TestCustomNodesAuthorization: """Tests that custom_nodes endpoints enforce AdminUserOrDefault. diff --git a/tests/app/routers/test_utilities_wildcards.py b/tests/app/routers/test_utilities_wildcards.py new file mode 100644 index 00000000000..656ec8b20be --- /dev/null +++ b/tests/app/routers/test_utilities_wildcards.py @@ -0,0 +1,201 @@ +import asyncio +import os + +from invokeai.app.api.routers.utilities import parse_dynamicprompts +from invokeai.app.util.wildcards import ( + find_missing_wildcard_references, + get_wildcard_references, + get_wildcard_values, + index_wildcards, +) + + +class _MockConfiguration: + def __init__(self, root_path): + self.root_path = root_path + + +class _MockServices: + def __init__(self, root_path): + self.configuration = _MockConfiguration(root_path) + + +class _MockInvoker: + def __init__(self, root_path): + self.services = _MockServices(root_path) + + +class _MockApiDependencies: + def __init__(self, root_path): + self.invoker = _MockInvoker(root_path) + + +def test_index_wildcards_supports_txt_json_and_nested_yaml(tmp_path): + wildcards_dir = tmp_path / "wildcards" + (wildcards_dir / "camera").mkdir(parents=True) + (wildcards_dir / "camera" / "lens.txt").write_text( + "# comment\n\n\ufeff50mm\n\u00ef\u00bb\u00bf85mm\n", encoding="utf-8" + ) + (wildcards_dir / "styles.json").write_text( + '{"lighting": ["soft light", "rim light"], "ignored": 123}', encoding="utf-8" + ) + (wildcards_dir / "packs.yaml").write_text( + "portrait:\n mood:\n - calm\n - intense\n ignored: 123\n", encoding="utf-8" + ) + + result = index_wildcards(wildcards_dir) + + assert result.errors == [] + indexed = {wildcard.path: wildcard for wildcard in result.wildcards} + assert indexed["camera/lens"].token == "__camera/lens__" + assert indexed["camera/lens"].value_count == 2 + assert indexed["camera/lens"].samples == ["50mm", "85mm"] + assert indexed["lighting"].file_type == "json" + assert indexed["portrait/mood"].file_type == "yaml" + + +def test_get_wildcard_values_resolves_txt_json_and_nested_yaml(tmp_path): + wildcards_dir = tmp_path / "wildcards" + (wildcards_dir / "camera").mkdir(parents=True) + (wildcards_dir / "camera" / "lens.txt").write_text("# comment\n35mm\n50mm\n85mm\n", encoding="utf-8") + (wildcards_dir / "styles.json").write_text('{"lighting": ["soft light", "rim light"]}', encoding="utf-8") + (wildcards_dir / "packs.yaml").write_text("portrait:\n mood:\n - calm\n - intense\n", encoding="utf-8") + + txt_values = get_wildcard_values(wildcards_dir, "camera/lens") + json_values = get_wildcard_values(wildcards_dir, "lighting") + yaml_values = get_wildcard_values(wildcards_dir, "portrait/mood") + + assert txt_values is not None + assert txt_values.values == ["35mm", "50mm", "85mm"] + assert txt_values.path == "camera/lens" + assert json_values is not None + assert json_values.values == ["soft light", "rim light"] + assert yaml_values is not None + assert yaml_values.values == ["calm", "intense"] + + +def test_get_wildcard_values_skips_malformed_structured_files_for_unrelated_wildcards(tmp_path): + wildcards_dir = tmp_path / "wildcards" + wildcards_dir.mkdir() + (wildcards_dir / "bad.json").write_text('{"broken": [', encoding="utf-8") + (wildcards_dir / "valid.txt").write_text("one\ntwo\n", encoding="utf-8") + + values = get_wildcard_values(wildcards_dir, "valid") + + assert values is not None + assert values.values == ["one", "two"] + + +def test_get_wildcard_values_rejects_traversal_and_unknown_paths(tmp_path): + wildcards_dir = tmp_path / "wildcards" + wildcards_dir.mkdir() + (wildcards_dir / "safe.txt").write_text("safe\n", encoding="utf-8") + + assert get_wildcard_values(wildcards_dir, "../safe") is None + assert get_wildcard_values(wildcards_dir, "/safe") is None + assert get_wildcard_values(wildcards_dir, "missing") is None + + +def test_get_wildcard_values_respects_limit(tmp_path): + wildcards_dir = tmp_path / "wildcards" + wildcards_dir.mkdir() + (wildcards_dir / "many.txt").write_text("one\ntwo\nthree\n", encoding="utf-8") + + values = get_wildcard_values(wildcards_dir, "many", limit=2) + + assert values is not None + assert values.values == ["one", "two"] + assert values.value_count == 3 + assert values.truncated is True + + +def test_index_wildcards_reports_invalid_structured_files(tmp_path): + wildcards_dir = tmp_path / "wildcards" + wildcards_dir.mkdir() + (wildcards_dir / "bad.yaml").write_text("portrait: [unterminated", encoding="utf-8") + (wildcards_dir / "bad.json").write_text('{"portrait": [', encoding="utf-8") + (wildcards_dir / "good.txt").write_text("usable\n", encoding="utf-8") + + result = index_wildcards(wildcards_dir) + + assert [wildcard.path for wildcard in result.wildcards] == ["good"] + assert {error.path for error in result.errors} == {"bad.json", "bad.yaml"} + + +def test_index_wildcards_reports_duplicate_structured_paths(tmp_path): + wildcards_dir = tmp_path / "wildcards" + wildcards_dir.mkdir() + (wildcards_dir / "a.json").write_text('{"k": ["x"]}', encoding="utf-8") + (wildcards_dir / "b.json").write_text('{"k": ["y"]}', encoding="utf-8") + + result = index_wildcards(wildcards_dir) + + assert [wildcard.path for wildcard in result.wildcards] == ["k"] + assert len(result.errors) == 1 + assert result.errors[0].path == "b.json" + assert "Duplicate wildcard path 'k'" in result.errors[0].message + + +def test_index_wildcards_blocks_symlinks_outside_root(tmp_path): + wildcards_dir = tmp_path / "wildcards" + outside_dir = tmp_path / "outside" + wildcards_dir.mkdir() + outside_dir.mkdir() + outside_file = outside_dir / "outside.txt" + outside_file.write_text("outside\n", encoding="utf-8") + link = wildcards_dir / "outside.txt" + + try: + os.symlink(outside_file, link) + except (OSError, NotImplementedError): + return + + result = index_wildcards(wildcards_dir) + + assert result.wildcards == [] + assert len(result.errors) == 1 + assert result.errors[0].message in { + "Wildcard file is outside the wildcards folder", + "Symlinks are not supported", + } + + +def test_get_wildcard_references_normalizes_sampler_and_template_args(): + refs = get_wildcard_references("a __~camera/lens__ and __@lighting*__ with __style(name=test)__") + + assert refs == ["camera/lens", "lighting*", "style"] + + +def test_find_missing_wildcard_references_supports_globs(tmp_path): + wildcards_dir = tmp_path / "wildcards" + (wildcards_dir / "camera").mkdir(parents=True) + (wildcards_dir / "camera" / "lens.txt").write_text("50mm\n", encoding="utf-8") + indexed = index_wildcards(wildcards_dir).wildcards + + missing = find_missing_wildcard_references("a __camera/*__ and __missing__", indexed) + + assert missing == ["missing"] + + +def test_parse_dynamicprompts_random_mode_returns_one_prompt(tmp_path, monkeypatch): + wildcards_dir = tmp_path / "wildcards" + wildcards_dir.mkdir() + (wildcards_dir / "colors.txt").write_text("red\nblue\n", encoding="utf-8") + monkeypatch.setattr("invokeai.app.api.routers.utilities.ApiDependencies", _MockApiDependencies(tmp_path)) + + result = asyncio.run(parse_dynamicprompts(None, prompt="__colors__", max_prompts=1, combinatorial=False, seed=1)) + + assert len(result.prompts) == 1 + assert result.prompts[0] in {"red", "blue"} + + +def test_parse_dynamicprompts_combinatorial_mode_respects_max_prompts(tmp_path, monkeypatch): + wildcards_dir = tmp_path / "wildcards" + wildcards_dir.mkdir() + (wildcards_dir / "colors.txt").write_text("\ufeffred\nblue\ngreen\n", encoding="utf-8") + monkeypatch.setattr("invokeai.app.api.routers.utilities.ApiDependencies", _MockApiDependencies(tmp_path)) + + result = asyncio.run(parse_dynamicprompts(None, prompt="__colors__", max_prompts=2, combinatorial=True)) + + assert len(result.prompts) == 2 + assert set(result.prompts).issubset({"red", "blue", "green"})