update ed1 mode

chain.bat

@@ -1,3 +1,5 @@
 python train.py --config chain0.json
 python train.py --config chain1.json
-python train.py --config chain2.json
+python train.py --config chain2.json
+
+pause

train.py

@@ -436,6 +436,9 @@ def main(args):
         unet.enable_gradient_checkpointing()
         text_encoder.gradient_checkpointing_enable()
 
+    if args.ed1_mode:
+        unet.set_attention_slice(4)
+    
     if not args.disable_xformers and is_xformers_available():
         try:
             unet.enable_xformers_memory_efficient_attention()

utils/patch_unet.py

@@ -27,7 +27,7 @@ def patch_unet(ckpt_path, force_sd1attn: bool = False):
 
 
     if force_sd1attn:
-        unet_cfg["attention_head_dim"] = [5, 8, 8, 8]
+        unet_cfg["attention_head_dim"] = [8, 8, 8, 8]
     else:     
         unet_cfg["attention_head_dim"] = [5, 10, 20, 20]