feat(ui): rough out imagen3 support for canvas

Author: psychedelicious

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

@@ -1322,6 +1322,7 @@
         "unableToCopyDesc": "Your browser does not support clipboard access. Firefox users may be able to fix this by following ",
         "unableToCopyDesc_theseSteps": "these steps",
         "fluxFillIncompatibleWithT2IAndI2I": "FLUX Fill is not compatible with Text to Image or Image to Image. Use other FLUX models for these tasks.",
+        "image3IncompatibleWithInpaintAndOutpaint": "Imagen3 is does not support Inpainting or Outpainting. Use other models for these tasks.",
         "problemUnpublishingWorkflow": "Problem Unpublishing Workflow",
         "problemUnpublishingWorkflowDescription": "There was a problem unpublishing the workflow. Please try again.",
         "workflowUnpublished": "Workflow Unpublished"

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedLinear.ts

@@ -8,6 +8,7 @@ import { $canvasManager } from 'features/controlLayers/store/ephemeral';
 import { prepareLinearUIBatch } from 'features/nodes/util/graph/buildLinearBatchConfig';
 import { buildCogView4Graph } from 'features/nodes/util/graph/generation/buildCogView4Graph';
 import { buildFLUXGraph } from 'features/nodes/util/graph/generation/buildFLUXGraph';
+import { buildImagen3Graph } from 'features/nodes/util/graph/generation/buildImagen3Graph';
 import { buildSD1Graph } from 'features/nodes/util/graph/generation/buildSD1Graph';
 import { buildSD3Graph } from 'features/nodes/util/graph/generation/buildSD3Graph';
 import { buildSDXLGraph } from 'features/nodes/util/graph/generation/buildSDXLGraph';
@@ -48,6 +49,8 @@ export const addEnqueueRequestedLinear = (startAppListening: AppStartListening)
             return await buildFLUXGraph(state, manager);
           case 'cogview4':
             return await buildCogView4Graph(state, manager);
+          case 'imagen3':
+            return await buildImagen3Graph(state, manager);
           default:
             assert(false, `No graph builders for base ${base}`);
         }
@@ -68,12 +71,20 @@ export const addEnqueueRequestedLinear = (startAppListening: AppStartListening)
         return;
       }
 
-      const { g, noise, posCond } = buildGraphResult.value;
+      const { g, seedFieldIdentifier, positivePromptFieldIdentifier } = buildGraphResult.value;
 
       const destination = state.canvasSettings.sendToCanvas ? 'canvas' : 'gallery';
 
       const prepareBatchResult = withResult(() =>
-        prepareLinearUIBatch(state, g, prepend, noise, posCond, 'canvas', destination)
+        prepareLinearUIBatch(
+          state,
+          g,
+          prepend,
+          seedFieldIdentifier,
+          positivePromptFieldIdentifier,
+          'canvas',
+          destination
+        )
       );
 
       if (prepareBatchResult.isErr()) {

invokeai/frontend/web/src/features/controlLayers/hooks/useIsEntityTypeEnabled.ts

@@ -1,5 +1,5 @@
 import { useAppSelector } from 'app/store/storeHooks';
-import { selectIsCogView4, selectIsSD3 } from 'features/controlLayers/store/paramsSlice';
+import { selectIsCogView4, selectIsImagen3, selectIsSD3 } from 'features/controlLayers/store/paramsSlice';
 import type { CanvasEntityType } from 'features/controlLayers/store/types';
 import { useMemo } from 'react';
 import type { Equals } from 'tsafe';
@@ -8,23 +8,24 @@ import { assert } from 'tsafe';
 export const useIsEntityTypeEnabled = (entityType: CanvasEntityType) => {
   const isSD3 = useAppSelector(selectIsSD3);
   const isCogView4 = useAppSelector(selectIsCogView4);
+  const isImagen3 = useAppSelector(selectIsImagen3);
 
   const isEntityTypeEnabled = useMemo<boolean>(() => {
     switch (entityType) {
       case 'reference_image':
-        return !isSD3 && !isCogView4;
+        return !isSD3 && !isCogView4 && !isImagen3;
       case 'regional_guidance':
-        return !isSD3 && !isCogView4;
+        return !isSD3 && !isCogView4 && !isImagen3;
       case 'control_layer':
-        return !isSD3 && !isCogView4;
+        return !isSD3 && !isCogView4 && !isImagen3;
       case 'inpaint_mask':
-        return true;
+        return !isImagen3;
       case 'raster_layer':
-        return true;
+        return !isImagen3;
       default:
         assert<Equals<typeof entityType, never>>(false);
     }
-  }, [entityType, isSD3, isCogView4]);
+  }, [entityType, isSD3, isCogView4, isImagen3]);
 
   return isEntityTypeEnabled;
 };

invokeai/frontend/web/src/features/controlLayers/konva/CanvasEntity/CanvasEntityTransformer.ts

@@ -346,7 +346,7 @@ export class CanvasEntityTransformer extends CanvasModuleBase {
 
     // If the user is not holding shift, the transform is retaining aspect ratio. It's not possible to snap to the grid
     // in this case, because that would change the aspect ratio. So, we only snap to the grid when shift is held.
-    const gridSize = this.manager.stateApi.$shiftKey.get() ? this.manager.stateApi.getGridSize() : 1;
+    const gridSize = this.manager.stateApi.$shiftKey.get() ? this.manager.stateApi.getPositionGridSize() : 1;
 
     // We need to snap the anchor to the selected grid size, but the positions provided to this callback are absolute,
     // scaled coordinates. They need to be converted to stage coordinates, snapped, then converted back to absolute
@@ -464,7 +464,7 @@ export class CanvasEntityTransformer extends CanvasModuleBase {
       return;
     }
     const { rect } = this.manager.stateApi.getBbox();
-    const gridSize = this.manager.stateApi.getGridSize();
+    const gridSize = this.manager.stateApi.getPositionGridSize();
     const width = this.konva.proxyRect.width();
     const height = this.konva.proxyRect.height();
     const scaleX = rect.width / width;
@@ -498,7 +498,7 @@ export class CanvasEntityTransformer extends CanvasModuleBase {
       return;
     }
     const { rect } = this.manager.stateApi.getBbox();
-    const gridSize = this.manager.stateApi.getGridSize();
+    const gridSize = this.manager.stateApi.getPositionGridSize();
     const width = this.konva.proxyRect.width();
     const height = this.konva.proxyRect.height();
     const scaleX = rect.width / width;
@@ -523,7 +523,7 @@ export class CanvasEntityTransformer extends CanvasModuleBase {
 
   onDragMove = () => {
     // Snap the interaction rect to the grid
-    const gridSize = this.manager.stateApi.getGridSize();
+    const gridSize = this.manager.stateApi.getPositionGridSize();
     this.konva.proxyRect.x(roundToMultiple(this.konva.proxyRect.x(), gridSize));
     this.konva.proxyRect.y(roundToMultiple(this.konva.proxyRect.y(), gridSize));
 

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

@@ -493,7 +493,7 @@ export class CanvasStateApiModule extends CanvasModuleBase {
    * Gets the _positional_ grid size for the current canvas. Note that this is not the same as bbox grid size, which is
    * based on the currently-selected model.
    */
-  getGridSize = (): number => {
+  getPositionGridSize = (): number => {
     const snapToGrid = this.getSettings().snapToGrid;
     if (!snapToGrid) {
       return 1;

invokeai/frontend/web/src/features/controlLayers/konva/CanvasTool/CanvasBboxToolModule.ts

@@ -4,8 +4,10 @@ import { CanvasModuleBase } from 'features/controlLayers/konva/CanvasModuleBase'
 import type { CanvasToolModule } from 'features/controlLayers/konva/CanvasTool/CanvasToolModule';
 import { fitRectToGrid, getKonvaNodeDebugAttrs, getPrefixedId } from 'features/controlLayers/konva/util';
 import { selectBboxOverlay } from 'features/controlLayers/store/canvasSettingsSlice';
+import { selectModel } from 'features/controlLayers/store/paramsSlice';
 import { selectBbox } from 'features/controlLayers/store/selectors';
-import type { Coordinate, Rect } from 'features/controlLayers/store/types';
+import type { Coordinate, Rect, Tool } from 'features/controlLayers/store/types';
+import type { ModelIdentifierField } from 'features/nodes/types/common';
 import Konva from 'konva';
 import { noop } from 'lodash-es';
 import { atom } from 'nanostores';
@@ -178,6 +180,9 @@ export class CanvasBboxToolModule extends CanvasModuleBase {
     // Listen for the bbox overlay setting to update the overlay's visibility
     this.subscriptions.add(this.manager.stateApi.createStoreSubscription(selectBboxOverlay, this.render));
 
+    // Listen for the model changing - some model types constraint the bbox to a certain size or aspect ratio.
+    this.subscriptions.add(this.manager.stateApi.createStoreSubscription(selectModel, this.render));
+
     // Update on busy state changes
     this.subscriptions.add(this.manager.$isBusy.listen(this.render));
   }
@@ -218,12 +223,25 @@ export class CanvasBboxToolModule extends CanvasModuleBase {
 
     this.syncOverlay();
 
+    const model = this.manager.stateApi.runSelector(selectModel);
+
     this.konva.transformer.setAttrs({
       listening: tool === 'bbox',
-      enabledAnchors: tool === 'bbox' ? ALL_ANCHORS : NO_ANCHORS,
+      enabledAnchors: this.getEnabledAnchors(tool, model),
     });
   };
 
+  getEnabledAnchors = (tool: Tool, model?: ModelIdentifierField | null): string[] => {
+    if (tool !== 'bbox') {
+      return NO_ANCHORS;
+    }
+    if (model?.base === 'imagen3') {
+      // The bbox is not resizable in imagen3 mode
+      return NO_ANCHORS;
+    }
+    return ALL_ANCHORS;
+  };
+
   syncOverlay = () => {
     const bboxOverlay = this.manager.stateApi.getSettings().bboxOverlay;
 
@@ -251,7 +269,7 @@ export class CanvasBboxToolModule extends CanvasModuleBase {
   onDragMove = () => {
     // The grid size here is the _position_ grid size, not the _dimension_ grid size - it is not constratined by the
     // currently-selected model.
-    const gridSize = this.manager.stateApi.getGridSize();
+    const gridSize = this.manager.stateApi.getPositionGridSize();
     const bbox = this.manager.stateApi.getBbox();
     const bboxRect: Rect = {
       ...bbox.rect,

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

@@ -67,7 +67,7 @@ import type {
   IPMethodV2,
   T2IAdapterConfig,
 } from './types';
-import { getEntityIdentifier, isRenderableEntity } from './types';
+import { getEntityIdentifier, isImagen3AspectRatioID, isRenderableEntity } from './types';
 import {
   converters,
   getControlLayerState,
@@ -1139,7 +1139,21 @@ export const canvasSlice = createSlice({
       syncScaledSize(state);
     },
     bboxChangedFromCanvas: (state, action: PayloadAction<IRect>) => {
-      state.bbox.rect = action.payload;
+      const newBboxRect = action.payload;
+      const oldBboxRect = state.bbox.rect;
+
+      state.bbox.rect = newBboxRect;
+
+      if (newBboxRect.width === oldBboxRect.width && newBboxRect.height === oldBboxRect.height) {
+        return;
+      }
+
+      const oldAspectRatio = state.bbox.aspectRatio.value;
+      const newAspectRatio = newBboxRect.width / newBboxRect.height;
+
+      if (oldAspectRatio === newAspectRatio) {
+        return;
+      }
 
       // TODO(psyche): Figure out a way to handle this without resetting the aspect ratio on every change.
       // This action is dispatched when the user resizes or moves the bbox from the canvas. For now, when the user
@@ -1198,6 +1212,26 @@ export const canvasSlice = createSlice({
       state.bbox.aspectRatio.id = id;
       if (id === 'Free') {
         state.bbox.aspectRatio.isLocked = false;
+      } else if (state.bbox.modelBase === 'imagen3' && isImagen3AspectRatioID(id)) {
+        // Imagen3 has specific output sizes that are not exactly the same as the aspect ratio. Need special handling.
+        if (id === '16:9') {
+          state.bbox.rect.width = 1408;
+          state.bbox.rect.height = 768;
+        } else if (id === '4:3') {
+          state.bbox.rect.width = 1280;
+          state.bbox.rect.height = 896;
+        } else if (id === '1:1') {
+          state.bbox.rect.width = 1024;
+          state.bbox.rect.height = 1024;
+        } else if (id === '3:4') {
+          state.bbox.rect.width = 896;
+          state.bbox.rect.height = 1280;
+        } else if (id === '9:16') {
+          state.bbox.rect.width = 768;
+          state.bbox.rect.height = 1408;
+        }
+        state.bbox.aspectRatio.value = state.bbox.rect.width / state.bbox.rect.height;
+        state.bbox.aspectRatio.isLocked = true;
       } else {
         state.bbox.aspectRatio.isLocked = true;
         state.bbox.aspectRatio.value = ASPECT_RATIO_MAP[id].ratio;
@@ -1670,6 +1704,13 @@ export const canvasSlice = createSlice({
       const base = model?.base;
       if (isMainModelBase(base) && state.bbox.modelBase !== base) {
         state.bbox.modelBase = base;
+        if (base === 'imagen3') {
+          state.bbox.aspectRatio.isLocked = true;
+          state.bbox.aspectRatio.value = 1;
+          state.bbox.aspectRatio.id = '1:1';
+          state.bbox.rect.width = 1024;
+          state.bbox.rect.height = 1024;
+        }
         syncScaledSize(state);
       }
     });
@@ -1802,6 +1843,10 @@ export const canvasPersistConfig: PersistConfig<CanvasState> = {
 };
 
 const syncScaledSize = (state: CanvasState) => {
+  if (state.bbox.modelBase === 'imagen3') {
+    // Imagen3 has fixed sizes. Scaled bbox is not supported.
+    return;
+  }
   if (state.bbox.scaleMethod === 'auto') {
     // Sync both aspect ratio and size
     const { width, height } = state.bbox.rect;

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

@@ -79,6 +79,7 @@ export type ParamsState = {
   clipLEmbedModel: ParameterCLIPLEmbedModel | null;
   clipGEmbedModel: ParameterCLIPGEmbedModel | null;
   controlLora: ParameterControlLoRAModel | null;
+  imagen3EnhancePrompt: boolean;
 };
 
 const initialState: ParamsState = {
@@ -128,6 +129,7 @@ const initialState: ParamsState = {
   clipLEmbedModel: null,
   clipGEmbedModel: null,
   controlLora: null,
+  imagen3EnhancePrompt: true,
 };
 
 export const paramsSlice = createSlice({
@@ -290,6 +292,9 @@ export const paramsSlice = createSlice({
     setCanvasCoherenceMinDenoise: (state, action: PayloadAction<number>) => {
       state.canvasCoherenceMinDenoise = action.payload;
     },
+    imagen3EnhancePromptChanged: (state, action: PayloadAction<boolean>) => {
+      state.imagen3EnhancePrompt = action.payload;
+    },
     paramsReset: (state) => resetState(state),
   },
   extraReducers(builder) {
@@ -357,6 +362,7 @@ export const {
   setRefinerStart,
   modelChanged,
   paramsReset,
+  imagen3EnhancePromptChanged,
 } = paramsSlice.actions;
 
 /* eslint-disable-next-line @typescript-eslint/no-explicit-any */
@@ -380,6 +386,7 @@ export const selectIsSDXL = createParamsSelector((params) => params.model?.base
 export const selectIsFLUX = createParamsSelector((params) => params.model?.base === 'flux');
 export const selectIsSD3 = createParamsSelector((params) => params.model?.base === 'sd-3');
 export const selectIsCogView4 = createParamsSelector((params) => params.model?.base === 'cogview4');
+export const selectIsImagen3 = createParamsSelector((params) => params.model?.base === 'imagen3');
 
 export const selectModel = createParamsSelector((params) => params.model);
 export const selectModelKey = createParamsSelector((params) => params.model?.key);
@@ -414,6 +421,7 @@ export const selectNegativePrompt = createParamsSelector((params) => params.nega
 export const selectPositivePrompt2 = createParamsSelector((params) => params.positivePrompt2);
 export const selectNegativePrompt2 = createParamsSelector((params) => params.negativePrompt2);
 export const selectShouldConcatPrompts = createParamsSelector((params) => params.shouldConcatPrompts);
+export const selectImagen3EnhancePrompt = createParamsSelector((params) => params.imagen3EnhancePrompt);
 export const selectScheduler = createParamsSelector((params) => params.scheduler);
 export const selectSeamlessXAxis = createParamsSelector((params) => params.seamlessXAxis);
 export const selectSeamlessYAxis = createParamsSelector((params) => params.seamlessYAxis);

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

@@ -388,8 +388,11 @@ export type StagingAreaImage = {
 };
 
 export const zAspectRatioID = z.enum(['Free', '16:9', '3:2', '4:3', '1:1', '3:4', '2:3', '9:16']);
+export const zImagen3AspectRatioID = z.enum(['16:9', '4:3', '1:1', '3:4', '9:16']);
+export const isImagen3AspectRatioID = (v: unknown): v is z.infer<typeof zImagen3AspectRatioID> =>
+  zImagen3AspectRatioID.safeParse(v).success;
 export type AspectRatioID = z.infer<typeof zAspectRatioID>;
-export const isAspectRatioID = (v: string): v is AspectRatioID => zAspectRatioID.safeParse(v).success;
+export const isAspectRatioID = (v: unknown): v is AspectRatioID => zAspectRatioID.safeParse(v).success;
 
 const zCanvasState = z.object({
   _version: z.literal(3),

invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelManagerPanel/ModelBaseBadge.tsx

@@ -16,6 +16,7 @@ export const BASE_COLOR_MAP: Record<BaseModelType, string> = {
   'sdxl-refiner': 'invokeBlue',
   flux: 'gold',
   cogview4: 'red',
+  imagen3: 'pink'
 };
 
 const ModelBaseBadge = ({ base }: Props) => {