diff --git a/scripts/daam/trace.py b/scripts/daam/trace.py
index 687ed82..5d1d819 100644
--- a/scripts/daam/trace.py
+++ b/scripts/daam/trace.py
@@ -46,10 +46,10 @@ def _forward(hk_self, self, *args, **kwargs):
 
 
 class HeatMap:
-    def __init__(self, prompt_analyzer: PromptAnalyzer, prompt: str, heat_maps: torch.Tensor):
-        self.prompt_analyzer = prompt_analyzer.create(prompt)
+    def __init__(self, prompt_analyzer: PromptAnalyzer, heat_maps: torch.Tensor):
+        self.prompt_analyzer = prompt_analyzer
         self.heat_maps = heat_maps
-        self.prompt = prompt
+        self.prompt = prompt_analyzer.prompt
 
     def compute_word_heat_map(self, word: str, word_idx: int = None) -> torch.Tensor:
         merge_idxs, _ = self.prompt_analyzer.calc_word_indecies(word)
@@ -182,7 +182,7 @@ def compute_global_heat_map(self, prompt_analyzer, prompt, batch_index, time_wei
         maps = torch.stack([torch.stack(x, 0) for x in all_merges], dim=0)
         maps = maps.sum(0).to(device).sum(2).sum(0)
 
-        return HeatMap(prompt_analyzer, prompt, maps)
+        return HeatMap(prompt_analyzer, maps)
 
 
 class UNetCrossAttentionHooker(ObjectHooker[CrossAttention]):
diff --git a/scripts/daam/utils.py b/scripts/daam/utils.py
index c271d52..0bda9b3 100644
--- a/scripts/daam/utils.py
+++ b/scripts/daam/utils.py
@@ -4,6 +4,7 @@
 from pathlib import Path
 import random
 import re
+import sys
 from typing import Union
 
 from PIL import Image, ImageFont, ImageDraw
@@ -276,6 +277,7 @@ def __init__(self, clip : Union[FrozenCLIPEmbedderWithCustomWordsBase, GeneralCo
         self.used_custom_terms = []
         self.hijack_comments = []
 
+        self.prompt = text
         chunks, token_count = self.tokenize_line(text)
 
         self.token_count = token_count