refactor(ui): focus/hotkey handling for new layout

Author: psychedelicious

invokeai/frontend/web/src/app/components/InvokeAIUI.tsx

@@ -23,6 +23,7 @@ import { $whatsNew } from 'app/store/nanostores/whatsNew';
 import { createStore } from 'app/store/store';
 import type { PartialAppConfig } from 'app/types/invokeai';
 import Loading from 'common/components/Loading/Loading';
+import { FocusApiProvider } from 'common/hooks/focus2';
 import type { WorkflowSortOption, WorkflowTagCategory } from 'features/nodes/store/workflowLibrarySlice';
 import {
   $workflowLibraryCategoriesOptions,
@@ -331,7 +332,9 @@ const InvokeAIUI = ({
       <Provider store={store}>
         <React.Suspense fallback={<Loading />}>
           <ThemeLocaleProvider>
-            <App config={config} studioInitAction={studioInitAction} />
+            <FocusApiProvider>
+              <App config={config} studioInitAction={studioInitAction} />
+            </FocusApiProvider>
           </ThemeLocaleProvider>
         </React.Suspense>
       </Provider>

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

@@ -9,48 +9,44 @@ interface FocusRegionWrapperProps extends BoxProps {
   focusOnMount?: boolean;
 }
 
-const FOCUS_REGION_STYLES: SystemStyleObject = {
+const BASE_FOCUS_REGION_STYLES: SystemStyleObject = {
   position: 'relative',
   '&[data-highlighted="true"]::after': {
-    borderColor: 'blue.700',
+    borderColor: 'blue.300',
   },
   '&::after': {
     content: '""',
     position: 'absolute',
-    inset: 0,
+    inset: '2px',
     zIndex: 1,
     borderRadius: 'base',
-    border: '2px solid',
+    border: '1px solid',
     borderColor: 'transparent',
     pointerEvents: 'none',
     transition: 'border-color 0.1s ease-in-out',
   },
 };
 
-export const FocusRegionWrapper = memo(
-  ({ region, focusOnMount = false, sx, children, ...boxProps }: FocusRegionWrapperProps) => {
-    const shouldHighlightFocusedRegions = useAppSelector(selectSystemShouldEnableHighlightFocusedRegions);
-
-    const ref = useRef<HTMLDivElement>(null);
-
-    const options = useMemo(() => ({ focusOnMount }), [focusOnMount]);
-
-    useFocusRegion(region, ref, options);
-    const isFocused = useIsRegionFocused(region);
-    const isHighlighted = isFocused && shouldHighlightFocusedRegions;
-
-    return (
-      <Box
-        ref={ref}
-        tabIndex={-1}
-        sx={useMemo(() => ({ ...FOCUS_REGION_STYLES, ...sx }), [sx])}
-        data-highlighted={isHighlighted}
-        {...boxProps}
-      >
-        {children}
-      </Box>
-    );
-  }
-);
+export const FocusRegionWrapper = memo((props: FocusRegionWrapperProps) => {
+  const { region, focusOnMount = false, sx: _sx, children, ...boxProps } = props;
+
+  const ref = useRef<HTMLDivElement>(null);
+
+  const sx = useMemo(() => ({ ...BASE_FOCUS_REGION_STYLES, ..._sx }), [_sx]);
+
+  const shouldHighlightFocusedRegions = useAppSelector(selectSystemShouldEnableHighlightFocusedRegions);
+
+  const options = useMemo(() => ({ focusOnMount }), [focusOnMount]);
+
+  useFocusRegion(region, ref, options);
+  const isFocused = useIsRegionFocused(region);
+  const isHighlighted = isFocused && shouldHighlightFocusedRegions;
+
+  return (
+    <Box ref={ref} tabIndex={-1} sx={sx} data-highlighted={isHighlighted} {...boxProps}>
+      {children}
+    </Box>
+  );
+});
 
 FocusRegionWrapper.displayName = 'FocusRegionWrapper';

invokeai/frontend/web/src/common/hooks/focus.ts

@@ -6,6 +6,7 @@ import { atom, computed } from 'nanostores';
 import type { RefObject } from 'react';
 import { useEffect } from 'react';
 import { objectKeys } from 'tsafe';
+import z from 'zod/v4';
 
 /**
  * We need to manage focus regions to conditionally enable hotkeys:
@@ -30,18 +31,29 @@ const log = logger('system');
 /**
  * The names of the focus regions.
  */
-export type FocusRegionName = 'gallery' | 'layers' | 'canvas' | 'workflows' | 'viewer';
+const zFocusRegionName = z.enum([
+  'launchpad',
+  'viewer',
+  'gallery',
+  'boards',
+  'layers',
+  'canvas',
+  'workflows',
+  'progress',
+  'settings',
+]);
+export type FocusRegionName = z.infer<typeof zFocusRegionName>;
 
 /**
  * A map of focus regions to the elements that are part of that region.
  */
-const REGION_TARGETS: Record<FocusRegionName, Set<HTMLElement>> = {
-  gallery: new Set<HTMLElement>(),
-  layers: new Set<HTMLElement>(),
-  canvas: new Set<HTMLElement>(),
-  workflows: new Set<HTMLElement>(),
-  viewer: new Set<HTMLElement>(),
-} as const;
+const REGION_TARGETS: Record<FocusRegionName, Set<HTMLElement>> = zFocusRegionName.options.values().reduce(
+  (acc, region) => {
+    acc[region] = new Set<HTMLElement>();
+    return acc;
+  },
+  {} as Record<FocusRegionName, Set<HTMLElement>>
+);
 
 /**
  * The currently-focused region or `null` if no region is focused.

invokeai/frontend/web/src/common/hooks/focus2.tsx

@@ -0,0 +1,102 @@
+import { useAppStore } from 'app/store/storeHooks';
+import { selectActiveTab } from 'features/ui/store/uiSelectors';
+import type { TabName } from 'features/ui/store/uiTypes';
+import type { Atom } from 'nanostores';
+import { atom } from 'nanostores';
+import type { PropsWithChildren } from 'react';
+import { createContext, memo, useContext, useEffect } from 'react';
+
+type FocusState = {
+  tab: TabName;
+  panels: Record<TabName, string | null>;
+};
+
+const $focusState = atom<FocusState>({
+  tab: 'generate',
+  panels: {
+    canvas: null,
+    generate: null,
+    upscaling: null,
+    workflows: null,
+    models: null,
+    queue: null,
+  },
+});
+
+type FocusedTabAndPanel = {
+  tab: TabName;
+  panelId: string | null;
+};
+
+const getFocusedTabAndPanel = (focusState: FocusState): FocusedTabAndPanel => {
+  const { tab, panels } = focusState;
+  return { tab, panelId: panels[tab] || null };
+};
+
+const $focusedTabAndPanel = atom<FocusedTabAndPanel>(getFocusedTabAndPanel($focusState.get()));
+
+$focusState.listen((focusState) => {
+  const prev = $focusedTabAndPanel.get();
+  const next = getFocusedTabAndPanel(focusState);
+  if (prev.tab === next.tab && prev.panelId === next.panelId) {
+    return;
+  }
+  $focusedTabAndPanel.set(next);
+});
+
+const _setFocusedTab = (tab: TabName) => {
+  const state = $focusState.get();
+  if (state.tab === tab) {
+    return;
+  }
+  $focusState.set({ ...state, tab });
+};
+const setFocusedPanel = (tab: TabName, panelId: string) => {
+  const state = $focusState.get();
+  if (state === null) {
+    return;
+  }
+  if (state.panels[tab] === panelId) {
+    return;
+  }
+  $focusState.set({ ...state, panels: { ...state.panels, [tab]: panelId } });
+};
+
+export type FocusApi = {
+  $focusState: Atom<FocusState>;
+  $focusedTabAndPanel: Atom<FocusedTabAndPanel>;
+  setFocusedPanel: (tab: TabName, panelId: string) => void;
+};
+
+export const focusApi: FocusApi = { $focusState, $focusedTabAndPanel, setFocusedPanel };
+
+const FocusApiContext = createContext<FocusApi | null>(null);
+
+export const FocusApiProvider = memo((props: PropsWithChildren) => {
+  const store = useAppStore();
+  useEffect(() => {
+    const cb = () => {
+      const tab = selectActiveTab(store.getState());
+      _setFocusedTab(tab);
+    };
+    const unsubscribe = store.subscribe(cb);
+    cb();
+    return () => {
+      unsubscribe();
+    };
+  }, [store]);
+  return <FocusApiContext.Provider value={focusApi}>{props.children}</FocusApiContext.Provider>;
+});
+FocusApiProvider.displayName = 'FocusApiProvider';
+
+export const useFocusApi = () => {
+  const context = useContext(FocusApiContext);
+  if (context === null) {
+    throw new Error('useFocusApi must be used within a FocusApiProvider');
+  }
+  return context;
+};
+
+$focusedTabAndPanel.subscribe((state) => {
+  console.log(state);
+});

invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/GenerateLaunchpadPanel.tsx

@@ -1,5 +1,6 @@
 import { Alert, Button, Flex, Grid, Heading, Text } from '@invoke-ai/ui-library';
 import { useAppDispatch } from 'app/store/storeHooks';
+import { FocusRegionWrapper } from 'common/components/FocusRegionWrapper';
 import { InitialStateMainModelPicker } from 'features/controlLayers/components/SimpleSession/InitialStateMainModelPicker';
 import { LaunchpadAddStyleReference } from 'features/controlLayers/components/SimpleSession/LaunchpadAddStyleReference';
 import { setActiveTab } from 'features/ui/store/uiSlice';
@@ -14,7 +15,15 @@ export const GenerateLaunchpadPanel = memo(() => {
   }, [dispatch]);
 
   return (
-    <Flex flexDir="column" h="full" w="full" alignItems="center" gap={2}>
+    <FocusRegionWrapper
+      region="launchpad"
+      display="flex"
+      flexDir="column"
+      h="full"
+      w="full"
+      alignItems="center"
+      gap={2}
+    >
       <Flex flexDir="column" w="full" gap={4} px={14} maxW={768} pt="20vh">
         <Heading mb={4}>Generate images from text prompts.</Heading>
         <Flex flexDir="column" gap={8}>
@@ -41,7 +50,7 @@ export const GenerateLaunchpadPanel = memo(() => {
           </Alert>
         </Flex>
       </Flex>
-    </Flex>
+    </FocusRegionWrapper>
   );
 });
 GenerateLaunchpadPanel.displayName = 'GenerateLaunchpad';

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

@@ -1,6 +1,7 @@
 import { Box, Button, Collapse, Divider, Flex, IconButton } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { FocusRegionWrapper } from 'common/components/FocusRegionWrapper';
 import { useDisclosure } from 'common/hooks/useBoolean';
 import { BoardsListWrapper } from 'features/gallery/components/Boards/BoardsList/BoardsListWrapper';
 import { BoardsSearch } from 'features/gallery/components/Boards/BoardsList/BoardsSearch';
@@ -45,7 +46,7 @@ export const BoardsPanel = memo(() => {
   }, [boardSearchText.length, searchDisclosure, collapsibleApi, dispatch]);
 
   return (
-    <Flex flexDir="column" w="full" h="full" p={2}>
+    <FocusRegionWrapper region="boards" display="flex" flexDir="column" w="full" h="full" p={2}>
       <Flex alignItems="center" justifyContent="space-between" w="full">
         <Flex flexGrow={1} flexBasis={0}>
           <Button
@@ -81,7 +82,7 @@ export const BoardsPanel = memo(() => {
       </Collapse>
       <Divider pt={2} />
       <BoardsListWrapper />
-    </Flex>
+    </FocusRegionWrapper>
   );
 });
 BoardsPanel.displayName = 'BoardsPanel';

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

@@ -2,6 +2,7 @@ import { Box, Button, ButtonGroup, Collapse, Divider, Flex, IconButton, Spacer }
 import { useStore } from '@nanostores/react';
 import { createSelector } from '@reduxjs/toolkit';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { FocusRegionWrapper } from 'common/components/FocusRegionWrapper';
 import { useDisclosure } from 'common/hooks/useBoolean';
 import { useGallerySearchTerm } from 'features/gallery/components/ImageGrid/useGallerySearchTerm';
 import { selectSelectedBoardId } from 'features/gallery/store/gallerySelectors';
@@ -67,7 +68,17 @@ export const GalleryPanel = memo(() => {
   const boardName = useBoardName(selectedBoardId);
 
   return (
-    <Flex flexDirection="column" alignItems="center" justifyContent="space-between" h="full" w="full" minH={0} p={2}>
+    <FocusRegionWrapper
+      region="gallery"
+      display="flex"
+      flexDirection="column"
+      alignItems="center"
+      justifyContent="space-between"
+      h="full"
+      w="full"
+      minH={0}
+      p={2}
+    >
       <Flex gap={2} fontSize="sm" alignItems="center" w="full">
         <Button
           size="sm"
@@ -123,7 +134,7 @@ export const GalleryPanel = memo(() => {
       <Flex w="full" h="full" pt={2}>
         <NewGallery />
       </Flex>
-    </Flex>
+    </FocusRegionWrapper>
   );
 });
 GalleryPanel.displayName = 'Gallery';

invokeai/frontend/web/src/features/gallery/components/ImageViewer/GenerationProgressPanel.tsx

@@ -1,13 +1,22 @@
-import { Flex } from '@invoke-ai/ui-library';
+import { FocusRegionWrapper } from 'common/components/FocusRegionWrapper';
 import { ProgressImage } from 'features/gallery/components/ImageViewer/ProgressImage';
 import { memo } from 'react';
 
 import { ProgressIndicator } from './ProgressIndicator';
 
 export const GenerationProgressPanel = memo(() => (
-  <Flex position="relative" flexDir="column" w="full" h="full" overflow="hidden" p={2}>
+  <FocusRegionWrapper
+    region="progress"
+    display="flex"
+    position="relative"
+    flexDir="column"
+    w="full"
+    h="full"
+    overflow="hidden"
+    p={2}
+  >
     <ProgressImage />
     <ProgressIndicator position="absolute" top={6} right={6} size={8} />
-  </Flex>
+  </FocusRegionWrapper>
 ));
 GenerationProgressPanel.displayName = 'GenerationProgressPanel';

invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageViewerPanel.tsx

@@ -1,4 +1,5 @@
-import { Divider, Flex } from '@invoke-ai/ui-library';
+import { Divider } from '@invoke-ai/ui-library';
+import { FocusRegionWrapper } from 'common/components/FocusRegionWrapper';
 import { ImageViewer } from 'features/gallery/components/ImageViewer/ImageViewer';
 import { ViewerToolbar } from 'features/gallery/components/ImageViewer/ViewerToolbar';
 import { memo } from 'react';
@@ -8,11 +9,20 @@ import { ImageViewerContextProvider } from './context';
 export const ImageViewerPanel = memo(() => {
   return (
     <ImageViewerContextProvider>
-      <Flex flexDir="column" w="full" h="full" overflow="hidden" p={2} gap={2}>
+      <FocusRegionWrapper
+        region="viewer"
+        display="flex"
+        flexDir="column"
+        w="full"
+        h="full"
+        overflow="hidden"
+        p={2}
+        gap={2}
+      >
         <ViewerToolbar />
         <Divider />
         <ImageViewer />
-      </Flex>
+      </FocusRegionWrapper>
     </ImageViewerContextProvider>
   );
 });