mirrors/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-09-12 12:37:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix memory usage issue with inspect

Browse Source
This commit is contained in:
Jedrzej Kosinski
2025-08-27 17:55:35 -07:00
parent 68b00e9c60
commit 29b7990dc2
1 changed files with 52 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

84
comfy/ldm/modules/attention.py
View File

@@ -142,38 +142,58 @@ def wrap_attn(func):
to_add = 1000 to_add = 1000
logged_stack = [] logged_stack = []
logged_stack_to_index = -1 logged_stack_to_index = -1
for frame_info in inspect.stack()[1:]:
if not continue_to_add: frame = inspect.currentframe()
break try:
if to_add == 0: # skip wrapper, start at actual wrapped function
break frame = frame.f_back
if frame_info.function == "_calc_cond_batch_outer":
break while frame and continue_to_add and to_add > 0:
if 'venv' in frame_info.filename: code = frame.f_code
continue filename = code.co_filename
elif 'ComfyUI' not in frame_info.filename: function = code.co_name
continue lineno = frame.f_lineno
elif 'execution.py' in frame_info.filename:
continue if function == "_calc_cond_batch_outer":
elif 'patcher_extension.py' in frame_info.filename: break
continue if 'venv' in filename:
to_add -= 1 frame = frame.f_back
cls_name = get_class_from_frame(frame_info.frame) continue
log_string = f"{frame_info.filename}:{frame_info.lineno}" elif 'ComfyUI' not in filename:
if cls_name: frame = frame.f_back
log_string += f":{cls_name}.{frame_info.function}" continue
else: elif 'execution.py' in filename:
log_string += f":{frame_info.function}" frame = frame.f_back
if has_transformer_options_passed(frame_info.frame): continue
log_string += ":✅" elif 'patcher_extension.py' in filename:
if logged_stack_to_index == -1: frame = frame.f_back
logged_stack_to_index = len(logged_stack) continue
else:
log_string += ":❌" to_add -= 1
logged_stack.append(log_string) cls_name = get_class_from_frame(frame)
# logging.info(f"Attn call stack: {logged_stack}") log_string = f"{filename}:{lineno}"
# logging.info(f"Logged stack to index: {logged_stack[:logged_stack_to_index+1]}") if cls_name:
LOG_CONTENTS["|".join(logged_stack)] = (logged_stack_to_index, logged_stack) log_string += f":{cls_name}.{function}"
else:
log_string += f":{function}"
if has_transformer_options_passed(frame):
log_string += ":✅"
if logged_stack_to_index == -1:
logged_stack_to_index = len(logged_stack)
else:
log_string += ":❌"
logged_stack.append(log_string)
# move up the stack
frame = frame.f_back
LOG_CONTENTS["|".join(logged_stack)] = (logged_stack_to_index, logged_stack)
finally:
# Important: break ref cycles so tensors aren't pinned
del frame
transformer_options = kwargs.pop("transformer_options", None) transformer_options = kwargs.pop("transformer_options", None)
if transformer_options is not None: if transformer_options is not None:
if "optimized_attention_override" in transformer_options: if "optimized_attention_override" in transformer_options: