Update

1 files changed, 3 insertions, 2 deletions

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