diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 3c9df9a00..804fd0df9 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -132,114 +132,12 @@ class FeedForward(nn.Module): def Normalize(in_channels, dtype=None, device=None): return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True, dtype=dtype, device=device) -import inspect -LOG_ATTN_CALLS = False -LOG_CONTENTS = {} - -def save_log_contents(): - import folder_paths - output_dir = folder_paths.get_output_directory() - - # Create attn_logs directory if it doesn't exist - attn_logs_dir = os.path.join(output_dir, "attn_logs") - os.makedirs(attn_logs_dir, exist_ok=True) - - # Generate timestamp filename (down to second) - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - filename = f"{timestamp}.json" - filepath = os.path.join(attn_logs_dir, filename) - - # Save LOG_CONTENTS as JSON file - try: - with open(filepath, 'w', encoding='utf-8') as f: - json.dump(list(LOG_CONTENTS.values()), f, indent=2, ensure_ascii=False) - logging.info(f"Saved attention log contents to {filepath}") - except Exception as e: - logging.error(f"Failed to save attention log contents: {e}") - -def get_class_from_frame(frame): - # Check for 'self' (instance method) or 'cls' (classmethod) - if 'self' in frame.f_locals: - return frame.f_locals['self'].__class__.__name__ - elif 'cls' in frame.f_locals: - return frame.f_locals['cls'].__name__ - return None - -def has_transformer_options_passed(frame): - if 'transformer_options' in frame.f_locals.keys(): - if frame.f_locals['transformer_options']: - return True - return False def wrap_attn(func): @functools.wraps(func) def wrapper(*args, **kwargs): remove_attn_wrapper_key = False try: - if LOG_ATTN_CALLS: - continue_to_add = True - to_add = 1000 - logged_stack = [] - logged_stack_to_index = -1 - - frame = inspect.currentframe() - try: - # skip wrapper, start at actual wrapped function - frame = frame.f_back - - while frame and continue_to_add and to_add > 0: - code = frame.f_code - filename = code.co_filename - function = code.co_name - lineno = frame.f_lineno - - if function == "_calc_cond_batch_outer": - break - if 'venv' in filename: - frame = frame.f_back - continue - elif 'ComfyUI' not in filename: - frame = frame.f_back - continue - elif 'execution.py' in filename: - frame = frame.f_back - continue - elif 'patcher_extension.py' in filename: - frame = frame.f_back - continue - - to_add -= 1 - cls_name = get_class_from_frame(frame) - log_string = f"{filename}:{lineno}" - if cls_name: - 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 - - # check if we get what we want from transformer_options - t_check = "❌❌❌" - transformer_options = kwargs.get("transformer_options", None) - if transformer_options is not None: - if "optimized_attention_override" in transformer_options: - t_check = "✅✅✅" - - LOG_CONTENTS["|".join(logged_stack)] = (t_check, logged_stack_to_index, logged_stack) - - finally: - # Important: break ref cycles so tensors aren't pinned - del frame if "_inside_attn_wrapper" not in kwargs: transformer_options = kwargs.get("transformer_options", None) remove_attn_wrapper_key = True diff --git a/comfy/samplers.py b/comfy/samplers.py index 7775bcada..b3202cec6 100755 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -1019,7 +1019,6 @@ class CFGGuider: self.conds[k] = list(map(lambda a: a.copy(), self.original_conds[k])) preprocess_conds_hooks(self.conds) - import comfy.ldm.modules.attention #TODO: Remove this $$$$$ try: orig_model_options = self.model_options self.model_options = comfy.model_patcher.create_model_options_clone(self.model_options) @@ -1034,23 +1033,12 @@ class CFGGuider: self, comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.OUTER_SAMPLE, self.model_options, is_model_options=True) ) - comfy.ldm.modules.attention.LOG_ATTN_CALLS = False #TODO: Remove this $$$$$ - comfy.ldm.modules.attention.LOG_CONTENTS = {} - if "optimized_attention_override" not in self.model_options["transformer_options"]: - def optimized_attention_override(func, *args, **kwargs): - return func(*args, **kwargs) - self.model_options["transformer_options"]["optimized_attention_override"] = optimized_attention_override - output = executor.execute(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed) finally: cast_to_load_options(self.model_options, device=self.model_patcher.offload_device) self.model_options = orig_model_options self.model_patcher.hook_mode = orig_hook_mode self.model_patcher.restore_hook_patches() - if comfy.ldm.modules.attention.LOG_ATTN_CALLS: - comfy.ldm.modules.attention.save_log_contents() - comfy.ldm.modules.attention.LOG_ATTN_CALLS = False #TODO: Remove this $$$$$ - comfy.ldm.modules.attention.LOG_CONTENTS = {} del self.conds return output