From e00fed5c46297d1b9499a8f1c258aff836e0a1f2 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Mon, 12 Jun 2023 16:54:01 +1000
Subject: [PATCH] feat(ui): support disabling controlnet models & processors

---
 .../frontend/web/src/app/types/invokeai.ts    |  2 ++
 .../parameters/ParamControlNetModel.tsx       | 17 ++++++++-----
 .../ParamControlNetProcessorSelect.tsx        | 24 +++++++++++++++++--
 .../src/features/system/store/configSlice.ts  |  2 ++
 4 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/invokeai/frontend/web/src/app/types/invokeai.ts b/invokeai/frontend/web/src/app/types/invokeai.ts
index 4931c498bf..b718c2f742 100644
--- a/invokeai/frontend/web/src/app/types/invokeai.ts
+++ b/invokeai/frontend/web/src/app/types/invokeai.ts
@@ -117,6 +117,8 @@ export type AppConfig = {
   canRestoreDeletedImagesFromBin: boolean;
   sd: {
     defaultModel?: string;
+    disabledControlNetModels: string[];
+    disabledControlNetProcessors: string[];
     iterations: {
       initial: number;
       min: number;
diff --git a/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetModel.tsx b/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetModel.tsx
index 6fd9cc26bd..cbdd727b42 100644
--- a/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetModel.tsx
+++ b/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetModel.tsx
@@ -1,4 +1,5 @@
-import { useAppDispatch } from 'app/store/storeHooks';
+import { createSelector } from '@reduxjs/toolkit';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import IAICustomSelect, {
   IAICustomSelectOption,
 } from 'common/components/IAICustomSelect';
@@ -9,6 +10,7 @@ import {
   ControlNetModelName,
 } from 'features/controlNet/store/constants';
 import { controlNetModelChanged } from 'features/controlNet/store/controlNetSlice';
+import { configSelector } from 'features/system/store/configSelectors';
 import { map } from 'lodash-es';
 import { ChangeEvent, memo, useCallback } from 'react';
 
@@ -17,10 +19,12 @@ type ParamControlNetModelProps = {
   model: ControlNetModelName;
 };
 
-const DATA = map(CONTROLNET_MODELS, (m) => ({
-  key: m.label,
-  value: m.type,
-}));
+const selector = createSelector(configSelector, (config) => {
+  return map(CONTROLNET_MODELS, (m) => ({
+    key: m.label,
+    value: m.type,
+  })).filter((d) => !config.sd.disabledControlNetModels.includes(d.value));
+});
 
 // const DATA: IAICustomSelectOption[] = map(CONTROLNET_MODELS, (m) => ({
 //   value: m.type,
@@ -30,6 +34,7 @@ const DATA = map(CONTROLNET_MODELS, (m) => ({
 
 const ParamControlNetModel = (props: ParamControlNetModelProps) => {
   const { controlNetId, model } = props;
+  const controlNetModels = useAppSelector(selector);
   const dispatch = useAppDispatch();
   const isReady = useIsReadyToInvoke();
 
@@ -55,7 +60,7 @@ const ParamControlNetModel = (props: ParamControlNetModelProps) => {
     <IAISelect
       tooltip={model}
       tooltipProps={{ placement: 'top', hasArrow: true }}
-      validValues={DATA}
+      validValues={controlNetModels}
       value={model}
       onChange={handleModelChanged}
       isDisabled={!isReady}
diff --git a/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetProcessorSelect.tsx b/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetProcessorSelect.tsx
index e35d0e5588..b48fa50188 100644
--- a/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetProcessorSelect.tsx
+++ b/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetProcessorSelect.tsx
@@ -7,11 +7,13 @@ import {
   ControlNetProcessorType,
 } from '../../store/types';
 import { controlNetProcessorTypeChanged } from '../../store/controlNetSlice';
-import { useAppDispatch } from 'app/store/storeHooks';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { CONTROLNET_PROCESSORS } from '../../store/constants';
 import { map } from 'lodash-es';
 import { useIsReadyToInvoke } from 'common/hooks/useIsReadyToInvoke';
 import IAISelect from 'common/components/IAISelect';
+import { createSelector } from '@reduxjs/toolkit';
+import { configSelector } from 'features/system/store/configSelectors';
 
 type ParamControlNetProcessorSelectProps = {
   controlNetId: string;
@@ -26,6 +28,22 @@ const CONTROLNET_PROCESSOR_TYPES = map(CONTROLNET_PROCESSORS, (p) => ({
   a.value === 'none' ? -1 : b.value === 'none' ? 1 : a.key.localeCompare(b.key)
 );
 
+const selector = createSelector(configSelector, (config) => {
+  return map(CONTROLNET_PROCESSORS, (p) => ({
+    value: p.type,
+    key: p.label,
+  }))
+    .sort((a, b) =>
+      // sort 'none' to the top
+      a.value === 'none'
+        ? -1
+        : b.value === 'none'
+        ? 1
+        : a.key.localeCompare(b.key)
+    )
+    .filter((d) => !config.sd.disabledControlNetProcessors.includes(d.value));
+});
+
 // const CONTROLNET_PROCESSOR_TYPES: IAICustomSelectOption[] = map(
 //   CONTROLNET_PROCESSORS,
 //   (p) => ({
@@ -48,6 +66,7 @@ const ParamControlNetProcessorSelect = (
   const { controlNetId, processorNode } = props;
   const dispatch = useAppDispatch();
   const isReady = useIsReadyToInvoke();
+  const controlNetProcessors = useAppSelector(selector);
 
   const handleProcessorTypeChanged = useCallback(
     (e: ChangeEvent<HTMLSelectElement>) => {
@@ -76,8 +95,9 @@ const ParamControlNetProcessorSelect = (
     <IAISelect
       label="Processor"
       value={processorNode.type ?? 'canny_image_processor'}
-      validValues={CONTROLNET_PROCESSOR_TYPES}
+      validValues={controlNetProcessors}
       onChange={handleProcessorTypeChanged}
+      isDisabled={!isReady}
     />
   );
   // return (
diff --git a/invokeai/frontend/web/src/features/system/store/configSlice.ts b/invokeai/frontend/web/src/features/system/store/configSlice.ts
index 5f4dd68959..02440894a2 100644
--- a/invokeai/frontend/web/src/features/system/store/configSlice.ts
+++ b/invokeai/frontend/web/src/features/system/store/configSlice.ts
@@ -10,6 +10,8 @@ export const initialConfigState: AppConfig = {
   disabledSDFeatures: [],
   canRestoreDeletedImagesFromBin: true,
   sd: {
+    disabledControlNetModels: [],
+    disabledControlNetProcessors: [],
     iterations: {
       initial: 1,
       min: 1,