feat(ui): genericizing picker

2026-04-23 03:00:31 -04:00 · 2025-04-16 14:46:21 +10:00
parent 2eb212f393
commit 9d8a71b362
3 changed files with 76 additions and 37 deletions
--- a/invokeai/frontend/web/src/common/components/Picker/Picker.tsx
+++ b/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,17 +37,18 @@ 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">
@@ -55,9 +56,9 @@ const DefaultNoOptionsFallback = typedMemo(() => {
    </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">
@@ -65,7 +66,7 @@ const DefaultNoMatchesFallback = typedMemo(() => {
    </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>
    );
  }
--- a/invokeai/frontend/web/src/common/util/fixedForwardRef.ts
+++ b/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;
+}
--- a/invokeai/frontend/web/src/features/settingsAccordions/components/GenerationSettingsAccordion/GenerationSettingsAccordion.tsx
+++ b/invokeai/frontend/web/src/features/settingsAccordions/components/GenerationSettingsAccordion/GenerationSettingsAccordion.tsx
@@ -1,4 +1,4 @@
-import type { FormLabelProps } from '@invoke-ai/ui-library';
+import type { FormLabelProps, InputProps } from '@invoke-ai/ui-library';
 import {
  Box,
  Button,
@@ -6,6 +6,7 @@ import {
  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 (