From 3d34bfafd5cba3814216b35d7a24071486b65ccd Mon Sep 17 00:00:00 2001
From: Dariusz L <dl122@wp.pl>
Date: Sun, 3 Aug 2025 22:14:39 +0200
Subject: [PATCH] Fix matting: refresh image after background removal

Fixed an issue where images were not immediately refreshed after background removal (matting). Now, the canvas updates instantly when the background is removed, ensuring correct display without requiring manual scaling or other actions.
---
 js/CanvasLayers.js  | 14 +++++---------
 js/CanvasView.js    |  2 ++
 src/CanvasLayers.ts | 18 +++++++-----------
 src/CanvasView.ts   |  4 ++++
 4 files changed, 18 insertions(+), 20 deletions(-)

diff --git a/js/CanvasLayers.js b/js/CanvasLayers.js
index 5e86ca9..425913e 100644
--- a/js/CanvasLayers.js
+++ b/js/CanvasLayers.js
@@ -961,17 +961,13 @@ export class CanvasLayers {
     }
     async getLayerImageData(layer) {
         try {
-            const { canvas: tempCanvas, ctx: tempCtx } = createCanvas(layer.width, layer.height, '2d', { willReadFrequently: true });
+            const width = layer.originalWidth || layer.width;
+            const height = layer.originalHeight || layer.height;
+            const { canvas: tempCanvas, ctx: tempCtx } = createCanvas(width, height, '2d', { willReadFrequently: true });
             if (!tempCtx)
                 throw new Error("Could not create canvas context");
-            // We need to draw the layer relative to the new canvas, so we "move" it to 0,0
-            // by creating a temporary layer object for drawing.
-            const layerToDraw = {
-                ...layer,
-                x: 0,
-                y: 0,
-            };
-            this._drawLayer(tempCtx, layerToDraw);
+            // Use original image directly to ensure full quality
+            tempCtx.drawImage(layer.image, 0, 0, width, height);
             const dataUrl = tempCanvas.toDataURL('image/png');
             if (!dataUrl.startsWith('data:image/png;base64,')) {
                 throw new Error("Invalid image data format");
diff --git a/js/CanvasView.js b/js/CanvasView.js
index 6ed769f..27a9a7a 100644
--- a/js/CanvasView.js
+++ b/js/CanvasView.js
@@ -430,6 +430,8 @@ async function createCanvasWidget(node, widget, app) {
                             delete newLayer.imageId;
                             canvas.layers[selectedLayerIndex] = newLayer;
                             canvas.canvasSelection.updateSelection([newLayer]);
+                            // Invalidate processed image cache when layer image changes (matting)
+                            canvas.canvasLayers.invalidateProcessedImageCache(newLayer.id);
                             canvas.render();
                             canvas.saveState();
                             showSuccessNotification("Background removed successfully!");
diff --git a/src/CanvasLayers.ts b/src/CanvasLayers.ts
index 632ffbe..dee496a 100644
--- a/src/CanvasLayers.ts
+++ b/src/CanvasLayers.ts
@@ -1095,19 +1095,15 @@ export class CanvasLayers {
 
     async getLayerImageData(layer: Layer): Promise<string> {
         try {
-            const { canvas: tempCanvas, ctx: tempCtx } = createCanvas(layer.width, layer.height, '2d', { willReadFrequently: true });
+            const width = layer.originalWidth || layer.width;
+            const height = layer.originalHeight || layer.height;
+
+            const { canvas: tempCanvas, ctx: tempCtx } = createCanvas(width, height, '2d', { willReadFrequently: true });
             if (!tempCtx) throw new Error("Could not create canvas context");
 
-            // We need to draw the layer relative to the new canvas, so we "move" it to 0,0
-            // by creating a temporary layer object for drawing.
-            const layerToDraw = {
-                ...layer,
-                x: 0,
-                y: 0,
-            };
-
-            this._drawLayer(tempCtx, layerToDraw);
-
+            // Use original image directly to ensure full quality
+            tempCtx.drawImage(layer.image, 0, 0, width, height);
+            
             const dataUrl = tempCanvas.toDataURL('image/png');
             if (!dataUrl.startsWith('data:image/png;base64,')) {
                 throw new Error("Invalid image data format");
diff --git a/src/CanvasView.ts b/src/CanvasView.ts
index 7aea28d..51868e0 100644
--- a/src/CanvasView.ts
+++ b/src/CanvasView.ts
@@ -515,6 +515,10 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp):
                             
                             canvas.layers[selectedLayerIndex] = newLayer;
                             canvas.canvasSelection.updateSelection([newLayer]);
+                            
+                            // Invalidate processed image cache when layer image changes (matting)
+                            canvas.canvasLayers.invalidateProcessedImageCache(newLayer.id);
+                            
                             canvas.render();
                             canvas.saveState();
                             showSuccessNotification("Background removed successfully!");