diff --git a/py/__init__.py b/py/__init__.py
index 0db041a..b3bdc1c 100644
--- a/py/__init__.py
+++ b/py/__init__.py
@@ -1,6 +1,6 @@
 import functools
 
-from .hook import pre_execute, pre_recursive_execute
+from .hook import pre_execute, pre_get_input_data
 import execution
 
 
@@ -19,6 +19,4 @@ execution.PromptExecutor.execute = prefix_function(
 )
 
 
-execution.recursive_execute = prefix_function(
-    execution.recursive_execute, pre_recursive_execute
-)
+execution.get_input_data = prefix_function(execution.get_input_data, pre_get_input_data)
diff --git a/py/hook.py b/py/hook.py
index 61323a8..84421f3 100644
--- a/py/hook.py
+++ b/py/hook.py
@@ -1,7 +1,9 @@
+from .nodes.node import SaveImageWithMetaData
+
 current_prompt = {}
 current_extra_data = {}
 prompt_executer = None
-current_node_id = -1
+current_save_image_node_id = -1
 
 
 def pre_execute(self, prompt, prompt_id, extra_data, execute_outputs):
@@ -14,17 +16,8 @@ def pre_execute(self, prompt, prompt_id, extra_data, execute_outputs):
     prompt_executer = self
 
 
-def pre_recursive_execute(
-    server,
-    prompt,
-    outputs,
-    current_item,
-    extra_data,
-    executed,
-    prompt_id,
-    outputs_ui,
-    object_storage,
-):
-    global current_node_id
+def pre_get_input_data(inputs, class_def, unique_id, *args):
+    global current_save_image_node_id
 
-    current_node_id = current_item
+    if class_def == SaveImageWithMetaData:
+        current_save_image_node_id = unique_id
diff --git a/py/nodes/node.py b/py/nodes/node.py
index b2a3f09..75b8a96 100644
--- a/py/nodes/node.py
+++ b/py/nodes/node.py
@@ -108,7 +108,7 @@ class SaveImageWithMetaData(BaseNode):
 
         # get sampler node before this node
         trace_tree_from_this_node = Trace.trace(
-            hook.current_node_id, hook.current_prompt
+            hook.current_save_image_node_id, hook.current_prompt
         )
         inputs_before_this_node = Trace.filter_inputs_by_trace_tree(
             inputs, trace_tree_from_this_node