feat(ui): genericizing picker

Author: psychedelicious

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

@@ -1,10 +1,10 @@
-import type { SystemStyleObject } from '@invoke-ai/ui-library';
+import type { InputProps, SystemStyleObject } from '@invoke-ai/ui-library';
 import { Box, Divider, Flex, Input, Text } from '@invoke-ai/ui-library';
 import { IAINoContentFallback } from 'common/components/IAIImageFallback';
 import ScrollableContent from 'common/components/OverlayScrollbars/ScrollableContent';
 import { useStateImperative } from 'common/hooks/useStateImperative';
+import { fixedForwardRef } from 'common/util/fixedForwardRef';
 import { typedMemo } from 'common/util/typedMemo';
-import { NavigateToModelManagerButton } from 'features/parameters/components/MainModel/NavigateToModelManagerButton';
 import type { ChangeEvent } from 'react';
 import {
   createContext,
@@ -37,35 +37,36 @@ export type ImperativeModelPickerHandle = {
   setSearchTerm: (searchTerm: string) => void;
 };
 
-const DefaultOptionComponent = typedMemo(({ id }: { id: string }) => {
-  return <Text fontWeight="bold">{id}</Text>;
+const DefaultOptionComponent = typedMemo(<T extends object>({ option }: { option: T }) => {
+  const { getOptionId } = usePickerContext();
+  return <Text fontWeight="bold">{getOptionId(option)}</Text>;
 });
 DefaultOptionComponent.displayName = 'DefaultOptionComponent';
 
-const DefaultGroupHeaderComponent = typedMemo(({ id }: { id: string }) => {
-  return <Text fontWeight="bold">{id}</Text>;
+const DefaultGroupHeaderComponent = typedMemo(<T extends object>({ group }: { group: Group<T> }) => {
+  return <Text fontWeight="bold">{group.id}</Text>;
 });
 DefaultGroupHeaderComponent.displayName = 'DefaultGroupHeaderComponent';
 
-const DefaultNoOptionsFallback = typedMemo(() => {
+const DefaultNoOptionsFallbackComponent = typedMemo(() => {
   const { t } = useTranslation();
   return (
     <Flex w="full" h="full" alignItems="center" justifyContent="center">
       <Text variant="subtext">{t('common.noOptions')}</Text>
     </Flex>
   );
 });
-DefaultNoOptionsFallback.displayName = 'DefaultNoOptionsFallback';
+DefaultNoOptionsFallbackComponent.displayName = 'DefaultNoOptionsFallbackComponent';
 
-const DefaultNoMatchesFallback = typedMemo(() => {
+const DefaultNoMatchesFallbackComponent = typedMemo(() => {
   const { t } = useTranslation();
   return (
     <Flex w="full" h="full" alignItems="center" justifyContent="center">
       <Text variant="subtext">{t('common.noMatches')}</Text>
     </Flex>
   );
 });
-DefaultNoMatchesFallback.displayName = 'DefaultNoMatchesFallback';
+DefaultNoMatchesFallbackComponent.displayName = 'DefaultNoMatchesFallbackComponent';
 
 export type PickerProps<T extends object> = {
   options: (T | Group<T>)[];
@@ -75,9 +76,10 @@ export type PickerProps<T extends object> = {
   selectedItem?: T;
   onSelect?: (option: T) => void;
   onClose?: () => void;
-  noOptionsFallback?: React.ReactNode;
-  noMatchesFallback?: React.ReactNode;
   handleRef?: React.Ref<ImperativeModelPickerHandle>;
+  SearchBarComponent?: ReturnType<typeof fixedForwardRef<HTMLInputElement, InputProps>>;
+  NoOptionsFallbackComponent?: React.ComponentType;
+  NoMatchesFallbackComponent?: React.ComponentType;
   OptionComponent?: React.ComponentType<{ option: T }>;
   GroupHeaderComponent?: React.ComponentType<{ group: Group<T> }>;
 };
@@ -88,10 +90,11 @@ type PickerContextState<T extends object> = {
   getIsDisabled?: (option: T) => boolean;
   setActiveOptionId: (id: string) => void;
   onSelectById: (id: string) => void;
-  noOptionsFallback?: React.ReactNode;
-  noMatchesFallback?: React.ReactNode;
-  OptionComponent?: React.ComponentType<{ option: T }>;
-  GroupHeaderComponent?: React.ComponentType<{ group: Group<T> }>;
+  SearchBarComponent: ReturnType<typeof fixedForwardRef<HTMLInputElement, InputProps>>;
+  NoOptionsFallbackComponent: React.ComponentType;
+  NoMatchesFallbackComponent: React.ComponentType;
+  OptionComponent: React.ComponentType<{ option: T }>;
+  GroupHeaderComponent: React.ComponentType<{ group: Group<T> }>;
 };
 
 /* eslint-disable-next-line @typescript-eslint/no-explicit-any */
@@ -180,13 +183,14 @@ export const Picker = typedMemo(<T extends object>(props: PickerProps<T>) => {
     handleRef,
     isMatch,
     getIsDisabled,
-    noMatchesFallback = <DefaultNoMatchesFallback />,
-    noOptionsFallback = <DefaultNoOptionsFallback />,
     onClose,
     onSelect,
     selectedItem,
-    OptionComponent,
-    GroupHeaderComponent,
+    SearchBarComponent = DefaultPickerSearchBarComponent,
+    NoMatchesFallbackComponent = DefaultNoMatchesFallbackComponent,
+    NoOptionsFallbackComponent = DefaultNoOptionsFallbackComponent,
+    OptionComponent = DefaultOptionComponent,
+    GroupHeaderComponent = DefaultGroupHeaderComponent,
   } = props;
   const [activeOptionId, setActiveOptionId, getActiveOptionId] = useStateImperative(() =>
     getFirstOptionId(options, getOptionId)
@@ -346,22 +350,24 @@ export const Picker = typedMemo(<T extends object>(props: PickerProps<T>) => {
         isMatch,
         getIsDisabled,
         onSelectById,
-        noOptionsFallback,
-        noMatchesFallback,
+        setActiveOptionId,
+        SearchBarComponent,
+        NoOptionsFallbackComponent,
+        NoMatchesFallbackComponent,
         OptionComponent,
         GroupHeaderComponent,
-        setActiveOptionId,
       }) satisfies PickerContextState<T>,
     [
-      GroupHeaderComponent,
-      OptionComponent,
-      getIsDisabled,
       getOptionId,
       isMatch,
-      noMatchesFallback,
-      noOptionsFallback,
+      getIsDisabled,
       onSelectById,
       setActiveOptionId,
+      SearchBarComponent,
+      NoOptionsFallbackComponent,
+      NoMatchesFallbackComponent,
+      OptionComponent,
+      GroupHeaderComponent,
     ]
   );
 
@@ -378,14 +384,11 @@ export const Picker = typedMemo(<T extends object>(props: PickerProps<T>) => {
         gap={2}
         onKeyDown={onKeyDown}
       >
-        <Flex gap={2} alignItems="center">
-          <Input ref={inputRef} value={searchTerm} onChange={onChangeSearchTerm} placeholder="Filter" />
-          <NavigateToModelManagerButton />
-        </Flex>
+        <SearchBarComponent ref={inputRef} value={searchTerm} onChange={onChangeSearchTerm} />
         <Divider />
         <Flex tabIndex={-1} w="full" flexGrow={1}>
-          {flattenedOptions.length === 0 && noOptionsFallback}
-          {flattenedOptions.length > 0 && flattenedFilteredOptions.length === 0 && noMatchesFallback}
+          {flattenedOptions.length === 0 && <NoOptionsFallbackComponent />}
+          {flattenedOptions.length > 0 && flattenedFilteredOptions.length === 0 && <NoMatchesFallbackComponent />}
           {flattenedOptions.length > 0 && flattenedFilteredOptions.length > 0 && (
             <ScrollableContent>
               <PickerList
@@ -402,6 +405,13 @@ export const Picker = typedMemo(<T extends object>(props: PickerProps<T>) => {
 });
 Picker.displayName = 'Picker';
 
+const DefaultPickerSearchBarComponent = typedMemo(
+  fixedForwardRef<HTMLInputElement, InputProps>((props, ref) => {
+    return <Input placeholder="Search" ref={ref} {...props} />;
+  })
+);
+DefaultPickerSearchBarComponent.displayName = 'DefaultPickerSearchBarComponent';
+
 const PickerList = typedMemo(
   <T extends object>({
     items,
@@ -473,7 +483,7 @@ const PickerOptionGroup = typedMemo(
 
     return (
       <Flex flexDir="column" gap={2} w="full">
-        {GroupHeaderComponent ? <GroupHeaderComponent group={group} /> : <DefaultGroupHeaderComponent id={group.id} />}
+        <GroupHeaderComponent group={group} />
         <Flex flexDir="column" gap={1} w="full">
           {group.options.map((item) => {
             const id = getOptionId(item);
@@ -535,7 +545,7 @@ const PickerOption = typedMemo(
         onPointerMove={isDisabled ? undefined : onPointerMove}
         onClick={isDisabled ? undefined : onClick}
       >
-        {OptionComponent ? <OptionComponent option={option} /> : <DefaultOptionComponent id={id} />}
+        <OptionComponent option={option} />
       </Box>
     );
   }

invokeai/frontend/web/src/common/util/fixedForwardRef.ts

@@ -0,0 +1,13 @@
+import { forwardRef } from 'react';
+
+/**
+ * A forwardRef that works with generics and doesn't require the use of `as` to cast the type.
+ * See: https://www.totaltypescript.com/forwardref-with-generic-components
+ */
+export function fixedForwardRef<T, P = object>(
+  render: (props: P, ref: React.Ref<T>) => React.ReactNode
+): (props: P & React.RefAttributes<T>) => React.ReactNode {
+  // @ts-expect-error: This is a workaround for forwardRef's crappy typing
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  return forwardRef(render) as any;
+}

invokeai/frontend/web/src/features/settingsAccordions/components/GenerationSettingsAccordion/GenerationSettingsAccordion.tsx

@@ -1,11 +1,12 @@
-import type { FormLabelProps } from '@invoke-ai/ui-library';
+import type { FormLabelProps, InputProps } from '@invoke-ai/ui-library';
 import {
   Box,
   Button,
   Expander,
   Flex,
   FormControlGroup,
   FormLabel,
+  Input,
   Popover,
   PopoverArrow,
   PopoverBody,
@@ -23,6 +24,7 @@ import { InformationalPopover } from 'common/components/InformationalPopover/Inf
 import type { Group, ImperativeModelPickerHandle } from 'common/components/Picker/Picker';
 import { getRegex, Picker } from 'common/components/Picker/Picker';
 import { useDisclosure } from 'common/hooks/useBoolean';
+import { fixedForwardRef } from 'common/util/fixedForwardRef';
 import { typedMemo } from 'common/util/typedMemo';
 import { selectLoRAsSlice } from 'features/controlLayers/store/lorasSlice';
 import { selectIsCogView4, selectIsFLUX, selectIsSD3 } from 'features/controlLayers/store/paramsSlice';
@@ -222,6 +224,7 @@ const MainModelPicker = memo(() => {
               isMatch={isMatch}
               OptionComponent={PickerItemComponent}
               GroupHeaderComponent={PickerGroupHeaderComponent}
+              SearchBarComponent={SearchBarComponent}
             />
           </PopoverBody>
         </PopoverContent>
@@ -231,6 +234,19 @@ const MainModelPicker = memo(() => {
 });
 MainModelPicker.displayName = 'MainModelPicker';
 
+const SearchBarComponent = typedMemo(
+  fixedForwardRef<HTMLInputElement, InputProps>((props, ref) => {
+    const { t } = useTranslation();
+    return (
+      <Flex gap={2} alignItems="center">
+        <Input ref={ref} {...props} placeholder={t('common.search')} />
+        <NavigateToModelManagerButton />
+      </Flex>
+    );
+  })
+);
+SearchBarComponent.displayName = 'SearchBarComponent';
+
 const PickerGroupHeaderComponent = memo(
   ({ group }: { group: Group<AnyModelConfig, { name: string; description: string }> }) => {
     return (