refactor: unify image cache contract

js/Canvas.js

@@ -11,6 +11,7 @@ import { CanvasRenderer } from "./CanvasRenderer.js";
 import { CanvasIO } from "./CanvasIO.js";
 import { ImageReferenceManager } from "./ImageReferenceManager.js";
 import { BatchPreviewManager } from "./BatchPreviewManager.js";
+import { ImageCache } from "./ImageCache.js";
 import { createModuleLogger } from "./log_system/log_funcs.js";
 import { debounce, createCanvas } from "./utils/CommonUtils.js";
 import { MaskEditorIntegration } from "./MaskEditorIntegration.js";
@@ -75,7 +76,7 @@ export class Canvas {
         this.pendingDataCheck = null;
         this.pendingInputDataCheck = null;
         this.inputDataLoaded = false;
-        this.imageCache = new Map();
+        this.imageCache = new ImageCache();
         this.requestSaveState = () => { };
         this.outputAreaShape = null;
         this.autoApplyShapeMask = false;

js/CanvasState.js

@@ -142,9 +142,8 @@ export class CanvasState {
         log.debug(`Layer ${index}: Loading image with id: ${layerData.imageId}`);
         if (this.canvas.imageCache.has(layerData.imageId)) {
             log.debug(`Layer ${index}: Image found in cache.`);
-            const imageData = this.canvas.imageCache.get(layerData.imageId);
-            if (imageData) {
-                const imageSrc = URL.createObjectURL(new Blob([imageData.data]));
+            const imageSrc = this.canvas.imageCache.get(layerData.imageId);
+            if (imageSrc) {
                 this._createLayerFromSrc(layerData, imageSrc, index, resolve);
             }
             else {

js/CanvasView.js

@@ -7,7 +7,6 @@ import { $el } from "../../scripts/ui.js";
 import { addStylesheet, getUrl, loadTemplate } from "./utils/ResourceManager.js";
 import { Canvas } from "./Canvas.js";
 import { clearAllCanvasStates, getCanvasState, setCanvasState } from "./db.js";
-import { ImageCache } from "./ImageCache.js";
 import { createCanvas } from "./utils/CommonUtils.js";
 import { createModuleLogger } from "./log_system/log_funcs.js";
 import { showErrorNotification, showSuccessNotification, showInfoNotification, showWarningNotification } from "./utils/NotificationUtils.js";
@@ -109,7 +108,6 @@ async function createCanvasWidget(node, widget, app) {
     const canvas = new Canvas(node, widget, {
         onStateChange: () => updateOutput(node, canvas)
     });
-    const imageCache = new ImageCache();
     /**
      * Helper function to update the icon of a switch component.
      * @param knobIconEl The HTML element for the switch's knob icon.

js/ImageCache.js

@@ -4,18 +4,21 @@ export class ImageCache {
     constructor() {
         this.cache = new Map();
     }
-    set(key, imageData) {
-        log.info("Caching image data for key:", key);
-        this.cache.set(key, imageData);
+    set(key, image) {
+        log.info("Caching image for key:", key);
+        this.cache.set(key, image);
     }
     get(key) {
-        const data = this.cache.get(key);
-        log.debug("Retrieved cached data for key:", key, !!data);
-        return data;
+        const image = this.cache.get(key);
+        log.debug("Retrieved cached image for key:", key, !!image);
+        return image;
     }
     has(key) {
         return this.cache.has(key);
     }
+    delete(key) {
+        return this.cache.delete(key);
+    }
     clear() {
         log.info("Clearing image cache");
         this.cache.clear();

src/Canvas.ts

@@ -17,6 +17,7 @@ import {CanvasRenderer} from "./CanvasRenderer.js";
 import {CanvasIO} from "./CanvasIO.js";
 import {ImageReferenceManager} from "./ImageReferenceManager.js";
 import {BatchPreviewManager} from "./BatchPreviewManager.js";
+import {ImageCache} from "./ImageCache.js";
 import {createModuleLogger} from "./log_system/log_funcs.js";
 import { debounce, createCanvas } from "./utils/CommonUtils.js";
 import {MaskEditorIntegration} from "./MaskEditorIntegration.js";
@@ -59,7 +60,7 @@ export class Canvas {
     ctx: CanvasRenderingContext2D;
     dataInitialized: boolean;
     height: number;
-    imageCache: Map<string, any>;
+    imageCache: ImageCache;
     imageReferenceManager: ImageReferenceManager;
     interaction: any;
     isMouseOver: boolean;
@@ -145,7 +146,7 @@ export class Canvas {
         this.pendingDataCheck = null;
         this.pendingInputDataCheck = null;
         this.inputDataLoaded = false;
-        this.imageCache = new Map();
+        this.imageCache = new ImageCache();
 
         this.requestSaveState = () => {};
         this.outputAreaShape = null;

src/CanvasState.ts

@@ -177,9 +177,8 @@ export class CanvasState {
 
         if (this.canvas.imageCache.has(layerData.imageId)) {
             log.debug(`Layer ${index}: Image found in cache.`);
-            const imageData = this.canvas.imageCache.get(layerData.imageId);
-            if (imageData) {
-                const imageSrc = URL.createObjectURL(new Blob([imageData.data]));
+            const imageSrc = this.canvas.imageCache.get(layerData.imageId);
+            if (imageSrc) {
                 this._createLayerFromSrc(layerData, imageSrc, index, resolve);
             } else {
                 resolve(null);

src/CanvasView.ts

@@ -13,7 +13,6 @@ import { addStylesheet, getUrl, loadTemplate } from "./utils/ResourceManager.js"
 
 import {Canvas} from "./Canvas.js";
 import {clearAllCanvasStates, getCanvasState, setCanvasState} from "./db.js";
-import {ImageCache} from "./ImageCache.js";
 import {generateUniqueFileName, createCanvas} from "./utils/CommonUtils.js";
 import {createModuleLogger} from "./log_system/log_funcs.js";
 import {showErrorNotification, showSuccessNotification, showInfoNotification, showWarningNotification} from "./utils/NotificationUtils.js";
@@ -160,7 +159,6 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp):
     const canvas = new Canvas(node, widget, {
         onStateChange: () => updateOutput(node, canvas)
     });
-    const imageCache = new ImageCache();
 
     /**
      * Helper function to update the icon of a switch component.

src/ImageCache.ts

@@ -1,30 +1,34 @@
 import {createModuleLogger} from "./log_system/log_funcs.js";
-import type { ImageDataPixel } from './types';
+import type { CachedImage, ImageCacheContract } from './types';
 
 const log = createModuleLogger('ImageCache');
 
-export class ImageCache {
-    private cache: Map<string, ImageDataPixel>;
+export class ImageCache implements ImageCacheContract {
+    private cache: Map<string, CachedImage>;
 
     constructor() {
         this.cache = new Map();
     }
 
-    set(key: string, imageData: ImageDataPixel): void {
-        log.info("Caching image data for key:", key);
-        this.cache.set(key, imageData);
+    set(key: string, image: CachedImage): void {
+        log.info("Caching image for key:", key);
+        this.cache.set(key, image);
     }
 
-    get(key: string): ImageDataPixel | undefined {
-        const data = this.cache.get(key);
-        log.debug("Retrieved cached data for key:", key, !!data);
-        return data;
+    get(key: string): CachedImage | undefined {
+        const image = this.cache.get(key);
+        log.debug("Retrieved cached image for key:", key, !!image);
+        return image;
     }
 
     has(key: string): boolean {
         return this.cache.has(key);
     }
 
+    delete(key: string): boolean {
+        return this.cache.delete(key);
+    }
+
     clear(): void {
         log.info("Clearing image cache");
         this.cache.clear();

src/types.ts

@@ -178,6 +178,16 @@ export interface Tensor {
     height: number;
 }
 
+export type CachedImage = string | ImageBitmap;
+
+export interface ImageCacheContract {
+    get(key: string): CachedImage | undefined;
+    set(key: string, image: CachedImage): void;
+    has(key: string): boolean;
+    delete(key: string): boolean;
+    clear(): void;
+}
+
 export interface ImageDataPixel {
     data: Uint8ClampedArray;
     width: number;

tests/test_image_cache_characterization.mjs

@@ -0,0 +1,44 @@
+import assert from 'node:assert/strict';
+import { readFile } from 'node:fs/promises';
+import test from 'node:test';
+
+import { ImageCache } from '../js/ImageCache.js';
+
+const sourceFiles = await Promise.all([
+  readFile(new URL('../src/ImageCache.ts', import.meta.url), 'utf8'),
+  readFile(new URL('../src/Canvas.ts', import.meta.url), 'utf8'),
+  readFile(new URL('../src/CanvasLayers.ts', import.meta.url), 'utf8'),
+  readFile(new URL('../src/CanvasState.ts', import.meta.url), 'utf8'),
+  readFile(new URL('../src/CanvasView.ts', import.meta.url), 'utf8'),
+]);
+
+test('image cache preserves key/value storage and clearing behavior', () => {
+  const cache = new ImageCache();
+  const imageSource = 'data:image/png;base64,encoded-image';
+
+  cache.set('image-1', imageSource);
+
+  assert.equal(cache.has('image-1'), true);
+  assert.equal(cache.get('image-1'), imageSource);
+  assert.equal(cache.delete('image-1'), true);
+  assert.equal(cache.has('image-1'), false);
+
+  cache.set('image-1', imageSource);
+
+  cache.clear();
+
+  assert.equal(cache.has('image-1'), false);
+  assert.equal(cache.get('image-1'), undefined);
+});
+
+test('layer persistence uses one image cache contract', () => {
+  const [cacheSource, canvasSource, layersSource, stateSource, viewSource] = sourceFiles;
+
+  assert.match(cacheSource, /private cache: Map<string, CachedImage>/);
+  assert.match(canvasSource, /imageCache: ImageCache/);
+  assert.match(canvasSource, /this\.imageCache = new ImageCache\(\)/);
+  assert.match(layersSource, /imageCache\.set\(imageId, image\.src\)/);
+  assert.match(stateSource, /const imageSrc = this\.canvas\.imageCache\.get\(layerData\.imageId\)/);
+  assert.doesNotMatch(stateSource, /imageData\.data/);
+  assert.doesNotMatch(viewSource, /new ImageCache\(\)/);
+});