Add Heun Karras Scheduler

invokeai/backend/args.py

@@ -113,6 +113,7 @@
     "lms",
     "pndm",
     "heun",
+    "heun_k",
     "euler",
     "euler_k",
     "euler_a",

invokeai/backend/stable_diffusion/schedulers/schedulers.py

@@ -9,7 +9,8 @@
     deis=(DEISMultistepScheduler, dict()),
     lms=(LMSDiscreteScheduler, dict()),
     pndm=(PNDMScheduler, dict()),
-    heun=(HeunDiscreteScheduler, dict()),
+    heun=(HeunDiscreteScheduler, dict(use_karras_sigmas=False)),
+    heun_k=(HeunDiscreteScheduler, dict(use_karras_sigmas=True)),
     euler=(EulerDiscreteScheduler, dict(use_karras_sigmas=False)),
     euler_k=(EulerDiscreteScheduler, dict(use_karras_sigmas=True)),
     euler_a=(EulerAncestralDiscreteScheduler, dict()),

invokeai/backend/web/modules/parameters.py

@@ -9,6 +9,7 @@
     "lms",
     "pndm",
     "heun",
+    'heun_k',
     "euler",
     "euler_k",
     "euler_a",

invokeai/frontend/web/package.json

@@ -62,11 +62,13 @@
     "@dagrejs/graphlib": "^2.1.12",
     "@emotion/react": "^11.10.6",
     "@emotion/styled": "^11.10.6",
+    "@floating-ui/react-dom": "^2.0.0",
     "@fontsource/inter": "^4.5.15",
     "@reduxjs/toolkit": "^1.9.5",
     "@roarr/browser-log-writer": "^1.1.5",
     "chakra-ui-contextmenu": "^1.0.5",
     "dateformat": "^5.0.3",
+    "downshift": "^7.6.0",
     "formik": "^2.2.9",
     "framer-motion": "^10.12.4",
     "fuse.js": "^6.6.2",

invokeai/frontend/web/public/locales/en.json

@@ -540,7 +540,10 @@
         "consoleLogLevel": "Log Level",
         "shouldLogToConsole": "Console Logging",
         "developer": "Developer",
-        "general": "General"
+        "general": "General",
+        "generation": "Generation",
+        "ui": "User Interface",
+        "availableSchedulers": "Available Schedulers"
     },
     "toast": {
         "serverError": "Server Error",

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

@@ -1,29 +1,24 @@
 // TODO: use Enums?
 
-export const DIFFUSERS_SCHEDULERS: Array<string> = [
+export const SCHEDULERS: Array<string> = [
   'ddim',
-  'ddpm',
-  'deis',
   'lms',
-  'pndm',
-  'heun',
   'euler',
   'euler_k',
   'euler_a',
-  'kdpm_2',
-  'kdpm_2_a',
   'dpmpp_2s',
   'dpmpp_2m',
   'dpmpp_2m_k',
+  'kdpm_2',
+  'kdpm_2_a',
+  'deis',
+  'ddpm',
+  'pndm',
+  'heun',
+  'heun_k',
   'unipc',
 ];
 
-export const IMG2IMG_DIFFUSERS_SCHEDULERS = DIFFUSERS_SCHEDULERS.filter(
-  (scheduler) => {
-    return scheduler !== 'dpmpp_2s';
-  }
-);
-
 // Valid image widths
 export const WIDTHS: Array<number> = Array.from(Array(64)).map(
   (_x, i) => (i + 1) * 64

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

@@ -52,6 +52,7 @@ export type CommonGeneratedImageMetadata = {
     | 'lms'
     | 'pndm'
     | 'heun'
+    | 'heun_k'
     | 'euler'
     | 'euler_k'
     | 'euler_a'

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

@@ -0,0 +1,172 @@
+import { CheckIcon } from '@chakra-ui/icons';
+import {
+  Box,
+  Flex,
+  FlexProps,
+  FormControl,
+  FormControlProps,
+  FormLabel,
+  Grid,
+  GridItem,
+  List,
+  ListItem,
+  Select,
+  Text,
+  Tooltip,
+  TooltipProps,
+} from '@chakra-ui/react';
+import { autoUpdate, offset, shift, useFloating } from '@floating-ui/react-dom';
+import { useSelect } from 'downshift';
+import { OverlayScrollbarsComponent } from 'overlayscrollbars-react';
+
+import { memo } from 'react';
+
+type IAICustomSelectProps = {
+  label?: string;
+  items: string[];
+  selectedItem: string;
+  setSelectedItem: (v: string | null | undefined) => void;
+  withCheckIcon?: boolean;
+  formControlProps?: FormControlProps;
+  buttonProps?: FlexProps;
+  tooltip?: string;
+  tooltipProps?: Omit<TooltipProps, 'children'>;
+};
+
+const IAICustomSelect = (props: IAICustomSelectProps) => {
+  const {
+    label,
+    items,
+    setSelectedItem,
+    selectedItem,
+    withCheckIcon,
+    formControlProps,
+    tooltip,
+    buttonProps,
+    tooltipProps,
+  } = props;
+
+  const {
+    isOpen,
+    getToggleButtonProps,
+    getLabelProps,
+    getMenuProps,
+    highlightedIndex,
+    getItemProps,
+  } = useSelect({
+    items,
+    selectedItem,
+    onSelectedItemChange: ({ selectedItem: newSelectedItem }) =>
+      setSelectedItem(newSelectedItem),
+  });
+
+  const { refs, floatingStyles } = useFloating<HTMLButtonElement>({
+    whileElementsMounted: autoUpdate,
+    middleware: [offset(4), shift({ crossAxis: true, padding: 8 })],
+  });
+
+  return (
+    <FormControl sx={{ w: 'full' }} {...formControlProps}>
+      {label && (
+        <FormLabel
+          {...getLabelProps()}
+          onClick={() => {
+            refs.floating.current && refs.floating.current.focus();
+          }}
+        >
+          {label}
+        </FormLabel>
+      )}
+      <Tooltip label={tooltip} {...tooltipProps}>
+        <Select
+          {...getToggleButtonProps({ ref: refs.setReference })}
+          {...buttonProps}
+          as={Flex}
+          sx={{
+            alignItems: 'center',
+            userSelect: 'none',
+            cursor: 'pointer',
+          }}
+        >
+          <Text sx={{ fontSize: 'sm', fontWeight: 500, color: 'base.100' }}>
+            {selectedItem}
+          </Text>
+        </Select>
+      </Tooltip>
+      <Box {...getMenuProps()}>
+        {isOpen && (
+          <List
+            as={Flex}
+            ref={refs.setFloating}
+            sx={{
+              ...floatingStyles,
+              width: 'max-content',
+              top: 0,
+              left: 0,
+              flexDirection: 'column',
+              zIndex: 1,
+              bg: 'base.800',
+              borderRadius: 'base',
+              border: '1px',
+              borderColor: 'base.700',
+              shadow: 'dark-lg',
+              py: 2,
+              px: 0,
+              h: 'fit-content',
+              maxH: 64,
+            }}
+          >
+            <OverlayScrollbarsComponent>
+              {items.map((item, index) => (
+                <ListItem
+                  sx={{
+                    bg: highlightedIndex === index ? 'base.700' : undefined,
+                    py: 1,
+                    paddingInlineStart: 3,
+                    paddingInlineEnd: 6,
+                    cursor: 'pointer',
+                    transitionProperty: 'common',
+                    transitionDuration: '0.15s',
+                  }}
+                  key={`${item}${index}`}
+                  {...getItemProps({ item, index })}
+                >
+                  {withCheckIcon ? (
+                    <Grid gridTemplateColumns="1.25rem auto">
+                      <GridItem>
+                        {selectedItem === item && <CheckIcon boxSize={2} />}
+                      </GridItem>
+                      <GridItem>
+                        <Text
+                          sx={{
+                            fontSize: 'sm',
+                            color: 'base.100',
+                            fontWeight: 500,
+                          }}
+                        >
+                          {item}
+                        </Text>
+                      </GridItem>
+                    </Grid>
+                  ) : (
+                    <Text
+                      sx={{
+                        fontSize: 'sm',
+                        color: 'base.100',
+                        fontWeight: 500,
+                      }}
+                    >
+                      {item}
+                    </Text>
+                  )}
+                </ListItem>
+              ))}
+            </OverlayScrollbarsComponent>
+          </List>
+        )}
+      </Box>
+    </FormControl>
+  );
+};
+
+export default memo(IAICustomSelect);

invokeai/frontend/web/src/features/gallery/components/CurrentImageDisplay.tsx

@@ -4,10 +4,10 @@ import { useAppSelector } from 'app/store/storeHooks';
 import { systemSelector } from 'features/system/store/systemSelectors';
 import { isEqual } from 'lodash-es';
 
-import { MdPhoto } from 'react-icons/md';
 import { selectedImageSelector } from '../store/gallerySelectors';
 import CurrentImageButtons from './CurrentImageButtons';
 import CurrentImagePreview from './CurrentImagePreview';
+import { FaImage } from 'react-icons/fa';
 
 export const currentImageDisplaySelector = createSelector(
   [systemSelector, selectedImageSelector],
@@ -61,7 +61,7 @@ const CurrentImageDisplay = () => {
           </>
         ) : (
           <Icon
-            as={MdPhoto}
+            as={FaImage}
             sx={{
               boxSize: 24,
               color: 'base.500',

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

@@ -1,9 +1,6 @@
-import {
-  DIFFUSERS_SCHEDULERS,
-  IMG2IMG_DIFFUSERS_SCHEDULERS,
-} from 'app/constants';
 import { RootState } from 'app/store/store';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import IAICustomSelect from 'common/components/IAICustomSelect';
 import IAISelect from 'common/components/IAISelect';
 import { setSampler } from 'features/parameters/store/generationSlice';
 import { activeTabNameSelector } from 'features/ui/store/uiSelectors';
@@ -17,25 +14,36 @@ const ParamSampler = () => {
 
   const activeTabName = useAppSelector(activeTabNameSelector);
 
+  const schedulers = useAppSelector((state: RootState) => state.ui.schedulers);
+
+  const img2imgSchedulers = schedulers.filter((scheduler) => {
+    return !['dpmpp_2s'].includes(scheduler);
+  });
+
   const dispatch = useAppDispatch();
   const { t } = useTranslation();
 
   const handleChange = useCallback(
-    (e: ChangeEvent<HTMLSelectElement>) => dispatch(setSampler(e.target.value)),
+    (v: string | null | undefined) => {
+      if (!v) {
+        return;
+      }
+      dispatch(setSampler(v));
+    },
     [dispatch]
   );
 
   return (
-    <IAISelect
+    <IAICustomSelect
       label={t('parameters.sampler')}
-      value={sampler}
-      onChange={handleChange}
-      validValues={
-        activeTabName === 'img2img' || activeTabName == 'unifiedCanvas'
-          ? IMG2IMG_DIFFUSERS_SCHEDULERS
-          : DIFFUSERS_SCHEDULERS
+      selectedItem={sampler}
+      setSelectedItem={handleChange}
+      items={
+        ['img2img', 'unifiedCanvas'].includes(activeTabName)
+          ? img2imgSchedulers
+          : schedulers
       }
-      minWidth={36}
+      withCheckIcon
     />
   );
 };

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

@@ -0,0 +1,19 @@
+import { Box, Flex } from '@chakra-ui/react';
+import { memo } from 'react';
+import ParamSampler from './ParamSampler';
+import ModelSelect from 'features/system/components/ModelSelect';
+
+const ParamSchedulerAndModel = () => {
+  return (
+    <Flex gap={3} w="full">
+      <Box w="16rem">
+        <ParamSampler />
+      </Box>
+      <Box w="full">
+        <ModelSelect />
+      </Box>
+    </Flex>
+  );
+};
+
+export default memo(ParamSchedulerAndModel);

invokeai/frontend/web/src/features/parameters/components/Parameters/ImageToImage/InitialImagePreview.tsx

@@ -1,4 +1,4 @@
-import { Flex, Image } from '@chakra-ui/react';
+import { Flex, Icon, Image } from '@chakra-ui/react';
 import { createSelector } from '@reduxjs/toolkit';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import SelectImagePlaceholder from 'common/components/SelectImagePlaceholder';
@@ -13,6 +13,7 @@ import { generationSelector } from 'features/parameters/store/generationSelector
 import { initialImageSelected } from 'features/parameters/store/actions';
 import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
 import ImageFallbackSpinner from 'features/gallery/components/ImageFallbackSpinner';
+import { FaImage } from 'react-icons/fa';
 
 const selector = createSelector(
   [generationSelector],
@@ -83,7 +84,15 @@ const InitialImagePreview = () => {
           <ImageMetadataOverlay image={initialImage} />
         </>
       )}
-      {!initialImage?.url && <SelectImagePlaceholder />}
+      {!initialImage?.url && (
+        <Icon
+          as={FaImage}
+          sx={{
+            boxSize: 24,
+            color: 'base.500',
+          }}
+        />
+      )}
     </Flex>
   );
 };