From 10fb99592801acc95c1f1af1539cbd1f19123e5a Mon Sep 17 00:00:00 2001
From: "Dr.Lt.Data" <dr.lt.data@gmail.com>
Date: Sat, 1 Jul 2023 10:43:59 +0900
Subject: [PATCH] object storage application fix2

---
 worklist_execution.py | 33 +++++++++++++++++++++------------
 1 file changed, 21 insertions(+), 12 deletions(-)

diff --git a/worklist_execution.py b/worklist_execution.py
index 8caccd842..9d9ad8940 100644
--- a/worklist_execution.py
+++ b/worklist_execution.py
@@ -125,6 +125,15 @@ def get_class_def(prompt, unique_id):
     return class_def
 
 
+def get_class_type(prompt, unique_id):
+    if 'class_type' in prompt[unique_id]:
+        class_type = prompt[unique_id]['class_type']
+    else:
+        class_type = "DummyNode"
+
+    return class_type
+
+
 def get_next_nodes_map(prompt):
     next_nodes = {}
     for key, value in prompt.items():
@@ -218,7 +227,7 @@ def worklist_execute(server, prompt, outputs, extra_data, prompt_id, outputs_ui,
         unique_id = get_work()
 
         inputs = prompt[unique_id]['inputs']
-        class_type = prompt[unique_id]['class_type']
+        class_type = get_class_type(prompt, unique_id)
         class_def = get_class_def(prompt, unique_id)
 
         print_dbg(lambda: f"work: {unique_id} ({class_def.__name__}) / worklist: {list(worklist.queue)}")
@@ -457,17 +466,17 @@ class PromptExecutor:
             d = self.outputs.pop(o)
             del d
 
-        # Next, remove the subsequent outputs since they will not be executed
-        to_delete = []
-        for o in self.outputs:
-            if (o not in current_outputs) and (o not in executed):
-                to_delete += [o]
-                if o in self.old_prompt:
-                    d = self.old_prompt.pop(o)
-                    del d
-        for o in to_delete:
-            d = self.outputs.pop(o)
-            del d
+            to_delete = []
+            for o in self.object_storage:
+                if o[0] not in prompt:
+                    to_delete += [o]
+                else:
+                    p = prompt[o[0]]
+                    if o[1] != p['class_type']:
+                        to_delete += [o]
+            for o in to_delete:
+                d = self.object_storage.pop(o)
+                del d
 
     def execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
         nodes.interrupt_processing(False)