feat(ui): wip raster layers

I meant to split this up into smaller commits and undo some of it, but I committed afterwards and it's tedious to undo.

invokeai/frontend/web/src/features/controlLayers/components/BrushColorPicker.tsx

@@ -1,4 +1,4 @@
-import { Flex, Popover, PopoverBody, PopoverContent, PopoverTrigger, Tooltip } from '@invoke-ai/ui-library';
+import { Flex, Popover, PopoverBody, PopoverContent, PopoverTrigger } from '@invoke-ai/ui-library';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import IAIColorPicker from 'common/components/IAIColorPicker';
 import { rgbaColorToString } from 'features/canvas/util/colorToString';
@@ -20,21 +20,17 @@ export const BrushColorPicker = memo(() => {
   return (
     <Popover isLazy>
       <PopoverTrigger>
-        <span>
+        <Flex
-          <Tooltip label={t('controlLayers.brushColor')}>
+          as="button"
-            <Flex
+          aria-label={t('controlLayers.brushColor')}
-              as="button"
+          borderRadius="full"
-              aria-label={t('controlLayers.brushColor')}
+          borderWidth={1}
-              borderRadius="full"
+          bg={rgbaColorToString(brushColor)}
-              borderWidth={1}
+          w={8}
-              bg={rgbaColorToString(brushColor)}
+          h={8}
-              w={8}
+          cursor="pointer"
-              h={8}
+          tabIndex={-1}
-              cursor="pointer"
+        />
-              tabIndex={-1}
-            />
-          </Tooltip>
-        </span>
       </PopoverTrigger>
       <PopoverContent>
         <PopoverBody minH={64}>

invokeai/frontend/web/src/features/controlLayers/components/BrushSize.tsx

@@ -28,7 +28,7 @@ export const BrushSize = memo(() => {
     [dispatch]
   );
   return (
-    <FormControl w="min-content">
+    <FormControl w="min-content" gap={2}>
       <FormLabel m={0}>{t('controlLayers.brushSize')}</FormLabel>
       <Popover isLazy>
         <PopoverTrigger>

invokeai/frontend/web/src/features/controlLayers/components/ControlLayersToolbar.tsx

@@ -1,31 +1,40 @@
 /* eslint-disable i18next/no-literal-string */
 import { Flex } from '@invoke-ai/ui-library';
+import { useStore } from '@nanostores/react';
 import { BrushColorPicker } from 'features/controlLayers/components/BrushColorPicker';
 import { BrushSize } from 'features/controlLayers/components/BrushSize';
 import ControlLayersSettingsPopover from 'features/controlLayers/components/ControlLayersSettingsPopover';
 import { ToolChooser } from 'features/controlLayers/components/ToolChooser';
 import { UndoRedoButtonGroup } from 'features/controlLayers/components/UndoRedoButtonGroup';
+import { $tool } from 'features/controlLayers/store/controlLayersSlice';
 import { ToggleProgressButton } from 'features/gallery/components/ImageViewer/ToggleProgressButton';
 import { ViewerToggleMenu } from 'features/gallery/components/ImageViewer/ViewerToggleMenu';
-import { memo } from 'react';
+import { memo, useMemo } from 'react';
 
 export const ControlLayersToolbar = memo(() => {
+  const tool = useStore($tool);
+  const withBrushSize = useMemo(() => {
+    return tool === 'brush' || tool === 'eraser';
+  }, [tool]);
+  const withBrushColor = useMemo(() => {
+    return tool === 'brush';
+  }, [tool]);
   return (
     <Flex w="full" gap={2}>
       <Flex flex={1} justifyContent="center">
         <Flex gap={2} marginInlineEnd="auto">
           <ToggleProgressButton />
+          <ToolChooser />
         </Flex>
       </Flex>
-      <Flex flex={1} gap={2} justifyContent="center">
+      <Flex flex={1} gap={2} justifyContent="center" alignItems="center">
-        <BrushSize />
+        {withBrushSize && <BrushSize />}
-        <BrushColorPicker />
+        {withBrushColor && <BrushColorPicker />}
-        <ToolChooser />
-        <UndoRedoButtonGroup />
-        <ControlLayersSettingsPopover />
       </Flex>
       <Flex flex={1} justifyContent="center">
         <Flex gap={2} marginInlineStart="auto">
+          <UndoRedoButtonGroup />
+          <ControlLayersSettingsPopover />
           <ViewerToggleMenu />
         </Flex>
       </Flex>

invokeai/frontend/web/src/features/controlLayers/konva/naming.ts

@@ -35,13 +35,15 @@ export const RASTER_LAYER_OBJECT_GROUP_NAME = 'raster_layer.object_group';
 export const RASTER_LAYER_BRUSH_LINE_NAME = 'raster_layer.brush_line';
 export const RASTER_LAYER_ERASER_LINE_NAME = 'raster_layer.eraser_line';
 export const RASTER_LAYER_RECT_SHAPE_NAME = 'raster_layer.rect_shape';
+export const RASTER_LAYER_IMAGE_NAME = 'raster_layer.image';
 
 // Getters for non-singleton layer and object IDs
 export const getRGLayerId = (layerId: string) => `${RG_LAYER_NAME}_${layerId}`;
 export const getRasterLayerId = (layerId: string) => `${RASTER_LAYER_NAME}_${layerId}`;
 export const getBrushLineId = (layerId: string, lineId: string) => `${layerId}.brush_line_${lineId}`;
 export const getEraserLineId = (layerId: string, lineId: string) => `${layerId}.eraser_line_${lineId}`;
-export const getRectId = (layerId: string, lineId: string) => `${layerId}.rect_${lineId}`;
+export const getRectShapeId = (layerId: string, lineId: string) => `${layerId}.rect_${lineId}`;
+export const getImageObjectId = (layerId: string, imageName: string) => `${layerId}.image_${imageName}`;
 export const getObjectGroupId = (layerId: string, groupId: string) => `${layerId}.objectGroup_${groupId}`;
 export const getLayerBboxId = (layerId: string) => `${layerId}.bbox`;
 export const getCALayerId = (layerId: string) => `control_adapter_layer_${layerId}`;

invokeai/frontend/web/src/features/controlLayers/konva/renderers/objects.ts

@@ -1,8 +1,9 @@
 import { rgbaColorToString } from 'features/canvas/util/colorToString';
 import { getObjectGroupId } from 'features/controlLayers/konva/naming';
-import type { BrushLine, EraserLine, RectShape } from 'features/controlLayers/store/types';
+import type { BrushLine, EraserLine, ImageObject, RectShape } from 'features/controlLayers/store/types';
 import { DEFAULT_RGBA_COLOR } from 'features/controlLayers/store/types';
 import Konva from 'konva';
+import { getImageDTO } from 'services/api/endpoints/images';
 import { v4 as uuidv4 } from 'uuid';
 
 /**
@@ -80,6 +81,34 @@ export const createRectShape = (rectShape: RectShape, layerObjectGroup: Konva.Gr
   return konvaRect;
 };
 
+export const createImageObject = async (
+  imageObject: ImageObject,
+  layerObjectGroup: Konva.Group,
+  name: string
+): Promise<Konva.Image | null> => {
+  const imageDTO = await getImageDTO(imageObject.image.name);
+  if (!imageDTO) {
+    return null;
+  }
+  return new Promise((resolve) => {
+    const imageEl = new Image();
+    imageEl.onload = () => {
+      const konvaImage = new Konva.Image({
+        id: imageObject.id,
+        name,
+        listening: false,
+        image: imageEl,
+      });
+      layerObjectGroup.add(konvaImage);
+      resolve(konvaImage);
+    };
+    imageEl.onerror = () => {
+      resolve(null);
+    };
+    imageEl.id = imageObject.id;
+    imageEl.src = imageDTO.image_url;
+  });
+};
 /**
  * Creates a konva group for a layer's objects.
  * @param konvaLayer The konva layer to add the object group to

invokeai/frontend/web/src/features/controlLayers/konva/renderers/rasterLayer.ts

@@ -1,6 +1,7 @@
 import {
   RASTER_LAYER_BRUSH_LINE_NAME,
   RASTER_LAYER_ERASER_LINE_NAME,
+  RASTER_LAYER_IMAGE_NAME,
   RASTER_LAYER_NAME,
   RASTER_LAYER_OBJECT_GROUP_NAME,
   RASTER_LAYER_RECT_SHAPE_NAME,
@@ -8,12 +9,14 @@ import {
 import {
   createBrushLine,
   createEraserLine,
+  createImageObject,
   createObjectGroup,
   createRectShape,
 } from 'features/controlLayers/konva/renderers/objects';
 import { getScaledFlooredCursorPosition, mapId, selectRasterObjects } from 'features/controlLayers/konva/util';
 import type { RasterLayer, Tool } from 'features/controlLayers/store/types';
 import Konva from 'konva';
+import { assert } from 'tsafe';
 
 /**
  * Logic for creating and rendering raster layers.
@@ -76,12 +79,12 @@ const createRasterLayer = (
  * @param tool The current tool
  * @param onLayerPosChanged Callback for when the layer's position changes
  */
-export const renderRasterLayer = (
+export const renderRasterLayer = async (
   stage: Konva.Stage,
   layerState: RasterLayer,
   tool: Tool,
   onLayerPosChanged?: (layerId: string, x: number, y: number) => void
-): void => {
+) => {
   const konvaLayer =
     stage.findOne<Konva.Layer>(`#${layerState.id}`) ?? createRasterLayer(stage, layerState, onLayerPosChanged);
 
@@ -106,26 +109,38 @@ export const renderRasterLayer = (
     }
   }
 
-  for (const obj of layerState.objects) {
+  for (let i = 0; i < layerState.objects.length; i++) {
+    const obj = layerState.objects[i];
+    assert(obj);
+    const zIndex = layerState.objects.length - i;
     if (obj.type === 'brush_line') {
       const konvaBrushLine =
-        stage.findOne<Konva.Line>(`#${obj.id}`) ?? createBrushLine(obj, konvaObjectGroup, RASTER_LAYER_BRUSH_LINE_NAME);
+        konvaObjectGroup.findOne<Konva.Line>(`#${obj.id}`) ??
+        createBrushLine(obj, konvaObjectGroup, RASTER_LAYER_BRUSH_LINE_NAME);
       // Only update the points if they have changed.
       if (konvaBrushLine.points().length !== obj.points.length) {
         konvaBrushLine.points(obj.points);
       }
+      konvaBrushLine.zIndex(zIndex);
     } else if (obj.type === 'eraser_line') {
       const konvaEraserLine =
-        stage.findOne<Konva.Line>(`#${obj.id}`) ??
+        konvaObjectGroup.findOne<Konva.Line>(`#${obj.id}`) ??
         createEraserLine(obj, konvaObjectGroup, RASTER_LAYER_ERASER_LINE_NAME);
       // Only update the points if they have changed.
       if (konvaEraserLine.points().length !== obj.points.length) {
         konvaEraserLine.points(obj.points);
       }
+      konvaEraserLine.zIndex(zIndex);
     } else if (obj.type === 'rect_shape') {
-      if (!stage.findOne<Konva.Rect>(`#${obj.id}`)) {
+      const konvaRect =
+        konvaObjectGroup.findOne<Konva.Rect>(`#${obj.id}`) ??
         createRectShape(obj, konvaObjectGroup, RASTER_LAYER_RECT_SHAPE_NAME);
-      }
+      konvaRect.zIndex(zIndex);
+    } else if (obj.type === 'image') {
+      const konvaImage =
+        konvaObjectGroup.findOne<Konva.Image>(`#${obj.id}`) ??
+        (await createImageObject(obj, konvaObjectGroup, RASTER_LAYER_IMAGE_NAME));
+      konvaImage?.zIndex(zIndex);
     }
   }
 

invokeai/frontend/web/src/features/controlLayers/konva/util.ts

@@ -3,6 +3,7 @@ import {
   INITIAL_IMAGE_LAYER_NAME,
   RASTER_LAYER_BRUSH_LINE_NAME,
   RASTER_LAYER_ERASER_LINE_NAME,
+  RASTER_LAYER_IMAGE_NAME,
   RASTER_LAYER_NAME,
   RASTER_LAYER_RECT_SHAPE_NAME,
   RG_LAYER_BRUSH_LINE_NAME,
@@ -115,5 +116,6 @@ export const selectVectorMaskObjects = (node: Konva.Node): boolean =>
 export const selectRasterObjects = (node: Konva.Node): boolean =>
   node.name() === RASTER_LAYER_BRUSH_LINE_NAME ||
   node.name() === RASTER_LAYER_ERASER_LINE_NAME ||
-  node.name() === RASTER_LAYER_RECT_SHAPE_NAME;
+  node.name() === RASTER_LAYER_RECT_SHAPE_NAME ||
+  node.name() === RASTER_LAYER_IMAGE_NAME;
 //#endregion

invokeai/frontend/web/src/features/controlLayers/store/controlLayersSlice.ts

@@ -11,7 +11,7 @@ import {
   getImageObjectId,
   getIPALayerId,
   getRasterLayerId,
-  getRectId,
+  getRectShapeId,
   getRGLayerId,
   INITIAL_IMAGE_LAYER_ID,
 } from 'features/controlLayers/konva/naming';