From 8ceb94497e335d1263d85110b7d2e4287fae7b38 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Thu, 2 May 2024 08:56:13 +1000
Subject: [PATCH] fix(ui): fix canvas rendering of control images

---
 .../controlLayers/store/controlLayersSlice.ts |  3 +++
 .../features/controlLayers/util/renderers.ts  | 25 +++++++++----------
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/invokeai/frontend/web/src/features/controlLayers/store/controlLayersSlice.ts b/invokeai/frontend/web/src/features/controlLayers/store/controlLayersSlice.ts
index d4c43fdc89..13b6188692 100644
--- a/invokeai/frontend/web/src/features/controlLayers/store/controlLayersSlice.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/store/controlLayersSlice.ts
@@ -305,6 +305,9 @@ export const controlLayersSlice = createSlice({
       const { layerId, processorConfig } = action.payload;
       const layer = selectCALayerOrThrow(state, layerId);
       layer.controlAdapter.processorConfig = processorConfig;
+      if (!processorConfig) {
+        layer.controlAdapter.processedImage = null;
+      }
     },
     caLayerIsFilterEnabledChanged: (state, action: PayloadAction<{ layerId: string; isFilterEnabled: boolean }>) => {
       const { layerId, isFilterEnabled } = action.payload;
diff --git a/invokeai/frontend/web/src/features/controlLayers/util/renderers.ts b/invokeai/frontend/web/src/features/controlLayers/util/renderers.ts
index c79278b03d..85f48baa6e 100644
--- a/invokeai/frontend/web/src/features/controlLayers/util/renderers.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/util/renderers.ts
@@ -431,29 +431,30 @@ const updateControlNetLayerImageSource = async (
   konvaLayer: Konva.Layer,
   reduxLayer: ControlAdapterLayer
 ) => {
-  if (reduxLayer.controlAdapter.image) {
-    const { imageName } = reduxLayer.controlAdapter.image;
+  const image = reduxLayer.controlAdapter.processedImage ?? reduxLayer.controlAdapter.image;
+  if (image) {
+    const { imageName } = image;
     const req = getStore().dispatch(imagesApi.endpoints.getImageDTO.initiate(imageName));
     const imageDTO = await req.unwrap();
     req.unsubscribe();
-    const image = new Image();
+    const imageEl = new Image();
     const imageId = getCALayerImageId(reduxLayer.id, imageName);
-    image.onload = () => {
+    imageEl.onload = () => {
       // Find the existing image or create a new one - must find using the name, bc the id may have just changed
       const konvaImage =
-        konvaLayer.findOne<Konva.Image>(`.${CA_LAYER_IMAGE_NAME}`) ?? createControlNetLayerImage(konvaLayer, image);
+        konvaLayer.findOne<Konva.Image>(`.${CA_LAYER_IMAGE_NAME}`) ?? createControlNetLayerImage(konvaLayer, imageEl);
 
       // Update the image's attributes
       konvaImage.setAttrs({
         id: imageId,
-        image,
+        image: imageEl,
       });
       updateControlNetLayerImageAttrs(stage, konvaImage, reduxLayer);
       // Must cache after this to apply the filters
       konvaImage.cache();
-      image.id = imageId;
+      imageEl.id = imageId;
     };
-    image.src = imageDTO.image_url;
+    imageEl.src = imageDTO.image_url;
   } else {
     konvaLayer.findOne(`.${CA_LAYER_IMAGE_NAME}`)?.destroy();
   }
@@ -499,12 +500,10 @@ const renderControlNetLayer = (stage: Konva.Stage, reduxLayer: ControlAdapterLay
   const canvasImageSource = konvaImage?.image();
   let imageSourceNeedsUpdate = false;
   if (canvasImageSource instanceof HTMLImageElement) {
-    if (
-      reduxLayer.controlAdapter.image &&
-      canvasImageSource.id !== getCALayerImageId(reduxLayer.id, reduxLayer.controlAdapter.image.imageName)
-    ) {
+    const image = reduxLayer.controlAdapter.processedImage ?? reduxLayer.controlAdapter.image;
+    if (image && canvasImageSource.id !== getCALayerImageId(reduxLayer.id, image.imageName)) {
       imageSourceNeedsUpdate = true;
-    } else if (!reduxLayer.controlAdapter.image) {
+    } else if (!image) {
       imageSourceNeedsUpdate = true;
     }
   } else if (!canvasImageSource) {