From 1120e3dee53916888312ed1e3c70bfa9654e42af Mon Sep 17 00:00:00 2001 From: shentanyue Date: Wed, 12 Mar 2025 14:44:33 +0800 Subject: [PATCH] fix: use weak references for _instances to prevent memory leaks --- attention_processor_advanced.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/attention_processor_advanced.py b/attention_processor_advanced.py index bf5fca4..1d24872 100644 --- a/attention_processor_advanced.py +++ b/attention_processor_advanced.py @@ -1,11 +1,12 @@ import torch import torch.nn as nn import torch.nn.functional as F +import weakref from diffusers.models.normalization import RMSNorm from einops import rearrange class IPAFluxAttnProcessor2_0Advanced(nn.Module): - _instances = set() + _instances = weakref.WeakSet() _global_call_count = 0 _last_timestep_printed = None _first_instance_for_timestep = None # Add this line @@ -61,6 +62,7 @@ class IPAFluxAttnProcessor2_0Advanced(nn.Module): # print(f"Steps and seen timesteps have been reset for this instance.") def __del__(self): + self.clear_memory() # Remove this instance from the set when it's deleted self.__class__._instances.remove(self)