3 files changed, 119 insertions, 57 deletions

diff --git a/models/attention/control.py b/models/attention/control.py
index 248bd9f..ec378c4 100644
--- a/models/attention/control.py
+++ b/models/attention/control.py
@@ -23,7 +23,7 @@ class AttentionControl(abc.ABC):
                 attn = self.forward(attn, is_cross, place_in_unet)
             else:
                 h = attn.shape[0]
-                attn[h // 2:] = self.forward(attn[h // 2:], is_cross, place_in_unet)
+                attn[h // 2 :] = self.forward(attn[h // 2 :], is_cross, place_in_unet)
         self.cur_att_layer += 1
         if self.cur_att_layer == self.num_att_layers + self.num_uncond_att_layers:
             self.cur_att_layer = 0
@@ -49,12 +49,18 @@ class EmptyControl(AttentionControl):
 class AttentionStore(AttentionControl):
     @staticmethod
     def get_empty_store():
-        return {"down_cross": [], "mid_cross": [], "up_cross": [],
-                "down_self": [],  "mid_self": [],  "up_self": []}
+        return {
+            "down_cross": [],
+            "mid_cross": [],
+            "up_cross": [],
+            "down_self": [],
+            "mid_self": [],
+            "up_self": [],
+        }
 
     def forward(self, attn, is_cross: bool, place_in_unet: str):
         key = f"{place_in_unet}_{'cross' if is_cross else 'self'}"
-        if attn.shape[1] <= 32 ** 2:  # avoid memory overhead
+        if attn.shape[1] <= 32**2:  # avoid memory overhead
             self.step_store[key].append(attn)
         return attn
 
@@ -68,8 +74,10 @@ class AttentionStore(AttentionControl):
         self.step_store = self.get_empty_store()
 
     def get_average_attention(self):
-        average_attention = {key: [item / self.cur_step for item in self.attention_store[key]]
-                             for key in self.attention_store}
+        average_attention = {
+            key: [item / self.cur_step for item in self.attention_store[key]]
+            for key in self.attention_store
+        }
         return average_attention
 
     def reset(self):
@@ -90,7 +98,7 @@ class AttentionControlEdit(AttentionStore, abc.ABC):
         return x_t
 
     def replace_self_attention(self, attn_base, att_replace):
-        if att_replace.shape[2] <= 16 ** 2:
+        if att_replace.shape[2] <= 16**2:
             return attn_base.unsqueeze(0).expand(att_replace.shape[0], *attn_base.shape)
         else:
             return att_replace
@@ -101,41 +109,62 @@ class AttentionControlEdit(AttentionStore, abc.ABC):
 
     def forward(self, attn, is_cross: bool, place_in_unet: str):
         super(AttentionControlEdit, self).forward(attn, is_cross, place_in_unet)
-        if is_cross or (self.num_self_replace[0] <= self.cur_step < self.num_self_replace[1]):
+        if is_cross or (
+            self.num_self_replace[0] <= self.cur_step < self.num_self_replace[1]
+        ):
             h = attn.shape[0] // (self.batch_size)
             attn = attn.reshape(self.batch_size, h, *attn.shape[1:])
             attn_base, attn_repalce = attn[0], attn[1:]
             if is_cross:
                 alpha_words = self.cross_replace_alpha[self.cur_step]
-                attn_repalce_new = self.replace_cross_attention(
-                    attn_base, attn_repalce) * alpha_words + (1 - alpha_words) * attn_repalce
+                attn_repalce_new = (
+                    self.replace_cross_attention(attn_base, attn_repalce) * alpha_words
+                    + (1 - alpha_words) * attn_repalce
+                )
                 attn[1:] = attn_repalce_new
             else:
                 attn[1:] = self.replace_self_attention(attn_base, attn_repalce)
             attn = attn.reshape(self.batch_size * h, *attn.shape[2:])
         return attn
 
-    def __init__(self, prompts, num_steps: int,
-                 cross_replace_steps: Union[float, Tuple[float, float], Dict[str, Tuple[float, float]]],
-                 self_replace_steps: Union[float, Tuple[float, float]],
-                 local_blend: Optional[LocalBlend]):
+    def __init__(
+        self,
+        prompts,
+        num_steps: int,
+        cross_replace_steps: Union[
+            float, Tuple[float, float], Dict[str, Tuple[float, float]]
+        ],
+        self_replace_steps: Union[float, Tuple[float, float]],
+        local_blend: Optional[LocalBlend],
+    ):
         super(AttentionControlEdit, self).__init__()
         self.batch_size = len(prompts)
         self.cross_replace_alpha = ptp_utils.get_time_words_attention_alpha(
-            prompts, num_steps, cross_replace_steps, tokenizer).to(device)
+            prompts, num_steps, cross_replace_steps, tokenizer
+        ).to(device)
         if type(self_replace_steps) is float:
             self_replace_steps = 0, self_replace_steps
-        self.num_self_replace = int(num_steps * self_replace_steps[0]), int(num_steps * self_replace_steps[1])
+        self.num_self_replace = int(num_steps * self_replace_steps[0]), int(
+            num_steps * self_replace_steps[1]
+        )
         self.local_blend = local_blend
 
 
 class AttentionReplace(AttentionControlEdit):
     def replace_cross_attention(self, attn_base, att_replace):
-        return torch.einsum('hpw,bwn->bhpn', attn_base, self.mapper)
-
-    def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float,
-                 local_blend: Optional[LocalBlend] = None):
-        super(AttentionReplace, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend)
+        return torch.einsum("hpw,bwn->bhpn", attn_base, self.mapper)
+
+    def __init__(
+        self,
+        prompts,
+        num_steps: int,
+        cross_replace_steps: float,
+        self_replace_steps: float,
+        local_blend: Optional[LocalBlend] = None,
+    ):
+        super(AttentionReplace, self).__init__(
+            prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend
+        )
         self.mapper = seq_aligner.get_replacement_mapper(prompts, tokenizer).to(device)
 
 
@@ -145,9 +174,17 @@ class AttentionRefine(AttentionControlEdit):
         attn_replace = attn_base_replace * self.alphas + att_replace * (1 - self.alphas)
         return attn_replace
 
-    def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float,
-                 local_blend: Optional[LocalBlend] = None):
-        super(AttentionRefine, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend)
+    def __init__(
+        self,
+        prompts,
+        num_steps: int,
+        cross_replace_steps: float,
+        self_replace_steps: float,
+        local_blend: Optional[LocalBlend] = None,
+    ):
+        super(AttentionRefine, self).__init__(
+            prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend
+        )
         self.mapper, alphas = seq_aligner.get_refinement_mapper(prompts, tokenizer)
         self.mapper, alphas = self.mapper.to(device), alphas.to(device)
         self.alphas = alphas.reshape(alphas.shape[0], 1, 1, alphas.shape[1])
@@ -156,13 +193,24 @@ class AttentionRefine(AttentionControlEdit):
 class AttentionReweight(AttentionControlEdit):
     def replace_cross_attention(self, attn_base, att_replace):
         if self.prev_controller is not None:
-            attn_base = self.prev_controller.replace_cross_attention(attn_base, att_replace)
+            attn_base = self.prev_controller.replace_cross_attention(
+                attn_base, att_replace
+            )
         attn_replace = attn_base[None, :, :, :] * self.equalizer[:, None, None, :]
         return attn_replace
 
-    def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float, equalizer,
-                 local_blend: Optional[LocalBlend] = None, controller: Optional[AttentionControlEdit] = None):
-        super(AttentionReweight, self).__init__(prompts, num_steps,
-                                                cross_replace_steps, self_replace_steps, local_blend)
+    def __init__(
+        self,
+        prompts,
+        num_steps: int,
+        cross_replace_steps: float,
+        self_replace_steps: float,
+        equalizer,
+        local_blend: Optional[LocalBlend] = None,
+        controller: Optional[AttentionControlEdit] = None,
+    ):
+        super(AttentionReweight, self).__init__(
+            prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend
+        )
         self.equalizer = equalizer.to(device)
         self.prev_controller = controller
diff --git a/models/attention/hook.py b/models/attention/hook.py
index 903de02..6b5fb68 100644
--- a/models/attention/hook.py
+++ b/models/attention/hook.py
@@ -3,6 +3,7 @@ import torch
 
 try:
     import xformers.ops
+
     xformers._is_functorch_available = True
     MEM_EFFICIENT_ATTN = True
 except ImportError:
@@ -42,10 +43,10 @@ def register_attention_control(model, controller):
         return forward
 
     def register_recr(net_, count, place_in_unet):
-        if net_.__class__.__name__ == 'CrossAttention':
+        if net_.__class__.__name__ == "CrossAttention":
             net_.forward = ca_forward(net_, place_in_unet)
             return count + 1
-        elif hasattr(net_, 'children'):
+        elif hasattr(net_, "children"):
             for net__ in net_.children():
                 count = register_recr(net__, count, place_in_unet)
         return count
diff --git a/models/attention/structured.py b/models/attention/structured.py
index 24d889f..5bbbc06 100644
--- a/models/attention/structured.py
+++ b/models/attention/structured.py
@@ -16,7 +16,9 @@ class StructuredAttentionControl(AttentionControl):
             if self.struct_attn:
                 out = self.struct_qkv(q, context, mask)
             else:
-                context = torch.cat([context[0], context[1]['k'][0]], dim=0)  # use key tensor for context
+                context = torch.cat(
+                    [context[0], context[1]["k"][0]], dim=0
+                )  # use key tensor for context
                 out = self.normal_qkv(q, context, mask)
         else:
             context = default(context, x)
@@ -29,11 +31,13 @@ class StructuredAttentionControl(AttentionControl):
         context: list of [uc, list of conditional context]
         """
         uc_context = context[0]
-        context_k, context_v = context[1]['k'], context[1]['v']
+        context_k, context_v = context[1]["k"], context[1]["v"]
 
         if isinstance(context_k, list) and isinstance(context_v, list):
             out = self.multi_qkv(q, uc_context, context_k, context_v, mask)
-        elif isinstance(context_k, torch.Tensor) and isinstance(context_v, torch.Tensor):
+        elif isinstance(context_k, torch.Tensor) and isinstance(
+            context_v, torch.Tensor
+        ):
             out = self.heterogeous_qkv(q, uc_context, context_k, context_v, mask)
         else:
             raise NotImplementedError
@@ -50,36 +54,45 @@ class StructuredAttentionControl(AttentionControl):
         k_c = [self.to_k(c_k) for c_k in context_k]
         v_c = [self.to_v(c_v) for c_v in context_v]
 
-        q = rearrange(q, 'b n (h d) -> (b h) n d', h=h)
+        q = rearrange(q, "b n (h d) -> (b h) n d", h=h)
 
-        k_uc = rearrange(k_uc, 'b n (h d) -> (b h) n d', h=h)
-        v_uc = rearrange(v_uc, 'b n (h d) -> (b h) n d', h=h)
+        k_uc = rearrange(k_uc, "b n (h d) -> (b h) n d", h=h)
+        v_uc = rearrange(v_uc, "b n (h d) -> (b h) n d", h=h)
 
-        k_c = [rearrange(k, 'b n (h d) -> (b h) n d', h=h) for k in k_c]  # NOTE: modification point
-        v_c = [rearrange(v, 'b n (h d) -> (b h) n d', h=h) for v in v_c]
+        k_c = [
+            rearrange(k, "b n (h d) -> (b h) n d", h=h) for k in k_c
+        ]  # NOTE: modification point
+        v_c = [rearrange(v, "b n (h d) -> (b h) n d", h=h) for v in v_c]
 
         # get composition
-        sim_uc = einsum('b i d, b j d -> b i j', q[:true_bs], k_uc) * self.scale
-        sim_c = [einsum('b i d, b j d -> b i j', q[true_bs:], k) * self.scale for k in k_c]
+        sim_uc = einsum("b i d, b j d -> b i j", q[:true_bs], k_uc) * self.scale
+        sim_c = [
+            einsum("b i d, b j d -> b i j", q[true_bs:], k) * self.scale for k in k_c
+        ]
 
         attn_uc = sim_uc.softmax(dim=-1)
         attn_c = [sim.softmax(dim=-1) for sim in sim_c]
 
         # get uc output
-        out_uc = einsum('b i j, b j d -> b i d', attn_uc, v_uc)
+        out_uc = einsum("b i j, b j d -> b i d", attn_uc, v_uc)
 
         # get c output
         if len(v_c) == 1:
             out_c_collect = []
             for attn in attn_c:
                 for v in v_c:
-                    out_c_collect.append(einsum('b i j, b j d -> b i d', attn, v))
+                    out_c_collect.append(einsum("b i j, b j d -> b i d", attn, v))
             out_c = sum(out_c_collect) / len(out_c_collect)
         else:
-            out_c = sum([einsum('b i j, b j d -> b i d', attn, v) for attn, v in zip(attn_c, v_c)]) / len(v_c)
+            out_c = sum(
+                [
+                    einsum("b i j, b j d -> b i d", attn, v)
+                    for attn, v in zip(attn_c, v_c)
+                ]
+            ) / len(v_c)
 
         out = torch.cat([out_uc, out_c], dim=0)
-        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
+        out = rearrange(out, "(b h) n d -> b n (h d)", h=h)
 
         return out
 
@@ -88,21 +101,21 @@ class StructuredAttentionControl(AttentionControl):
         k = self.to_k(context)
         v = self.to_v(context)
 
-        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v))
 
-        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
+        sim = einsum("b i d, b j d -> b i j", q, k) * self.scale
 
         if exists(mask):
-            mask = rearrange(mask, 'b ... -> b (...)')
+            mask = rearrange(mask, "b ... -> b (...)")
             max_neg_value = -torch.finfo(sim.dtype).max
-            mask = repeat(mask, 'b j -> (b h) () j', h=h)
+            mask = repeat(mask, "b j -> (b h) () j", h=h)
             sim.masked_fill_(~mask, max_neg_value)
 
         # attention, what we cannot get enough of
         attn = sim.softmax(dim=-1)
 
-        out = einsum('b i j, b j d -> b i d', attn, v)
-        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
+        out = einsum("b i j, b j d -> b i d", attn, v)
+        out = rearrange(out, "(b h) n d -> b n (h d)", h=h)
 
         return out
 
@@ -111,21 +124,21 @@ class StructuredAttentionControl(AttentionControl):
         k = self.to_k(torch.cat([uc_context, context_k], dim=0))
         v = self.to_v(torch.cat([uc_context, context_v], dim=0))
 
-        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v))
 
-        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
+        sim = einsum("b i d, b j d -> b i j", q, k) * self.scale
 
         if exists(mask):
-            mask = rearrange(mask, 'b ... -> b (...)')
+            mask = rearrange(mask, "b ... -> b (...)")
             max_neg_value = -torch.finfo(sim.dtype).max
-            mask = repeat(mask, 'b j -> (b h) () j', h=h)
+            mask = repeat(mask, "b j -> (b h) () j", h=h)
             sim.masked_fill_(~mask, max_neg_value)
 
         # attention, what we cannot get enough of
         attn = sim.softmax(dim=-1)
 
-        out = einsum('b i j, b j d -> b i d', attn, v)
-        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
+        out = einsum("b i j, b j d -> b i d", attn, v)
+        out = rearrange(out, "(b h) n d -> b n (h d)", h=h)
         return out
 
     def get_kv(self, context):