diff --git a/invokeai/app/invocations/compel.py b/invokeai/app/invocations/compel.py index 5494c3261f..49c62cff56 100644 --- a/invokeai/app/invocations/compel.py +++ b/invokeai/app/invocations/compel.py @@ -1,4 +1,3 @@ -import re from dataclasses import dataclass from typing import List, Optional, Union @@ -17,6 +16,7 @@ from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ( from ...backend.model_management.lora import ModelPatcher from ...backend.model_management.models import ModelNotFoundException, ModelType from ...backend.util.devices import torch_dtype +from ..util.ti_utils import extract_ti_triggers_from_prompt from .baseinvocation import ( BaseInvocation, BaseInvocationOutput, @@ -87,7 +87,7 @@ class CompelInvocation(BaseInvocation): # loras = [(context.services.model_manager.get_model(**lora.dict(exclude={"weight"})).context.model, lora.weight) for lora in self.clip.loras] ti_list = [] - for trigger in re.findall(r"<[a-zA-Z0-9., _-]+>", self.prompt): + for trigger in extract_ti_triggers_from_prompt(self.prompt): name = trigger[1:-1] try: ti_list.append( @@ -210,7 +210,7 @@ class SDXLPromptInvocationBase: # loras = [(context.services.model_manager.get_model(**lora.dict(exclude={"weight"})).context.model, lora.weight) for lora in self.clip.loras] ti_list = [] - for trigger in re.findall(r"<[a-zA-Z0-9., _-]+>", prompt): + for trigger in extract_ti_triggers_from_prompt(prompt): name = trigger[1:-1] try: ti_list.append( diff --git a/invokeai/app/invocations/onnx.py b/invokeai/app/invocations/onnx.py index 9eca5e083e..759cfde700 100644 --- a/invokeai/app/invocations/onnx.py +++ b/invokeai/app/invocations/onnx.py @@ -1,7 +1,6 @@ # Copyright (c) 2023 Borisov Sergey (https://github.com/StAlKeR7779) import inspect -import re # from contextlib import ExitStack from typing import List, Literal, Union @@ -21,6 +20,7 @@ from invokeai.backend import BaseModelType, ModelType, SubModelType from ...backend.model_management import ONNXModelPatcher from ...backend.stable_diffusion import PipelineIntermediateState from ...backend.util import choose_torch_device +from ..util.ti_utils import extract_ti_triggers_from_prompt from .baseinvocation import ( BaseInvocation, BaseInvocationOutput, @@ -78,7 +78,7 @@ class ONNXPromptInvocation(BaseInvocation): ] ti_list = [] - for trigger in re.findall(r"<[a-zA-Z0-9., _-]+>", self.prompt): + for trigger in extract_ti_triggers_from_prompt(self.prompt): name = trigger[1:-1] try: ti_list.append( diff --git a/invokeai/app/util/ti_utils.py b/invokeai/app/util/ti_utils.py new file mode 100644 index 0000000000..a66a832b42 --- /dev/null +++ b/invokeai/app/util/ti_utils.py @@ -0,0 +1,8 @@ +import re + + +def extract_ti_triggers_from_prompt(prompt: str) -> list[str]: + ti_triggers = [] + for trigger in re.findall(r"<[a-zA-Z0-9., _-]+>", prompt): + ti_triggers.append(trigger) + return ti_triggers