From fb214d7b2f1d32905676f89d5100f1bdcfa44fcb Mon Sep 17 00:00:00 2001
From: CMonk <tianlang0704@gmail.com>
Date: Tue, 23 Apr 2024 05:56:01 +0800
Subject: [PATCH] Add history check to improve auto queue compatibility

---
 comfy/__init__.py                       | 33 +++++++++++++++++++++----
 photoshop/dist/index.js                 | 17 +++++++++++++
 photoshop/src/system/ComfyConnection.js | 21 ++++++++++++++++
 3 files changed, 66 insertions(+), 5 deletions(-)

diff --git a/comfy/__init__.py b/comfy/__init__.py
index d457b9b..6026eb0 100644
--- a/comfy/__init__.py
+++ b/comfy/__init__.py
@@ -61,6 +61,9 @@ class PhotoshopInstance:
         self.wsCallsManager = WSCallsManager(ws)
         self.destroyed = False
         self.layers = []
+        self.get_img_state_id = None
+        self.last_get_img_id = None
+        self.send_img_state_id = None
         
     async def poll_layers(self):
         await asyncio.sleep(1)
@@ -111,13 +114,33 @@ class PhotoshopInstance:
         return layers
         
     async def get_image(self, layer_id, use_layer_bounds=False):
+        history_state_id = await self.get_active_history_state_id()
+        if self.get_img_state_id == history_state_id:
+            return self.last_get_img_id
         result = await self.wsCallsManager.call('get_image', {'layer_id': layer_id, 'use_layer_bounds': use_layer_bounds}, timeout=60)
         id = result['upload_name']
+        self.get_img_state_id = history_state_id
+        self.last_get_img_id = id
         return id
     
+    def is_get_image_changed(self):
+        current_id = _invoke_async(self.get_active_history_state_id())
+        return self.get_img_state_id != current_id
+
     async def send_images(self, image_ids, layer_name=""):
+        history_state_id = await self.get_active_history_state_id()
+        if self.send_img_state_id == history_state_id:
+            return {}
         result = await self.wsCallsManager.call('send_images', {'image_ids': image_ids, 'layer_name': layer_name})
+        history_state_id = await self.get_active_history_state_id() # need to get new id after operation, it causes history change
+        self.send_img_state_id = history_state_id
+        self.get_img_state_id = history_state_id # it gets into a loop if get img state is not updated
         return result
+    
+    async def get_active_history_state_id(self):
+        result = await self.wsCallsManager.call('get_active_history_state_id', {})
+        id = result['history_state_id']
+        return id
             
     async def destroy(self):
         await self.wsCallsManager.ws.close()
@@ -231,8 +254,11 @@ class GetImageFromPhotoshopLayerNode:
         return True
     
     @classmethod
-    def IS_CHANGED(s, layer):
-        return np.random.rand()
+    def IS_CHANGED(self, layer, use_layer_bounds):
+        if (PhotoshopInstance.instance is None):
+            return np.random.rand()
+        else:
+            return PhotoshopInstance.instance.is_get_image_changed()
     
     @classmethod
     def INPUT_TYPES(cls):
@@ -278,9 +304,6 @@ class GetImageFromPhotoshopLayerNode:
         loadImage = LoadImage()
         (output_image, output_mask) = loadImage.load_image(image_id)
         
-        # image_out = ImageCache.data[id]
-        # ImageCache.data.pop(id)
-        
         return (output_image,)
     
 
diff --git a/photoshop/dist/index.js b/photoshop/dist/index.js
index 24fe7c6..6cdbbf5 100644
--- a/photoshop/dist/index.js
+++ b/photoshop/dist/index.js
@@ -413,6 +413,23 @@ class ComfyConnection {
         }, {
           commandName: "get content of layer " + layerID
         });
+      } else if (payload.action == 'get_active_history_state_id') {
+        try {
+          const historyStates = photoshop__WEBPACK_IMPORTED_MODULE_0__.app.activeDocument.historyStates;
+          const historyState = historyStates[historyStates.length - 1];
+          this.socket.send(JSON.stringify({
+            call_id: payload.call_id,
+            result: {
+              history_state_id: historyState.id
+            }
+          }));
+        } catch (e) {
+          this.socket.send(JSON.stringify({
+            call_id: payload.call_id,
+            error: e.message
+          }));
+          throw e;
+        }
       }
     } catch (e) {
       console.error(e);
diff --git a/photoshop/src/system/ComfyConnection.js b/photoshop/src/system/ComfyConnection.js
index 011ec57..f542272 100644
--- a/photoshop/src/system/ComfyConnection.js
+++ b/photoshop/src/system/ComfyConnection.js
@@ -278,6 +278,27 @@ class ComfyConnection {
                     )
                 }, { commandName: "get content of layer " + layerID })
 
+            } else if (payload.action == 'get_active_history_state_id') {
+                try {
+                    const historyStates = app.activeDocument.historyStates;
+                    const historyState = historyStates[historyStates.length - 1];
+                    this.socket.send(
+                        JSON.stringify({
+                            call_id: payload.call_id,
+                            result: { 
+                                history_state_id: historyState.id,
+                            }
+                        })
+                    )
+                } catch (e) {
+                    this.socket.send(
+                        JSON.stringify({
+                            call_id: payload.call_id,
+                            error: e.message
+                        })
+                    )
+                    throw e
+                }
             }
         } catch (e) { console.error(e) }
     }