alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-03-25 07:05:43 -03:00
Code Issues Packages Projects Releases Wiki Activity

feat: enhance async metadata collection by updating function signature and preserving all parameters. Fixes #328 #327

Browse Source
This commit is contained in:
Will Miao
2025-08-01 21:47:52 +08:00
parent 848c1741fe
commit c0eff2bb5e
1 changed files with 9 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
py/metadata_collector/metadata_hook.py
View File

@@ -146,52 +146,40 @@ class MetadataHook:
# Store the original _async_map_node_over_list function # Store the original _async_map_node_over_list function
original_map_node_over_list = getattr(execution, map_node_func_name) original_map_node_over_list = getattr(execution, map_node_func_name)
# Define the wrapped async function - NOTE: Updated signature with prompt_id and unique_id! # Wrapped async function, compatible with both stable and nightly
async def async_map_node_over_list_with_metadata(prompt_id, unique_id, obj, input_data_all, func, allow_interrupt=False, execution_block_cb=None, pre_execute_cb=None): async def async_map_node_over_list_with_metadata(prompt_id, unique_id, obj, input_data_all, func, allow_interrupt=False, execution_block_cb=None, pre_execute_cb=None, *args, **kwargs):
hidden_inputs = kwargs.get('hidden_inputs', None)
# Only collect metadata when calling the main function of nodes # Only collect metadata when calling the main function of nodes
if func == obj.FUNCTION and hasattr(obj, '__class__'): if func == obj.FUNCTION and hasattr(obj, '__class__'):
try: try:
# Get the current prompt_id from the registry
registry = MetadataRegistry() registry = MetadataRegistry()
# We now have prompt_id directly from the function parameters
if prompt_id is not None: if prompt_id is not None:
# Get node class type
class_type = obj.__class__.__name__ class_type = obj.__class__.__name__
# Use the passed unique_id parameter instead of trying to extract it
node_id = unique_id node_id = unique_id
# Record inputs before execution
if node_id is not None: if node_id is not None:
registry.record_node_execution(node_id, class_type, input_data_all, None) registry.record_node_execution(node_id, class_type, input_data_all, None)
except Exception as e: except Exception as e:
print(f"Error collecting metadata (pre-execution): {str(e)}") print(f"Error collecting metadata (pre-execution): {str(e)}")
# Execute the original async function with ALL parameters in the correct order # Call original function with all args/kwargs
results = await original_map_node_over_list(prompt_id, unique_id, obj, input_data_all, func, allow_interrupt, execution_block_cb, pre_execute_cb) results = await original_map_node_over_list(
prompt_id, unique_id, obj, input_data_all, func,
allow_interrupt, execution_block_cb, pre_execute_cb, *args, **kwargs
)
# After execution, collect outputs for relevant nodes
if func == obj.FUNCTION and hasattr(obj, '__class__'): if func == obj.FUNCTION and hasattr(obj, '__class__'):
try: try:
# Get the current prompt_id from the registry
registry = MetadataRegistry() registry = MetadataRegistry()
if prompt_id is not None: if prompt_id is not None:
# Get node class type
class_type = obj.__class__.__name__ class_type = obj.__class__.__name__
# Use the passed unique_id parameter
node_id = unique_id node_id = unique_id
# Record outputs after execution
if node_id is not None: if node_id is not None:
registry.update_node_execution(node_id, class_type, results) registry.update_node_execution(node_id, class_type, results)
except Exception as e: except Exception as e:
print(f"Error collecting metadata (post-execution): {str(e)}") print(f"Error collecting metadata (post-execution): {str(e)}")
return results return results
# Also hook the execute function to track the current prompt_id # Also hook the execute function to track the current prompt_id
original_execute = execution.execute original_execute = execution.execute