castorini / daam

Diffusion attentive attribution maps for interpreting Stable Diffusion.
MIT License
689 stars 63 forks source link

heatmap for num_inference_steps=1 using callable #38

Closed NamHyelin closed 1 year ago

NamHyelin commented 1 year ago

Hi, I want to compare heatmaps for operating diffusion steps for one time and 20 times. How should I add code in your colab? I think I should use callable and callback_steps options in 'model().images[0]', but I don't know what to do after when plotting heatmaps. I would greatly appreciate it if you could provide me with guidance thanks!

daemon commented 1 year ago

Hello, are you trying to get intermediate heat maps? If so, I would recommend calling compute_global_heat_map based on the value of time_idx here: https://github.com/castorini/daam/blob/f2466b8e066494d0f6817c29f7a4c9cf88d656ba/daam/trace.py#L38

For a quick and dirty solution, that logic could sit in https://github.com/castorini/daam/blob/f2466b8e066494d0f6817c29f7a4c9cf88d656ba/daam/trace.py#L59

E.g.,

if self.time_idx == 1:
    # Do something if it's the first time step
    map1 = self.compute_global_heat_map(prompt=...)  
elif self.time_idx == 20:
    # Do something if it's the twentieth time step
    map20 = self.compute_global_heat_map(prompt=...)