mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
feat(backend): allow/deny nodes
Allow denying and explicitly allowing nodes. When a not-allowed node is used, a pydantic `ValidationError` will be raised. - When collecting all invocations, check against the allowlist and denylist first. When pydantic constructs any unions related to nodes, the denied nodes will be omitted - Add `allow_nodes` and `deny_nodes` to `InvokeAIAppConfig`. These are `Union[list[str], None]`, and may be populated with the `type` of invocations. - When `allow_nodes` is `None`, allow all nodes, else if it is `list[str]`, only allow nodes in the list - When `deny_nodes` is `None`, deny no nodes, else if it is `list[str]`, deny nodes in the list - `deny_nodes` overrides `allow_nodes`
This commit is contained in:
parent
dccf291f64
commit
dc771d9645
@ -28,6 +28,8 @@ from pydantic.fields import Undefined, ModelField
|
|||||||
from pydantic.typing import NoArgAnyCallable
|
from pydantic.typing import NoArgAnyCallable
|
||||||
import semver
|
import semver
|
||||||
|
|
||||||
|
from invokeai.app.services.config.invokeai_config import InvokeAIAppConfig
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from ..services.invocation_services import InvocationServices
|
from ..services.invocation_services import InvocationServices
|
||||||
|
|
||||||
@ -470,6 +472,8 @@ class BaseInvocation(ABC, BaseModel):
|
|||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def get_all_subclasses(cls):
|
def get_all_subclasses(cls):
|
||||||
|
app_config = InvokeAIAppConfig.get_config()
|
||||||
|
app_config.parse_args()
|
||||||
subclasses = []
|
subclasses = []
|
||||||
toprocess = [cls]
|
toprocess = [cls]
|
||||||
while len(toprocess) > 0:
|
while len(toprocess) > 0:
|
||||||
@ -477,7 +481,23 @@ class BaseInvocation(ABC, BaseModel):
|
|||||||
next_subclasses = next.__subclasses__()
|
next_subclasses = next.__subclasses__()
|
||||||
subclasses.extend(next_subclasses)
|
subclasses.extend(next_subclasses)
|
||||||
toprocess.extend(next_subclasses)
|
toprocess.extend(next_subclasses)
|
||||||
return subclasses
|
allowed_invocations = []
|
||||||
|
for sc in subclasses:
|
||||||
|
is_in_allowlist = (
|
||||||
|
sc.__fields__.get("type").default in app_config.allow_nodes
|
||||||
|
if isinstance(app_config.allow_nodes, list)
|
||||||
|
else True
|
||||||
|
)
|
||||||
|
|
||||||
|
is_in_denylist = (
|
||||||
|
sc.__fields__.get("type").default in app_config.deny_nodes
|
||||||
|
if isinstance(app_config.deny_nodes, list)
|
||||||
|
else False
|
||||||
|
)
|
||||||
|
|
||||||
|
if is_in_allowlist and not is_in_denylist:
|
||||||
|
allowed_invocations.append(sc)
|
||||||
|
return allowed_invocations
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def get_invocations(cls):
|
def get_invocations(cls):
|
||||||
|
@ -254,6 +254,10 @@ class InvokeAIAppConfig(InvokeAISettings):
|
|||||||
attention_slice_size: Literal[tuple(["auto", "balanced", "max", 1, 2, 3, 4, 5, 6, 7, 8])] = Field(default="auto", description='Slice size, valid when attention_type=="sliced"', category="Generation", )
|
attention_slice_size: Literal[tuple(["auto", "balanced", "max", 1, 2, 3, 4, 5, 6, 7, 8])] = Field(default="auto", description='Slice size, valid when attention_type=="sliced"', category="Generation", )
|
||||||
force_tiled_decode: bool = Field(default=False, description="Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty)", category="Generation",)
|
force_tiled_decode: bool = Field(default=False, description="Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty)", category="Generation",)
|
||||||
|
|
||||||
|
# NODES
|
||||||
|
allow_nodes : Optional[List[str]] = Field(default=None, description="List of nodes to allow. Omit to allow all.", category="Nodes")
|
||||||
|
deny_nodes : Optional[List[str]] = Field(default=None, description="List of nodes to deny. Omit to deny none.", category="Nodes")
|
||||||
|
|
||||||
# DEPRECATED FIELDS - STILL HERE IN ORDER TO OBTAN VALUES FROM PRE-3.1 CONFIG FILES
|
# DEPRECATED FIELDS - STILL HERE IN ORDER TO OBTAN VALUES FROM PRE-3.1 CONFIG FILES
|
||||||
always_use_cpu : bool = Field(default=False, description="If true, use the CPU for rendering even if a GPU is available.", category='Memory/Performance')
|
always_use_cpu : bool = Field(default=False, description="If true, use the CPU for rendering even if a GPU is available.", category='Memory/Performance')
|
||||||
free_gpu_mem : Optional[bool] = Field(default=None, description="If true, purge model from GPU after each generation.", category='Memory/Performance')
|
free_gpu_mem : Optional[bool] = Field(default=None, description="If true, purge model from GPU after each generation.", category='Memory/Performance')
|
||||||
|
Loading…
Reference in New Issue
Block a user