diff --git a/.pylintrc b/.pylintrc
index 5bd9493dd..7b6cdb7cb 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -44,6 +44,7 @@ ignore-paths=/usr/lib/.*$,
              pipelines/segmoe,
              scripts/consistory,
              scripts/ctrlx,
+             scripts/daam,
              scripts/demofusion,
              scripts/freescale,
              scripts/infiniteyou,
diff --git a/.ruff.toml b/.ruff.toml
index 74b734878..297515402 100644
--- a/.ruff.toml
+++ b/.ruff.toml
@@ -24,6 +24,7 @@ exclude = [
     "pipelines/segmoe",
 
     "scripts/lbm",
+    "scripts/daam",
     "scripts/xadapter",
     "scripts/pulid",
     "scripts/instantir",
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 05680394d..bcfa7def8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -21,10 +21,14 @@
     enable in *settings -> pipeline modifers -> cache-dit*  
   - [Nunchaku Flux.1 PulID](https://nunchaku.tech/docs/nunchaku/python_api/nunchaku.pipeline.pipeline_flux_pulid.html)  
     automatically enabled if loaded model is FLUX.1 with Nunchaku engine enabled and when PulID script is enabled  
-- **Extensions**
-  - **agent-scheduler** was a high-value built-in extension, but it has not been maintained for 1.5 years  
+- **Extensions**  
+  - [Agent-Scheduler](https://github.com/SipherAGI/sd-webui-agent-scheduler)  
+    was a high-value built-in extension, but it has not been maintained for 1.5 years  
     it also does not work with control and video tabs which are the core of sdnext nowadays  
     so it has been removed from built-in extensions: manual installation is still possible  
+  - [DAAM: Diffusion Attentive Attribution Maps](https://github.com/castorini/daam)  
+    create heatmap visualizations of which parts of the prompt influenced which parts of the image  
+    available in scripts for sdxl text-to-image workflows  
 - **Offloading**
   - improve offloading for pipelines with multiple stages such as *wan-2.2-14b*  
   - add timers to measure onload/offload times during generate  
diff --git a/modules/modular.py b/modules/modular.py
index b2b6a8124..595b860f7 100644
--- a/modules/modular.py
+++ b/modules/modular.py
@@ -57,3 +57,4 @@ def restore_standard(modular_pipe):
     if hasattr(modular_pipe, 'original_pipe'):
         shared.log.debug(f'Modular: source={modular_pipe.__class__.__name__} target={modular_pipe.original_pipe.__class__.__name__}')
         return modular_pipe.original_pipe
+    return modular_pipe
diff --git a/modules/processing_args.py b/modules/processing_args.py
index 7770f210d..c1f0b9cdb 100644
--- a/modules/processing_args.py
+++ b/modules/processing_args.py
@@ -18,8 +18,8 @@ debug_log = shared.log.trace if debug_enabled else lambda *args, **kwargs: None
 disable_pbar = os.environ.get('SD_DISABLE_PBAR', None) is not None
 
 
-def task_modular_kwargs(p, model):
-    model_cls = model.__class__.__name__ # pylint: disable=unused-variable
+def task_modular_kwargs(p, model): # pylint: disable=unused-argument
+    # model_cls = model.__class__.__name__
     task_args = {}
     p.ops.append('modular')
 
diff --git a/scripts/daam/__init__.py b/scripts/daam/__init__.py
new file mode 100644
index 000000000..a407cfd4d
--- /dev/null
+++ b/scripts/daam/__init__.py
@@ -0,0 +1,5 @@
+from .experiment import *
+from .heatmap import *
+from .hook import *
+from .utils import *
+from .trace import *
diff --git a/scripts/daam/evaluate.py b/scripts/daam/evaluate.py
new file mode 100644
index 000000000..ec9e1567b
--- /dev/null
+++ b/scripts/daam/evaluate.py
@@ -0,0 +1,122 @@
+from collections import defaultdict
+from typing import List, Union
+
+from scipy.optimize import linear_sum_assignment
+import PIL.Image as Image
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+__all__ = ['compute_iou', 'MeanEvaluator', 'load_mask', 'compute_ioa']
+
+
+def compute_iou(a: torch.Tensor, b: torch.Tensor) -> float:
+    if a.shape[0] != b.shape[0]:
+        a = F.interpolate(a.unsqueeze(0).unsqueeze(0).float(), size=b.shape, mode='bicubic').squeeze()
+        a[a < 1] = 0
+        a[a >= 1] = 1
+
+    intersection = (a * b).sum()
+    union = a.sum() + b.sum() - intersection
+
+    return (intersection / (union + 1e-8)).item()
+
+
+def compute_ioa(a: torch.Tensor, b: torch.Tensor) -> float:
+    if a.shape[0] != b.shape[0]:
+        a = F.interpolate(a.unsqueeze(0).unsqueeze(0).float(), size=b.shape, mode='bicubic').squeeze()
+        a[a < 1] = 0
+        a[a >= 1] = 1
+
+    intersection = (a * b).sum()
+    area = a.sum()
+
+    return (intersection / (area + 1e-8)).item()
+
+
+def load_mask(path: str) -> torch.Tensor:
+    mask = np.array(Image.open(path))
+    mask = torch.from_numpy(mask).float()[:, :, 3]  # use alpha channel
+    mask = (mask > 0).float()
+
+    return mask
+
+
+class UnsupervisedEvaluator:
+    def __init__(self, name: str = 'UnsupervisedEvaluator'):
+        self.name = name
+        self.ious = defaultdict(list)
+        self.num_samples = 0
+
+    def log_iou(self, preds: Union[torch.Tensor, List[torch.Tensor]], truth: torch.Tensor, gt_idx: int = 0, pred_idx: int = 0):
+        if not isinstance(preds, list):
+            preds = [preds]
+
+        iou = max(compute_iou(pred, truth) for pred in preds)
+        self.ious[gt_idx].append((pred_idx, iou))
+
+    @property
+    def mean_iou(self) -> float:
+        n = max(max(self.ious), max([y[0] for x in self.ious.values() for y in x])) + 1
+        iou_matrix = np.zeros((n, n))
+        count_matrix = np.zeros((n, n))
+
+        for gt_idx, ious in self.ious.items():
+            for pred_idx, iou in ious:
+                iou_matrix[gt_idx, pred_idx] += iou
+                count_matrix[gt_idx, pred_idx] += 1
+
+        row_ind, col_ind = linear_sum_assignment(iou_matrix, maximize=True)
+        return iou_matrix[row_ind, col_ind].sum() / count_matrix[row_ind, col_ind].sum()
+
+    def increment(self):
+        self.num_samples += 1
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def __str__(self):
+        return f'{self.name}<{self.mean_iou:.4f} (mIoU) {len(self)} samples>'
+
+
+class MeanEvaluator:
+    def __init__(self, name: str = 'MeanEvaluator'):
+        self.ious: List[float] = []
+        self.intensities: List[float] = []
+        self.name = name
+
+    def log_iou(self, preds: Union[torch.Tensor, List[torch.Tensor]], truth: torch.Tensor):
+        if not isinstance(preds, list):
+            preds = [preds]
+
+        self.ious.append(max(compute_iou(pred, truth) for pred in preds))
+        return self
+
+    def log_intensity(self, pred: torch.Tensor):
+        self.intensities.append(pred.mean().item())
+        return self
+
+    @property
+    def mean_iou(self) -> float:
+        return np.mean(self.ious)
+
+    @property
+    def mean_intensity(self) -> float:
+        return np.mean(self.intensities)
+
+    @property
+    def ci95_miou(self) -> float:
+        return 1.96 * np.std(self.ious) / np.sqrt(len(self.ious))
+
+    def __len__(self) -> int:
+        return max(len(self.ious), len(self.intensities))
+
+    def __str__(self):
+        return f'{self.name}<{self.mean_iou:.4f} (±{self.ci95_miou:.3f} mIoU) {self.mean_intensity:.4f} (mInt) {len(self)} samples>'
+
+
+if __name__ == '__main__':
+    mask = load_mask('truth/output/452/sink.gt.png')
+
+    print(MeanEvaluator().log_iou(mask, mask))
diff --git a/scripts/daam/experiment.py b/scripts/daam/experiment.py
new file mode 100644
index 000000000..4465a3054
--- /dev/null
+++ b/scripts/daam/experiment.py
@@ -0,0 +1,344 @@
+from pathlib import Path
+from typing import List, Optional, Dict, Any, Union
+from dataclasses import dataclass
+import json
+
+from transformers import PreTrainedTokenizer, AutoTokenizer
+import PIL.Image
+import numpy as np
+import torch
+
+from .utils import auto_autocast
+from .evaluate import load_mask
+
+
+__all__ = ['GenerationExperiment', 'COCO80_LABELS', 'COCOSTUFF27_LABELS', 'COCO80_INDICES', 'build_word_list_coco80']
+
+
+COCO80_LABELS: List[str] = [
+    'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
+    'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
+    'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
+    'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
+    'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
+    'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
+    'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
+    'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
+    'hair drier', 'toothbrush'
+]
+
+COCO80_INDICES: Dict[str, int] = {x: i for i, x in enumerate(COCO80_LABELS)}
+
+UNUSED_LABELS: List[str] = [f'__unused_{i}__' for i in range(1, 200)]
+
+COCOSTUFF27_LABELS: List[str] = [
+    'electronic', 'appliance', 'food', 'furniture', 'indoor', 'kitchen', 'accessory', 'animal', 'outdoor', 'person',
+    'sports', 'vehicle', 'ceiling', 'floor', 'food', 'furniture', 'rawmaterial', 'textile', 'wall', 'window',
+    'building', 'ground', 'plant', 'sky', 'solid', 'structural', 'water'
+]
+
+COCO80_ONTOLOGY = {
+    'two-wheeled vehicle': ['bicycle', 'motorcycle'],
+    'vehicle': ['two-wheeled vehicle', 'four-wheeled vehicle'],
+    'four-wheeled vehicle': ['bus', 'truck', 'car'],
+    'four-legged animals': ['livestock', 'pets', 'wild animals'],
+    'livestock': ['cow', 'horse', 'sheep'],
+    'pets': ['cat', 'dog'],
+    'wild animals': ['elephant', 'bear', 'zebra', 'giraffe'],
+    'bags': ['backpack', 'handbag', 'suitcase'],
+    'sports boards': ['snowboard', 'surfboard', 'skateboard'],
+    'utensils': ['fork', 'knife', 'spoon'],
+    'receptacles': ['bowl', 'cup'],
+    'fruits': ['banana', 'apple', 'orange'],
+    'foods': ['fruits', 'meals', 'desserts'],
+    'meals': ['sandwich', 'hot dog', 'pizza'],
+    'desserts': ['cake', 'donut'],
+    'furniture': ['chair', 'couch', 'bench'],
+    'electronics': ['monitors', 'appliances'],
+    'monitors': ['tv', 'cell phone', 'laptop'],
+    'appliances': ['oven', 'toaster', 'refrigerator']
+}
+
+COCO80_TO_27 = {
+    'bicycle': 'vehicle', 'car': 'vehicle', 'motorcycle': 'vehicle', 'airplane': 'vehicle', 'bus': 'vehicle',
+    'train': 'vehicle', 'truck': 'vehicle', 'boat': 'vehicle', 'traffic light': 'accessory', 'fire hydrant': 'accessory',
+    'stop sign': 'accessory', 'parking meter': 'accessory', 'bench': 'furniture', 'bird': 'animal', 'cat': 'animal',
+    'dog': 'animal', 'horse': 'animal', 'sheep': 'animal', 'cow': 'animal', 'elephant': 'animal', 'bear': 'animal',
+    'zebra': 'animal', 'giraffe': 'animal', 'backpack': 'accessory', 'umbrella': 'accessory', 'handbag': 'accessory',
+    'tie': 'accessory', 'suitcase': 'accessory', 'frisbee': 'sports', 'skis': 'sports', 'snowboard': 'sports',
+    'sports ball': 'sports', 'kite': 'sports', 'baseball bat': 'sports', 'baseball glove': 'sports',
+    'skateboard': 'sports', 'surfboard': 'sports', 'tennis racket': 'sports', 'bottle': 'food', 'wine glass': 'food',
+    'cup': 'food', 'fork': 'food', 'knife': 'food', 'spoon': 'food', 'bowl': 'food', 'banana': 'food', 'apple': 'food',
+    'sandwich': 'food', 'orange': 'food', 'broccoli': 'food', 'carrot': 'food', 'hot dog': 'food', 'pizza': 'food',
+    'donut': 'food', 'cake': 'food', 'chair': 'furniture', 'couch': 'furniture', 'potted plant': 'plant',
+    'bed': 'furniture', 'dining table': 'furniture', 'toilet': 'furniture', 'tv': 'electronic', 'laptop': 'electronic',
+    'mouse': 'electronic', 'remote': 'electronic', 'keyboard': 'electronic', 'cell phone': 'electronic',
+    'microwave': 'appliance', 'oven': 'appliance', 'toaster': 'appliance', 'sink': 'appliance',
+    'refrigerator': 'appliance', 'book': 'indoor', 'clock': 'indoor', 'vase': 'indoor', 'scissors': 'indoor',
+    'teddy bear': 'indoor', 'hair drier': 'indoor', 'toothbrush': 'indoor'
+}
+
+
+def build_word_list_coco80() -> Dict[str, List[str]]:
+    words_map = COCO80_ONTOLOGY.copy()
+    words_map = {k: v for k, v in words_map.items() if not any(item in COCO80_ONTOLOGY for item in v)}
+
+    return words_map
+
+
+def _add_mask(masks: Dict[str, torch.Tensor], word: str, mask: torch.Tensor, simplify80: bool = False) -> Dict[str, torch.Tensor]:
+    if simplify80:
+        word = COCO80_TO_27.get(word, word)
+
+    if word in masks:
+        masks[word] = masks[word.lower()] + mask
+        masks[word].clamp_(0, 1)
+    else:
+        masks[word] = mask
+
+    return masks
+
+
+@dataclass
+class GenerationExperiment:
+    """Class to hold experiment parameters. Pickleable."""
+    image: PIL.Image.Image
+    global_heat_map: torch.Tensor
+    prompt: str
+
+    seed: int = None
+    id: str = '.'
+    path: Optional[Path] = None
+
+    truth_masks: Optional[Dict[str, torch.Tensor]] = None
+    prediction_masks: Optional[Dict[str, torch.Tensor]] = None
+    annotations: Optional[Dict[str, Any]] = None
+    subtype: Optional[str] = '.'
+    tokenizer: AutoTokenizer = None
+
+    def __post_init__(self):
+        if isinstance(self.path, str):
+            self.path = Path(self.path)
+
+        self.path = None if self.path is None else self.path / self.id
+
+    def nsfw(self) -> bool:
+        return np.sum(np.array(self.image)) == 0
+
+    def heat_map(self, tokenizer: AutoTokenizer = None):
+        if tokenizer is None:
+            tokenizer = self.tokenizer
+
+        from daam import GlobalHeatMap
+        return GlobalHeatMap(tokenizer, self.prompt, self.global_heat_map)
+
+    def clear_checkpoint(self):
+        path = self if isinstance(self, Path) else self.path
+
+        (path / 'generation.pt').unlink(missing_ok=True)
+
+    def save(self, path: str = None, heat_maps: bool = True, tokenizer: AutoTokenizer = None):
+        if path is None:
+            path = self.path
+        else:
+            path = Path(path) / self.id
+
+        if tokenizer is None:
+            tokenizer = self.tokenizer
+
+        (path / self.subtype).mkdir(parents=True, exist_ok=True)
+        torch.save(self, path / self.subtype / 'generation.pt')
+        self.image.save(path / self.subtype / 'output.png')
+
+        with (path / 'prompt.txt').open('w') as f:
+            f.write(self.prompt)
+
+        with (path / 'seed.txt').open('w') as f:
+            f.write(str(self.seed))
+
+        if self.truth_masks is not None:
+            for name, mask in self.truth_masks.items():
+                im = PIL.Image.fromarray((mask * 255).unsqueeze(-1).expand(-1, -1, 4).byte().numpy())
+                im.save(path / f'{name.lower()}.gt.png')
+
+        if heat_maps and tokenizer is not None:
+            self.save_all_heat_maps(tokenizer)
+
+        self.save_annotations()
+
+    def save_annotations(self, path: Path = None):
+        if path is None:
+            path = self.path
+
+        if self.annotations is not None:
+            with (path / 'annotations.json').open('w') as f:
+                json.dump(self.annotations, f)
+
+    def _load_truth_masks(self, simplify80: bool = False) -> Dict[str, torch.Tensor]:
+        masks = {}
+
+        for mask_path in self.path.glob('*.gt.png'):
+            word = mask_path.name.split('.gt.png')[0].lower()
+            mask = load_mask(str(mask_path))
+            _add_mask(masks, word, mask, simplify80)
+
+        return masks
+
+    def _load_pred_masks(self, pred_prefix, composite=False, simplify80=False, vocab=None):
+        # type: (str, bool, bool, List[str] | None) -> Dict[str, torch.Tensor]
+        masks = {}
+
+        if vocab is None:
+            vocab = UNUSED_LABELS
+
+        if composite:
+            try:
+                im = PIL.Image.open(self.path / self.subtype / f'composite.{pred_prefix}.pred.png')
+                im = np.array(im)
+
+                for mask_idx in np.unique(im):
+                    mask = torch.from_numpy((im == mask_idx).astype(np.float32))
+                    _add_mask(masks, vocab[mask_idx], mask, simplify80)
+            except FileNotFoundError:
+                pass
+        else:
+            for mask_path in (self.path / self.subtype).glob(f'*.{pred_prefix}.pred.png'):
+                mask = load_mask(str(mask_path))
+                word = mask_path.name.split(f'.{pred_prefix}.pred')[0].lower()
+                _add_mask(masks, word, mask, simplify80)
+
+        return masks
+
+    def clear_prediction_masks(self, name: str):
+        path = self if isinstance(self, Path) else self.path
+        path = path / self.subtype
+
+        for mask_path in path.glob(f'*.{name}.pred.png'):
+            mask_path.unlink()
+
+    def save_prediction_mask(self, mask: torch.Tensor, word: str, name: str):
+        path = self if isinstance(self, Path) else self.path
+        im = PIL.Image.fromarray((mask * 255).unsqueeze(-1).expand(-1, -1, 4).cpu().byte().numpy())
+        im.save(path / self.subtype / f'{word.lower()}.{name}.pred.png')
+
+    def save_heat_map(
+            self,
+            word: str,
+            tokenizer: PreTrainedTokenizer = None,
+            crop: int = None,
+            output_prefix: str = '',
+            absolute: bool = False
+    ) -> Path:
+        from .trace import GlobalHeatMap  # because of cyclical import
+
+        if tokenizer is None:
+            tokenizer = self.tokenizer
+
+        with auto_autocast(dtype=torch.float32):
+            path = self.path / self.subtype / f'{output_prefix}{word.lower()}.heat_map.png'
+            heat_map = GlobalHeatMap(tokenizer, self.prompt, self.global_heat_map)
+            heat_map.compute_word_heat_map(word).expand_as(self.image, color_normalize=not absolute, out_file=path, plot=True)
+
+        return path
+
+    def save_all_heat_maps(self, tokenizer: PreTrainedTokenizer = None, crop: int = None) -> Dict[str, Path]:
+        path_map = {}
+
+        if tokenizer is None:
+            tokenizer = self.tokenizer
+
+        for word in self.prompt.split(' '):
+            try:
+                path = self.save_heat_map(word, tokenizer, crop=crop)
+                path_map[word] = path
+            except:
+                pass
+
+        return path_map
+
+    @staticmethod
+    def contains_truth_mask(path: Union[str, Path], prompt_id: str = None) -> bool:
+        if prompt_id is None:
+            return any(Path(path).glob('*.gt.png'))
+        else:
+            return any((Path(path) / prompt_id).glob('*.gt.png'))
+
+    @staticmethod
+    def read_seed(path: Union[str, Path], prompt_id: str = None) -> int:
+        if prompt_id is None:
+            return int(Path(path).joinpath('seed.txt').read_text())
+        else:
+            return int(Path(path).joinpath(prompt_id).joinpath('seed.txt').read_text())
+
+    @staticmethod
+    def has_annotations(path: Union[str, Path]) -> bool:
+        return Path(path).joinpath('annotations.json').exists()
+
+    @staticmethod
+    def has_experiment(path: Union[str, Path], prompt_id: str) -> bool:
+        return (Path(path) / prompt_id / 'generation.pt').exists()
+
+    @staticmethod
+    def read_prompt(path: Union[str, Path], prompt_id: str = None) -> str:
+        if prompt_id is None:
+            prompt_id = '.'
+
+        with (Path(path) / prompt_id / 'prompt.txt').open('r') as f:
+            return f.read().strip()
+
+    def _try_load_annotations(self):
+        if not (self.path / 'annotations.json').exists():
+            return None
+
+        return json.load((self.path / 'annotations.json').open())
+
+    def annotate(self, key: str, value: Any) -> 'GenerationExperiment':
+        if self.annotations is None:
+            self.annotations = {}
+
+        self.annotations[key] = value
+
+        return self
+
+    @classmethod
+    def load(
+            cls,
+            path,
+            pred_prefix='daam',
+            composite=False,
+            simplify80=False,
+            vocab=None,
+            subtype='.',
+            all_subtypes=False
+    ):
+        # type: (str, str, bool, bool, List[str] | None, str, bool) -> GenerationExperiment | List[GenerationExperiment]
+        if all_subtypes:
+            experiments = []
+
+            for directory in Path(path).iterdir():
+                if not directory.is_dir():
+                    continue
+
+                try:
+                    experiments.append(cls.load(
+                        path,
+                        pred_prefix=pred_prefix,
+                        composite=composite,
+                        simplify80=simplify80,
+                        vocab=vocab,
+                        subtype=directory.name
+                    ))
+                except:
+                    pass
+
+            return experiments
+
+        path = Path(path)
+        exp = torch.load(path / subtype / 'generation.pt')
+        exp.subtype = subtype
+        exp.path = path
+        exp.truth_masks = exp._load_truth_masks(simplify80=simplify80)
+        exp.prediction_masks = exp._load_pred_masks(pred_prefix, composite=composite, simplify80=simplify80, vocab=vocab)
+        exp.annotations = exp._try_load_annotations()
+
+        return exp
diff --git a/scripts/daam/heatmap.py b/scripts/daam/heatmap.py
new file mode 100644
index 000000000..0f7a311f0
--- /dev/null
+++ b/scripts/daam/heatmap.py
@@ -0,0 +1,185 @@
+import io
+from collections import defaultdict
+from dataclasses import dataclass
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Dict, Tuple, Set, Iterable
+
+from matplotlib import pyplot as plt
+import numpy as np
+import PIL.Image
+import spacy.tokens
+import torch
+import torch.nn.functional as F
+
+from .evaluate import compute_ioa
+from .utils import compute_token_merge_indices, cached_nlp, auto_autocast
+
+__all__ = ['GlobalHeatMap', 'RawHeatMapCollection', 'WordHeatMap', 'ParsedHeatMap', 'SyntacticHeatMapPair']
+
+
+def plot_overlay_heat_map(im, heat_map, word=None, out_file=None, crop=None, color_normalize=True, ax=None, cmap='jet'):
+    # type: (PIL.Image.Image | np.ndarray, torch.Tensor, str, Path, int, bool, plt.Axes) -> None
+    if ax is None:
+        plt.rcParams['font.size'] = 16
+        plt.rcParams['figure.facecolor'] = 'black'
+        plt.rcParams['text.color'] = 'white'
+        plt.rcParams['axes.labelcolor'] = 'white'
+        plt.rcParams['xtick.color'] = 'black'
+        plt.rcParams['ytick.color'] = 'black'
+        plt.clf()
+        plt_ = plt
+    else:
+        plt_ = ax
+
+    with auto_autocast(dtype=torch.float32):
+        im = np.array(im)
+
+        if crop is not None:
+            heat_map = heat_map.squeeze()[crop:-crop, crop:-crop]
+            im = im[crop:-crop, crop:-crop]
+
+        if color_normalize:
+            plt_.imshow(heat_map.squeeze().cpu().numpy(), cmap=cmap)
+        else:
+            heat_map = heat_map.clamp_(min=0, max=1)
+            plt_.imshow(heat_map.squeeze().cpu().numpy(), cmap=cmap, vmin=0.0, vmax=1.0)
+
+        im = torch.from_numpy(im).float() / 255
+        im = torch.cat((im, (1 - heat_map.unsqueeze(-1))), dim=-1)
+        plt_.imshow(im)
+
+        if word is not None:
+            if ax is None:
+                plt.title(word)
+            else:
+                ax.set_title(word)
+
+        if out_file is not None:
+            plt.savefig(out_file)
+
+        buf = io.BytesIO()
+        plt.savefig(buf, format='png', bbox_inches='tight')
+        buf.seek(0)
+        image = PIL.Image.open(buf)
+        return image
+
+
+class WordHeatMap:
+    def __init__(self, heatmap: torch.Tensor, word: str = None, word_idx: int = None):
+        self.word = word
+        self.word_idx = word_idx
+        self.heatmap = heatmap
+
+    @property
+    def value(self):
+        return self.heatmap
+
+    def plot_overlay(self, image, out_file=None, color_normalize=True, ax=None, cmap='jet', **expand_kwargs):
+        # type: (PIL.Image.Image | np.ndarray, Path, bool, plt.Axes, Dict[str, Any]) -> None
+        return plot_overlay_heat_map(
+            image,
+            self.expand_as(image, **expand_kwargs),
+            word=self.word,
+            out_file=out_file,
+            color_normalize=color_normalize,
+            ax=ax,
+            cmap=cmap,
+        )
+
+    def expand_as(self, image, absolute=False, threshold=None, plot=False, **plot_kwargs):
+        # type: (PIL.Image.Image, bool, float, bool, Dict[str, Any]) -> torch.Tensor
+        im = self.heatmap.unsqueeze(0).unsqueeze(0)
+        im = F.interpolate(im.float().detach(), size=(image.size[0], image.size[1]), mode='bicubic')
+
+        if not absolute:
+            im = (im - im.min()) / (im.max() - im.min() + 1e-8)
+
+        if threshold:
+            im = (im > threshold).float()
+
+        im = im.cpu().detach().squeeze()
+
+        if plot:
+            self.plot_overlay(image, **plot_kwargs)
+
+        return im
+
+    def compute_ioa(self, other: 'WordHeatMap'):
+        return compute_ioa(self.heatmap, other.heatmap)
+
+
+@dataclass
+class SyntacticHeatMapPair:
+    head_heat_map: WordHeatMap
+    dep_heat_map: WordHeatMap
+    head_text: str
+    dep_text: str
+    relation: str
+
+
+@dataclass
+class ParsedHeatMap:
+    word_heat_map: WordHeatMap
+    token: spacy.tokens.Token
+
+
+class GlobalHeatMap:
+    def __init__(self, tokenizer: Any, prompt: str, heat_maps: torch.Tensor):
+        self.tokenizer = tokenizer
+        self.heat_maps = heat_maps
+        self.prompt = prompt
+        self.compute_word_heat_map = lru_cache(maxsize=50)(self.compute_word_heat_map)
+
+    def compute_word_heat_map(self, word: str, word_idx: int = None, offset_idx: int = 0) -> WordHeatMap:
+        merge_idxs, word_idx = compute_token_merge_indices(self.tokenizer, self.prompt, word, word_idx, offset_idx)
+        return WordHeatMap(self.heat_maps[merge_idxs].mean(0), word, word_idx)
+
+    def parsed_heat_maps(self) -> Iterable[ParsedHeatMap]:
+        for token in cached_nlp(self.prompt):
+            try:
+                heat_map = self.compute_word_heat_map(token.text)
+                yield ParsedHeatMap(heat_map, token)
+            except ValueError:
+                pass
+
+    def dependency_relations(self) -> Iterable[SyntacticHeatMapPair]:
+        for token in cached_nlp(self.prompt):
+            if token.dep_ != 'ROOT':
+                try:
+                    dep_heat_map = self.compute_word_heat_map(token.text)
+                    head_heat_map = self.compute_word_heat_map(token.head.text)
+
+                    yield SyntacticHeatMapPair(head_heat_map, dep_heat_map, token.head.text, token.text, token.dep_)
+                except ValueError:
+                    pass
+
+
+RawHeatMapKey = Tuple[int, int, int]  # factor, layer, head
+
+
+class RawHeatMapCollection:
+    def __init__(self):
+        self.ids_to_heatmaps: Dict[RawHeatMapKey, torch.Tensor] = defaultdict(lambda: 0.0)
+        self.ids_to_num_maps: Dict[RawHeatMapKey, int] = defaultdict(lambda: 0)
+
+    def update(self, factor: int, layer_idx: int, head_idx: int, heatmap: torch.Tensor):
+        with auto_autocast(dtype=torch.float32):
+            key = (factor, layer_idx, head_idx)
+            self.ids_to_heatmaps[key] = self.ids_to_heatmaps[key] + heatmap
+
+    def factors(self) -> Set[int]:
+        return set(key[0] for key in self.ids_to_heatmaps.keys())
+
+    def layers(self) -> Set[int]:
+        return set(key[1] for key in self.ids_to_heatmaps.keys())
+
+    def heads(self) -> Set[int]:
+        return set(key[2] for key in self.ids_to_heatmaps.keys())
+
+    def __iter__(self):
+        return iter(self.ids_to_heatmaps.items())
+
+    def clear(self):
+        self.ids_to_heatmaps.clear()
+        self.ids_to_num_maps.clear()
diff --git a/scripts/daam/hook.py b/scripts/daam/hook.py
new file mode 100644
index 000000000..f82762c4e
--- /dev/null
+++ b/scripts/daam/hook.py
@@ -0,0 +1,127 @@
+from typing import List, Generic, TypeVar, Callable, Union, Any
+import functools
+import itertools
+
+from diffusers import UNet2DConditionModel
+from diffusers.models.attention_processor import Attention
+import torch.nn as nn
+
+
+__all__ = ['ObjectHooker', 'ModuleLocator', 'AggregateHooker', 'UNetCrossAttentionLocator']
+
+
+ModuleType = TypeVar('ModuleType')
+ModuleListType = TypeVar('ModuleListType', bound=List)
+
+
+class ModuleLocator(Generic[ModuleType]):
+    def locate(self, model: nn.Module) -> List[ModuleType]:
+        raise NotImplementedError
+
+
+class ObjectHooker(Generic[ModuleType]):
+    def __init__(self, module: ModuleType):
+        self.module: ModuleType = module
+        self.hooked = False
+        self.old_state = dict()
+
+    def __enter__(self):
+        self.hook()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.unhook()
+
+    def hook(self):
+        if self.hooked:
+            raise RuntimeError('Already hooked module')
+
+        self.old_state = dict()
+        self.hooked = True
+        self._hook_impl()
+
+        return self
+
+    def unhook(self):
+        if not self.hooked:
+            raise RuntimeError('Module is not hooked')
+
+        for k, v in self.old_state.items():
+            if k.startswith('old_fn_'):
+                setattr(self.module, k[7:], v)
+
+        self.hooked = False
+        self._unhook_impl()
+
+        return self
+
+    def monkey_patch(self, fn_name, fn, strict: bool = True):
+        try:
+            self.old_state[f'old_fn_{fn_name}'] = getattr(self.module, fn_name)
+            setattr(self.module, fn_name, functools.partial(fn, self.module))
+        except AttributeError:
+            if strict:
+                raise
+
+    def monkey_super(self, fn_name, *args, **kwargs):
+        return self.old_state[f'old_fn_{fn_name}'](*args, **kwargs)
+
+    def _hook_impl(self):
+        raise NotImplementedError
+
+    def _unhook_impl(self):
+        pass
+
+
+class AggregateHooker(ObjectHooker[ModuleListType]):
+    def _hook_impl(self):
+        for h in self.module:
+            h.hook()
+
+    def _unhook_impl(self):
+        for h in self.module:
+            h.unhook()
+
+    def register_hook(self, hook: ObjectHooker):
+        self.module.append(hook)
+
+
+class UNetCrossAttentionLocator(ModuleLocator[Attention]):
+    def __init__(self, restrict: bool = None, locate_middle_block: bool = False):
+        self.restrict = restrict
+        self.layer_names = []
+        self.locate_middle_block = locate_middle_block
+
+    def locate(self, model: UNet2DConditionModel) -> List[Attention]:
+        """
+        Locate all cross-attention modules in a UNet2DConditionModel.
+
+        Args:
+            model (`UNet2DConditionModel`): The model to locate the cross-attention modules in.
+
+        Returns:
+            `List[Attention]`: The list of cross-attention modules.
+        """
+        self.layer_names.clear()
+        blocks_list = []
+        up_names = ['up'] * len(model.up_blocks)
+        down_names = ['down'] * len(model.down_blocks)
+
+        for unet_block, name in itertools.chain(
+                zip(model.up_blocks, up_names),
+                zip(model.down_blocks, down_names),
+                zip([model.mid_block], ['mid']) if self.locate_middle_block else [],
+        ):
+            if 'CrossAttn' in unet_block.__class__.__name__:
+                blocks = []
+
+                for spatial_transformer in unet_block.attentions:
+                    for transformer_block in spatial_transformer.transformer_blocks:
+                        blocks.append(transformer_block.attn2)
+
+                blocks = [b for idx, b in enumerate(blocks) if self.restrict is None or idx in self.restrict]
+                names = [f'{name}-attn-{i}' for i in range(len(blocks)) if self.restrict is None or i in self.restrict]
+                blocks_list.extend(blocks)
+                self.layer_names.extend(names)
+
+        return blocks_list
diff --git a/scripts/daam/trace.py b/scripts/daam/trace.py
new file mode 100644
index 000000000..1b6a6cd3f
--- /dev/null
+++ b/scripts/daam/trace.py
@@ -0,0 +1,323 @@
+from pathlib import Path
+from typing import List, Type, Any, Dict, Tuple, Union
+import math
+
+from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.attention_processor import Attention
+import numpy as np
+import PIL.Image as Image
+import torch
+import torch.nn.functional as F
+
+from .utils import cache_dir, auto_autocast
+from .experiment import GenerationExperiment
+from .heatmap import RawHeatMapCollection, GlobalHeatMap
+from .hook import ObjectHooker, AggregateHooker, UNetCrossAttentionLocator
+
+
+__all__ = ['trace', 'DiffusionHeatMapHooker', 'GlobalHeatMap']
+
+
+class DiffusionHeatMapHooker(AggregateHooker):
+    def __init__(
+            self,
+            pipeline: Union[StableDiffusionPipeline, StableDiffusionXLPipeline],
+            low_memory: bool = False,
+            load_heads: bool = False,
+            save_heads: bool = False,
+            data_dir: str = None
+    ):
+        self.all_heat_maps = RawHeatMapCollection()
+        h = (pipeline.unet.config.sample_size * pipeline.vae_scale_factor)
+        self.latent_hw = 4096 if h == 512 or h == 1024 else 9216  # 64x64 or 96x96 depending on if it's 2.0-v or 2.0
+        locate_middle = load_heads or save_heads
+        self.locator = UNetCrossAttentionLocator(restrict={0} if low_memory else None, locate_middle_block=locate_middle)
+        self.last_prompt: str = ''
+        self.last_image: Image = None
+        self.time_idx = 0
+        self._gen_idx = 0
+
+        modules = [
+            UNetCrossAttentionHooker(
+                x,
+                self,
+                layer_idx=idx,
+                latent_hw=self.latent_hw,
+                load_heads=load_heads,
+                save_heads=save_heads,
+                data_dir=data_dir
+            ) for idx, x in enumerate(self.locator.locate(pipeline.unet))
+        ]
+
+        modules.append(PipelineHooker(pipeline, self))
+
+        if type(pipeline) == StableDiffusionXLPipeline:
+            modules.append(ImageProcessorHooker(pipeline.image_processor, self))
+
+        super().__init__(modules)
+        self.pipe = pipeline
+
+    def time_callback(self, *args, **kwargs):
+        self.time_idx += 1
+
+    @property
+    def layer_names(self):
+        return self.locator.layer_names
+
+    def to_experiment(self, path, seed=None, id='.', subtype='.', **compute_kwargs):
+        # type: (Union[Path, str], int, str, str, Dict[str, Any]) -> GenerationExperiment
+        """Exports the last generation call to a serializable generation experiment."""
+
+        return GenerationExperiment(
+            self.last_image,
+            self.compute_global_heat_map(**compute_kwargs).heat_maps,
+            self.last_prompt,
+            seed=seed,
+            id=id,
+            subtype=subtype,
+            path=path,
+            tokenizer=self.pipe.tokenizer,
+        )
+
+    def compute_global_heat_map(self, prompt=None, factors=None, head_idx=None, layer_idx=None, normalize=False):
+        # type: (str, List[float], int, int, bool) -> GlobalHeatMap
+        """
+        Compute the global heat map for the given prompt, aggregating across time (inference steps) and space (different
+        spatial transformer block heat maps).
+
+        Args:
+            prompt: The prompt to compute the heat map for. If none, uses the last prompt that was used for generation.
+            factors: Restrict the application to heat maps with spatial factors in this set. If `None`, use all sizes.
+            head_idx: Restrict the application to heat maps with this head index. If `None`, use all heads.
+            layer_idx: Restrict the application to heat maps with this layer index. If `None`, use all layers.
+
+        Returns:
+            A heat map object for computing word-level heat maps.
+        """
+        heat_maps = self.all_heat_maps
+
+        if prompt is None:
+            prompt = self.last_prompt
+
+        if factors is None:
+            factors = {0, 1, 2, 4, 8, 16, 32, 64}
+        else:
+            factors = set(factors)
+
+        all_merges = []
+        x = int(np.sqrt(self.latent_hw))
+
+        with auto_autocast(dtype=torch.float32):
+            for (factor, layer, head), heat_map in heat_maps:
+                if factor in factors and (head_idx is None or head_idx == head) and (layer_idx is None or layer_idx == layer):
+                    heat_map = heat_map.unsqueeze(1)
+                    # The clamping fixes undershoot.
+                    all_merges.append(F.interpolate(heat_map, size=(x, x), mode='bicubic').clamp_(min=0))
+
+            try:
+                maps = torch.stack(all_merges, dim=0)
+            except RuntimeError:
+                if head_idx is not None or layer_idx is not None:
+                    raise RuntimeError('No heat maps found for the given parameters.')
+                else:
+                    raise RuntimeError('No heat maps found. Did you forget to call `with trace(...)` during generation?')
+
+            maps = maps.mean(0)[:, 0]
+            maps = maps[:len(self.pipe.tokenizer.tokenize(prompt)) + 2]  # 1 for SOS and 1 for padding
+
+            if normalize:
+                maps = maps / (maps[1:-1].sum(0, keepdim=True) + 1e-6)  # drop out [SOS] and [PAD] for proper probabilities
+
+        return GlobalHeatMap(self.pipe.tokenizer, prompt, maps)
+
+
+class ImageProcessorHooker(ObjectHooker[VaeImageProcessor]):
+    def __init__(self, processor: VaeImageProcessor, parent_trace: 'trace'):
+        super().__init__(processor)
+        self.parent_trace = parent_trace
+
+    def _hooked_postprocess(hk_self, _: VaeImageProcessor, *args, **kwargs):
+        images = hk_self.monkey_super('postprocess', *args, **kwargs)
+        hk_self.parent_trace.last_image = images[0]
+
+        return images
+
+    def _hook_impl(self):
+        self.monkey_patch('postprocess', self._hooked_postprocess)
+
+
+class PipelineHooker(ObjectHooker[StableDiffusionPipeline]):
+    def __init__(self, pipeline: StableDiffusionPipeline, parent_trace: 'trace'):
+        super().__init__(pipeline)
+        self.heat_maps = parent_trace.all_heat_maps
+        self.parent_trace = parent_trace
+
+    def _hooked_run_safety_checker(hk_self, self: StableDiffusionPipeline, image, *args, **kwargs):
+        image, has_nsfw = hk_self.monkey_super('run_safety_checker', image, *args, **kwargs)
+
+        if self.image_processor:
+            if torch.is_tensor(image):
+                images = self.image_processor.postprocess(image, output_type='pil')
+            else:
+                images = self.image_processor.numpy_to_pil(image)
+        else:
+            images = self.numpy_to_pil(image)
+
+        hk_self.parent_trace.last_image = images[len(images)-1]
+
+        return image, has_nsfw
+
+    def _hooked_check_inputs(hk_self, _: StableDiffusionPipeline, prompt: Union[str, List[str]], *args, **kwargs):
+        if not isinstance(prompt, str) and len(prompt) > 1:
+            raise ValueError('Only single prompt generation is supported for heat map computation.')
+        elif not isinstance(prompt, str):
+            last_prompt = prompt[0]
+        else:
+            last_prompt = prompt
+
+        hk_self.heat_maps.clear()
+        hk_self.parent_trace.last_prompt = last_prompt
+
+        return hk_self.monkey_super('check_inputs', prompt, *args, **kwargs)
+
+    def _hook_impl(self):
+        self.monkey_patch('run_safety_checker', self._hooked_run_safety_checker, strict=False)  # not present in SDXL
+        self.monkey_patch('check_inputs', self._hooked_check_inputs)
+
+
+class UNetCrossAttentionHooker(ObjectHooker[Attention]):
+    def __init__(
+            self,
+            module: Attention,
+            parent_trace: 'trace',
+            context_size: int = 77,
+            layer_idx: int = 0,
+            latent_hw: int = 9216,
+            load_heads: bool = False,
+            save_heads: bool = False,
+            data_dir: Union[str, Path] = None,
+    ):
+        super().__init__(module)
+        self.heat_maps = parent_trace.all_heat_maps
+        self.context_size = context_size
+        self.layer_idx = layer_idx
+        self.latent_hw = latent_hw
+
+        self.load_heads = load_heads
+        self.save_heads = save_heads
+        self.trace = parent_trace
+
+        if data_dir is not None:
+            data_dir = Path(data_dir)
+        else:
+            data_dir = cache_dir() / 'heads'
+
+        self.data_dir = data_dir
+        self.data_dir.mkdir(parents=True, exist_ok=True)
+
+    @torch.no_grad()
+    def _unravel_attn(self, x):
+        # type: (torch.Tensor) -> torch.Tensor
+        # x shape: (heads, height * width, tokens)
+        """
+        Unravels the attention, returning it as a collection of heat maps.
+
+        Args:
+            x (`torch.Tensor`): cross attention slice/map between the words and the tokens.
+            value (`torch.Tensor`): the value tensor.
+
+        Returns:
+            `List[Tuple[int, torch.Tensor]]`: the list of heat maps across heads.
+        """
+        h = w = int(math.sqrt(x.size(1)))
+        maps = []
+        x = x.permute(2, 0, 1)
+
+        with auto_autocast(dtype=torch.float32):
+            for map_ in x:
+                map_ = map_.view(map_.size(0), h, w)
+                # For Instruct Pix2Pix, divide the map into three parts: text condition, image condition and unconditional,
+                # and only keep the text condition part, which is first of the three parts(as per diffusers implementation).
+                if map_.size(0) == 24:
+                    map_ = map_[:((map_.size(0) // 3)+1)]  # Filter out unconditional and image condition
+                else:
+                    map_ = map_[map_.size(0) // 2:] #  # Filter out unconditional
+                maps.append(map_)
+
+        maps = torch.stack(maps, 0)  # shape: (tokens, heads, height, width)
+        return maps.permute(1, 0, 2, 3).contiguous()  # shape: (heads, tokens, height, width)
+
+    def _save_attn(self, attn_slice: torch.Tensor):
+        torch.save(attn_slice, self.data_dir / f'{self.trace._gen_idx}.pt')
+
+    def _load_attn(self) -> torch.Tensor:
+        return torch.load(self.data_dir / f'{self.trace._gen_idx}.pt')
+
+    def __call__(
+            self,
+            attn: Attention,
+            hidden_states,
+            encoder_hidden_states=None,
+            attention_mask=None,
+    ):
+        """Capture attentions and aggregate them."""
+        batch_size, sequence_length, _ = hidden_states.shape
+        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross is not None:
+            encoder_hidden_states = attn.norm_cross(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        query = attn.head_to_batch_dim(query)
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+
+        # DAAM save heads
+        if self.save_heads:
+            self._save_attn(attention_probs)
+        elif self.load_heads:
+            attention_probs = self._load_attn()
+
+        # compute shape factor
+        factor = int(math.sqrt(self.latent_hw // attention_probs.shape[1]))
+        self.trace._gen_idx += 1
+
+        # skip if too large
+        if attention_probs.shape[-1] == self.context_size and factor != 8:
+            # shape: (batch_size, 64 // factor, 64 // factor, 77)
+            maps = self._unravel_attn(attention_probs)
+
+            for head_idx, heatmap in enumerate(maps):
+                self.heat_maps.update(factor, self.layer_idx, head_idx, heatmap)
+
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        return hidden_states
+
+    def _hook_impl(self):
+        self.original_processor = self.module.processor
+        self.module.set_processor(self)
+
+    def _unhook_impl(self):
+        self.module.set_processor(self.original_processor)
+
+    @property
+    def num_heat_maps(self):
+        return len(next(iter(self.heat_maps.values())))
+
+
+trace: Type[DiffusionHeatMapHooker] = DiffusionHeatMapHooker
diff --git a/scripts/daam/utils.py b/scripts/daam/utils.py
new file mode 100644
index 000000000..8cfde13f2
--- /dev/null
+++ b/scripts/daam/utils.py
@@ -0,0 +1,109 @@
+from functools import lru_cache
+from pathlib import Path
+import os
+import sys
+import random
+from typing import TypeVar
+
+import PIL.Image
+import matplotlib.pyplot as plt
+import numpy as np
+import spacy
+import torch
+import torch.nn.functional as F
+
+
+__all__ = ['set_seed', 'compute_token_merge_indices', 'plot_mask_heat_map', 'cached_nlp', 'cache_dir', 'auto_device', 'auto_autocast']
+
+
+T = TypeVar('T')
+
+
+def auto_device(obj: T = torch.device('cpu')) -> T:
+    if isinstance(obj, torch.device):
+        return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    if torch.cuda.is_available():
+        return obj.to('cuda')
+
+    return obj
+
+
+def auto_autocast(*args, **kwargs):
+    if not torch.cuda.is_available():
+        kwargs['enabled'] = False
+
+    return torch.cuda.amp.autocast(*args, **kwargs)
+
+
+def plot_mask_heat_map(im: PIL.Image.Image, heat_map: torch.Tensor, threshold: float = 0.4):
+    im = torch.from_numpy(np.array(im)).float() / 255
+    mask = (heat_map.squeeze() > threshold).float()
+    im = im * mask.unsqueeze(-1)
+    plt.imshow(im)
+
+
+def set_seed(seed: int) -> torch.Generator:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    gen = torch.Generator(device=auto_device())
+    gen.manual_seed(seed)
+
+    return gen
+
+
+def cache_dir() -> Path:
+    # *nix
+    if os.name == 'posix' and sys.platform != 'darwin':
+        xdg = os.environ.get('XDG_CACHE_HOME', os.path.expanduser('~/.cache'))
+        return Path(xdg, 'daam')
+    elif sys.platform == 'darwin':
+        # Mac OS
+        return Path(os.path.expanduser('~'), 'Library/Caches/daam')
+    else:
+        # Windows
+        local = os.environ.get('LOCALAPPDATA', None) \
+                or os.path.expanduser('~\\AppData\\Local')
+        return Path(local, 'daam')
+
+
+def compute_token_merge_indices(tokenizer, prompt: str, word: str, word_idx: int = None, offset_idx: int = 0):
+    merge_idxs = []
+    tokens = tokenizer.tokenize(prompt.lower())
+    tokens = [x.replace('</w>', '') for x in tokens]  # New tokenizer uses wordpiece markers.
+
+    if word_idx is None:
+        word = word.lower()
+        search_tokens = [x.replace('</w>', '') for x in tokenizer.tokenize(word)]  # New tokenizer uses wordpiece markers.
+        start_indices = [x + offset_idx for x in range(len(tokens)) if tokens[x:x + len(search_tokens)] == search_tokens]
+
+        for indice in start_indices:
+            merge_idxs += [i + indice for i in range(0, len(search_tokens))]
+
+        if not merge_idxs:
+            raise ValueError(f'Search word {word} not found in prompt!')
+    else:
+        merge_idxs.append(word_idx)
+
+    return [x + 1 for x in merge_idxs], word_idx  # Offset by 1.
+
+
+nlp = None
+
+
+@lru_cache(maxsize=100000)
+def cached_nlp(prompt: str, type='en_core_web_md'):
+    global nlp
+
+    if nlp is None:
+        try:
+            nlp = spacy.load(type)
+        except OSError:
+            import os
+            os.system(f'python -m spacy download {type}')
+            nlp = spacy.load(type)
+
+    return nlp(prompt)
diff --git a/scripts/daam_ext.py b/scripts/daam_ext.py
new file mode 100644
index 000000000..5ae16a855
--- /dev/null
+++ b/scripts/daam_ext.py
@@ -0,0 +1,57 @@
+# https://github.com/genforce/ctrl-x
+
+import gradio as gr
+from installer import install
+from modules import shared, scripts_manager, processing
+
+
+COLORMAP = ['autumn', 'bone', 'jet', 'winter', 'rainbow', 'ocean', 'summer', 'spring', 'cool', 'hsv', 'pink', 'hot', 'parula', 'magma', 'inferno', 'plasma', 'viridis', 'cividis', 'twilight', 'shifted', 'turbo', 'deepgreen']
+
+
+class Script(scripts_manager.Script):
+    def title(self):
+        return 'DAAM: Diffusion Attentive Attribution Maps'
+
+    def show(self, is_img2img):
+        return not is_img2img
+
+    def ui(self, _is_img2img):
+        with gr.Row():
+            gr.HTML('<a href="https://github.com/castorini/daam">&nbsp DAAM: Diffusion Attentive Attribution Maps</a><br>')
+        with gr.Row():
+            append_images = gr.Checkbox(label='Append heatmaps to results', value=True, elem_id='daam_append_images')
+            colormap = gr.Dropdown(label='Colormap', choices=COLORMAP, value='jet', type='value', elem_id='daam_colormap')
+        return append_images, colormap
+
+    def run(self, p: processing.StableDiffusionProcessing, append_images, colormap): # pylint: disable=arguments-differ
+        c = shared.sd_model.__class__.__name__ if shared.sd_loaded else ''
+        if shared.sd_model_type != 'sdxl':
+            shared.log.warning(f'DAAM: pipeline={c} required=StableDiffusionXLPipeline')
+            return None
+
+        install('thinc==8.3.4')
+        install('spacy==3.8.4')
+
+        from scripts import daam # pylint: disable=no-name-in-module
+        orig_prompt_attention = shared.opts.prompt_attention
+        shared.opts.data['prompt_attention'] = 'fixed'
+
+        # process
+        with daam.trace(shared.sd_model) as tc:
+            processed: processing.Processed = processing.process_images(p)
+            global_heat_map = tc.compute_global_heat_map()
+            shared.log.info(f'DAAM: prompt="{global_heat_map.prompt}" heatmaps={global_heat_map.heat_maps.shape}')
+
+            # word_heat_map = global_heat_map.compute_word_heat_map('woman')
+            parsed_heat_maps = global_heat_map.parsed_heat_maps()
+            if append_images:
+                image = processed.images[0]
+                for parsed_heat_map in parsed_heat_maps:
+                    if len(parsed_heat_map.token.text) > 1:
+                        shared.log.debug(f'DAAM: token="{parsed_heat_map.token.text}"')
+                        overlay = parsed_heat_map.word_heat_map.plot_overlay(image=image, color_normalize=True, cmap=colormap)
+                        processed.images.append(overlay)
+
+        # restore and return
+        shared.opts.data['prompt_attention'] = orig_prompt_attention
+        return processed