Update utils.py

utils.py

@@ -174,7 +174,7 @@ def save_attention_maps(attn_maps, tokenizer, prompts, base_dir='attn_maps', unc
             
     
     total_attn_map /= total_attn_map_number
-    final_attn_map = {}
+    final_attn_map = []
     for batch, (attn_map, tokens) in enumerate(zip(total_attn_map, total_tokens)):
         batch_dir = os.path.join(base_dir, f'batch-{batch}')
         if not os.path.exists(batch_dir):
@@ -198,6 +198,6 @@ def save_attention_maps(attn_maps, tokenizer, prompts, base_dir='attn_maps', unc
                     token = '-' + token + '-'
             
 
-            final_attn_map[f'{i}-{token}.png'] = to_pil(a.to(torch.float32))
+            final_attn_map.append((to_pil(a.to(torch.float32)), f'{i}-{token}'))
 
     return final_attn_map