feat: Port Schedulers to Mantine

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

@@ -1,6 +1,28 @@
-// TODO: use Enums?
+import { SelectItem } from '@mantine/core';
 
-export const SCHEDULERS = [
+// TODO: use Enums?
+export const SCHEDULERS: SelectItem[] = [
+  { label: 'euler', value: 'euler', group: 'Standard' },
+  { label: 'deis', value: 'deis', group: 'Standard' },
+  { label: 'ddim', value: 'ddim', group: 'Standard' },
+  { label: 'ddpm', value: 'ddpm', group: 'Standard' },
+  { label: 'dpmpp_2s', value: 'dpmpp_2s', group: 'Standard' },
+  { label: 'dpmpp_2m', value: 'dpmpp_2m', group: 'Standard' },
+  { label: 'heun', value: 'heun', group: 'Standard' },
+  { label: 'kdpm_2', value: 'kdpm_2', group: 'Standard' },
+  { label: 'lms', value: 'lms', group: 'Standard' },
+  { label: 'pndm', value: 'pndm', group: 'Standard' },
+  { label: 'unipc', value: 'unipc', group: 'Standard' },
+  { label: 'euler_k', value: 'euler_k', group: 'Karras' },
+  { label: 'dpmpp_2s_k', value: 'dpmpp_2s_k', group: 'Karras' },
+  { label: 'dpmpp_2m_k', value: 'dpmpp_2m_k', group: 'Karras' },
+  { label: 'heun_k', value: 'heun_k', group: 'Karras' },
+  { label: 'lms_k', value: 'lms_k', group: 'Karras' },
+  { label: 'euler_a', value: 'euler_a', group: 'Ancestral' },
+  { label: 'kdpm_2_a', value: 'kdpm_2_a', group: 'Ancestral' },
+];
+
+export const SCHEDULER_ITEMS = [
   'ddim',
   'lms',
   'lms_k',
@@ -19,9 +41,9 @@ export const SCHEDULERS = [
   'heun',
   'heun_k',
   'unipc',
-] as const;
+];
 
-export type Scheduler = (typeof SCHEDULERS)[number];
+export type Scheduler = typeof SCHEDULERS;
 
 // Valid upscaling levels
 export const UPSCALING_LEVELS: Array<{ label: string; value: string }> = [

invokeai/frontend/web/src/common/components/IAIMantineMultiSelect.tsx

@@ -0,0 +1,79 @@
+import { Tooltip } from '@chakra-ui/react';
+import { MultiSelect, MultiSelectProps } from '@mantine/core';
+import { memo } from 'react';
+
+type IAIMultiSelectProps = MultiSelectProps & {
+  tooltip?: string;
+};
+
+const IAIMantineMultiSelect = (props: IAIMultiSelectProps) => {
+  const { searchable = true, tooltip, ...rest } = props;
+  return (
+    <Tooltip label={tooltip} placement="top" hasArrow>
+      <MultiSelect
+        searchable={searchable}
+        styles={() => ({
+          label: {
+            color: 'var(--invokeai-colors-base-300)',
+            fontWeight: 'normal',
+          },
+          input: {
+            backgroundColor: 'var(--invokeai-colors-base-900)',
+            borderWidth: '2px',
+            borderColor: 'var(--invokeai-colors-base-800)',
+            color: 'var(--invokeai-colors-base-100)',
+            padding: 10,
+            paddingRight: 24,
+            fontWeight: 600,
+            '&:hover': { borderColor: 'var(--invokeai-colors-base-700)' },
+            '&:focus': {
+              borderColor: 'var(--invokeai-colors-accent-600)',
+            },
+          },
+          value: {
+            backgroundColor: 'var(--invokeai-colors-base-800)',
+            color: 'var(--invokeai-colors-base-100)',
+            '&:hover': {
+              backgroundColor: 'var(--invokeai-colors-base-700)',
+              cursor: 'pointer',
+            },
+          },
+          dropdown: {
+            backgroundColor: 'var(--invokeai-colors-base-800)',
+            borderColor: 'var(--invokeai-colors-base-700)',
+          },
+          item: {
+            backgroundColor: 'var(--invokeai-colors-base-800)',
+            color: 'var(--invokeai-colors-base-200)',
+            padding: 6,
+            '&[data-hovered]': {
+              color: 'var(--invokeai-colors-base-100)',
+              backgroundColor: 'var(--invokeai-colors-base-750)',
+            },
+            '&[data-active]': {
+              backgroundColor: 'var(--invokeai-colors-base-750)',
+              '&:hover': {
+                color: 'var(--invokeai-colors-base-100)',
+                backgroundColor: 'var(--invokeai-colors-base-750)',
+              },
+            },
+            '&[data-selected]': {
+              color: 'var(--invokeai-colors-base-50)',
+              backgroundColor: 'var(--invokeai-colors-accent-650)',
+              fontWeight: 600,
+              '&:hover': {
+                backgroundColor: 'var(--invokeai-colors-accent-600)',
+              },
+            },
+          },
+          rightSection: {
+            width: 24,
+          },
+        })}
+        {...rest}
+      />
+    </Tooltip>
+  );
+};
+
+export default memo(IAIMantineMultiSelect);

invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamScheduler.tsx

@@ -1,43 +1,44 @@
-import { createSelector } from '@reduxjs/toolkit';
+import { SCHEDULER_ITEMS } from 'app/constants';
-import { Scheduler } from 'app/constants';
+import { RootState } from 'app/store/store';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
+import IAIMantineSelect from 'common/components/IAIMantineSelect';
-import IAIMantineSelect, {
-  IAISelectDataType,
-} from 'common/components/IAIMantineSelect';
-import { generationSelector } from 'features/parameters/store/generationSelectors';
 import { setScheduler } from 'features/parameters/store/generationSlice';
-import { uiSelector } from 'features/ui/store/uiSelectors';
+import { setSelectedSchedulers } from 'features/ui/store/uiSlice';
-import { memo, useCallback } from 'react';
+import { memo, useCallback, useEffect } from 'react';
 import { useTranslation } from 'react-i18next';
 
-const selector = createSelector(
+const ParamScheduler = () => {
-  [uiSelector, generationSelector],
+  const scheduler = useAppSelector(
-  (ui, generation) => {
+    (state: RootState) => state.generation.scheduler
-    const allSchedulers: string[] = ui.schedulers
-      .slice()
-      .sort((a, b) => a.localeCompare(b));
-
-    return {
-      scheduler: generation.scheduler,
-      allSchedulers,
-    };
-  },
-  defaultSelectorOptions
   );
 
-const ParamScheduler = () => {
+  const selectedSchedulers = useAppSelector(
-  const { allSchedulers, scheduler } = useAppSelector(selector);
+    (state: RootState) => state.ui.selectedSchedulers
+  );
+
+  const activeSchedulers = useAppSelector(
+    (state: RootState) => state.ui.activeSchedulers
+  );
 
   const dispatch = useAppDispatch();
   const { t } = useTranslation();
 
+  useEffect(() => {
+    if (selectedSchedulers.length === 0)
+      dispatch(setSelectedSchedulers(SCHEDULER_ITEMS));
+
+    const schedulerFound = activeSchedulers.find(
+      (activeSchedulers) => activeSchedulers.label === scheduler
+    );
+    if (!schedulerFound) dispatch(setScheduler(activeSchedulers[0].value));
+  }, [dispatch, selectedSchedulers, scheduler, activeSchedulers]);
+
   const handleChange = useCallback(
     (v: string | null) => {
       if (!v) {
         return;
       }
-      dispatch(setScheduler(v as Scheduler));
+      dispatch(setScheduler(v));
     },
     [dispatch]
   );
@@ -46,7 +47,7 @@ const ParamScheduler = () => {
     <IAIMantineSelect
       label={t('parameters.scheduler')}
       value={scheduler}
-      data={allSchedulers}
+      data={activeSchedulers}
       onChange={handleChange}
     />
   );

invokeai/frontend/web/src/features/parameters/store/generationSlice.ts

@@ -1,10 +1,10 @@
 import type { PayloadAction } from '@reduxjs/toolkit';
 import { createSlice } from '@reduxjs/toolkit';
-import { clamp, sortBy } from 'lodash-es';
-import { receivedModels } from 'services/thunks/model';
-import { Scheduler } from 'app/constants';
-import { ImageDTO } from 'services/api';
 import { configChanged } from 'features/system/store/configSlice';
+import { clamp, sortBy } from 'lodash-es';
+import { ImageDTO } from 'services/api';
+import { imageUrlsReceived } from 'services/thunks/image';
+import { receivedModels } from 'services/thunks/model';
 import {
   CfgScaleParam,
   HeightParam,
@@ -17,7 +17,6 @@ import {
   StrengthParam,
   WidthParam,
 } from './parameterZodSchemas';
-import { imageUrlsReceived } from 'services/thunks/image';
 
 export interface GenerationState {
   cfgScale: CfgScaleParam;
@@ -133,7 +132,7 @@ export const generationSlice = createSlice({
     setWidth: (state, action: PayloadAction<number>) => {
       state.width = action.payload;
     },
-    setScheduler: (state, action: PayloadAction<Scheduler>) => {
+    setScheduler: (state, action: PayloadAction<string>) => {
       state.scheduler = action.payload;
     },
     setSeed: (state, action: PayloadAction<number>) => {

invokeai/frontend/web/src/features/parameters/store/parameterZodSchemas.ts

@@ -1,4 +1,4 @@
-import { NUMPY_RAND_MAX, SCHEDULERS } from 'app/constants';
+import { NUMPY_RAND_MAX } from 'app/constants';
 import { z } from 'zod';
 
 /**
@@ -73,7 +73,7 @@ export const isValidCfgScale = (val: unknown): val is CfgScaleParam =>
 /**
  * Zod schema for scheduler parameter
  */
-export const zScheduler = z.enum(SCHEDULERS);
+export const zScheduler = z.string();
 /**
  * Type alias for scheduler parameter, inferred from its zod schema
  */

invokeai/frontend/web/src/features/system/components/SettingsModal/SettingsSchedulers.tsx

@@ -1,47 +1,33 @@
-import {
-  Menu,
-  MenuButton,
-  MenuItemOption,
-  MenuList,
-  MenuOptionGroup,
-} from '@chakra-ui/react';
 import { SCHEDULERS } from 'app/constants';
-
 import { RootState } from 'app/store/store';
+
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import IAIButton from 'common/components/IAIButton';
+import IAIMantineMultiSelect from 'common/components/IAIMantineMultiSelect';
-import { setSchedulers } from 'features/ui/store/uiSlice';
+import { setSelectedSchedulers } from 'features/ui/store/uiSlice';
-import { isArray } from 'lodash-es';
 import { useTranslation } from 'react-i18next';
 
 export default function SettingsSchedulers() {
-  const schedulers = useAppSelector((state: RootState) => state.ui.schedulers);
-
   const dispatch = useAppDispatch();
+
+  const selectedSchedulers = useAppSelector(
+    (state: RootState) => state.ui.selectedSchedulers
+  );
+
   const { t } = useTranslation();
 
-  const schedulerSettingsHandler = (v: string | string[]) => {
+  const schedulerSettingsHandler = (v: string[]) => {
-    if (isArray(v)) dispatch(setSchedulers(v.sort()));
+    dispatch(setSelectedSchedulers(v));
   };
 
   return (
-    <Menu closeOnSelect={false}>
+    <IAIMantineMultiSelect
-      <MenuButton as={IAIButton}>
+      label={t('settings.availableSchedulers')}
-        {t('settings.availableSchedulers')}
+      value={selectedSchedulers}
-      </MenuButton>
+      data={SCHEDULERS}
-      <MenuList maxHeight={64} overflowY="scroll">
-        <MenuOptionGroup
-          value={schedulers}
-          type="checkbox"
       onChange={schedulerSettingsHandler}
-        >
+      clearable
-          {SCHEDULERS.map((scheduler) => (
+      searchable
-            <MenuItemOption key={scheduler} value={scheduler}>
+      maxSelectedValues={99}
-              {scheduler}
+    />
-            </MenuItemOption>
-          ))}
-        </MenuOptionGroup>
-      </MenuList>
-    </Menu>
   );
 }

invokeai/frontend/web/src/features/ui/store/uiSlice.ts

@@ -1,10 +1,11 @@
+import { SelectItem } from '@mantine/core';
 import type { PayloadAction } from '@reduxjs/toolkit';
 import { createSlice } from '@reduxjs/toolkit';
+import { SCHEDULERS } from 'app/constants';
+import { initialImageChanged } from 'features/parameters/store/generationSlice';
 import { setActiveTabReducer } from './extraReducers';
 import { InvokeTabName } from './tabMap';
 import { AddNewModelType, UIState } from './uiTypes';
-import { initialImageChanged } from 'features/parameters/store/generationSlice';
-import { SCHEDULERS } from 'app/constants';
 
 export const initialUIState: UIState = {
   activeTab: 0,
@@ -20,7 +21,8 @@ export const initialUIState: UIState = {
   shouldShowGallery: true,
   shouldHidePreview: false,
   shouldShowProgressInViewer: true,
-  schedulers: SCHEDULERS,
+  activeSchedulers: [],
+  selectedSchedulers: [],
 };
 
 export const uiSlice = createSlice({
@@ -94,9 +96,20 @@ export const uiSlice = createSlice({
     setShouldShowProgressInViewer: (state, action: PayloadAction<boolean>) => {
       state.shouldShowProgressInViewer = action.payload;
     },
-    setSchedulers: (state, action: PayloadAction<string[]>) => {
+    setSelectedSchedulers: (state, action: PayloadAction<string[]>) => {
-      state.schedulers = [];
+      const selectedSchedulerData: SelectItem[] = [];
-      state.schedulers = action.payload;
+
+      if (action.payload.length === 0) action.payload = [SCHEDULERS[0].value];
+
+      action.payload.forEach((item) => {
+        const schedulerData = SCHEDULERS.find(
+          (scheduler) => scheduler.value === item
+        );
+        if (schedulerData) selectedSchedulerData.push(schedulerData);
+      });
+
+      state.activeSchedulers = selectedSchedulerData;
+      state.selectedSchedulers = action.payload;
     },
   },
   extraReducers(builder) {
@@ -124,7 +137,7 @@ export const {
   toggleParametersPanel,
   toggleGalleryPanel,
   setShouldShowProgressInViewer,
-  setSchedulers,
+  setSelectedSchedulers,
 } = uiSlice.actions;
 
 export default uiSlice.reducer;

invokeai/frontend/web/src/features/ui/store/uiTypes.ts

@@ -1,3 +1,5 @@
+import { SelectItem } from '@mantine/core';
+
 export type AddNewModelType = 'ckpt' | 'diffusers' | null;
 
 export type Coordinates = {
@@ -26,5 +28,6 @@ export interface UIState {
   shouldPinGallery: boolean;
   shouldShowGallery: boolean;
   shouldShowProgressInViewer: boolean;
-  schedulers: string[];
+  activeSchedulers: SelectItem[];
+  selectedSchedulers: string[];
 }