Maryhipp/chatgpt UI (#7969)

* add GPTimage1 as allowed base model

* fix for non-disabled inpaint layers

* lots of boilerplate for adding gpt-image base model and disabling things along with imagen

* handle gpt-image dimensions

* build graph for gpt-image

* lint

* feat(ui): make chatgpt model naming consistent

* feat(ui): graph builder naming

* feat(ui): disable img2img for imagen3

* feat(ui): more naming

* feat(ui): support presigned url prefetch

* feat(ui): disable neg prompt for chatgpt

* docs(ui): update docstring

* feat(ui): fix graph building issues for chatgpt

* fix(ui): node ids for chatgpt/imagen

* chore(ui): typegen

---------

Co-authored-by: Mary Hipp <maryhipp@Marys-MacBook-Air.local>
Co-authored-by: psychedelicious <4822129+psychedelicious@users.noreply.github.com>

Author: 3 people

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

@@ -1322,7 +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 does not support Inpainting or Outpainting. Use other models for these tasks.",
+        "imagen3IncompatibleGenerationMode": "Imagen3 only supports Text to Image. Use other models for Image to Image, Inpainting and Outpainting 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

@@ -6,6 +6,7 @@ import { withResult, withResultAsync } from 'common/util/result';
 import { parseify } from 'common/util/serialize';
 import { $canvasManager } from 'features/controlLayers/store/ephemeral';
 import { prepareLinearUIBatch } from 'features/nodes/util/graph/buildLinearBatchConfig';
+import { buildChatGPT4oGraph } from 'features/nodes/util/graph/generation/buildChatGPT4oGraph';
 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';
@@ -51,6 +52,8 @@ export const addEnqueueRequestedLinear = (startAppListening: AppStartListening)
             return await buildCogView4Graph(state, manager);
           case 'imagen3':
             return await buildImagen3Graph(state, manager);
+          case 'chatgpt-4o':
+            return await buildChatGPT4oGraph(state, manager);
           default:
             assert(false, `No graph builders for base ${base}`);
         }
@@ -76,15 +79,15 @@ export const addEnqueueRequestedLinear = (startAppListening: AppStartListening)
       const destination = state.canvasSettings.sendToCanvas ? 'canvas' : 'gallery';
 
       const prepareBatchResult = withResult(() =>
-        prepareLinearUIBatch(
+        prepareLinearUIBatch({
           state,
           g,
           prepend,
           seedFieldIdentifier,
           positivePromptFieldIdentifier,
-          'canvas',
-          destination
-        )
+          origin: 'canvas',
+          destination,
+        })
       );
 
       if (prepareBatchResult.isErr()) {

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

@@ -20,15 +20,15 @@ export const addEnqueueRequestedUpscale = (startAppListening: AppStartListening)
 
       const { g, seedFieldIdentifier, positivePromptFieldIdentifier } = await buildMultidiffusionUpscaleGraph(state);
 
-      const batchConfig = prepareLinearUIBatch(
+      const batchConfig = prepareLinearUIBatch({
         state,
         g,
         prepend,
         seedFieldIdentifier,
         positivePromptFieldIdentifier,
-        'upscaling',
-        'gallery'
-      );
+        origin: 'upscaling',
+        destination: 'gallery',
+      });
 
       const req = dispatch(queueApi.endpoints.enqueueBatch.initiate(batchConfig, enqueueMutationFixedCacheKeyOptions));
       try {

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

@@ -24,6 +24,7 @@ export const CanvasAddEntityButtons = memo(() => {
   const isReferenceImageEnabled = useIsEntityTypeEnabled('reference_image');
   const isRegionalGuidanceEnabled = useIsEntityTypeEnabled('regional_guidance');
   const isControlLayerEnabled = useIsEntityTypeEnabled('control_layer');
+  const isInpaintLayerEnabled = useIsEntityTypeEnabled('inpaint_mask');
 
   return (
     <Flex w="full" h="full" justifyContent="center" gap={4}>
@@ -52,6 +53,7 @@ export const CanvasAddEntityButtons = memo(() => {
               justifyContent="flex-start"
               leftIcon={<PiPlusBold />}
               onClick={addInpaintMask}
+              isDisabled={!isInpaintLayerEnabled}
             >
               {t('controlLayers.inpaintMask')}
             </Button>

invokeai/frontend/web/src/features/controlLayers/components/CanvasEntityList/EntityListGlobalActionBarAddLayerMenu.tsx

@@ -25,6 +25,7 @@ export const EntityListGlobalActionBarAddLayerMenu = memo(() => {
   const isReferenceImageEnabled = useIsEntityTypeEnabled('reference_image');
   const isRegionalGuidanceEnabled = useIsEntityTypeEnabled('regional_guidance');
   const isControlLayerEnabled = useIsEntityTypeEnabled('control_layer');
+  const isInpaintLayerEnabled = useIsEntityTypeEnabled('inpaint_mask');
 
   return (
     <Menu>
@@ -46,7 +47,7 @@ export const EntityListGlobalActionBarAddLayerMenu = memo(() => {
           </MenuItem>
         </MenuGroup>
         <MenuGroup title={t('controlLayers.regional')}>
-          <MenuItem icon={<PiPlusBold />} onClick={addInpaintMask}>
+          <MenuItem icon={<PiPlusBold />} onClick={addInpaintMask} isDisabled={!isInpaintLayerEnabled}>
             {t('controlLayers.inpaintMask')}
           </MenuItem>
           <MenuItem icon={<PiPlusBold />} onClick={addRegionalGuidance} isDisabled={!isRegionalGuidanceEnabled}>

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

@@ -1,5 +1,10 @@
 import { useAppSelector } from 'app/store/storeHooks';
-import { selectIsCogView4, selectIsImagen3, selectIsSD3 } from 'features/controlLayers/store/paramsSlice';
+import {
+  selectIsChatGTP4o,
+  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';
@@ -9,23 +14,24 @@ export const useIsEntityTypeEnabled = (entityType: CanvasEntityType) => {
   const isSD3 = useAppSelector(selectIsSD3);
   const isCogView4 = useAppSelector(selectIsCogView4);
   const isImagen3 = useAppSelector(selectIsImagen3);
+  const isChatGPT4o = useAppSelector(selectIsChatGTP4o);
 
   const isEntityTypeEnabled = useMemo<boolean>(() => {
     switch (entityType) {
       case 'reference_image':
-        return !isSD3 && !isCogView4 && !isImagen3;
+        return !isSD3 && !isCogView4 && !isImagen3 && !isChatGPT4o;
       case 'regional_guidance':
-        return !isSD3 && !isCogView4 && !isImagen3;
+        return !isSD3 && !isCogView4 && !isImagen3 && !isChatGPT4o;
       case 'control_layer':
-        return !isSD3 && !isCogView4 && !isImagen3;
+        return !isSD3 && !isCogView4 && !isImagen3 && !isChatGPT4o;
       case 'inpaint_mask':
-        return !isImagen3;
+        return !isImagen3 && !isChatGPT4o;
       case 'raster_layer':
-        return !isImagen3;
+        return !isImagen3 && !isChatGPT4o;
       default:
         assert<Equals<typeof entityType, never>>(false);
     }
-  }, [entityType, isSD3, isCogView4, isImagen3]);
+  }, [entityType, isSD3, isCogView4, isImagen3, isChatGPT4o]);
 
   return isEntityTypeEnabled;
 };

invokeai/frontend/web/src/features/controlLayers/konva/CanvasObject/CanvasObjectImage.ts

@@ -112,7 +112,7 @@ export class CanvasObjectImage extends CanvasModuleBase {
       return;
     }
 
-    const imageElementResult = await withResultAsync(() => loadImage(imageDTO.image_url));
+    const imageElementResult = await withResultAsync(() => loadImage(imageDTO.image_url, true));
     if (imageElementResult.isErr()) {
       // Image loading failed (e.g. the URL to the "physical" image is invalid)
       this.onFailedToLoadImage(t('controlLayers.unableToLoadImage', 'Unable to load image'));

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

@@ -235,8 +235,8 @@ export class CanvasBboxToolModule extends CanvasModuleBase {
     if (tool !== 'bbox') {
       return NO_ANCHORS;
     }
-    if (model?.base === 'imagen3') {
-      // The bbox is not resizable in imagen3 mode
+    if (model?.base === 'imagen3' || model?.base === 'chatgpt-4o') {
+      // The bbox is not resizable in these modes
       return NO_ANCHORS;
     }
     return ALL_ANCHORS;

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

@@ -476,15 +476,24 @@ export function getImageDataTransparency(imageData: ImageData): Transparency {
 /**
  * Loads an image from a URL and returns a promise that resolves with the loaded image element.
  * @param src The image source URL
+ * @param fetchUrlFirst Whether to fetch the image's URL first, assuming the provided `src` will redirect to a different URL. This addresses an issue where CORS headers are dropped during a redirect.
  * @returns A promise that resolves with the loaded image element
  */
-export function loadImage(src: string): Promise<HTMLImageElement> {
+export async function loadImage(src: string, fetchUrlFirst?: boolean): Promise<HTMLImageElement> {
+  const authToken = $authToken.get();
+  let url = src;
+  if (authToken && fetchUrlFirst) {
+    const response = await fetch(`${src}?url_only=true`, { credentials: 'include' });
+    const data = await response.json();
+    url = data.url;
+  }
+
   return new Promise((resolve, reject) => {
     const imageElement = new Image();
     imageElement.onload = () => resolve(imageElement);
     imageElement.onerror = (error) => reject(error);
     imageElement.crossOrigin = $authToken.get() ? 'use-credentials' : 'anonymous';
-    imageElement.src = src;
+    imageElement.src = url;
   });
 }
 

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

@@ -67,7 +67,7 @@ import type {
   IPMethodV2,
   T2IAdapterConfig,
 } from './types';
-import { getEntityIdentifier, isImagen3AspectRatioID, isRenderableEntity } from './types';
+import { getEntityIdentifier, isChatGPT4oAspectRatioID, isImagen3AspectRatioID, isRenderableEntity } from './types';
 import {
   converters,
   getControlLayerState,
@@ -1232,6 +1232,20 @@ export const canvasSlice = createSlice({
         }
         state.bbox.aspectRatio.value = state.bbox.rect.width / state.bbox.rect.height;
         state.bbox.aspectRatio.isLocked = true;
+      } else if (state.bbox.modelBase === 'chatgpt-4o' && isChatGPT4oAspectRatioID(id)) {
+        // gpt-image has specific output sizes that are not exactly the same as the aspect ratio. Need special handling.
+        if (id === '3:2') {
+          state.bbox.rect.width = 1536;
+          state.bbox.rect.height = 1024;
+        } else if (id === '1:1') {
+          state.bbox.rect.width = 1024;
+          state.bbox.rect.height = 1024;
+        } else if (id === '2:3') {
+          state.bbox.rect.width = 1024;
+          state.bbox.rect.height = 1536;
+        }
+        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;
@@ -1704,7 +1718,7 @@ export const canvasSlice = createSlice({
       const base = model?.base;
       if (isMainModelBase(base) && state.bbox.modelBase !== base) {
         state.bbox.modelBase = base;
-        if (base === 'imagen3') {
+        if (base === 'imagen3' || base === 'chatgpt-4o') {
           state.bbox.aspectRatio.isLocked = true;
           state.bbox.aspectRatio.value = 1;
           state.bbox.aspectRatio.id = '1:1';
@@ -1843,7 +1857,7 @@ export const canvasPersistConfig: PersistConfig<CanvasState> = {
 };
 
 const syncScaledSize = (state: CanvasState) => {
-  if (state.bbox.modelBase === 'imagen3') {
+  if (state.bbox.modelBase === 'imagen3' || state.bbox.modelBase === 'chatgpt-4o') {
     // Imagen3 has fixed sizes. Scaled bbox is not supported.
     return;
   }

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

@@ -381,6 +381,7 @@ export const selectIsFLUX = createParamsSelector((params) => params.model?.base
 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 selectIsChatGTP4o = createParamsSelector((params) => params.model?.base === 'chatgpt-4o');
 
 export const selectModel = createParamsSelector((params) => params.model);
 export const selectModelKey = createParamsSelector((params) => params.model?.key);

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

@@ -388,9 +388,15 @@ 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 const zChatGPT4oAspectRatioID = z.enum(['3:2', '1:1', '2:3']);
+export const isChatGPT4oAspectRatioID = (v: unknown): v is z.infer<typeof zChatGPT4oAspectRatioID> =>
+  zChatGPT4oAspectRatioID.safeParse(v).success;
+
 export type AspectRatioID = z.infer<typeof zAspectRatioID>;
 export const isAspectRatioID = (v: unknown): v is AspectRatioID => zAspectRatioID.safeParse(v).success;
 

invokeai/frontend/web/src/features/dynamicPrompts/store/dynamicPromptsSlice.ts

@@ -4,7 +4,7 @@ import type { PersistConfig, RootState } from 'app/store/store';
 import { z } from 'zod';
 
 const zSeedBehaviour = z.enum(['PER_ITERATION', 'PER_PROMPT']);
-type SeedBehaviour = z.infer<typeof zSeedBehaviour>;
+export type SeedBehaviour = z.infer<typeof zSeedBehaviour>;
 export const isSeedBehaviour = (v: unknown): v is SeedBehaviour => zSeedBehaviour.safeParse(v).success;
 
 export interface DynamicPromptsState {

invokeai/frontend/web/src/features/nodes/components/sidePanel/workflow/publish.ts

@@ -123,6 +123,8 @@ const NODE_TYPE_PUBLISH_DENYLIST = [
   'metadata_to_t2i_adapters',
   'google_imagen3_generate',
   'google_imagen3_edit',
+  'chatgpt_create_image',
+  'chatgpt_edit_image',
 ];
 
 export const selectHasUnpublishableNodes = createSelector(selectNodes, (nodes) => {

invokeai/frontend/web/src/features/nodes/util/graph/buildLinearBatchConfig.ts

@@ -2,35 +2,57 @@ import { NUMPY_RAND_MAX, NUMPY_RAND_MIN } from 'app/constants';
 import type { RootState } from 'app/store/store';
 import { generateSeeds } from 'common/util/generateSeeds';
 import randomInt from 'common/util/randomInt';
+import type { SeedBehaviour } from 'features/dynamicPrompts/store/dynamicPromptsSlice';
+import type { ModelIdentifierField } from 'features/nodes/types/common';
 import type { FieldIdentifier } from 'features/nodes/types/field';
 import type { Graph } from 'features/nodes/util/graph/generation/Graph';
 import { range } from 'lodash-es';
 import type { components } from 'services/api/schema';
 import type { Batch, EnqueueBatchArg } from 'services/api/types';
+import { assert } from 'tsafe';
 
-export const prepareLinearUIBatch = (
-  state: RootState,
-  g: Graph,
-  prepend: boolean,
-  seedFieldIdentifier: FieldIdentifier,
-  positivePromptFieldIdentifier: FieldIdentifier,
-  origin: 'canvas' | 'workflows' | 'upscaling',
-  destination: 'canvas' | 'gallery'
-): EnqueueBatchArg => {
+const getExtendedPrompts = (arg: {
+  seedBehaviour: SeedBehaviour;
+  iterations: number;
+  prompts: string[];
+  model: ModelIdentifierField;
+}): string[] => {
+  const { seedBehaviour, iterations, prompts, model } = arg;
+  // Normally, the seed behaviour implicity determines the batch size. But when we use models without seeds (like
+  // ChatGPT 4o) in conjunction with the per-prompt seed behaviour, we lose out on that implicit batch size. To rectify
+  // this, we need to create a batch of the right size by repeating the prompts.
+  if (seedBehaviour === 'PER_PROMPT' || model.base === 'chatgpt-4o') {
+    return range(iterations).flatMap(() => prompts);
+  }
+  return prompts;
+};
+
+export const prepareLinearUIBatch = (arg: {
+  state: RootState;
+  g: Graph;
+  prepend: boolean;
+  seedFieldIdentifier?: FieldIdentifier;
+  positivePromptFieldIdentifier: FieldIdentifier;
+  origin: 'canvas' | 'workflows' | 'upscaling';
+  destination: 'canvas' | 'gallery';
+}): EnqueueBatchArg => {
+  const { state, g, prepend, seedFieldIdentifier, positivePromptFieldIdentifier, origin, destination } = arg;
   const { iterations, model, shouldRandomizeSeed, seed, shouldConcatPrompts } = state.params;
   const { prompts, seedBehaviour } = state.dynamicPrompts;
 
+  assert(model, 'No model found in state when preparing batch');
+
   const data: Batch['data'] = [];
   const firstBatchDatumList: components['schemas']['BatchDatum'][] = [];
   const secondBatchDatumList: components['schemas']['BatchDatum'][] = [];
 
   // add seeds first to ensure the output order groups the prompts
-  if (seedBehaviour === 'PER_PROMPT') {
+  if (seedFieldIdentifier && seedBehaviour === 'PER_PROMPT') {
     const seeds = generateSeeds({
       count: prompts.length * iterations,
-      // Imagen3's support for seeded generation is iffy, we are just not going too use in in linear UI generations.
+      // Imagen3's support for seeded generation is iffy, we are just not going too use it in linear UI generations.
       start:
-        model?.base === 'imagen3' ? randomInt(NUMPY_RAND_MIN, NUMPY_RAND_MAX) : shouldRandomizeSeed ? undefined : seed,
+        model.base === 'imagen3' ? randomInt(NUMPY_RAND_MIN, NUMPY_RAND_MAX) : shouldRandomizeSeed ? undefined : seed,
     });
 
     firstBatchDatumList.push({
@@ -46,13 +68,13 @@ export const prepareLinearUIBatch = (
       field_name: 'seed',
       items: seeds,
     });
-  } else {
+  } else if (seedFieldIdentifier && seedBehaviour === 'PER_ITERATION') {
     // seedBehaviour = SeedBehaviour.PerRun
     const seeds = generateSeeds({
       count: iterations,
       // Imagen3's support for seeded generation is iffy, we are just not going too use in in linear UI generations.
       start:
-        model?.base === 'imagen3' ? randomInt(NUMPY_RAND_MIN, NUMPY_RAND_MAX) : shouldRandomizeSeed ? undefined : seed,
+        model.base === 'imagen3' ? randomInt(NUMPY_RAND_MIN, NUMPY_RAND_MAX) : shouldRandomizeSeed ? undefined : seed,
     });
 
     secondBatchDatumList.push({
@@ -71,7 +93,7 @@ export const prepareLinearUIBatch = (
     data.push(secondBatchDatumList);
   }
 
-  const extendedPrompts = seedBehaviour === 'PER_PROMPT' ? range(iterations).flatMap(() => prompts) : prompts;
+  const extendedPrompts = getExtendedPrompts({ seedBehaviour, iterations, prompts, model });
 
   // zipped batch of prompts
   firstBatchDatumList.push({
@@ -88,7 +110,7 @@ export const prepareLinearUIBatch = (
     items: extendedPrompts,
   });
 
-  if (shouldConcatPrompts && model?.base === 'sdxl') {
+  if (shouldConcatPrompts && model.base === 'sdxl') {
     firstBatchDatumList.push({
       node_path: positivePromptFieldIdentifier.nodeId,
       field_name: 'style',