reset fix

models/region_diffusion.py

@@ -285,8 +285,10 @@ class RegionDiffusion(nn.Module):
         We reset attention maps because we append them while getting hooks
         to visualize attention maps for every step.
         """
-        for key in self.attention_maps:
-            self.attention_maps[key] = []
+        for key in self.selfattn_maps:
+            self.selfattn_maps[key] = []
+        for key in self.crossattn_maps:
+            self.crossattn_maps[key] = []
 
     def register_evaluation_hooks(self):
         r"""Function for registering hooks during evaluation.

utils/attention_utils.py

@@ -123,6 +123,7 @@ def plot_attention_maps(atten_map_list, obj_tokens, save_dir, seed, tokens_vis=N
                             dtype='uint8').reshape((height, width, 3))
 
         fig.tight_layout()
+        plt.close()
     return img