├── LICENSE
├── README.md
├── bounded_attention.py
├── images
    ├── example.jpg
    └── teaser.jpg
├── injection_utils.py
├── pipeline_stable_diffusion_opt.py
├── pipeline_stable_diffusion_xl_opt.py
├── requirements.txt
├── run_sd.py
├── run_xl.py
└── utils.py


/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2024 Omer Dahary
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy
 6 | of this software and associated documentation files (the "Software"), to deal
 7 | in the Software without restriction, including without limitation the rights
 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 | 
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 | 
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
  1 | # Be Yourself: Bounded Attention for Multi-Subject Text-to-Image Generation (ECCV 2024)
  2 | 
  3 | > **Omer Dahary, Or Patashnik, Kfir Aberman, Daniel Cohen-Or**
  4 | > 
  5 | > Text-to-image diffusion models have an unprecedented ability to generate diverse and high-quality images. However, they often struggle to faithfully capture the intended semantics of complex input prompts that include multiple subjects. Recently, numerous layout-to-image extensions have been introduced to improve user control, aiming to localize subjects represented by specific tokens. Yet, these methods often produce semantically inaccurate images, especially when dealing with multiple semantically or visually similar subjects. In this work, we study and analyze the causes of these limitations. Our exploration reveals that the primary issue stems from inadvertent semantic leakage between subjects in the denoising process. This leakage is attributed to the diffusion model’s attention layers, which tend to blend the visual features of different subjects. To address these issues, we introduce Bounded Attention, a training-free method for bounding the information flow in the sampling process. Bounded Attention prevents detrimental leakage among subjects and enables guiding the generation to promote each subject's individuality, even with complex multi-subject conditioning. Through extensive experimentation, we demonstrate that our method empowers the generation of multiple subjects that better align with given prompts and layouts.
  6 | 
  7 | <a href="https://omer11a.github.io/bounded-attention/"><img src="https://img.shields.io/static/v1?label=Project&message=Website&color=red" height=20.5></a> 
  8 | <a href="https://arxiv.org/abs/2403.16990"><img src="https://img.shields.io/badge/arXiv-BA-b31b1b.svg" height=20.5></a>
  9 | [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/omer11a/bounded-attention)
 10 | 
 11 | <p align="center">
 12 | <img src="images/teaser.jpg" width="800px"/>
 13 | </p>
 14 | 
 15 | ## Description  
 16 | Official implementation of our "Be Yourself: Bounded Attention for Multi-Subject Text-to-Image Generation" paper.
 17 | 
 18 | ## Setup
 19 | 
 20 | ### Environment
 21 | 
 22 | To set up the environment, run:
 23 | 
 24 | ```
 25 | conda create --name bounded-attention python=3.11.4
 26 | conda activate bounded-attention
 27 | pip install -r requirements.txt
 28 | ```
 29 | 
 30 | Then, run in Python:
 31 | 
 32 | ```
 33 | import nltk
 34 | nltk.download('averaged_perceptron_tagger')
 35 | ```
 36 | 
 37 | ### Demo
 38 | 
 39 | This project has a gradio [demo](https://huggingface.co/spaces/omer11a/bounded-attention) deployed in HuggingFace.
 40 | To run the demo locally, run the following: 
 41 | ```shell
 42 | gradio app.py
 43 | ```
 44 | Then, you can connect to the local demo by browsing to `http://localhost:7860/`.
 45 | 
 46 | ## Usage
 47 | 
 48 | <p align="center">
 49 | <img src="images/example.jpg" width="800px"/>  
 50 | <br>
 51 | Example generations by SDXL with and without Bounded Attention.
 52 | </p>
 53 | 
 54 | ### Basics
 55 | 
 56 | To generate images, you can run `run_xl.py` for our SDXL version, and `run_sd.py` for our Stable Diffusion version.
 57 | In each script, we call the `run` function to generate the images. E.g.,
 58 | ```
 59 | boxes = [
 60 |     [0.35, 0.4, 0.65, 0.9],
 61 |     [0, 0.6, 0.3, 0.9],
 62 |     [0.7, 0.55, 1, 0.85],
 63 | ]
 64 | 
 65 | prompt = "3 D Pixar animation of a cute unicorn and a pink hedgehog and a nerdy owl traveling in a magical forest"
 66 | subject_token_indices = [[7, 8, 17], [11, 12, 17], [15, 16, 17]]
 67 | 
 68 | run(boxes, prompt, subject_token_indices, init_step_size=25, final_step_size=10)
 69 | ```
 70 | 
 71 | The `run` function receives the following parameters:
 72 | - boxes: the bounding box of each subject in the format [(x0, y0, x1, y1), ...], where x=0,y=0 represents the top-left corner of the image, and x=1,y=1 represents the bottom-right corner.
 73 | - prompt: the textual prompt.
 74 | - subject_token_indices: The indices of each token belonging to each subject, where the indices start from 1. Tokens can be shared between subjects.
 75 | - out_dir: The output directory. Defaults to "out".
 76 | - seed: The random seed.
 77 | - batch_size: The number of generated images.
 78 | - filter_token_indices: The indices of the tokens to ignore. This is automatically inferred, but we recommend explicitly ignoring prepositions, numbers and positional relations.
 79 | - eos_token_index: The index of the EOS token (the first padding token appended to the end of the prompt). This is automatically inferred, but we recommend explicitly passing it, as we use it to verify you have correctly counted the number of tokens.
 80 | 
 81 | ### Advanced options
 82 | 
 83 | The `run` function also supports the following optional hyperparameters:
 84 | 
 85 | - init_step_size: The initial step size of the linear step size scheduler when performing guidance.
 86 | - final_step_size: The final step size of the linear step size scheduler when performing guidance.
 87 | - num_clusters_per_subject: The number of clusters computed when clustering the self-attention maps (#clusters = #subject x #clusters_per_subject). Changing this value might improve semantics (adherence to the prompt), especially when the subjects exceed their bounding boxes.
 88 | - cross_loss_scale: The scale factor of the cross-attention loss term. Increasing it will improve semantic control (adherence to the prompt), but may reduce image quality.
 89 | - self_loss_scale: The scale factor of the self-attention loss term. Increasing it will improve layout control (adherence to the bounding boxes), but may reduce image quality.
 90 | - classifier_free_guidance_scale: The scale factor of classifier-free guidance.
 91 | - num_guidance_steps: The number of timesteps in which to perform guidance. Decreasing this also decreases the runtime.
 92 | - first_refinement_step: The timestep from which subject mask refinement is performed.
 93 | - num_gd_iterations: The number of Gradient Descent iterations for each timestep when performing guidance.
 94 | - loss_threshold: If the loss is below the threshold, Gradient Descent stops for that timestep.
 95 | 
 96 | ## Acknowledgements 
 97 | 
 98 | This code was built using the code from the following repositories:
 99 | - [diffusers](https://github.com/huggingface/diffusers)
100 | - [Prompt-to-Prompt](https://github.com/google/prompt-to-prompt/)
101 | - [MasaCtrl](https://github.com/TencentARC/MasaCtrl)
102 | 
103 | ## Citation
104 | 
105 | If you use this code for your research, please cite our paper:
106 | 
107 | ```
108 | @misc{dahary2024yourself,
109 |     title={Be Yourself: Bounded Attention for Multi-Subject Text-to-Image Generation},
110 |     author={Omer Dahary and Or Patashnik and Kfir Aberman and Daniel Cohen-Or},
111 |     year={2024},
112 |     eprint={2403.16990},
113 |     archivePrefix={arXiv},
114 |     primaryClass={cs.CV}
115 |  }
116 | ```
117 | 


--------------------------------------------------------------------------------
/bounded_attention.py:
--------------------------------------------------------------------------------
  1 | import nltk
  2 | import einops
  3 | import torch
  4 | import torch.nn.functional as F
  5 | import torchvision.utils
  6 | from torch_kmeans import KMeans
  7 | 
  8 | import os
  9 | 
 10 | import injection_utils
 11 | import utils
 12 | 
 13 | 
 14 | class BoundedAttention(injection_utils.AttentionBase):
 15 |     EPSILON = 1e-5
 16 |     FILTER_TAGS = {
 17 |         'CC', 'CD', 'DT', 'EX', 'IN', 'LS', 'MD', 'PDT', 'POS', 'PRP', 'PRP$', 'RP', 'TO', 'UH', 'WDT', 'WP', 'WRB'}
 18 |     TAG_RULES = {'left': 'IN', 'right': 'IN', 'top': 'IN', 'bottom': 'IN'}
 19 | 
 20 |     def __init__(
 21 |         self,
 22 |         boxes,
 23 |         prompts,
 24 |         subject_token_indices,
 25 |         cross_loss_layers,
 26 |         self_loss_layers,
 27 |         cross_mask_layers=None,
 28 |         self_mask_layers=None,
 29 |         eos_token_index=None,
 30 |         filter_token_indices=None,
 31 |         leading_token_indices=None,
 32 |         mask_cross_during_guidance=True,
 33 |         mask_eos=True,
 34 |         cross_loss_coef=1.5,
 35 |         self_loss_coef=0.5,
 36 |         max_guidance_iter=15,
 37 |         max_guidance_iter_per_step=5,
 38 |         start_step_size=18,
 39 |         end_step_size=5,
 40 |         loss_stopping_value=0.2,
 41 |         min_clustering_step=15,
 42 |         cross_mask_threshold=0.2,
 43 |         self_mask_threshold=0.2,
 44 |         delta_refine_mask_steps=5,
 45 |         pca_rank=None,
 46 |         num_clusters=None,
 47 |         num_clusters_per_box=3,
 48 |         max_resolution=None,
 49 |         map_dir=None,
 50 |         debug=False,
 51 |         delta_debug_attention_steps=20,
 52 |         delta_debug_mask_steps=5,
 53 |         debug_layers=None,
 54 |         saved_resolution=64,
 55 |     ):
 56 |         super().__init__()
 57 |         self.boxes = boxes
 58 |         self.prompts = prompts
 59 |         self.subject_token_indices = subject_token_indices
 60 |         self.cross_loss_layers = set(cross_loss_layers)
 61 |         self.self_loss_layers = set(self_loss_layers)
 62 |         self.cross_mask_layers = self.cross_loss_layers if cross_mask_layers is None else set(cross_mask_layers)
 63 |         self.self_mask_layers = self.self_loss_layers if self_mask_layers is None else set(self_mask_layers)
 64 | 
 65 |         self.eos_token_index = eos_token_index
 66 |         self.filter_token_indices = filter_token_indices
 67 |         self.leading_token_indices = leading_token_indices
 68 |         self.mask_cross_during_guidance = mask_cross_during_guidance
 69 |         self.mask_eos = mask_eos
 70 |         self.cross_loss_coef = cross_loss_coef
 71 |         self.self_loss_coef = self_loss_coef
 72 |         self.max_guidance_iter = max_guidance_iter
 73 |         self.max_guidance_iter_per_step = max_guidance_iter_per_step
 74 |         self.start_step_size = start_step_size
 75 |         self.step_size_coef = (end_step_size - start_step_size) / max_guidance_iter
 76 |         self.loss_stopping_value = loss_stopping_value
 77 |         self.min_clustering_step = min_clustering_step
 78 |         self.cross_mask_threshold = cross_mask_threshold
 79 |         self.self_mask_threshold = self_mask_threshold
 80 | 
 81 |         self.delta_refine_mask_steps = delta_refine_mask_steps
 82 |         self.pca_rank = pca_rank
 83 |         num_clusters = len(boxes) * num_clusters_per_box if num_clusters is None else num_clusters
 84 |         self.clustering = KMeans(n_clusters=num_clusters, num_init=100)
 85 |         self.centers = None
 86 | 
 87 |         self.max_resolution = max_resolution
 88 |         self.map_dir = map_dir
 89 |         self.debug = debug
 90 |         self.delta_debug_attention_steps = delta_debug_attention_steps
 91 |         self.delta_debug_mask_steps = delta_debug_mask_steps
 92 |         self.debug_layers = self.cross_loss_layers | self.self_loss_layers if debug_layers is None else debug_layers
 93 |         self.saved_resolution = saved_resolution
 94 | 
 95 |         self.optimized = False
 96 |         self.cross_foreground_values = []
 97 |         self.self_foreground_values = []
 98 |         self.cross_background_values = []
 99 |         self.self_background_values = []
100 |         self.mean_cross_map = 0
101 |         self.num_cross_maps = 0
102 |         self.mean_self_map = 0
103 |         self.num_self_maps = 0
104 |         self.self_masks = None
105 | 
106 |     def clear_values(self, include_maps=False):
107 |         lists = (
108 |             self.cross_foreground_values,
109 |             self.self_foreground_values,
110 |             self.cross_background_values,
111 |             self.self_background_values,
112 |         )
113 | 
114 |         for values in lists:
115 |             values.clear()
116 | 
117 |         if include_maps:
118 |             self.mean_cross_map = 0
119 |             self.num_cross_maps = 0
120 |             self.mean_self_map = 0
121 |             self.num_self_maps = 0
122 | 
123 |     def before_step(self):
124 |         self.clear_values()
125 |         if self.cur_step == 0:
126 |             self._determine_tokens()
127 | 
128 |     def reset(self):
129 |         self.clear_values(include_maps=True)
130 |         super().reset()
131 | 
132 |     def forward(self, q, k, v, is_cross, place_in_unet, num_heads, **kwargs):
133 |         batch_size = q.size(0) // num_heads
134 |         n = q.size(1)
135 |         d = k.size(1)
136 |         dtype = q.dtype
137 |         device = q.device
138 |         if is_cross:
139 |             masks = self._hide_other_subjects_from_tokens(batch_size // 2, n, d, dtype, device)
140 |         else:
141 |             masks = self._hide_other_subjects_from_subjects(batch_size // 2, n, dtype, device)
142 | 
143 |         resolution = int(n ** 0.5)
144 |         if (self.max_resolution is not None) and (resolution > self.max_resolution):
145 |             return super().forward(q, k, v, is_cross, place_in_unet, num_heads, mask=masks)
146 | 
147 |         sim = torch.einsum('b i d, b j d -> b i j', q, k) * kwargs['scale']
148 |         attn = sim.softmax(-1)
149 |         self._display_attention_maps(attn, is_cross, num_heads)
150 |         sim = sim.reshape(batch_size, num_heads, n, d) + masks
151 |         attn = sim.reshape(-1, n, d).softmax(-1)
152 |         self._save(attn, is_cross, num_heads)
153 |         self._display_attention_maps(attn, is_cross, num_heads, prefix='masked')
154 |         self._debug_hook(q, k, v, sim, attn, is_cross, place_in_unet, num_heads, **kwargs)
155 |         out = torch.bmm(attn, v)
156 |         return einops.rearrange(out, '(b h) n d -> b n (h d)', h=num_heads)
157 | 
158 |     def update_loss(self, forward_pass, latents, i):
159 |         if i >= self.max_guidance_iter:
160 |             return latents
161 | 
162 |         step_size = self.start_step_size + self.step_size_coef * i
163 | 
164 |         self.optimized = True
165 |         normalized_loss = torch.tensor(10000)
166 |         with torch.enable_grad():
167 |             latents = latents.clone().detach().requires_grad_(True)
168 |             for guidance_iter in range(self.max_guidance_iter_per_step):
169 |                 if normalized_loss < self.loss_stopping_value:
170 |                     break
171 | 
172 |                 latent_model_input = torch.cat([latents] * 2)
173 |                 cur_step = self.cur_step
174 |                 forward_pass(latent_model_input)
175 |                 self.cur_step = cur_step
176 | 
177 |                 loss, normalized_loss = self._compute_loss()
178 |                 grad_cond = torch.autograd.grad(loss, [latents])[0]
179 |                 latents = latents - step_size * grad_cond
180 |                 if self.debug:
181 |                     print(f'Loss at step={i}, iter={guidance_iter}: {normalized_loss}')
182 |                     grad_norms = grad_cond.flatten(start_dim=2).norm(dim=1)
183 |                     grad_norms = grad_norms / grad_norms.max(dim=1, keepdim=True)[0]
184 |                     self._save_maps(grad_norms, 'grad_norms')
185 | 
186 |         self.optimized = False
187 |         return latents
188 | 
189 |     def _tokenize(self):
190 |         ids = self.model.tokenizer.encode(self.prompts[0])
191 |         tokens = self.model.tokenizer.convert_ids_to_tokens(ids, skip_special_tokens=True)
192 |         return [token[:-4] for token in tokens]  # remove ending </w>
193 | 
194 |     def _tag_tokens(self):
195 |         tagged_tokens = nltk.pos_tag(self._tokenize())
196 |         return [type(self).TAG_RULES.get(token, tag) for token, tag in tagged_tokens]
197 | 
198 |     def _determine_eos_token(self):
199 |         tokens = self._tokenize()
200 |         eos_token_index = len(tokens) + 1
201 |         if self.eos_token_index is None:
202 |             self.eos_token_index = eos_token_index
203 |         elif eos_token_index != self.eos_token_index:
204 |             raise ValueError(f'Wrong EOS token index. Tokens are: {tokens}.')
205 | 
206 |     def _determine_filter_tokens(self):
207 |         if self.filter_token_indices is not None:
208 |             return
209 | 
210 |         tags = self._tag_tokens()
211 |         self.filter_token_indices = [i + 1 for i, tag in enumerate(tags) if tag in type(self).FILTER_TAGS]
212 | 
213 |     def _determine_leading_tokens(self):
214 |         if self.leading_token_indices is not None:
215 |             return
216 | 
217 |         tags = self._tag_tokens()
218 |         leading_token_indices = []
219 |         for indices in self.subject_token_indices:
220 |             subject_noun_indices = [i for i in indices if tags[i - 1].startswith('NN')]
221 |             leading_token_candidates = subject_noun_indices if len(subject_noun_indices) > 0 else indices
222 |             leading_token_indices.append(leading_token_candidates[-1])
223 | 
224 |         self.leading_token_indices = leading_token_indices
225 | 
226 |     def _determine_tokens(self):
227 |         self._determine_eos_token()
228 |         self._determine_filter_tokens()
229 |         self._determine_leading_tokens()
230 | 
231 |     def _split_references(self, tensor, num_heads):
232 |         tensor = tensor.reshape(-1, num_heads, *tensor.shape[1:])
233 |         unconditional, conditional = tensor.chunk(2)
234 | 
235 |         num_subjects = len(self.boxes)
236 |         batch_unconditional = unconditional[:-num_subjects]
237 |         references_unconditional = unconditional[-num_subjects:]
238 |         batch_conditional = conditional[:-num_subjects]
239 |         references_conditional = conditional[-num_subjects:]
240 | 
241 |         batch = torch.cat((batch_unconditional, batch_conditional))
242 |         references = torch.cat((references_unconditional, references_conditional))
243 |         batch = batch.reshape(-1, *batch_unconditional.shape[2:])
244 |         references = references.reshape(-1, *references_unconditional.shape[2:])
245 |         return batch, references
246 | 
247 |     def _hide_other_subjects_from_tokens(self, batch_size, n, d, dtype, device):  # b h i j
248 |         resolution = int(n ** 0.5)
249 |         subject_masks, background_masks = self._obtain_masks(resolution, batch_size=batch_size, device=device)  # b s n
250 |         include_background = self.optimized or (not self.mask_cross_during_guidance and self.cur_step < self.max_guidance_iter_per_step)
251 |         subject_masks = torch.logical_or(subject_masks, background_masks.unsqueeze(1)) if include_background else subject_masks
252 |         min_value = torch.finfo(dtype).min
253 |         sim_masks = torch.zeros((batch_size, n, d), dtype=dtype, device=device)  # b i j
254 |         for token_indices in (*self.subject_token_indices, self.filter_token_indices):
255 |             sim_masks[:, :, token_indices] = min_value
256 | 
257 |         for batch_index in range(batch_size):
258 |             for subject_mask, token_indices in zip(subject_masks[batch_index], self.subject_token_indices):
259 |                 for token_index in token_indices:
260 |                     sim_masks[batch_index, subject_mask, token_index] = 0
261 | 
262 |         if self.mask_eos and not include_background:
263 |             for batch_index, background_mask in zip(range(batch_size), background_masks):
264 |                 sim_masks[batch_index, background_mask, self.eos_token_index] = min_value
265 | 
266 |         return torch.cat((torch.zeros_like(sim_masks), sim_masks)).unsqueeze(1)
267 | 
268 |     def _hide_other_subjects_from_subjects(self, batch_size, n, dtype, device):  # b h i j
269 |         resolution = int(n ** 0.5)
270 |         subject_masks, background_masks = self._obtain_masks(resolution, batch_size=batch_size, device=device)  # b s n
271 |         min_value = torch.finfo(dtype).min
272 |         sim_masks = torch.zeros((batch_size, n, n), dtype=dtype, device=device)  # b i j
273 |         for batch_index, background_mask in zip(range(batch_size), background_masks):
274 |             all_subject_mask = ~background_mask.unsqueeze(0) * ~background_mask.unsqueeze(1)
275 |             sim_masks[batch_index, all_subject_mask] = min_value
276 | 
277 |         for batch_index in range(batch_size):
278 |             for subject_mask in subject_masks[batch_index]:
279 |                 subject_sim_mask = sim_masks[batch_index, subject_mask]
280 |                 condition = torch.logical_or(subject_sim_mask == 0, subject_mask.unsqueeze(0))
281 |                 sim_masks[batch_index, subject_mask] = torch.where(condition, 0, min_value).to(dtype=dtype)
282 | 
283 |         return torch.cat((sim_masks, sim_masks)).unsqueeze(1)
284 | 
285 |     def _save(self, attn, is_cross, num_heads):
286 |         _, attn = attn.chunk(2)
287 |         attn = attn.reshape(-1, num_heads, *attn.shape[-2:])  # b h n k
288 | 
289 |         self._save_mask_maps(attn, is_cross)
290 |         self._save_loss_values(attn, is_cross)
291 | 
292 |     def _save_mask_maps(self, attn, is_cross):
293 |         if (
294 |             (self.optimized) or
295 |             (is_cross and self.cur_att_layer not in self.cross_mask_layers) or
296 |             ((not is_cross) and (self.cur_att_layer not in self.self_mask_layers))
297 |         ):
298 |             return
299 | 
300 |         if is_cross:
301 |             attn = attn[..., self.leading_token_indices]
302 |             mean_map = self.mean_cross_map
303 |             num_maps = self.num_cross_maps
304 |         else:
305 |             mean_map = self.mean_self_map
306 |             num_maps = self.num_self_maps
307 | 
308 |         num_maps += 1
309 |         attn = attn.mean(dim=1)  # mean over heads
310 |         mean_map = ((num_maps - 1) / num_maps) * mean_map +  (1 / num_maps) * attn
311 |         if is_cross:
312 |             self.mean_cross_map = mean_map
313 |             self.num_cross_maps = num_maps
314 |         else:
315 |             self.mean_self_map = mean_map
316 |             self.num_self_maps = num_maps
317 | 
318 |     def _save_loss_values(self, attn, is_cross):
319 |         if (
320 |             (not self.optimized) or
321 |             (is_cross and (self.cur_att_layer not in self.cross_loss_layers)) or
322 |             ((not is_cross) and (self.cur_att_layer not in self.self_loss_layers))
323 |         ):
324 |             return
325 | 
326 |         resolution = int(attn.size(2) ** 0.5)
327 |         boxes = self._convert_boxes_to_masks(resolution, device=attn.device)  # s n
328 |         background_mask = boxes.sum(dim=0) == 0
329 | 
330 |         if is_cross:
331 |             saved_foreground_values = self.cross_foreground_values
332 |             saved_background_values = self.cross_background_values
333 |             contexts = [indices + [self.eos_token_index] for indices in self.subject_token_indices]  # TODO: fix EOS loss term
334 |         else:
335 |             saved_foreground_values = self.self_foreground_values
336 |             saved_background_values = self.self_background_values
337 |             contexts = boxes
338 | 
339 |         foreground_values = []
340 |         background_values = []
341 |         for i, (box, context) in enumerate(zip(boxes, contexts)):
342 |             context_attn = attn[:, :, :, context]
343 |             
344 |             # sum over heads, pixels and contexts
345 |             foreground_values.append(context_attn[:, :, box].sum(dim=(1, 2, 3)))
346 |             background_values.append(context_attn[:, :, background_mask].sum(dim=(1, 2, 3)))
347 | 
348 |         saved_foreground_values.append(torch.stack(foreground_values, dim=1))
349 |         saved_background_values.append(torch.stack(background_values, dim=1))
350 | 
351 |     def _compute_loss(self):
352 |         cross_losses = self._compute_loss_term(self.cross_foreground_values, self.cross_background_values)
353 |         self_losses = self._compute_loss_term(self.self_foreground_values, self.self_background_values)
354 |         b, s = cross_losses.shape
355 | 
356 |         # sum over samples and subjects
357 |         total_cross_loss = cross_losses.sum()
358 |         total_self_loss = self_losses.sum()
359 | 
360 |         loss = self.cross_loss_coef * total_cross_loss + self.self_loss_coef * total_self_loss
361 |         normalized_loss = loss / b / s
362 |         return loss, normalized_loss
363 | 
364 |     def _compute_loss_term(self, foreground_values, background_values):
365 |         # mean over layers
366 |         mean_foreground = torch.stack(foreground_values).mean(dim=0)
367 |         mean_background = torch.stack(background_values).mean(dim=0)
368 |         iou = mean_foreground / (mean_foreground + len(self.boxes) * mean_background)
369 |         return (1 - iou) ** 2
370 | 
371 |     def _obtain_masks(self, resolution, return_boxes=False, return_existing=False, batch_size=None, device=None):
372 |         return_boxes = return_boxes or (return_existing and self.self_masks is None)
373 |         if return_boxes or self.cur_step < self.min_clustering_step:
374 |             masks = self._convert_boxes_to_masks(resolution, device=device).unsqueeze(0)
375 |             if batch_size is not None:
376 |                 masks = masks.expand(batch_size, *masks.shape[1:])
377 |         else:
378 |             masks = self._obtain_self_masks(resolution, return_existing=return_existing)
379 |             if device is not None:
380 |                 masks = masks.to(device=device)
381 | 
382 |         background_mask = masks.sum(dim=1) == 0
383 |         return masks, background_mask
384 | 
385 |     def _convert_boxes_to_masks(self, resolution, device=None):  # s n
386 |         boxes = torch.zeros(len(self.boxes), resolution, resolution, dtype=bool, device=device)
387 |         for i, box in enumerate(self.boxes):
388 |             x0, x1 = box[0] * resolution, box[2] * resolution
389 |             y0, y1 = box[1] * resolution, box[3] * resolution
390 | 
391 |             boxes[i, round(y0) : round(y1), round(x0) : round(x1)] = True
392 | 
393 |         return boxes.flatten(start_dim=1)
394 | 
395 |     def _obtain_self_masks(self, resolution, return_existing=False):
396 |         if (
397 |             (self.self_masks is None) or
398 |             (
399 |                 (self.cur_step % self.delta_refine_mask_steps == 0) and
400 |                 (self.cur_att_layer == 0) and
401 |                 (not return_existing)
402 |             )
403 |         ):
404 |             self.self_masks = self._fix_zero_masks(self._build_self_masks())
405 | 
406 |         b, s, n = self.self_masks.shape
407 |         mask_resolution = int(n ** 0.5)
408 |         self_masks = self.self_masks.reshape(b, s, mask_resolution, mask_resolution).float()
409 |         self_masks = F.interpolate(self_masks, resolution, mode='nearest-exact')
410 |         return self_masks.flatten(start_dim=2).bool()
411 | 
412 |     def _build_self_masks(self):
413 |         c, clusters = self._cluster_self_maps()  # b n
414 |         cluster_masks = torch.stack([(clusters == cluster_index) for cluster_index in range(c)], dim=2)  # b n c
415 |         cluster_area = cluster_masks.sum(dim=1, keepdim=True)  # b 1 c
416 | 
417 |         n = clusters.size(1)
418 |         resolution = int(n ** 0.5)
419 |         cross_masks = self._obtain_cross_masks(resolution)  # b s n
420 |         cross_mask_area = cross_masks.sum(dim=2, keepdim=True)  # b s 1
421 | 
422 |         intersection = torch.bmm(cross_masks.float(), cluster_masks.float())  # b s c
423 |         min_area = torch.minimum(cross_mask_area, cluster_area)  # b s c
424 |         score_per_cluster, subject_per_cluster = torch.max(intersection / min_area, dim=1)  # b c
425 |         subjects = torch.gather(subject_per_cluster, 1, clusters)  # b n
426 |         scores = torch.gather(score_per_cluster, 1, clusters)  # b n
427 | 
428 |         s = cross_masks.size(1)
429 |         self_masks = torch.stack([(subjects == subject_index) for subject_index in range(s)], dim=1)  # b s n
430 |         scores = scores.unsqueeze(1).expand(-1 ,s, n)  # b s n
431 |         self_masks[scores < self.self_mask_threshold] = False
432 |         self._save_maps(self_masks, 'self_masks')
433 |         return self_masks
434 | 
435 |     def _cluster_self_maps(self):  # b s n
436 |         self_maps = self._compute_maps(self.mean_self_map)  # b n m
437 |         if self.pca_rank is not None:
438 |             dtype = self_maps.dtype
439 |             _, _, eigen_vectors = torch.pca_lowrank(self_maps.float(), self.pca_rank)
440 |             self_maps = torch.matmul(self_maps, eigen_vectors.to(dtype=dtype))
441 | 
442 |         clustering_results = self.clustering(self_maps, centers=self.centers)
443 |         self.clustering.num_init = 1  # clustering is deterministic after the first time
444 |         self.centers = clustering_results.centers
445 |         clusters = clustering_results.labels
446 |         num_clusters = self.clustering.n_clusters
447 |         self._save_maps(clusters / num_clusters, f'clusters')
448 |         return num_clusters, clusters
449 | 
450 |     def _obtain_cross_masks(self, resolution, scale=10):
451 |         maps = self._compute_maps(self.mean_cross_map, resolution=resolution)  # b n k
452 |         maps = F.sigmoid(scale * (maps - self.cross_mask_threshold))
453 |         maps = self._normalize_maps(maps, reduce_min=True)
454 |         maps = maps.transpose(1, 2)  # b k n
455 |         existing_masks, _ = self._obtain_masks(
456 |             resolution, return_existing=True, batch_size=maps.size(0), device=maps.device)
457 |         maps = maps * existing_masks.to(dtype=maps.dtype)
458 |         self._save_maps(maps, 'cross_masks')
459 |         return maps
460 | 
461 |     def _fix_zero_masks(self, masks):
462 |         b, s, n = masks.shape
463 |         resolution = int(n ** 0.5)
464 |         boxes = self._convert_boxes_to_masks(resolution, device=masks.device)  # s n
465 | 
466 |         for i in range(b):
467 |             for j in range(s):
468 |                 if masks[i, j].sum() == 0:
469 |                     print('******Found a zero mask!******')
470 |                     for k in range(s):
471 |                         masks[i, k] = boxes[j] if (k == j) else masks[i, k].logical_and(~boxes[j])
472 | 
473 |         return masks
474 | 
475 |     def _compute_maps(self, maps, resolution=None):  # b n k
476 |         if resolution is not None:
477 |             b, n, k = maps.shape
478 |             original_resolution = int(n ** 0.5)
479 |             maps = maps.transpose(1, 2).reshape(b, k, original_resolution, original_resolution)
480 |             maps = F.interpolate(maps, resolution, mode='bilinear', antialias=True)
481 |             maps = maps.reshape(b, k, -1).transpose(1, 2)
482 | 
483 |         maps = self._normalize_maps(maps)
484 |         return maps
485 | 
486 |     @classmethod
487 |     def _normalize_maps(cls, maps, reduce_min=False):  # b n k
488 |         max_values = maps.max(dim=1, keepdim=True)[0]
489 |         min_values = maps.min(dim=1, keepdim=True)[0] if reduce_min else 0
490 |         numerator = maps - min_values
491 |         denominator = max_values - min_values + cls.EPSILON
492 |         return numerator / denominator
493 | 
494 |     def _save_maps(self, maps, prefix):
495 |         if self.map_dir is None or self.cur_step % self.delta_debug_mask_steps != 0:
496 |             return
497 | 
498 |         resolution = int(maps.size(-1) ** 0.5)
499 |         maps = maps.reshape(-1, 1, resolution, resolution).float()
500 |         maps = F.interpolate(maps, self.saved_resolution, mode='bilinear', antialias=True)
501 |         path = os.path.join(self.map_dir, f'map_{prefix}_{self.cur_step}_{self.cur_att_layer}.png')
502 |         torchvision.utils.save_image(maps, path)
503 | 
504 |     def _display_attention_maps(self, attention_maps, is_cross, num_heads, prefix=None):
505 |         if (not self.debug) or (self.cur_step == 0) or (self.cur_step % self.delta_debug_attention_steps > 0) or (self.cur_att_layer not in self.debug_layers):
506 |             return
507 | 
508 |         dir_name = self.map_dir
509 |         if prefix is not None:
510 |             splits = list(os.path.split(dir_name))
511 |             splits[-1] = '_'.join((prefix, splits[-1]))
512 |             dir_name = os.path.join(*splits)
513 | 
514 |         resolution = int(attention_maps.size(-2) ** 0.5)
515 |         if is_cross:
516 |             attention_maps = einops.rearrange(attention_maps, 'b (r1 r2) k -> b k r1 r2', r1=resolution)
517 |             attention_maps = F.interpolate(attention_maps, self.saved_resolution, mode='bilinear', antialias=True)
518 |             attention_maps = einops.rearrange(attention_maps, 'b k r1 r2 -> b (r1 r2) k')
519 | 
520 |         utils.display_attention_maps(
521 |             attention_maps,
522 |             is_cross,
523 |             num_heads,
524 |             self.model.tokenizer,
525 |             self.prompts,
526 |             dir_name,
527 |             self.cur_step,
528 |             self.cur_att_layer,
529 |             resolution,
530 |         )
531 | 
532 |     def _debug_hook(self, q, k, v, sim, attn, is_cross, place_in_unet, num_heads, **kwargs):
533 |         pass
534 | 


--------------------------------------------------------------------------------
/images/example.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/omer11a/bounded-attention/aeb6580e2f203f5f78e532b9c97f103219849e5c/images/example.jpg


--------------------------------------------------------------------------------
/images/teaser.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/omer11a/bounded-attention/aeb6580e2f203f5f78e532b9c97f103219849e5c/images/teaser.jpg


--------------------------------------------------------------------------------
/injection_utils.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | import numpy as np
  3 | import torch
  4 | import torch.nn as nn
  5 | import torch.nn.functional as F
  6 | 
  7 | from typing import Optional, Union, Tuple, List, Callable, Dict
  8 | 
  9 | from torchvision.utils import save_image
 10 | from einops import rearrange, repeat
 11 | 
 12 | 
 13 | class AttentionBase:
 14 |     def __init__(self):
 15 |         self.cur_step = 0
 16 |         self.num_att_layers = -1
 17 |         self.cur_att_layer = 0
 18 | 
 19 |     def before_step(self):
 20 |         pass
 21 | 
 22 |     def after_step(self):
 23 |         pass
 24 | 
 25 |     def __call__(self, q, k, v, is_cross, place_in_unet, num_heads, **kwargs):
 26 |         if self.cur_att_layer == 0:
 27 |             self.before_step()
 28 | 
 29 |         out = self.forward(q, k, v, is_cross, place_in_unet, num_heads, **kwargs)
 30 |         self.cur_att_layer += 1
 31 |         if self.cur_att_layer == self.num_att_layers:
 32 |             self.cur_att_layer = 0
 33 |             self.cur_step += 1
 34 |             self.after_step()
 35 | 
 36 |         return out
 37 | 
 38 |     def forward(self, q, k, v, is_cross, place_in_unet, num_heads, **kwargs):
 39 |         batch_size = q.size(0) // num_heads
 40 |         n = q.size(1)
 41 |         d = k.size(1)
 42 | 
 43 |         q = q.reshape(batch_size, num_heads, n, -1)
 44 |         k = k.reshape(batch_size, num_heads, d, -1)
 45 |         v = v.reshape(batch_size, num_heads, d, -1)
 46 |         out = F.scaled_dot_product_attention(q, k, v, attn_mask=kwargs['mask'])
 47 |         out = out.reshape(batch_size * num_heads, n, -1)
 48 |         out = rearrange(out, '(b h) n d -> b n (h d)', h=num_heads)
 49 |         return out
 50 | 
 51 |     def reset(self):
 52 |         self.cur_step = 0
 53 |         self.cur_att_layer = 0
 54 | 
 55 | 
 56 | def register_attention_editor_diffusers(model, editor: AttentionBase):
 57 |     """
 58 |     Register a attention editor to Diffuser Pipeline, refer from [Prompt-to-Prompt]
 59 |     """
 60 |     def ca_forward(self, place_in_unet):
 61 |         def forward(x, encoder_hidden_states=None, attention_mask=None, context=None, mask=None):
 62 |             """
 63 |             The attention is similar to the original implementation of LDM CrossAttention class
 64 |             except adding some modifications on the attention
 65 |             """
 66 |             if encoder_hidden_states is not None:
 67 |                 context = encoder_hidden_states
 68 |             if attention_mask is not None:
 69 |                 mask = attention_mask
 70 | 
 71 |             to_out = self.to_out
 72 |             if isinstance(to_out, nn.modules.container.ModuleList):
 73 |                 to_out = self.to_out[0]
 74 |             else:
 75 |                 to_out = self.to_out
 76 | 
 77 |             h = self.heads
 78 |             q = self.to_q(x)
 79 |             is_cross = context is not None
 80 |             context = context if is_cross else x
 81 |             k = self.to_k(context)
 82 |             v = self.to_v(context)
 83 |             q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
 84 |             out = editor(
 85 |                 q, k, v, is_cross, place_in_unet,
 86 |                 self.heads, scale=self.scale, mask=mask)
 87 | 
 88 |             return to_out(out)
 89 | 
 90 |         return forward
 91 | 
 92 |     def register_editor(net, count, place_in_unet):
 93 |         for name, subnet in net.named_children():
 94 |             if net.__class__.__name__ == 'Attention':  # spatial Transformer layer
 95 |                 net.original_forward = net.forward
 96 |                 net.forward = ca_forward(net, place_in_unet)
 97 |                 return count + 1
 98 |             elif hasattr(net, 'children'):
 99 |                 count = register_editor(subnet, count, place_in_unet)
100 |         return count
101 | 
102 |     cross_att_count = 0
103 |     for net_name, net in model.unet.named_children():
104 |         if "down" in net_name:
105 |             cross_att_count += register_editor(net, 0, "down")
106 |         elif "mid" in net_name:
107 |             cross_att_count += register_editor(net, 0, "mid")
108 |         elif "up" in net_name:
109 |             cross_att_count += register_editor(net, 0, "up")
110 | 
111 |     editor.num_att_layers = cross_att_count
112 |     editor.model = model
113 |     model.editor = editor
114 | 
115 | 
116 | def unregister_attention_editor_diffusers(model):
117 |     def unregister_editor(net):
118 |         for name, subnet in net.named_children():
119 |             if net.__class__.__name__ == 'Attention':  # spatial Transformer layer
120 |                 net.forward = net.original_forward
121 |                 net.original_forward = None
122 |             elif hasattr(net, 'children'):
123 |                 unregister_editor(subnet)
124 | 
125 |     for net_name, net in model.unet.named_children():
126 |         if "down" in net_name:
127 |             unregister_editor(net)
128 |         elif "mid" in net_name:
129 |             unregister_editor(net)
130 |         elif "up" in net_name:
131 |             unregister_editor(net)
132 | 
133 |     editor.model = None
134 |     model.editor = None
135 | 


--------------------------------------------------------------------------------
/pipeline_stable_diffusion_opt.py:
--------------------------------------------------------------------------------
  1 | """
  2 | Util functions based on Diffuser framework.
  3 | """
  4 | 
  5 | 
  6 | import os
  7 | import torch
  8 | import numpy as np
  9 | 
 10 | import torch.nn.functional as F
 11 | from tqdm import tqdm
 12 | from PIL import Image
 13 | from torchvision.utils import save_image
 14 | from torchvision.io import read_image
 15 | 
 16 | from diffusers import StableDiffusionPipeline
 17 | 
 18 | from pytorch_lightning import seed_everything
 19 | 
 20 | 
 21 | class StableDiffusionPipeline(StableDiffusionPipeline):
 22 |     def next_step(
 23 |         self,
 24 |         model_output: torch.FloatTensor,
 25 |         timestep: int,
 26 |         x: torch.FloatTensor,
 27 |         eta=0.,
 28 |         verbose=False
 29 |     ):
 30 |         """
 31 |         Inverse sampling for DDIM Inversion
 32 |         """
 33 |         if verbose:
 34 |             print("timestep: ", timestep)
 35 |         next_step = timestep
 36 |         timestep = min(timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps, 999)
 37 |         alpha_prod_t = self.scheduler.alphas_cumprod[timestep] if timestep >= 0 else self.scheduler.final_alpha_cumprod
 38 |         alpha_prod_t_next = self.scheduler.alphas_cumprod[next_step]
 39 |         beta_prod_t = 1 - alpha_prod_t
 40 |         pred_x0 = (x - beta_prod_t**0.5 * model_output) / alpha_prod_t**0.5
 41 |         pred_dir = (1 - alpha_prod_t_next)**0.5 * model_output
 42 |         x_next = alpha_prod_t_next**0.5 * pred_x0 + pred_dir
 43 |         return x_next, pred_x0
 44 | 
 45 |     def step(
 46 |         self,
 47 |         model_output: torch.FloatTensor,
 48 |         timestep: int,
 49 |         x: torch.FloatTensor,
 50 |         eta: float=0.0,
 51 |         verbose=False,
 52 |     ):
 53 |         """
 54 |         predict the sampe the next step in the denoise process.
 55 |         """
 56 |         prev_timestep = timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
 57 |         alpha_prod_t = self.scheduler.alphas_cumprod[timestep]
 58 |         alpha_prod_t_prev = self.scheduler.alphas_cumprod[prev_timestep] if prev_timestep > 0 else self.scheduler.final_alpha_cumprod
 59 |         beta_prod_t = 1 - alpha_prod_t
 60 |         pred_x0 = (x - beta_prod_t**0.5 * model_output) / alpha_prod_t**0.5
 61 |         pred_dir = (1 - alpha_prod_t_prev)**0.5 * model_output
 62 |         x_prev = alpha_prod_t_prev**0.5 * pred_x0 + pred_dir
 63 |         return x_prev, pred_x0
 64 | 
 65 |     @torch.no_grad()
 66 |     def image2latent(self, image):
 67 |         DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 68 |         if type(image) is Image:
 69 |             image = np.array(image)
 70 |             image = torch.from_numpy(image).float() / 127.5 - 1
 71 |             image = image.permute(2, 0, 1).unsqueeze(0).to(DEVICE)
 72 |         # input image density range [-1, 1]
 73 |         latents = self.vae.encode(image)['latent_dist'].mean
 74 |         latents = latents * 0.18215
 75 |         return latents
 76 | 
 77 |     @torch.no_grad()
 78 |     def latent2image(self, latents, return_type='np'):
 79 |         latents = latents.half()
 80 |         latents = 1 / 0.18215 * latents.detach()
 81 |         image = self.vae.decode(latents)['sample']
 82 |         if return_type == 'np':
 83 |             image = (image / 2 + 0.5).clamp(0, 1)
 84 |             image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
 85 |             image = (image * 255).astype(np.uint8)
 86 |         elif return_type == "pt":
 87 |             image = (image / 2 + 0.5).clamp(0, 1)
 88 | 
 89 |         return image
 90 | 
 91 |     def latent2image_grad(self, latents):
 92 |         latents = 1 / 0.18215 * latents
 93 |         image = self.vae.decode(latents)['sample']
 94 | 
 95 |         return image  # range [-1, 1]
 96 |     
 97 |     def update_loss(self, latents, text_embeddings, i, t):
 98 |         def forward_pass(latent_model_input):
 99 |             self.unet(
100 |                 latent_model_input,
101 |                 t,
102 |                 encoder_hidden_states=text_embeddings,
103 |             )
104 |             self.unet.zero_grad()
105 | 
106 |         return self.editor.update_loss(forward_pass, latents, i)
107 | 
108 |     @torch.no_grad()
109 |     def __call__(
110 |         self,
111 |         prompt,
112 |         batch_size=1,
113 |         height=512,
114 |         width=512,
115 |         num_inference_steps=50,
116 |         guidance_scale=7.5,
117 |         eta=0.0,
118 |         latents=None,
119 |         unconditioning=None,
120 |         neg_prompt=None,
121 |         ref_intermediate_latents=None,
122 |         return_intermediates=False,
123 |         **kwds):
124 |         DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
125 |         if isinstance(prompt, list):
126 |             batch_size = len(prompt)
127 |         elif isinstance(prompt, str):
128 |             if batch_size > 1:
129 |                 prompt = [prompt] * batch_size
130 | 
131 |         # text embeddings
132 |         text_input = self.tokenizer(
133 |             prompt,
134 |             padding="max_length",
135 |             max_length=77,
136 |             return_tensors="pt"
137 |         )
138 | 
139 |         text_embeddings = self.text_encoder(text_input.input_ids.to(DEVICE))[0]
140 |         print("input text embeddings :", text_embeddings.shape)
141 |         if kwds.get("dir"):
142 |             dir = text_embeddings[-2] - text_embeddings[-1]
143 |             u, s, v = torch.pca_lowrank(dir.transpose(-1, -2), q=1, center=True)
144 |             text_embeddings[-1] = text_embeddings[-1] + kwds.get("dir") * v
145 |             print(u.shape)
146 |             print(v.shape)
147 | 
148 |         # define initial latents
149 |         latents_shape = (batch_size, self.unet.in_channels, height//8, width//8)
150 |         if latents is None:
151 |             latents = torch.randn(latents_shape, device=DEVICE)
152 |         else:
153 |             assert latents.shape == latents_shape, f"The shape of input latent tensor {latents.shape} should equal to predefined one."
154 | 
155 |         # unconditional embedding for classifier free guidance
156 |         if guidance_scale > 1.:
157 |             max_length = text_input.input_ids.shape[-1]
158 |             if neg_prompt:
159 |                 uc_text = neg_prompt
160 |             else:
161 |                 uc_text = ""
162 |             # uc_text = "ugly, tiling, poorly drawn hands, poorly drawn feet, body out of frame, cut off, low contrast, underexposed, distorted face"
163 |             unconditional_input = self.tokenizer(
164 |                 [uc_text] * batch_size,
165 |                 padding="max_length",
166 |                 max_length=77,
167 |                 return_tensors="pt"
168 |             )
169 |             # unconditional_input.input_ids = unconditional_input.input_ids[:, 1:]
170 |             unconditional_embeddings = self.text_encoder(unconditional_input.input_ids.to(DEVICE))[0]
171 |             text_embeddings = torch.cat([unconditional_embeddings, text_embeddings], dim=0)
172 | 
173 |         print("latents shape: ", latents.shape)
174 |         # iterative sampling
175 |         self.scheduler.set_timesteps(num_inference_steps)
176 |         # print("Valid timesteps: ", reversed(self.scheduler.timesteps))
177 |         latents_list = [latents]
178 |         pred_x0_list = [latents]
179 |         latents = latents.half()
180 |         text_embeddings = text_embeddings.half()
181 |         for i, t in enumerate(tqdm(self.scheduler.timesteps, desc="DDIM Sampler")):
182 |             if ref_intermediate_latents is not None:
183 |                 # note that the batch_size >= 2
184 |                 latents_ref = ref_intermediate_latents[-1 - i]
185 |                 _, latents_cur = latents.chunk(2)
186 |                 latents = torch.cat([latents_ref, latents_cur])
187 | 
188 |             if unconditioning is not None and isinstance(unconditioning, list):
189 |                 _, text_embeddings = text_embeddings.chunk(2)
190 |                 text_embeddings = torch.cat([unconditioning[i].expand(*text_embeddings.shape), text_embeddings]) 
191 | 
192 |             #with torch.enable_grad():
193 |             latents = self.update_loss(latents, text_embeddings, i, t)
194 | 
195 |             if guidance_scale > 1.:
196 |                 model_inputs = torch.cat([latents] * 2)
197 |             else:
198 |                 model_inputs = latents
199 |             #model_inputs = self.scheduler.scale_model_input(model_inputs, t)
200 | 
201 |             # predict tghe noise
202 |             noise_pred = self.unet(model_inputs, t, encoder_hidden_states=text_embeddings).sample
203 |             if guidance_scale > 1.:
204 |                 noise_pred_uncon, noise_pred_con = noise_pred.chunk(2, dim=0)
205 |                 noise_pred = noise_pred_uncon + guidance_scale * (noise_pred_con - noise_pred_uncon)
206 |             # compute the previous noise sample x_t -> x_t-1
207 |             #print('noise_pred.shape', noise_pred.shape)
208 |             #print('latents.shape', latents.shape)
209 |             latents, pred_x0 = self.step(noise_pred, t, latents)
210 |             #print('new_latents.shape', latents.shape)
211 |             #print('pred_x0.shape', pred_x0.shape)
212 |             latents_list.append(latents)
213 |             pred_x0_list.append(pred_x0)
214 | 
215 |         image = self.latent2image(latents, return_type="pt")
216 |         if return_intermediates:
217 |             pred_x0_list = [self.latent2image(img, return_type="pt") for img in pred_x0_list]
218 |             latents_list = [self.latent2image(img, return_type="pt") for img in latents_list]
219 |             return image, pred_x0_list, latents_list
220 |         return image
221 | 
222 |     @torch.no_grad()
223 |     def invert(
224 |         self,
225 |         image: torch.Tensor,
226 |         prompt,
227 |         num_inference_steps=50,
228 |         guidance_scale=7.5,
229 |         eta=0.0,
230 |         return_intermediates=False,
231 |         **kwds):
232 |         """
233 |         invert a real image into noise map with determinisc DDIM inversion
234 |         """
235 |         DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
236 |         batch_size = image.shape[0]
237 |         if isinstance(prompt, list):
238 |             if batch_size == 1:
239 |                 image = image.expand(len(prompt), -1, -1, -1)
240 |         elif isinstance(prompt, str):
241 |             if batch_size > 1:
242 |                 prompt = [prompt] * batch_size
243 | 
244 |         # text embeddings
245 |         text_input = self.tokenizer(
246 |             prompt,
247 |             padding="max_length",
248 |             max_length=77,
249 |             return_tensors="pt"
250 |         )
251 |         text_embeddings = self.text_encoder(text_input.input_ids.to(DEVICE))[0]
252 |         print("input text embeddings :", text_embeddings.shape)
253 |         # define initial latents
254 |         latents = self.image2latent(image)
255 |         start_latents = latents
256 |         # print(latents)
257 |         # exit()
258 |         # unconditional embedding for classifier free guidance
259 |         if guidance_scale > 1.:
260 |             max_length = text_input.input_ids.shape[-1]
261 |             unconditional_input = self.tokenizer(
262 |                 [""] * batch_size,
263 |                 padding="max_length",
264 |                 max_length=77,
265 |                 return_tensors="pt"
266 |             )
267 |             unconditional_embeddings = self.text_encoder(unconditional_input.input_ids.to(DEVICE))[0]
268 |             text_embeddings = torch.cat([unconditional_embeddings, text_embeddings], dim=0)
269 | 
270 |         print("latents shape: ", latents.shape)
271 |         # interative sampling
272 |         self.scheduler.set_timesteps(num_inference_steps)
273 |         print("Valid timesteps: ", reversed(self.scheduler.timesteps))
274 |         # print("attributes: ", self.scheduler.__dict__)
275 |         latents_list = [latents]
276 |         pred_x0_list = [latents]
277 |         for i, t in enumerate(tqdm(reversed(self.scheduler.timesteps), desc="DDIM Inversion")):
278 |             if guidance_scale > 1.:
279 |                 model_inputs = torch.cat([latents] * 2)
280 |             else:
281 |                 model_inputs = latents
282 | 
283 |             # predict the noise
284 |             noise_pred = self.unet(model_inputs, t, encoder_hidden_states=text_embeddings).sample
285 |             if guidance_scale > 1.:
286 |                 noise_pred_uncon, noise_pred_con = noise_pred.chunk(2, dim=0)
287 |                 noise_pred = noise_pred_uncon + guidance_scale * (noise_pred_con - noise_pred_uncon)
288 |             # compute the previous noise sample x_t-1 -> x_t
289 |             latents, pred_x0 = self.next_step(noise_pred, t, latents)
290 |             latents_list.append(latents)
291 |             pred_x0_list.append(pred_x0)
292 | 
293 |         if return_intermediates:
294 |             # return the intermediate laters during inversion
295 |             # pred_x0_list = [self.latent2image(img, return_type="pt") for img in pred_x0_list]
296 |             return latents, latents_list
297 |         return latents, start_latents
298 | 


--------------------------------------------------------------------------------
/pipeline_stable_diffusion_xl_opt.py:
--------------------------------------------------------------------------------
  1 | # Copyright 2023 The HuggingFace Team. All rights reserved.
  2 | #
  3 | # Licensed under the Apache License, Version 2.0 (the "License");
  4 | # you may not use this file except in compliance with the License.
  5 | # You may obtain a copy of the License at
  6 | #
  7 | #     http://www.apache.org/licenses/LICENSE-2.0
  8 | #
  9 | # Unless required by applicable law or agreed to in writing, software
 10 | # distributed under the License is distributed on an "AS IS" BASIS,
 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12 | # See the License for the specific language governing permissions and
 13 | # limitations under the License.
 14 | 
 15 | import inspect
 16 | import os
 17 | from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 18 | 
 19 | import torch
 20 | from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
 21 | 
 22 | from diffusers.image_processor import VaeImageProcessor
 23 | from diffusers.loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin
 24 | from diffusers.models import AutoencoderKL, UNet2DConditionModel
 25 | from diffusers.models.attention_processor import (
 26 |     AttnProcessor2_0,
 27 |     LoRAAttnProcessor2_0,
 28 |     LoRAXFormersAttnProcessor,
 29 |     XFormersAttnProcessor,
 30 | )
 31 | from diffusers.schedulers import KarrasDiffusionSchedulers
 32 | from diffusers.utils import (
 33 |     is_accelerate_available,
 34 |     is_accelerate_version,
 35 |     is_invisible_watermark_available,
 36 |     logging,
 37 |     randn_tensor,
 38 |     replace_example_docstring,
 39 | )
 40 | from diffusers.pipeline_utils import DiffusionPipeline
 41 | from diffusers.pipelines.stable_diffusion_xl import StableDiffusionXLPipelineOutput
 42 | 
 43 | 
 44 | if is_invisible_watermark_available():
 45 |     from .watermark import StableDiffusionXLWatermarker
 46 | 
 47 | 
 48 | logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 49 | 
 50 | EXAMPLE_DOC_STRING = """
 51 |     Examples:
 52 |         ```py
 53 |         >>> import torch
 54 |         >>> from diffusers import StableDiffusionXLPipeline
 55 | 
 56 |         >>> pipe = StableDiffusionXLPipeline.from_pretrained(
 57 |         ...     "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
 58 |         ... )
 59 |         >>> pipe = pipe.to("cuda")
 60 | 
 61 |         >>> prompt = "a photo of an astronaut riding a horse on mars"
 62 |         >>> image = pipe(prompt).images[0]
 63 |         ```
 64 | """
 65 | 
 66 | 
 67 | # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg
 68 | def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
 69 |     """
 70 |     Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
 71 |     Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
 72 |     """
 73 |     std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
 74 |     std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
 75 |     # rescale the results from guidance (fixes overexposure)
 76 |     noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
 77 |     # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
 78 |     noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
 79 |     return noise_cfg
 80 | 
 81 | 
 82 | class StableDiffusionXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoaderMixin):
 83 |     r"""
 84 |     Pipeline for text-to-image generation using Stable Diffusion XL.
 85 | 
 86 |     This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
 87 |     library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
 88 | 
 89 |     In addition the pipeline inherits the following loading methods:
 90 |         - *LoRA*: [`StableDiffusionXLPipeline.load_lora_weights`]
 91 |         - *Ckpt*: [`loaders.FromSingleFileMixin.from_single_file`]
 92 | 
 93 |     as well as the following saving methods:
 94 |         - *LoRA*: [`loaders.StableDiffusionXLPipeline.save_lora_weights`]
 95 | 
 96 |     Args:
 97 |         vae ([`AutoencoderKL`]):
 98 |             Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
 99 |         text_encoder ([`CLIPTextModel`]):
100 |             Frozen text-encoder. Stable Diffusion XL uses the text portion of
101 |             [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
102 |             the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
103 |         text_encoder_2 ([` CLIPTextModelWithProjection`]):
104 |             Second frozen text-encoder. Stable Diffusion XL uses the text and pool portion of
105 |             [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection),
106 |             specifically the
107 |             [laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)
108 |             variant.
109 |         tokenizer (`CLIPTokenizer`):
110 |             Tokenizer of class
111 |             [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
112 |         tokenizer_2 (`CLIPTokenizer`):
113 |             Second Tokenizer of class
114 |             [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
115 |         unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
116 |         scheduler ([`SchedulerMixin`]):
117 |             A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
118 |             [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
119 |     """
120 | 
121 |     def __init__(
122 |         self,
123 |         vae: AutoencoderKL,
124 |         text_encoder: CLIPTextModel,
125 |         text_encoder_2: CLIPTextModelWithProjection,
126 |         tokenizer: CLIPTokenizer,
127 |         tokenizer_2: CLIPTokenizer,
128 |         unet: UNet2DConditionModel,
129 |         scheduler: KarrasDiffusionSchedulers,
130 |         force_zeros_for_empty_prompt: bool = True,
131 |         add_watermarker: Optional[bool] = None,
132 |     ):
133 |         super().__init__()
134 | 
135 |         self.register_modules(
136 |             vae=vae,
137 |             text_encoder=text_encoder,
138 |             text_encoder_2=text_encoder_2,
139 |             tokenizer=tokenizer,
140 |             tokenizer_2=tokenizer_2,
141 |             unet=unet,
142 |             scheduler=scheduler,
143 |         )
144 |         self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt)
145 |         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
146 |         self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
147 |         self.default_sample_size = self.unet.config.sample_size
148 | 
149 |         add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available()
150 | 
151 |         if add_watermarker:
152 |             self.watermark = StableDiffusionXLWatermarker()
153 |         else:
154 |             self.watermark = None
155 | 
156 |     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_slicing
157 |     def enable_vae_slicing(self):
158 |         r"""
159 |         Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
160 |         compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
161 |         """
162 |         self.vae.enable_slicing()
163 | 
164 |     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_slicing
165 |     def disable_vae_slicing(self):
166 |         r"""
167 |         Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
168 |         computing decoding in one step.
169 |         """
170 |         self.vae.disable_slicing()
171 | 
172 |     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_tiling
173 |     def enable_vae_tiling(self):
174 |         r"""
175 |         Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
176 |         compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
177 |         processing larger images.
178 |         """
179 |         self.vae.enable_tiling()
180 | 
181 |     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_tiling
182 |     def disable_vae_tiling(self):
183 |         r"""
184 |         Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
185 |         computing decoding in one step.
186 |         """
187 |         self.vae.disable_tiling()
188 | 
189 |     def enable_model_cpu_offload(self, gpu_id=0):
190 |         r"""
191 |         Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
192 |         to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
193 |         method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with
194 |         `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.
195 |         """
196 |         if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
197 |             from accelerate import cpu_offload_with_hook
198 |         else:
199 |             raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.")
200 | 
201 |         device = torch.device(f"cuda:{gpu_id}")
202 | 
203 |         if self.device.type != "cpu":
204 |             self.to("cpu", silence_dtype_warnings=True)
205 |             torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
206 | 
207 |         model_sequence = (
208 |             [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2]
209 |         )
210 |         model_sequence.extend([self.unet, self.vae])
211 | 
212 |         hook = None
213 |         for cpu_offloaded_model in model_sequence:
214 |             _, hook = cpu_offload_with_hook(cpu_offloaded_model, device, prev_module_hook=hook)
215 | 
216 |         # We'll offload the last model manually.
217 |         self.final_offload_hook = hook
218 | 
219 |     def encode_prompt(
220 |         self,
221 |         prompt: str,
222 |         prompt_2: Optional[str] = None,
223 |         device: Optional[torch.device] = None,
224 |         num_images_per_prompt: int = 1,
225 |         do_classifier_free_guidance: bool = True,
226 |         negative_prompt: Optional[str] = None,
227 |         negative_prompt_2: Optional[str] = None,
228 |         prompt_embeds: Optional[torch.FloatTensor] = None,
229 |         negative_prompt_embeds: Optional[torch.FloatTensor] = None,
230 |         pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
231 |         negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
232 |         lora_scale: Optional[float] = None,
233 |     ):
234 |         r"""
235 |         Encodes the prompt into text encoder hidden states.
236 | 
237 |         Args:
238 |             prompt (`str` or `List[str]`, *optional*):
239 |                 prompt to be encoded
240 |             prompt_2 (`str` or `List[str]`, *optional*):
241 |                 The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
242 |                 used in both text-encoders
243 |             device: (`torch.device`):
244 |                 torch device
245 |             num_images_per_prompt (`int`):
246 |                 number of images that should be generated per prompt
247 |             do_classifier_free_guidance (`bool`):
248 |                 whether to use classifier free guidance or not
249 |             negative_prompt (`str` or `List[str]`, *optional*):
250 |                 The prompt or prompts not to guide the image generation. If not defined, one has to pass
251 |                 `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
252 |                 less than `1`).
253 |             negative_prompt_2 (`str` or `List[str]`, *optional*):
254 |                 The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
255 |                 `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
256 |             prompt_embeds (`torch.FloatTensor`, *optional*):
257 |                 Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
258 |                 provided, text embeddings will be generated from `prompt` input argument.
259 |             negative_prompt_embeds (`torch.FloatTensor`, *optional*):
260 |                 Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
261 |                 weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
262 |                 argument.
263 |             pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
264 |                 Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
265 |                 If not provided, pooled text embeddings will be generated from `prompt` input argument.
266 |             negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
267 |                 Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
268 |                 weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
269 |                 input argument.
270 |             lora_scale (`float`, *optional*):
271 |                 A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
272 |         """
273 |         device = device or self._execution_device
274 | 
275 |         # set lora scale so that monkey patched LoRA
276 |         # function of text encoder can correctly access it
277 |         if lora_scale is not None and isinstance(self, LoraLoaderMixin):
278 |             self._lora_scale = lora_scale
279 | 
280 |         if prompt is not None and isinstance(prompt, str):
281 |             batch_size = 1
282 |         elif prompt is not None and isinstance(prompt, list):
283 |             batch_size = len(prompt)
284 |         else:
285 |             batch_size = prompt_embeds.shape[0]
286 | 
287 |         # Define tokenizers and text encoders
288 |         tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2]
289 |         text_encoders = (
290 |             [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2]
291 |         )
292 | 
293 |         if prompt_embeds is None:
294 |             prompt_2 = prompt_2 or prompt
295 |             # textual inversion: procecss multi-vector tokens if necessary
296 |             prompt_embeds_list = []
297 |             prompts = [prompt, prompt_2]
298 |             for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
299 |                 if isinstance(self, TextualInversionLoaderMixin):
300 |                     prompt = self.maybe_convert_prompt(prompt, tokenizer)
301 | 
302 |                 text_inputs = tokenizer(
303 |                     prompt,
304 |                     padding="max_length",
305 |                     max_length=tokenizer.model_max_length,
306 |                     truncation=True,
307 |                     return_tensors="pt",
308 |                 )
309 | 
310 |                 text_input_ids = text_inputs.input_ids
311 |                 untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
312 | 
313 |                 if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
314 |                     text_input_ids, untruncated_ids
315 |                 ):
316 |                     removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1 : -1])
317 |                     logger.warning(
318 |                         "The following part of your input was truncated because CLIP can only handle sequences up to"
319 |                         f" {tokenizer.model_max_length} tokens: {removed_text}"
320 |                     )
321 | 
322 |                 prompt_embeds = text_encoder(
323 |                     text_input_ids.to(device),
324 |                     output_hidden_states=True,
325 |                 )
326 | 
327 |                 # We are only ALWAYS interested in the pooled output of the final text encoder
328 |                 pooled_prompt_embeds = prompt_embeds[0]
329 |                 ### TODO: remove
330 |                 null_text_inputs = tokenizer(
331 |                     ['a realistic photo of an empty background'] * batch_size,
332 |                     padding="max_length",
333 |                     max_length=tokenizer.model_max_length,
334 |                     truncation=True,
335 |                     return_tensors="pt",
336 |                 )
337 |                 null_input_ids = null_text_inputs.input_ids
338 |                 null_prompt_embeds = text_encoder(
339 |                     null_input_ids.to(device),
340 |                     output_hidden_states=True,
341 |                 )
342 |                 pooled_prompt_embeds = null_prompt_embeds[0]
343 |                 ### TODO: remove
344 |                 prompt_embeds = prompt_embeds.hidden_states[-2]
345 | 
346 |                 prompt_embeds_list.append(prompt_embeds)
347 | 
348 |             prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
349 | 
350 |         # get unconditional embeddings for classifier free guidance
351 |         zero_out_negative_prompt = negative_prompt is None and self.config.force_zeros_for_empty_prompt
352 |         if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt:
353 |             negative_prompt_embeds = torch.zeros_like(prompt_embeds)
354 |             negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
355 |         elif do_classifier_free_guidance and negative_prompt_embeds is None:
356 |             negative_prompt = negative_prompt or ""
357 |             negative_prompt_2 = negative_prompt_2 or negative_prompt
358 | 
359 |             uncond_tokens: List[str]
360 |             if prompt is not None and type(prompt) is not type(negative_prompt):
361 |                 raise TypeError(
362 |                     f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
363 |                     f" {type(prompt)}."
364 |                 )
365 |             elif isinstance(negative_prompt, str):
366 |                 uncond_tokens = [negative_prompt, negative_prompt_2]
367 |             elif batch_size != len(negative_prompt):
368 |                 raise ValueError(
369 |                     f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
370 |                     f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
371 |                     " the batch size of `prompt`."
372 |                 )
373 |             else:
374 |                 uncond_tokens = [negative_prompt, negative_prompt_2]
375 | 
376 |             negative_prompt_embeds_list = []
377 |             for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders):
378 |                 if isinstance(self, TextualInversionLoaderMixin):
379 |                     negative_prompt = self.maybe_convert_prompt(negative_prompt, tokenizer)
380 | 
381 |                 max_length = prompt_embeds.shape[1]
382 |                 uncond_input = tokenizer(
383 |                     negative_prompt,
384 |                     padding="max_length",
385 |                     max_length=max_length,
386 |                     truncation=True,
387 |                     return_tensors="pt",
388 |                 )
389 | 
390 |                 negative_prompt_embeds = text_encoder(
391 |                     uncond_input.input_ids.to(device),
392 |                     output_hidden_states=True,
393 |                 )
394 |                 # We are only ALWAYS interested in the pooled output of the final text encoder
395 |                 negative_pooled_prompt_embeds = negative_prompt_embeds[0]
396 |                 negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2]
397 | 
398 |                 negative_prompt_embeds_list.append(negative_prompt_embeds)
399 | 
400 |             negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1)
401 | 
402 |         prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
403 |         bs_embed, seq_len, _ = prompt_embeds.shape
404 |         # duplicate text embeddings for each generation per prompt, using mps friendly method
405 |         prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
406 |         prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
407 | 
408 |         if do_classifier_free_guidance:
409 |             # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
410 |             seq_len = negative_prompt_embeds.shape[1]
411 |             negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
412 |             negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
413 |             negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
414 | 
415 |         pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
416 |             bs_embed * num_images_per_prompt, -1
417 |         )
418 |         if do_classifier_free_guidance:
419 |             negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
420 |                 bs_embed * num_images_per_prompt, -1
421 |             )
422 | 
423 |         return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
424 | 
425 |     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
426 |     def prepare_extra_step_kwargs(self, generator, eta):
427 |         # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
428 |         # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
429 |         # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
430 |         # and should be between [0, 1]
431 | 
432 |         accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
433 |         extra_step_kwargs = {}
434 |         if accepts_eta:
435 |             extra_step_kwargs["eta"] = eta
436 | 
437 |         # check if the scheduler accepts generator
438 |         accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
439 |         if accepts_generator:
440 |             extra_step_kwargs["generator"] = generator
441 |         return extra_step_kwargs
442 | 
443 |     def check_inputs(
444 |         self,
445 |         prompt,
446 |         prompt_2,
447 |         height,
448 |         width,
449 |         callback_steps,
450 |         negative_prompt=None,
451 |         negative_prompt_2=None,
452 |         prompt_embeds=None,
453 |         negative_prompt_embeds=None,
454 |         pooled_prompt_embeds=None,
455 |         negative_pooled_prompt_embeds=None,
456 |     ):
457 |         if height % 8 != 0 or width % 8 != 0:
458 |             raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
459 | 
460 |         if (callback_steps is None) or (
461 |             callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
462 |         ):
463 |             raise ValueError(
464 |                 f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
465 |                 f" {type(callback_steps)}."
466 |             )
467 | 
468 |         if prompt is not None and prompt_embeds is not None:
469 |             raise ValueError(
470 |                 f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
471 |                 " only forward one of the two."
472 |             )
473 |         elif prompt_2 is not None and prompt_embeds is not None:
474 |             raise ValueError(
475 |                 f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
476 |                 " only forward one of the two."
477 |             )
478 |         elif prompt is None and prompt_embeds is None:
479 |             raise ValueError(
480 |                 "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
481 |             )
482 |         elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
483 |             raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
484 |         elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
485 |             raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
486 | 
487 |         if negative_prompt is not None and negative_prompt_embeds is not None:
488 |             raise ValueError(
489 |                 f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
490 |                 f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
491 |             )
492 |         elif negative_prompt_2 is not None and negative_prompt_embeds is not None:
493 |             raise ValueError(
494 |                 f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:"
495 |                 f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
496 |             )
497 | 
498 |         if prompt_embeds is not None and negative_prompt_embeds is not None:
499 |             if prompt_embeds.shape != negative_prompt_embeds.shape:
500 |                 raise ValueError(
501 |                     "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
502 |                     f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
503 |                     f" {negative_prompt_embeds.shape}."
504 |                 )
505 | 
506 |         if prompt_embeds is not None and pooled_prompt_embeds is None:
507 |             raise ValueError(
508 |                 "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
509 |             )
510 | 
511 |         if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:
512 |             raise ValueError(
513 |                 "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`."
514 |             )
515 | 
516 |     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
517 |     def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
518 |         shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
519 |         if isinstance(generator, list) and len(generator) != batch_size:
520 |             raise ValueError(
521 |                 f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
522 |                 f" size of {batch_size}. Make sure the batch size matches the length of the generators."
523 |             )
524 | 
525 |         if latents is None:
526 |             latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
527 |         else:
528 |             latents = latents.to(device)
529 | 
530 |         # scale the initial noise by the standard deviation required by the scheduler
531 |         latents = latents * self.scheduler.init_noise_sigma
532 |         return latents
533 | 
534 |     def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, dtype):
535 |         add_time_ids = list(original_size + crops_coords_top_left + target_size)
536 | 
537 |         passed_add_embed_dim = (
538 |             self.unet.config.addition_time_embed_dim * len(add_time_ids) + self.text_encoder_2.config.projection_dim
539 |         )
540 |         expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
541 | 
542 |         if expected_add_embed_dim != passed_add_embed_dim:
543 |             raise ValueError(
544 |                 f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
545 |             )
546 | 
547 |         add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
548 |         return add_time_ids
549 | 
550 |     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae
551 |     def upcast_vae(self):
552 |         dtype = self.vae.dtype
553 |         self.vae.to(dtype=torch.float32)
554 |         use_torch_2_0_or_xformers = isinstance(
555 |             self.vae.decoder.mid_block.attentions[0].processor,
556 |             (
557 |                 AttnProcessor2_0,
558 |                 XFormersAttnProcessor,
559 |                 LoRAXFormersAttnProcessor,
560 |                 LoRAAttnProcessor2_0,
561 |             ),
562 |         )
563 |         # if xformers or torch_2_0 is used attention block does not need
564 |         # to be in float32 which can save lots of memory
565 |         if use_torch_2_0_or_xformers:
566 |             self.vae.post_quant_conv.to(dtype)
567 |             self.vae.decoder.conv_in.to(dtype)
568 |             self.vae.decoder.mid_block.to(dtype)
569 | 
570 |     def update_loss(self, latents, i, t, prompt_embeds, cross_attention_kwargs, add_text_embeds, add_time_ids):
571 |         def forward_pass(latent_model_input):
572 |             latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
573 |             added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
574 |             self.unet(
575 |                 latent_model_input,
576 |                 t,
577 |                 encoder_hidden_states=prompt_embeds,
578 |                 cross_attention_kwargs=cross_attention_kwargs,
579 |                 added_cond_kwargs=added_cond_kwargs,
580 |                 return_dict=False,
581 |             )
582 |             self.unet.zero_grad()
583 | 
584 |         return self.editor.update_loss(forward_pass, latents, i)
585 | 
586 |     @torch.no_grad()
587 |     @replace_example_docstring(EXAMPLE_DOC_STRING)
588 |     def __call__(
589 |         self,
590 |         prompt: Union[str, List[str]] = None,
591 |         prompt_2: Optional[Union[str, List[str]]] = None,
592 |         height: Optional[int] = None,
593 |         width: Optional[int] = None,
594 |         num_inference_steps: int = 50,
595 |         denoising_end: Optional[float] = None,
596 |         guidance_scale: float = 5.0,
597 |         negative_prompt: Optional[Union[str, List[str]]] = None,
598 |         negative_prompt_2: Optional[Union[str, List[str]]] = None,
599 |         num_images_per_prompt: Optional[int] = 1,
600 |         eta: float = 0.0,
601 |         generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
602 |         latents: Optional[torch.FloatTensor] = None,
603 |         prompt_embeds: Optional[torch.FloatTensor] = None,
604 |         negative_prompt_embeds: Optional[torch.FloatTensor] = None,
605 |         pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
606 |         negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
607 |         output_type: Optional[str] = "pil",
608 |         return_dict: bool = True,
609 |         callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
610 |         callback_steps: int = 1,
611 |         cross_attention_kwargs: Optional[Dict[str, Any]] = None,
612 |         guidance_rescale: float = 0.0,
613 |         original_size: Optional[Tuple[int, int]] = None,
614 |         crops_coords_top_left: Tuple[int, int] = (0, 0),
615 |         target_size: Optional[Tuple[int, int]] = None,
616 |     ):
617 |         r"""
618 |         Function invoked when calling the pipeline for generation.
619 | 
620 |         Args:
621 |             prompt (`str` or `List[str]`, *optional*):
622 |                 The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
623 |                 instead.
624 |             prompt_2 (`str` or `List[str]`, *optional*):
625 |                 The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
626 |                 used in both text-encoders
627 |             height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
628 |                 The height in pixels of the generated image.
629 |             width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
630 |                 The width in pixels of the generated image.
631 |             num_inference_steps (`int`, *optional*, defaults to 50):
632 |                 The number of denoising steps. More denoising steps usually lead to a higher quality image at the
633 |                 expense of slower inference.
634 |             denoising_end (`float`, *optional*):
635 |                 When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be
636 |                 completed before it is intentionally prematurely terminated. As a result, the returned sample will
637 |                 still retain a substantial amount of noise as determined by the discrete timesteps selected by the
638 |                 scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a
639 |                 "Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image
640 |                 Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output)
641 |             guidance_scale (`float`, *optional*, defaults to 5.0):
642 |                 Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
643 |                 `guidance_scale` is defined as `w` of equation 2. of [Imagen
644 |                 Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
645 |                 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
646 |                 usually at the expense of lower image quality.
647 |             negative_prompt (`str` or `List[str]`, *optional*):
648 |                 The prompt or prompts not to guide the image generation. If not defined, one has to pass
649 |                 `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
650 |                 less than `1`).
651 |             negative_prompt_2 (`str` or `List[str]`, *optional*):
652 |                 The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
653 |                 `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
654 |             num_images_per_prompt (`int`, *optional*, defaults to 1):
655 |                 The number of images to generate per prompt.
656 |             eta (`float`, *optional*, defaults to 0.0):
657 |                 Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
658 |                 [`schedulers.DDIMScheduler`], will be ignored for others.
659 |             generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
660 |                 One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
661 |                 to make generation deterministic.
662 |             latents (`torch.FloatTensor`, *optional*):
663 |                 Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
664 |                 generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
665 |                 tensor will ge generated by sampling using the supplied random `generator`.
666 |             prompt_embeds (`torch.FloatTensor`, *optional*):
667 |                 Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
668 |                 provided, text embeddings will be generated from `prompt` input argument.
669 |             negative_prompt_embeds (`torch.FloatTensor`, *optional*):
670 |                 Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
671 |                 weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
672 |                 argument.
673 |             pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
674 |                 Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
675 |                 If not provided, pooled text embeddings will be generated from `prompt` input argument.
676 |             negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
677 |                 Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
678 |                 weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
679 |                 input argument.
680 |             output_type (`str`, *optional*, defaults to `"pil"`):
681 |                 The output format of the generate image. Choose between
682 |                 [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
683 |             return_dict (`bool`, *optional*, defaults to `True`):
684 |                 Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead
685 |                 of a plain tuple.
686 |             callback (`Callable`, *optional*):
687 |                 A function that will be called every `callback_steps` steps during inference. The function will be
688 |                 called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
689 |             callback_steps (`int`, *optional*, defaults to 1):
690 |                 The frequency at which the `callback` function will be called. If not specified, the callback will be
691 |                 called at every step.
692 |             cross_attention_kwargs (`dict`, *optional*):
693 |                 A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
694 |                 `self.processor` in
695 |                 [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
696 |             guidance_rescale (`float`, *optional*, defaults to 0.7):
697 |                 Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
698 |                 Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of
699 |                 [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
700 |                 Guidance rescale factor should fix overexposure when using zero terminal SNR.
701 |             original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
702 |                 If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
703 |                 `original_size` defaults to `(width, height)` if not specified. Part of SDXL's micro-conditioning as
704 |                 explained in section 2.2 of
705 |                 [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
706 |             crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
707 |                 `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
708 |                 `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
709 |                 `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
710 |                 [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
711 |             target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
712 |                 For most cases, `target_size` should be set to the desired height and width of the generated image. If
713 |                 not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in
714 |                 section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
715 | 
716 |         Examples:
717 | 
718 |         Returns:
719 |             [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`:
720 |             [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a
721 |             `tuple`. When returning a tuple, the first element is a list with the generated images.
722 |         """
723 |         # 0. Default height and width to unet
724 |         height = height or self.default_sample_size * self.vae_scale_factor
725 |         width = width or self.default_sample_size * self.vae_scale_factor
726 | 
727 |         original_size = original_size or (height, width)
728 |         target_size = target_size or (height, width)
729 | 
730 |         # 1. Check inputs. Raise error if not correct
731 |         self.check_inputs(
732 |             prompt,
733 |             prompt_2,
734 |             height,
735 |             width,
736 |             callback_steps,
737 |             negative_prompt,
738 |             negative_prompt_2,
739 |             prompt_embeds,
740 |             negative_prompt_embeds,
741 |             pooled_prompt_embeds,
742 |             negative_pooled_prompt_embeds,
743 |         )
744 | 
745 |         # 2. Define call parameters
746 |         if prompt is not None and isinstance(prompt, str):
747 |             batch_size = 1
748 |         elif prompt is not None and isinstance(prompt, list):
749 |             batch_size = len(prompt)
750 |         else:
751 |             batch_size = prompt_embeds.shape[0]
752 | 
753 |         device = self._execution_device
754 | 
755 |         # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
756 |         # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
757 |         # corresponds to doing no classifier free guidance.
758 |         do_classifier_free_guidance = guidance_scale > 1.0
759 | 
760 |         # 3. Encode input prompt
761 |         text_encoder_lora_scale = (
762 |             cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
763 |         )
764 |         (
765 |             prompt_embeds,
766 |             negative_prompt_embeds,
767 |             pooled_prompt_embeds,
768 |             negative_pooled_prompt_embeds,
769 |         ) = self.encode_prompt(
770 |             prompt=prompt,
771 |             prompt_2=prompt_2,
772 |             device=device,
773 |             num_images_per_prompt=num_images_per_prompt,
774 |             do_classifier_free_guidance=do_classifier_free_guidance,
775 |             negative_prompt=negative_prompt,
776 |             negative_prompt_2=negative_prompt_2,
777 |             prompt_embeds=prompt_embeds,
778 |             negative_prompt_embeds=negative_prompt_embeds,
779 |             pooled_prompt_embeds=pooled_prompt_embeds,
780 |             negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
781 |             lora_scale=text_encoder_lora_scale,
782 |         )
783 | 
784 |         # 4. Prepare timesteps
785 |         self.scheduler.set_timesteps(num_inference_steps, device=device)
786 | 
787 |         timesteps = self.scheduler.timesteps
788 | 
789 |         # 5. Prepare latent variables
790 |         num_channels_latents = self.unet.config.in_channels
791 |         latents = self.prepare_latents(
792 |             batch_size * num_images_per_prompt,
793 |             num_channels_latents,
794 |             height,
795 |             width,
796 |             prompt_embeds.dtype,
797 |             device,
798 |             generator,
799 |             latents,
800 |         )
801 | 
802 |         # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
803 |         extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
804 | 
805 |         # 7. Prepare added time ids & embeddings
806 |         add_text_embeds = pooled_prompt_embeds
807 |         add_time_ids = self._get_add_time_ids(
808 |             original_size, crops_coords_top_left, target_size, dtype=prompt_embeds.dtype
809 |         )
810 | 
811 |         if do_classifier_free_guidance:
812 |             prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
813 |             add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
814 |             add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0)
815 | 
816 |         prompt_embeds = prompt_embeds.to(device)
817 |         add_text_embeds = add_text_embeds.to(device)
818 |         add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
819 | 
820 |         # 8. Denoising loop
821 |         num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
822 | 
823 |         # 7.1 Apply denoising_end
824 |         if denoising_end is not None and type(denoising_end) == float and denoising_end > 0 and denoising_end < 1:
825 |             discrete_timestep_cutoff = int(
826 |                 round(
827 |                     self.scheduler.config.num_train_timesteps
828 |                     - (denoising_end * self.scheduler.config.num_train_timesteps)
829 |                 )
830 |             )
831 |             num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))
832 |             timesteps = timesteps[:num_inference_steps]
833 | 
834 |         latents = latents.half()
835 |         prompt_embeds = prompt_embeds.half()
836 |         with self.progress_bar(total=num_inference_steps) as progress_bar:
837 |             for i, t in enumerate(timesteps):
838 |                 latents = self.update_loss(latents, i, t, prompt_embeds, cross_attention_kwargs, add_text_embeds, add_time_ids)
839 | 
840 |                 # expand the latents if we are doing classifier free guidance
841 |                 latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
842 | 
843 |                 latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
844 | 
845 |                 # predict the noise residual
846 |                 added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
847 |                 noise_pred = self.unet(
848 |                     latent_model_input,
849 |                     t,
850 |                     encoder_hidden_states=prompt_embeds,
851 |                     cross_attention_kwargs=cross_attention_kwargs,
852 |                     added_cond_kwargs=added_cond_kwargs,
853 |                     return_dict=False,
854 |                 )[0]
855 | 
856 |                 # perform guidance
857 |                 if do_classifier_free_guidance:
858 |                     noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
859 |                     noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
860 | 
861 |                 if do_classifier_free_guidance and guidance_rescale > 0.0:
862 |                     # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
863 |                     noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
864 | 
865 |                 # compute the previous noisy sample x_t -> x_t-1
866 |                 latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
867 | 
868 |                 # call the callback, if provided
869 |                 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
870 |                     progress_bar.update()
871 |                     if callback is not None and i % callback_steps == 0:
872 |                         callback(i, t, latents)
873 | 
874 |         # make sure the VAE is in float32 mode, as it overflows in float16
875 |         if self.vae.dtype == torch.float16 and self.vae.config.force_upcast:
876 |             self.upcast_vae()
877 |             latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
878 | 
879 |         if not output_type == "latent":
880 |             image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
881 |         else:
882 |             image = latents
883 |             return StableDiffusionXLPipelineOutput(images=image)
884 | 
885 |         # apply watermark if available
886 |         if self.watermark is not None:
887 |             image = self.watermark.apply_watermark(image)
888 | 
889 |         image = self.image_processor.postprocess(image, output_type=output_type)
890 | 
891 |         # Offload last model to CPU
892 |         if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
893 |             self.final_offload_hook.offload()
894 | 
895 |         if not return_dict:
896 |             return (image,)
897 | 
898 |         return StableDiffusionXLPipelineOutput(images=image)
899 | 
900 |     # Overrride to properly handle the loading and unloading of the additional text encoder.
901 |     def load_lora_weights(self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], **kwargs):
902 |         # We could have accessed the unet config from `lora_state_dict()` too. We pass
903 |         # it here explicitly to be able to tell that it's coming from an SDXL
904 |         # pipeline.
905 |         state_dict, network_alphas = self.lora_state_dict(
906 |             pretrained_model_name_or_path_or_dict,
907 |             unet_config=self.unet.config,
908 |             **kwargs,
909 |         )
910 |         self.load_lora_into_unet(state_dict, network_alphas=network_alphas, unet=self.unet)
911 | 
912 |         text_encoder_state_dict = {k: v for k, v in state_dict.items() if "text_encoder." in k}
913 |         if len(text_encoder_state_dict) > 0:
914 |             self.load_lora_into_text_encoder(
915 |                 text_encoder_state_dict,
916 |                 network_alphas=network_alphas,
917 |                 text_encoder=self.text_encoder,
918 |                 prefix="text_encoder",
919 |                 lora_scale=self.lora_scale,
920 |             )
921 | 
922 |         text_encoder_2_state_dict = {k: v for k, v in state_dict.items() if "text_encoder_2." in k}
923 |         if len(text_encoder_2_state_dict) > 0:
924 |             self.load_lora_into_text_encoder(
925 |                 text_encoder_2_state_dict,
926 |                 network_alphas=network_alphas,
927 |                 text_encoder=self.text_encoder_2,
928 |                 prefix="text_encoder_2",
929 |                 lora_scale=self.lora_scale,
930 |             )
931 | 
932 |     @classmethod
933 |     def save_lora_weights(
934 |         self,
935 |         save_directory: Union[str, os.PathLike],
936 |         unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
937 |         text_encoder_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
938 |         text_encoder_2_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
939 |         is_main_process: bool = True,
940 |         weight_name: str = None,
941 |         save_function: Callable = None,
942 |         safe_serialization: bool = True,
943 |     ):
944 |         state_dict = {}
945 | 
946 |         def pack_weights(layers, prefix):
947 |             layers_weights = layers.state_dict() if isinstance(layers, torch.nn.Module) else layers
948 |             layers_state_dict = {f"{prefix}.{module_name}": param for module_name, param in layers_weights.items()}
949 |             return layers_state_dict
950 | 
951 |         state_dict.update(pack_weights(unet_lora_layers, "unet"))
952 | 
953 |         if text_encoder_lora_layers and text_encoder_2_lora_layers:
954 |             state_dict.update(pack_weights(text_encoder_lora_layers, "text_encoder"))
955 |             state_dict.update(pack_weights(text_encoder_2_lora_layers, "text_encoder_2"))
956 | 
957 |         self.write_lora_layers(
958 |             state_dict=state_dict,
959 |             save_directory=save_directory,
960 |             is_main_process=is_main_process,
961 |             weight_name=weight_name,
962 |             save_function=save_function,
963 |             safe_serialization=safe_serialization,
964 |         )
965 | 
966 |     def _remove_text_encoder_monkey_patch(self):
967 |         self._remove_text_encoder_monkey_patch_classmethod(self.text_encoder)
968 |         self._remove_text_encoder_monkey_patch_classmethod(self.text_encoder_2)
969 | 


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
 1 | accelerate==0.21.0
 2 | diffusers==0.20.0
 3 | einops==0.6.1
 4 | lightning-utilities==0.9.0
 5 | matplotlib==3.7.3
 6 | nltk==3.8.1
 7 | numpy
 8 | opencv-python==4.8.1.78
 9 | Pillow==9.4.0
10 | pytorch-lightning==2.0.7
11 | scikit-image==0.22.0
12 | scikit-learn==1.3.1
13 | torch==2.0.1
14 | torch-kmeans==0.2.0
15 | torchvision==0.15.2
16 | transformers==4.32.0
17 | spaces==0.24.2
18 | xformers==0.0.21
19 | 


--------------------------------------------------------------------------------
/run_sd.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | import torch
  3 | import torchvision.transforms.functional as F
  4 | 
  5 | from diffusers import DDIMScheduler
  6 | from pipeline_stable_diffusion_opt import StableDiffusionPipeline
  7 | from pytorch_lightning import seed_everything
  8 | 
  9 | from injection_utils import register_attention_editor_diffusers
 10 | from bounded_attention import BoundedAttention
 11 | import utils
 12 | 
 13 | 
 14 | def load_model(device):
 15 |     scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
 16 |     model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", scheduler=scheduler, cross_attention_kwargs={"scale": 0.5}, torch_dtype=torch.float16, use_safetensors=True).to(device)
 17 |     model.enable_xformers_memory_efficient_attention()
 18 |     model.enable_sequential_cpu_offload()
 19 |     return model
 20 | 
 21 | 
 22 | def run(
 23 |     boxes,
 24 |     prompt,
 25 |     subject_token_indices,
 26 |     out_dir='out',
 27 |     seed=160,
 28 |     batch_size=1,
 29 |     filter_token_indices=None,
 30 |     eos_token_index=None,
 31 |     init_step_size=8,
 32 |     final_step_size=2,
 33 |     first_refinement_step=15,
 34 |     num_clusters_per_subject=3,
 35 |     cross_loss_scale=1.5,
 36 |     self_loss_scale=0.5,
 37 |     classifier_free_guidance_scale=7.5,
 38 |     num_gd_iterations=5,
 39 |     loss_threshold=0.2,
 40 |     num_guidance_steps=15,
 41 | ):
 42 |     device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 43 |     model = load_model(device)
 44 | 
 45 |     seed_everything(seed)
 46 |     prompts = [prompt] * batch_size
 47 |     start_code = torch.randn([len(prompts), 4, 64, 64], device=device)
 48 | 
 49 |     os.makedirs(out_dir, exist_ok=True)
 50 |     sample_count = len(os.listdir(out_dir))
 51 |     out_dir = os.path.join(out_dir, f"sample_{sample_count}")
 52 |     os.makedirs(out_dir)
 53 | 
 54 |     editor = BoundedAttention(
 55 |         boxes,
 56 |         prompts,
 57 |         subject_token_indices,
 58 |         list(range(12, 20)),
 59 |         list(range(12, 20)),
 60 |         cross_mask_layers=list(range(14, 20)),
 61 |         self_mask_layers=list(range(14, 20)),
 62 |         filter_token_indices=filter_token_indices,
 63 |         eos_token_index=eos_token_index,
 64 |         cross_loss_coef=cross_loss_scale,
 65 |         self_loss_coef=self_loss_scale,
 66 |         max_guidance_iter=num_guidance_steps,
 67 |         max_guidance_iter_per_step=num_gd_iterations,
 68 |         start_step_size=init_step_size,
 69 |         end_step_size=final_step_size,
 70 |         loss_stopping_value=loss_threshold,
 71 |         min_clustering_step=first_refinement_step,
 72 |         num_clusters_per_box=num_clusters_per_subject,
 73 |     )
 74 | 
 75 |     register_attention_editor_diffusers(model, editor)
 76 |     images = model(prompts, latents=start_code, guidance_scale=classifier_free_guidance_scale)
 77 | 
 78 |     for i, image in enumerate(images):
 79 |         image = F.to_pil_image(image)
 80 |         image.save(os.path.join(out_dir, f'{seed}_{i}.png'))
 81 |         utils.draw_box(image, boxes).save(os.path.join(out_dir, f'{seed}_{i}_boxes.png'))
 82 | 
 83 |     print("Syntheiszed images are saved in", out_dir)
 84 | 
 85 | 
 86 | def main():
 87 |     boxes = [
 88 |         [0.05, 0.2, 0.45, 0.8],
 89 |         [0.55, 0.2, 0.95, 0.8],
 90 |     ]
 91 | 
 92 |     prompt = "A ginger kitten and a gray puppy in a yard"
 93 |     subject_token_indices = [[2, 3], [6, 7]]
 94 | 
 95 |     run(boxes, prompt, subject_token_indices, init_step_size=8, final_step_size=2)
 96 | 
 97 | 
 98 | if __name__ == "__main__":
 99 |     main()
100 | 


--------------------------------------------------------------------------------
/run_xl.py:
--------------------------------------------------------------------------------
 1 | import os
 2 | import torch
 3 | 
 4 | from diffusers import DDIMScheduler
 5 | from pipeline_stable_diffusion_xl_opt import StableDiffusionXLPipeline
 6 | from pytorch_lightning import seed_everything
 7 | 
 8 | from injection_utils import register_attention_editor_diffusers
 9 | from bounded_attention import BoundedAttention
10 | import utils
11 | 
12 | 
13 | def load_model(device):
14 |     scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
15 |     model = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", scheduler=scheduler, torch_dtype=torch.float16).to(device)
16 |     model.enable_xformers_memory_efficient_attention()
17 |     model.enable_sequential_cpu_offload()
18 |     return model
19 | 
20 | 
21 | def run(
22 |     boxes,
23 |     prompt,
24 |     subject_token_indices,
25 |     out_dir='out',
26 |     seed=160,
27 |     batch_size=1,
28 |     filter_token_indices=None,
29 |     eos_token_index=None,
30 |     init_step_size=18,
31 |     final_step_size=5,
32 |     first_refinement_step=15,
33 |     num_clusters_per_subject=3,
34 |     cross_loss_scale=1.5,
35 |     self_loss_scale=0.5,
36 |     classifier_free_guidance_scale=7.5,
37 |     num_gd_iterations=5,
38 |     loss_threshold=0.2,
39 |     num_guidance_steps=15,
40 | ):
41 |     device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
42 |     model = load_model(device)
43 | 
44 |     seed_everything(seed)
45 |     prompts = [prompt] * batch_size
46 |     start_code = torch.randn([len(prompts), 4, 128, 128], device=device)
47 | 
48 |     os.makedirs(out_dir, exist_ok=True)
49 |     sample_count = len(os.listdir(out_dir))
50 |     out_dir = os.path.join(out_dir, f"sample_{sample_count}")
51 |     os.makedirs(out_dir)
52 | 
53 |     editor = BoundedAttention(
54 |         boxes,
55 |         prompts,
56 |         subject_token_indices,
57 |         list(range(70, 82)),
58 |         list(range(70, 82)),
59 |         filter_token_indices=filter_token_indices,
60 |         eos_token_index=eos_token_index,
61 |         cross_loss_coef=cross_loss_scale,
62 |         self_loss_coef=self_loss_scale,
63 |         max_guidance_iter=num_guidance_steps,
64 |         max_guidance_iter_per_step=num_gd_iterations,
65 |         start_step_size=init_step_size,
66 |         end_step_size=final_step_size,
67 |         loss_stopping_value=loss_threshold,
68 |         min_clustering_step=first_refinement_step,
69 |         num_clusters_per_box=num_clusters_per_subject,
70 |     )
71 | 
72 |     register_attention_editor_diffusers(model, editor)
73 |     images = model(prompts, latents=start_code, guidance_scale=classifier_free_guidance_scale).images
74 | 
75 |     for i, image in enumerate(images):
76 |         image.save(os.path.join(out_dir, f'{seed}_{i}.png'))
77 |         utils.draw_box(image, boxes).save(os.path.join(out_dir, f'{seed}_{i}_boxes.png'))
78 | 
79 |     print("Syntheiszed images are saved in", out_dir)
80 | 
81 | 
82 | def main():
83 |     boxes = [
84 |         [0, 0.5, 0.2, 0.8],
85 |         [0.2, 0.2, 0.4, 0.5],
86 |         [0.4, 0.5, 0.6, 0.8],
87 |         [0.6, 0.2, 0.8, 0.5],
88 |         [0.8, 0.5, 1, 0.8],
89 |     ]
90 | 
91 |     prompt = "a golden retriever and a german shepherd and a boston terrier and an english bulldog and a border collie in a pool"
92 |     subject_token_indices = [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]]
93 | 
94 |     run(boxes, prompt, subject_token_indices, init_step_size=18, final_step_size=5)
95 | 
96 | 
97 | if __name__ == "__main__":
98 |     main()
99 | 


--------------------------------------------------------------------------------
/utils.py:
--------------------------------------------------------------------------------
  1 | import numpy as np
  2 | from PIL import Image, ImageDraw, ImageFont
  3 | import cv2
  4 | from sklearn.decomposition import PCA
  5 | from torchvision import transforms
  6 | import matplotlib.pyplot as plt
  7 | import torch
  8 | 
  9 | import os
 10 | 
 11 | 
 12 | def display_attention_maps(
 13 |     attention_maps,
 14 |     is_cross,
 15 |     num_heads,
 16 |     tokenizer,
 17 |     prompts,
 18 |     dir_name,
 19 |     step,
 20 |     layer,
 21 |     resolution,
 22 |     is_query=False,
 23 |     is_key=False,
 24 |     points=None,
 25 |     image_path=None,
 26 | ):
 27 |     attention_maps = attention_maps.reshape(-1, num_heads, attention_maps.size(-2), attention_maps.size(-1))
 28 |     num_samples = len(attention_maps) // 2
 29 |     attention_type = 'cross' if is_cross else 'self'
 30 |     for i, attention_map in enumerate(attention_maps):
 31 |         if is_query:
 32 |             attention_type = f'{attention_type}_queries'
 33 |         elif is_key:
 34 |             attention_type = f'{attention_type}_keys'
 35 | 
 36 |         cond = 'uncond' if i < num_samples else 'cond'
 37 |         i = i % num_samples
 38 |         cur_dir_name = f'{dir_name}/{resolution}/{attention_type}/{layer}/{cond}/{i}'
 39 |         os.makedirs(cur_dir_name, exist_ok=True)
 40 | 
 41 |         if is_cross and not is_query:
 42 |             fig = show_cross_attention(attention_map, tokenizer, prompts[i % num_samples])
 43 |         else:
 44 |             fig = show_self_attention(attention_map)
 45 |             if points is not None:
 46 |                 point_dir_name = f'{cur_dir_name}/points'
 47 |                 os.makedirs(point_dir_name, exist_ok=True)
 48 |                 for j, point in enumerate(points):
 49 |                     specific_point_dir_name = f'{point_dir_name}/{j}'
 50 |                     os.makedirs(specific_point_dir_name, exist_ok=True)
 51 |                     point_path = f'{specific_point_dir_name}/{step}.png'
 52 |                     point_fig = show_individual_self_attention(attention_map, point, image_path=image_path)
 53 |                     point_fig.save(point_path)
 54 |                     point_fig.close()
 55 | 
 56 |         fig.save(f'{cur_dir_name}/{step}.png')
 57 |         fig.close()
 58 | 
 59 | 
 60 | def text_under_image(image: np.ndarray, text: str, text_color: tuple[int, int, int] = (0, 0, 0)):
 61 |     h, w, c = image.shape
 62 |     offset = int(h * .2)
 63 |     font = cv2.FONT_HERSHEY_SIMPLEX
 64 |     # font = ImageFont.truetype("/usr/share/fonts/truetype/noto/NotoMono-Regular.ttf", font_size)
 65 |     text_size = cv2.getTextSize(text, font, 1, 2)[0]
 66 |     lines = text.splitlines()
 67 |     img = np.ones((h + offset + (text_size[1] + 2) * len(lines) - 2, w, c), dtype=np.uint8) * 255
 68 |     img[:h, :w] = image
 69 | 
 70 |     for i, line in enumerate(lines):
 71 |         text_size = cv2.getTextSize(line, font, 1, 2)[0]
 72 |         text_x, text_y = ((w - text_size[0]) // 2, h + offset + i * (text_size[1] + 2))
 73 |         cv2.putText(img, line, (text_x, text_y), font, 1, text_color, 2)
 74 | 
 75 |     return img
 76 | 
 77 | def view_images(images, num_rows=1, offset_ratio=0.02):
 78 |     if type(images) is list:
 79 |         num_empty = len(images) % num_rows
 80 |     elif images.ndim == 4:
 81 |         num_empty = images.shape[0] % num_rows
 82 |     else:
 83 |         images = [images]
 84 |         num_empty = 0
 85 | 
 86 |     empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255
 87 |     images = [image.astype(np.uint8) for image in images] + [empty_images] * num_empty
 88 |     num_items = len(images)
 89 | 
 90 |     h, w, c = images[0].shape
 91 |     offset = int(h * offset_ratio)
 92 |     num_cols = num_items // num_rows
 93 |     image_ = np.ones((h * num_rows + offset * (num_rows - 1),
 94 |                       w * num_cols + offset * (num_cols - 1), 3), dtype=np.uint8) * 255
 95 |     for i in range(num_rows):
 96 |         for j in range(num_cols):
 97 |             image_[i * (h + offset): i * (h + offset) + h:, j * (w + offset): j * (w + offset) + w] = images[
 98 |                 i * num_cols + j]
 99 | 
100 |     return Image.fromarray(image_)
101 | 
102 | 
103 | def show_cross_attention(attention_maps, tokenizer, prompt, k_norms=None, v_norms=None):
104 |     attention_maps = attention_maps.mean(dim=0)
105 |     res = int(attention_maps.size(-2) ** 0.5)
106 |     attention_maps = attention_maps.reshape(res, res, -1)
107 |     tokens = tokenizer.encode(prompt)
108 |     decoder = tokenizer.decode
109 |     if k_norms is not None:
110 |         k_norms = k_norms.round(decimals=1)
111 |     if v_norms is not None:
112 |         v_norms = v_norms.round(decimals=1)
113 |     images = []
114 |     for i in range(len(tokens) + 5):
115 |         image = attention_maps[:, :, i]
116 |         image = 255 * image / image.max()
117 |         image = image.unsqueeze(-1).expand(*image.shape, 3)
118 |         image = image.detach().cpu().numpy().astype(np.uint8)
119 |         image = np.array(Image.fromarray(image).resize((256, 256)))
120 |         token = tokens[i] if i < len(tokens) else tokens[-1]
121 |         text = decoder(int(token))
122 |         if k_norms is not None and v_norms is not None:
123 |             text += f'\n{k_norms[i]}\n{v_norms[i]})'
124 |         image = text_under_image(image, text)
125 |         images.append(image)
126 |     return view_images(np.stack(images, axis=0))
127 | 
128 | 
129 | def show_queries_keys(queries, keys, colors, labels):  # [h ni d]
130 |     num_queries = [query.size(1) for query in queries]
131 |     num_keys = [key.size(1) for key in keys]
132 |     h, _, d = queries[0].shape
133 | 
134 |     data = torch.cat((*queries, *keys), dim=1)  # h n d
135 |     data = data.permute(1, 0, 2)  # n h d
136 |     data = data.reshape(-1, h * d).detach().cpu().numpy()
137 |     pca = PCA(n_components=2)
138 |     data = pca.fit_transform(data)  # n 2
139 | 
140 |     query_indices = np.array(num_queries).cumsum()
141 |     total_num_queries = query_indices[-1]
142 |     queries = np.split(data[:total_num_queries], query_indices[:-1])
143 |     if len(num_keys) == 0:
144 |         keys = [None, ] * len(labels)
145 |     else:
146 |         key_indices = np.array(num_keys).cumsum()
147 |         keys = np.split(data[total_num_queries:], key_indices[:-1])
148 | 
149 |     fig, ax = plt.subplots()
150 |     marker_size = plt.rcParams['lines.markersize'] ** 2
151 |     query_size = int(1.25 * marker_size)
152 |     key_size = int(2 * marker_size)
153 |     for query, key, color, label in zip(queries, keys, colors, labels):
154 |         print(f'# queries of {label}', query.shape[0])
155 |         ax.scatter(query[:, 0], query[:, 1], s=query_size, color=color, marker='o', label=f'"{label}" queries')
156 | 
157 |         if key is None:
158 |             continue
159 | 
160 |         print(f'# keys of {label}', key.shape[0])
161 |         keys_label = f'"{label}" key'
162 |         if key.shape[0] > 1:
163 |             keys_label += 's'
164 |         ax.scatter(key[:, 0], key[:, 1], s=key_size, color=color, marker='x', label=keys_label)
165 | 
166 |     ax.set_axis_off()
167 |     #ax.set_xlabel('X-axis')
168 |     #ax.set_ylabel('Y-axis')
169 |     #ax.set_title('Scatter Plot with Circles and Crosses')
170 | 
171 |     #ax.legend()
172 |     return fig
173 | 
174 | 
175 | def show_self_attention(attention_maps):  # h n m
176 |     attention_maps = attention_maps.transpose(0, 1).flatten(start_dim=1).detach().cpu().numpy()
177 |     pca = PCA(n_components=3)
178 |     pca_img = pca.fit_transform(attention_maps)  # N X 3
179 |     h = w = int(pca_img.shape[0] ** 0.5)
180 |     pca_img = pca_img.reshape(h, w, 3)
181 |     pca_img_min = pca_img.min(axis=(0, 1))
182 |     pca_img_max = pca_img.max(axis=(0, 1))
183 |     pca_img = (pca_img - pca_img_min) / (pca_img_max - pca_img_min)
184 |     pca_img = Image.fromarray((pca_img * 255).astype(np.uint8))
185 |     pca_img = transforms.Resize(256, interpolation=transforms.InterpolationMode.NEAREST)(pca_img)
186 |     return pca_img
187 | 
188 | 
189 | def draw_box(pil_img, bboxes, colors=None, width=5):
190 |     draw = ImageDraw.Draw(pil_img)
191 |     #font = ImageFont.truetype('./FreeMono.ttf', 25)
192 |     w, h = pil_img.size
193 |     colors = ['red'] * len(bboxes) if colors is None else colors
194 |     for obj_bbox, color in zip(bboxes, colors):
195 |         x_0, y_0, x_1, y_1 = obj_bbox[0], obj_bbox[1], obj_bbox[2], obj_bbox[3]
196 |         draw.rectangle([int(x_0 * w), int(y_0 * h), int(x_1 * w), int(y_1 * h)], outline=color, width=width)
197 |     return pil_img
198 | 
199 | 
200 | def show_individual_self_attention(attn, point, image_path=None):
201 |     resolution = int(attn.size(-1) ** 0.5)
202 |     attn = attn.mean(dim=0).reshape(resolution, resolution, resolution, resolution)
203 |     attn = attn[round(point[1] * resolution), round(point[0] * resolution)]
204 |     attn = (attn - attn.min()) / (attn.max() - attn.min())
205 |     image = None if image_path is None else Image.open(image_path).convert('RGB')
206 |     image = show_image_relevance(attn, image=image)
207 |     return Image.fromarray(image)
208 | 
209 | 
210 | def show_image_relevance(image_relevance, image: Image.Image = None, relevnace_res=16):
211 |     # create heatmap from mask on image
212 |     def show_cam_on_image(img, mask):
213 |         img = img.resize((relevnace_res ** 2, relevnace_res ** 2))
214 |         img = np.array(img)
215 |         img = (img - img.min()) / (img.max() - img.min())
216 |         heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)
217 |         heatmap = np.float32(heatmap) / 255
218 |         cam = heatmap + np.float32(img)
219 |         cam = cam / np.max(cam)
220 |         return cam
221 | 
222 |     image_relevance = image_relevance.reshape(1, 1, image_relevance.shape[-1], image_relevance.shape[-1])
223 |     image_relevance = image_relevance.cuda() # because float16 precision interpolation is not supported on cpu
224 |     image_relevance = torch.nn.functional.interpolate(image_relevance, size=relevnace_res ** 2, mode='bilinear')
225 |     image_relevance = image_relevance.cpu() # send it back to cpu
226 |     image_relevance = (image_relevance - image_relevance.min()) / (image_relevance.max() - image_relevance.min())
227 |     image_relevance = image_relevance.reshape(relevnace_res ** 2, relevnace_res ** 2)
228 |     vis = image_relevance if image is None else show_cam_on_image(image, image_relevance)
229 |     vis = np.uint8(255 * vis)
230 |     vis = cv2.cvtColor(np.array(vis), cv2.COLOR_RGB2BGR)
231 |     return vis
232 | 


--------------------------------------------------------------------------------