feat: nodes denylist

invokeai/app/invocations/baseinvocation.py

@@ -28,6 +28,8 @@
 from pydantic.typing import NoArgAnyCallable
 import semver
 
+from invokeai.app.services.config.invokeai_config import InvokeAIAppConfig
+
 if TYPE_CHECKING:
     from ..services.invocation_services import InvocationServices
 
@@ -470,14 +472,32 @@ class BaseInvocation(ABC, BaseModel):
 
     @classmethod
     def get_all_subclasses(cls):
+        app_config = InvokeAIAppConfig.get_config()
+        app_config.parse_args()
         subclasses = []
         toprocess = [cls]
         while len(toprocess) > 0:
             next = toprocess.pop(0)
             next_subclasses = next.__subclasses__()
             subclasses.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
     def get_invocations(cls):

invokeai/app/services/config/invokeai_config.py

@@ -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", )
     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
     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')

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/receivedOpenAPISchema.ts

@@ -9,13 +9,16 @@ import { startAppListening } from '..';
 export const addReceivedOpenAPISchemaListener = () => {
   startAppListening({
     actionCreator: receivedOpenAPISchema.fulfilled,
-    effect: (action, { dispatch }) => {
+    effect: (action, { dispatch, getState }) => {
       const log = logger('system');
       const schemaJSON = action.payload;
 
       log.debug({ schemaJSON }, 'Received OpenAPI schema');
 
-      const nodeTemplates = parseSchema(schemaJSON);
+      const nodeTemplates = parseSchema(
+        schemaJSON,
+        getState().config.nodesDenylist
+      );
 
       log.debug(
         { nodeTemplates: parseify(nodeTemplates) },

invokeai/frontend/web/src/app/types/invokeai.ts

@@ -49,6 +49,7 @@ export type AppConfig = {
   disabledFeatures: AppFeature[];
   disabledSDFeatures: SDFeature[];
   canRestoreDeletedImagesFromBin: boolean;
+  nodesDenylist: string[];
   sd: {
     defaultModel?: string;
     disabledControlNetModels: string[];

invokeai/frontend/web/src/features/nodes/util/parseSchema.ts

@@ -60,11 +60,15 @@ const isNotInDenylist = (schema: InvocationSchemaObject) =>
   !invocationDenylist.includes(schema.properties.type.default);
 
 export const parseSchema = (
-  openAPI: OpenAPIV3.Document
+  openAPI: OpenAPIV3.Document,
+  nodesDenylistExtra: string[] = []
 ): Record<string, InvocationTemplate> => {
   const filteredSchemas = Object.values(openAPI.components?.schemas ?? {})
     .filter(isInvocationSchemaObject)
-    .filter(isNotInDenylist);
+    .filter(isNotInDenylist)
+    .filter(
+      (schema) => !nodesDenylistExtra.includes(schema.properties.type.default)
+    );
 
   const invocations = filteredSchemas.reduce<
     Record<string, InvocationTemplate>

invokeai/frontend/web/src/features/system/store/configSlice.ts

@@ -14,6 +14,7 @@ export const initialConfigState: AppConfig = {
     'perlinNoise',
     'noiseThreshold',
   ],
+  nodesDenylist: [],
   canRestoreDeletedImagesFromBin: true,
   sd: {
     disabledControlNetModels: [],