Improved memory efficiency

bounded_attention.py

@@ -44,7 +44,7 @@ class BoundedAttention(injection_utils.AttentionBase):
         pca_rank=None,
         num_clusters=None,
         num_clusters_per_box=3,
-        max_resolution=32,
+        max_resolution=None,
         map_dir=None,
         debug=False,
         delta_debug_attention_steps=20,
@@ -95,8 +95,10 @@ class BoundedAttention(injection_utils.AttentionBase):
         self.self_foreground_values = []
         self.cross_background_values = []
         self.self_background_values = []
-        self.cross_maps = []
-        self.self_maps = []
+        self.mean_cross_map = 0
+        self.num_cross_maps = 0
+        self.mean_self_map = 0
+        self.num_self_maps = 0
         self.self_masks = None
 
     def clear_values(self, include_maps=False):
@@ -107,16 +109,15 @@ class BoundedAttention(injection_utils.AttentionBase):
             self.self_background_values,
         )
 
-        if include_maps:
-            lists = (
-                *all_values,
-                self.cross_maps,
-                self.self_maps,
-            )
-
         for values in lists:
             values.clear()
 
+        if include_maps:
+            self.mean_cross_map = 0
+            self.num_cross_maps = 0
+            self.mean_self_map = 0
+            self.num_self_maps = 0
+
     def before_step(self):
         self.clear_values()
         if self.cur_step == 0:
@@ -137,37 +138,31 @@ class BoundedAttention(injection_utils.AttentionBase):
         else:
             masks = self._hide_other_subjects_from_subjects(batch_size // 2, n, dtype, device)
 
-        if int(n ** 0.5) > self.max_resolution:
-            q = q.reshape(batch_size, num_heads, n, -1)
-            k = k.reshape(batch_size, num_heads, d, -1)
-            v = v.reshape(batch_size, num_heads, d, -1)
-            out = F.scaled_dot_product_attention(q, k, v, attn_mask=masks)
-            out = out.reshape(batch_size * num_heads, n, -1)
-        else:
-            sim = torch.einsum('b i d, b j d -> b i j', q, k) * kwargs['scale']
-            attn = sim.softmax(-1)
-            self._display_attention_maps(attn, is_cross, num_heads)
-            sim = sim.reshape(batch_size, num_heads, n, d) + masks
-            attn = sim.reshape(-1, n, d).softmax(-1)
-            self._save(attn, is_cross, num_heads)
-            self._display_attention_maps(attn, is_cross, num_heads, prefix='masked')
-            self._debug_hook(q, k, v, sim, attn, is_cross, place_in_unet, num_heads, **kwargs)
-            out = torch.bmm(attn, v)
-
-        out = einops.rearrange(out, '(b h) n d -> b n (h d)', h=num_heads)
-        return out
+        resolution = int(n ** 0.5)
+        if (self.max_resolution is not None) and (resolution > self.max_resolution):
+            return super().forward(q, k, v, is_cross, place_in_unet, num_heads, mask=masks)
+
+        sim = torch.einsum('b i d, b j d -> b i j', q, k) * kwargs['scale']
+        attn = sim.softmax(-1)
+        self._display_attention_maps(attn, is_cross, num_heads)
+        sim = sim.reshape(batch_size, num_heads, n, d) + masks
+        attn = sim.reshape(-1, n, d).softmax(-1)
+        self._save(attn, is_cross, num_heads)
+        self._display_attention_maps(attn, is_cross, num_heads, prefix='masked')
+        self._debug_hook(q, k, v, sim, attn, is_cross, place_in_unet, num_heads, **kwargs)
+        out = torch.bmm(attn, v)
+        return einops.rearrange(out, '(b h) n d -> b n (h d)', h=num_heads)
 
     def update_loss(self, forward_pass, latents, i):
         if i >= self.max_guidance_iter:
             return latents
 
         step_size = self.start_step_size + self.step_size_coef * i
-        updated_latents = latents
 
         self.optimized = True
         normalized_loss = torch.tensor(10000)
         with torch.enable_grad():
-            latents = updated_latents = updated_latents.clone().detach().requires_grad_(True)
+            latents = latents.clone().detach().requires_grad_(True)
             for guidance_iter in range(self.max_guidance_iter_per_step):
                 if normalized_loss < self.loss_stopping_value:
                     break
@@ -178,8 +173,8 @@ class BoundedAttention(injection_utils.AttentionBase):
                 self.cur_step = cur_step
 
                 loss, normalized_loss = self._compute_loss()
-                grad_cond = torch.autograd.grad(loss, [updated_latents])[0]
-                latents = updated_latents = updated_latents - step_size * grad_cond
+                grad_cond = torch.autograd.grad(loss, [latents])[0]
+                latents = latents - step_size * grad_cond
                 if self.debug:
                     print(f'Loss at step={i}, iter={guidance_iter}: {normalized_loss}')
                     grad_norms = grad_cond.flatten(start_dim=2).norm(dim=1)
@@ -301,13 +296,21 @@ class BoundedAttention(injection_utils.AttentionBase):
 
         if is_cross:
             attn = attn[..., self.leading_token_indices]
-            mask_maps = self.cross_maps
+            mean_map = self.mean_cross_map
+            num_maps = self.num_cross_maps
         else:
-            mask_maps = self.self_maps
+            mean_map = self.mean_self_map
+            num_maps = self.num_self_maps
 
-        mask_maps.append(attn.mean(dim=1))  # mean over heads
-        if self.cur_step > 0:
-            mask_maps.pop(0)  # throw away old maps
+        num_maps += 1
+        attn = attn.mean(dim=1)  # mean over heads
+        mean_map = ((num_maps - 1) / num_maps) * mean_map +  (1 / num_maps) * attn
+        if is_cross:
+            self.mean_cross_map = mean_map
+            self.num_cross_maps = num_maps
+        else:
+            self.mean_self_map = mean_map
+            self.num_self_maps = num_maps
 
     def _save_loss_values(self, attn, is_cross):
         if (
@@ -404,7 +407,7 @@ class BoundedAttention(injection_utils.AttentionBase):
         return self_masks.flatten(start_dim=2).bool()
 
     def _cluster_self_maps(self):  # b s n
-        self_maps = self._aggregate_maps(self.self_maps)  # b n m
+        self_maps = self._compute_maps(self.mean_self_map)  # b n m
         if self.pca_rank is not None:
             dtype = self_maps.dtype
             _, _, eigen_vectors = torch.pca_lowrank(self_maps.float(), self.pca_rank)
@@ -442,7 +445,7 @@ class BoundedAttention(injection_utils.AttentionBase):
         return self_masks
 
     def _obtain_cross_masks(self, resolution, scale=10):
-        maps = self._aggregate_maps(self.cross_maps, resolution=resolution)  # b n k
+        maps = self._compute_maps(self.mean_cross_map, resolution=resolution)  # b n k
         maps = F.sigmoid(scale * (maps - self.cross_mask_threshold))
         maps = self._normalize_maps(maps, reduce_min=True)
         maps = maps.transpose(1, 2)  # b k n
@@ -466,8 +469,7 @@ class BoundedAttention(injection_utils.AttentionBase):
 
         return masks
 
-    def _aggregate_maps(self, maps, resolution=None):  # b n k
-        maps = torch.stack(maps).mean(0)  # mean over layers
+    def _compute_maps(self, maps, resolution=None):  # b n k
         if resolution is not None:
             b, n, k = maps.shape
             original_resolution = int(n ** 0.5)

injection_utils.py

@@ -35,7 +35,7 @@ class AttentionBase:
 
         return out
 
-    def forward(self, q, k, v, sim, attn, is_cross, place_in_unet, num_heads, **kwargs):
+    def forward(self, q, k, v, is_cross, place_in_unet, num_heads, **kwargs):
         batch_size = q.size(0) // num_heads
         n = q.size(1)
         d = k.size(1)