Fix Gradient Mask (Canvas) Inpaint Model Support

invokeai/app/invocations/latent.py

@@ -51,6 +51,7 @@
 from invokeai.backend.ip_adapter.ip_adapter import IPAdapter, IPAdapterPlus
 from invokeai.backend.lora import LoRAModelRaw
 from invokeai.backend.model_manager import BaseModelType, LoadedModel
+from invokeai.backend.model_manager.config import MainConfigBase, ModelVariantType
 from invokeai.backend.model_patcher import ModelPatcher
 from invokeai.backend.stable_diffusion import PipelineIntermediateState, set_seamless
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
@@ -185,7 +186,7 @@ class GradientMaskOutput(BaseInvocationOutput):
     title="Create Gradient Mask",
     tags=["mask", "denoise"],
     category="latents",
-    version="1.0.0",
+    version="1.1.0",
 )
 class CreateGradientMaskInvocation(BaseInvocation):
     """Creates mask for denoising model run."""
@@ -198,6 +199,32 @@ class CreateGradientMaskInvocation(BaseInvocation):
     minimum_denoise: float = InputField(
         default=0.0, ge=0, le=1, description="Minimum denoise level for the coherence region", ui_order=4
     )
+    image: Optional[ImageField] = InputField(
+        default=None,
+        description="OPTIONAL: Only connect for specialized Inpainting models, masked_latents will be generated from the image with the VAE",
+        title="[OPTIONAL] Image",
+        ui_order=6,
+    )
+    unet: Optional[UNetField] = InputField(
+        description="OPTIONAL: If the Unet is a specialized Inpainting model, masked_latents will be generated from the image with the VAE",
+        default=None,
+        input=Input.Connection,
+        title="[OPTIONAL] UNet",
+        ui_order=5,
+    )
+    vae: Optional[VAEField] = InputField(
+        default=None,
+        description="OPTIONAL: Only connect for specialized Inpainting models, masked_latents will be generated from the image with the VAE",
+        title="[OPTIONAL] VAE",
+        input=Input.Connection,
+        ui_order=7,
+    )
+    tiled: bool = InputField(default=False, description=FieldDescriptions.tiled, ui_order=8)
+    fp32: bool = InputField(
+        default=DEFAULT_PRECISION == "float32",
+        description=FieldDescriptions.fp32,
+        ui_order=9,
+    )
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> GradientMaskOutput:
@@ -233,8 +260,27 @@ def invoke(self, context: InvocationContext) -> GradientMaskOutput:
         expanded_mask_image = Image.fromarray((expanded_mask.squeeze(0).numpy() * 255).astype(np.uint8), mode="L")
         expanded_image_dto = context.images.save(expanded_mask_image)
 
+        masked_latents_name = None
+        if self.unet is not None and self.vae is not None and self.image is not None:
+            # all three fields must be present at the same time
+            main_model_config = context.models.get_config(self.unet.unet.key)
+            assert isinstance(main_model_config, MainConfigBase)
+            if main_model_config.variant is ModelVariantType.Inpaint:
+                mask = blur_tensor
+                vae_info: LoadedModel = context.models.load(self.vae.vae)
+                image = context.images.get_pil(self.image.image_name)
+                image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB"))
+                if image_tensor.dim() == 3:
+                    image_tensor = image_tensor.unsqueeze(0)
+                img_mask = tv_resize(mask, image_tensor.shape[-2:], T.InterpolationMode.BILINEAR, antialias=False)
+                masked_image = image_tensor * torch.where(img_mask < 0.5, 0.0, 1.0)
+                masked_latents = ImageToLatentsInvocation.vae_encode(
+                    vae_info, self.fp32, self.tiled, masked_image.clone()
+                )
+                masked_latents_name = context.tensors.save(tensor=masked_latents)
+
         return GradientMaskOutput(
-            denoise_mask=DenoiseMaskField(mask_name=mask_name, masked_latents_name=None, gradient=True),
+            denoise_mask=DenoiseMaskField(mask_name=mask_name, masked_latents_name=masked_latents_name, gradient=True),
             expanded_mask_area=ImageField(image_name=expanded_image_dto.image_name),
         )
 

invokeai/backend/model_manager/config.py

@@ -301,12 +301,12 @@ class MainConfigBase(ModelConfigBase):
     default_settings: Optional[MainModelDefaultSettings] = Field(
         description="Default settings for this model", default=None
     )
+    variant: ModelVariantType = ModelVariantType.Normal
 
 
 class MainCheckpointConfig(CheckpointConfigBase, MainConfigBase):
     """Model config for main checkpoint models."""
 
-    variant: ModelVariantType = ModelVariantType.Normal
     prediction_type: SchedulerPredictionType = SchedulerPredictionType.Epsilon
     upcast_attention: bool = False
 

invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ControlNetOrT2IAdapterDefaultSettings/ControlNetOrT2IAdapterDefaultSettings.tsx

@@ -1,4 +1,4 @@
-import { Button, Flex, Heading, Text } from '@invoke-ai/ui-library';
+import { Button, Flex, Heading, SimpleGrid, Text } from '@invoke-ai/ui-library';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { useControlNetOrT2IAdapterDefaultSettings } from 'features/modelManagerV2/hooks/useControlNetOrT2IAdapterDefaultSettings';
 import { DefaultPreprocessor } from 'features/modelManagerV2/subpanels/ModelPanel/ControlNetOrT2IAdapterDefaultSettings/DefaultPreprocessor';
@@ -92,13 +92,9 @@ export const ControlNetOrT2IAdapterDefaultSettings = () => {
         </Button>
       </Flex>
 
-      <Flex flexDir="column" gap={8}>
-        <Flex gap={8}>
-          <Flex gap={4} w="full">
-            <DefaultPreprocessor control={control} name="preprocessor" />
-          </Flex>
-        </Flex>
-      </Flex>
+      <SimpleGrid columns={2} gap={8}>
+        <DefaultPreprocessor control={control} name="preprocessor" />
+      </SimpleGrid>
     </>
   );
 };

invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/MainModelDefaultSettings/MainModelDefaultSettings.tsx

@@ -1,4 +1,4 @@
-import { Button, Flex, Heading, Text } from '@invoke-ai/ui-library';
+import { Button, Flex, Heading, SimpleGrid, Text } from '@invoke-ai/ui-library';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { useMainModelDefaultSettings } from 'features/modelManagerV2/hooks/useMainModelDefaultSettings';
 import { DefaultHeight } from 'features/modelManagerV2/subpanels/ModelPanel/MainModelDefaultSettings/DefaultHeight';
@@ -122,40 +122,16 @@ export const MainModelDefaultSettings = () => {
         </Button>
       </Flex>
 
-      <Flex flexDir="column" gap={8}>
-        <Flex gap={8}>
-          <Flex gap={4} w="full">
-            <DefaultVae control={control} name="vae" />
-          </Flex>
-          <Flex gap={4} w="full">
-            <DefaultVaePrecision control={control} name="vaePrecision" />
-          </Flex>
-        </Flex>
-        <Flex gap={8}>
-          <Flex gap={4} w="full">
-            <DefaultScheduler control={control} name="scheduler" />
-          </Flex>
-          <Flex gap={4} w="full">
-            <DefaultSteps control={control} name="steps" />
-          </Flex>
-        </Flex>
-        <Flex gap={8}>
-          <Flex gap={4} w="full">
-            <DefaultCfgScale control={control} name="cfgScale" />
-          </Flex>
-          <Flex gap={4} w="full">
-            <DefaultCfgRescaleMultiplier control={control} name="cfgRescaleMultiplier" />
-          </Flex>
-        </Flex>
-        <Flex gap={8}>
-          <Flex gap={4} w="full">
-            <DefaultWidth control={control} optimalDimension={optimalDimension} />
-          </Flex>
-          <Flex gap={4} w="full">
-            <DefaultHeight control={control} optimalDimension={optimalDimension} />
-          </Flex>
-        </Flex>
-      </Flex>
+      <SimpleGrid columns={2} gap={8}>
+        <DefaultVae control={control} name="vae" />
+        <DefaultVaePrecision control={control} name="vaePrecision" />
+        <DefaultScheduler control={control} name="scheduler" />
+        <DefaultSteps control={control} name="steps" />
+        <DefaultCfgScale control={control} name="cfgScale" />
+        <DefaultCfgRescaleMultiplier control={control} name="cfgRescaleMultiplier" />
+        <DefaultWidth control={control} optimalDimension={optimalDimension} />
+        <DefaultHeight control={control} optimalDimension={optimalDimension} />
+      </SimpleGrid>
     </>
   );
 };

invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelEdit.tsx

@@ -6,6 +6,7 @@ import {
   FormLabel,
   Heading,
   Input,
+  SimpleGrid,
   Text,
   Textarea,
 } from '@invoke-ai/ui-library';
@@ -66,25 +67,21 @@ export const ModelEdit = ({ form }: Props) => {
           <Heading as="h3" fontSize="md" mt="4">
             {t('modelManager.modelSettings')}
           </Heading>
-          <Flex gap={4}>
+          <SimpleGrid columns={2} gap={4}>
             <FormControl flexDir="column" alignItems="flex-start" gap={1}>
               <FormLabel>{t('modelManager.baseModel')}</FormLabel>
               <BaseModelSelect control={form.control} />
             </FormControl>
-          </Flex>
-          {data.type === 'main' && data.format === 'checkpoint' && (
-            <>
-              <Flex gap={4}>
+            <FormControl flexDir="column" alignItems="flex-start" gap={1}>
+              <FormLabel>{t('modelManager.variant')}</FormLabel>
+              <ModelVariantSelect control={form.control} />
+            </FormControl>
+            {data.type === 'main' && data.format === 'checkpoint' && (
+              <>
                 <FormControl flexDir="column" alignItems="flex-start" gap={1}>
                   <FormLabel>{t('modelManager.pathToConfig')}</FormLabel>
                   <Input {...form.register('config_path', stringFieldOptions)} />
                 </FormControl>
-                <FormControl flexDir="column" alignItems="flex-start" gap={1}>
-                  <FormLabel>{t('modelManager.variant')}</FormLabel>
-                  <ModelVariantSelect control={form.control} />
-                </FormControl>
-              </Flex>
-              <Flex gap={4}>
                 <FormControl flexDir="column" alignItems="flex-start" gap={1}>
                   <FormLabel>{t('modelManager.predictionType')}</FormLabel>
                   <PredictionTypeSelect control={form.control} />
@@ -93,9 +90,9 @@ export const ModelEdit = ({ form }: Props) => {
                   <FormLabel>{t('modelManager.upcastAttention')}</FormLabel>
                   <Checkbox {...form.register('upcast_attention')} />
                 </FormControl>
-              </Flex>
-            </>
-          )}
+              </>
+            )}
+          </SimpleGrid>
         </Flex>
       </form>
     </Flex>

invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelView.tsx

@@ -1,4 +1,4 @@
-import { Box, Flex, Text } from '@invoke-ai/ui-library';
+import { Box, Flex, SimpleGrid, Text } from '@invoke-ai/ui-library';
 import { skipToken } from '@reduxjs/toolkit/query';
 import { useAppSelector } from 'app/store/storeHooks';
 import { ControlNetOrT2IAdapterDefaultSettings } from 'features/modelManagerV2/subpanels/ModelPanel/ControlNetOrT2IAdapterDefaultSettings/ControlNetOrT2IAdapterDefaultSettings';
@@ -24,57 +24,32 @@ export const ModelView = () => {
   return (
     <Flex flexDir="column" h="full" gap={4}>
       <Box layerStyle="second" borderRadius="base" p={4}>
-        <Flex flexDir="column" gap={4}>
-          <Flex gap={2}>
-            <ModelAttrView label={t('modelManager.baseModel')} value={data.base} />
-            <ModelAttrView label={t('modelManager.modelType')} value={data.type} />
-          </Flex>
-          <Flex gap={2}>
-            <ModelAttrView label={t('common.format')} value={data.format} />
-            <ModelAttrView label={t('modelManager.path')} value={data.path} />
-          </Flex>
-
+        <SimpleGrid columns={2} gap={4}>
+          <ModelAttrView label={t('modelManager.baseModel')} value={data.base} />
+          <ModelAttrView label={t('modelManager.modelType')} value={data.type} />
+          <ModelAttrView label={t('common.format')} value={data.format} />
+          <ModelAttrView label={t('modelManager.path')} value={data.path} />
+          {data.type === 'main' && <ModelAttrView label={t('modelManager.variant')} value={data.variant} />}
           {data.type === 'main' && data.format === 'diffusers' && data.repo_variant && (
-            <Flex gap={2}>
-              <ModelAttrView label={t('modelManager.repoVariant')} value={data.repo_variant} />
-            </Flex>
+            <ModelAttrView label={t('modelManager.repoVariant')} value={data.repo_variant} />
           )}
-
           {data.type === 'main' && data.format === 'checkpoint' && (
             <>
-              <Flex gap={2}>
-                <ModelAttrView label={t('modelManager.pathToConfig')} value={data.config_path} />
-                <ModelAttrView label={t('modelManager.variant')} value={data.variant} />
-              </Flex>
-              <Flex gap={2}>
-                <ModelAttrView label={t('modelManager.predictionType')} value={data.prediction_type} />
-                <ModelAttrView label={t('modelManager.upcastAttention')} value={`${data.upcast_attention}`} />
-              </Flex>
+              <ModelAttrView label={t('modelManager.pathToConfig')} value={data.config_path} />
+              <ModelAttrView label={t('modelManager.predictionType')} value={data.prediction_type} />
+              <ModelAttrView label={t('modelManager.upcastAttention')} value={`${data.upcast_attention}`} />
             </>
           )}
-
           {data.type === 'ip_adapter' && data.format === 'invokeai' && (
-            <Flex gap={2}>
-              <ModelAttrView label={t('modelManager.imageEncoderModelId')} value={data.image_encoder_model_id} />
-            </Flex>
+            <ModelAttrView label={t('modelManager.imageEncoderModelId')} value={data.image_encoder_model_id} />
           )}
-        </Flex>
+        </SimpleGrid>
+      </Box>
+      <Box layerStyle="second" borderRadius="base" p={4}>
+        {data.type === 'main' && data.base !== 'sdxl-refiner' && <MainModelDefaultSettings />}
+        {(data.type === 'controlnet' || data.type === 't2i_adapter') && <ControlNetOrT2IAdapterDefaultSettings />}
+        {(data.type === 'main' || data.type === 'lora') && <TriggerPhrases />}
       </Box>
-      {data.type === 'main' && data.base !== 'sdxl-refiner' && (
-        <Box layerStyle="second" borderRadius="base" p={4}>
-          <MainModelDefaultSettings />
-        </Box>
-      )}
-      {(data.type === 'controlnet' || data.type === 't2i_adapter') && (
-        <Box layerStyle="second" borderRadius="base" p={4}>
-          <ControlNetOrT2IAdapterDefaultSettings />
-        </Box>
-      )}
-      {(data.type === 'main' || data.type === 'lora') && (
-        <Box layerStyle="second" borderRadius="base" p={4}>
-          <TriggerPhrases />
-        </Box>
-      )}
     </Flex>
   );
 };

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

@@ -9,6 +9,7 @@ import {
   CANVAS_TEXT_TO_IMAGE_GRAPH,
   IMAGE_TO_IMAGE_GRAPH,
   IMAGE_TO_LATENTS,
+  INPAINT_CREATE_MASK,
   INPAINT_IMAGE,
   LATENTS_TO_IMAGE,
   MAIN_MODEL_LOADER,
@@ -145,6 +146,16 @@ export const addVAEToGraph = async (
           field: 'vae',
         },
       },
+      {
+        source: {
+          node_id: isSeamlessEnabled ? SEAMLESS : isAutoVae ? modelLoaderNodeId : VAE_LOADER,
+          field: 'vae',
+        },
+        destination: {
+          node_id: INPAINT_CREATE_MASK,
+          field: 'vae',
+        },
+      },
 
       {
         source: {

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

@@ -133,6 +133,8 @@ export const buildCanvasInpaintGraph = async (
         coherence_mode: canvasCoherenceMode,
         minimum_denoise: canvasCoherenceMinDenoise,
         edge_radius: canvasCoherenceEdgeSize,
+        tiled: false,
+        fp32: fp32,
       },
       [DENOISE_LATENTS]: {
         type: 'denoise_latents',
@@ -182,6 +184,16 @@ export const buildCanvasInpaintGraph = async (
           field: 'clip',
         },
       },
+      {
+        source: {
+          node_id: modelLoaderNodeId,
+          field: 'unet',
+        },
+        destination: {
+          node_id: INPAINT_CREATE_MASK,
+          field: 'unet',
+        },
+      },
       // Connect CLIP Skip to Conditioning
       {
         source: {
@@ -331,6 +343,16 @@ export const buildCanvasInpaintGraph = async (
           field: 'mask',
         },
       },
+      {
+        source: {
+          node_id: INPAINT_IMAGE_RESIZE_UP,
+          field: 'image',
+        },
+        destination: {
+          node_id: INPAINT_CREATE_MASK,
+          field: 'image',
+        },
+      },
       // Resize Down
       {
         source: {