1 files changed, 269 insertions, 220 deletions

diff --git a/train_lora.py b/train_lora.py
index c74dd8f..fccf48d 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -16,6 +16,7 @@ from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import LoggerType, set_seed
 from peft import LoraConfig, get_peft_model
+
 # from diffusers.models.attention_processor import AttnProcessor
 from diffusers.utils.import_utils import is_xformers_available
 import transformers
@@ -34,15 +35,20 @@ from util.files import load_config, load_embeddings_from_dir
 
 # https://github.com/huggingface/peft/blob/main/examples/lora_dreambooth/train_dreambooth.py
 UNET_TARGET_MODULES_ORIG = ["to_q", "to_v", "query", "value"]
-UNET_TARGET_MODULES = UNET_TARGET_MODULES_ORIG + ["to_out.0"] # ["to_k", "key"]
+UNET_TARGET_MODULES = UNET_TARGET_MODULES_ORIG + ["to_out.0", "to_k", "key"]  # []
 TEXT_ENCODER_TARGET_MODULES_ORIG = ["q_proj", "v_proj"]
-TEXT_ENCODER_TARGET_MODULES = TEXT_ENCODER_TARGET_MODULES_ORIG + ["out_proj"] # ["k_proj"]
-TEXT_ENCODER_TARGET_MODULES_WITH_EMBEDDING = TEXT_ENCODER_TARGET_MODULES + ["token_embedding"]
+TEXT_ENCODER_TARGET_MODULES = TEXT_ENCODER_TARGET_MODULES_ORIG + [
+    "out_proj",
+    "k_proj",
+]  # []
+TEXT_ENCODER_TARGET_MODULES_WITH_EMBEDDING = TEXT_ENCODER_TARGET_MODULES + [
+    "token_embedding"
+]
 
 
 logger = get_logger(__name__)
 
-warnings.filterwarnings('ignore')
+warnings.filterwarnings("ignore")
 
 
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -55,20 +61,27 @@ hidet.torch.dynamo_config.use_tensor_core(True)
 hidet.torch.dynamo_config.search_space(0)
 
 
-if is_xformers_available():
-    import xformers
-    import xformers.ops
-    
-    orig_xformers_memory_efficient_attention = xformers.ops.memory_efficient_attention
-    def xformers_memory_efficient_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, **kwargs):
-        return orig_xformers_memory_efficient_attention(query.to(key.dtype), key, value.to(key.dtype), **kwargs)
-    xformers.ops.memory_efficient_attention = xformers_memory_efficient_attention
+def patch_xformers(dtype):
+    if is_xformers_available():
+        import xformers
+        import xformers.ops
+
+        orig_xformers_memory_efficient_attention = (
+            xformers.ops.memory_efficient_attention
+        )
+
+        def xformers_memory_efficient_attention(
+            query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, **kwargs
+        ):
+            return orig_xformers_memory_efficient_attention(
+                query.to(dtype), key.to(dtype), value.to(dtype), **kwargs
+            )
+
+        xformers.ops.memory_efficient_attention = xformers_memory_efficient_attention
 
 
 def parse_args():
-    parser = argparse.ArgumentParser(
-        description="Simple example of a training script."
-    )
+    parser = argparse.ArgumentParser(description="Simple example of a training script.")
     parser.add_argument(
         "--pretrained_model_name_or_path",
         type=str,
@@ -85,7 +98,7 @@ def parse_args():
         "--train_data_file",
         type=str,
         default=None,
-        help="A folder containing the training data."
+        help="A folder containing the training data.",
     )
     parser.add_argument(
         "--train_data_template",
@@ -96,13 +109,13 @@ def parse_args():
         "--train_set_pad",
         type=int,
         default=None,
-        help="The number to fill train dataset items up to."
+        help="The number to fill train dataset items up to.",
     )
     parser.add_argument(
         "--valid_set_pad",
         type=int,
         default=None,
-        help="The number to fill validation dataset items up to."
+        help="The number to fill validation dataset items up to.",
     )
     parser.add_argument(
         "--project",
@@ -111,64 +124,52 @@ def parse_args():
         help="The name of the current project.",
     )
     parser.add_argument(
-        "--auto_cycles",
-        type=str,
-        default="o",
-        help="Cycles to run automatically."
+        "--auto_cycles", type=str, default="o", help="Cycles to run automatically."
     )
     parser.add_argument(
-        "--cycle_decay",
-        type=float,
-        default=1.0,
-        help="Learning rate decay per cycle."
+        "--cycle_decay", type=float, default=1.0, help="Learning rate decay per cycle."
     )
     parser.add_argument(
         "--placeholder_tokens",
         type=str,
-        nargs='*',
+        nargs="*",
         help="A token to use as a placeholder for the concept.",
     )
     parser.add_argument(
         "--initializer_tokens",
         type=str,
-        nargs='*',
-        help="A token to use as initializer word."
+        nargs="*",
+        help="A token to use as initializer word.",
     )
     parser.add_argument(
-        "--filter_tokens",
-        type=str,
-        nargs='*',
-        help="Tokens to filter the dataset by."
+        "--filter_tokens", type=str, nargs="*", help="Tokens to filter the dataset by."
     )
     parser.add_argument(
         "--initializer_noise",
         type=float,
         default=0,
-        help="Noise to apply to the initializer word"
+        help="Noise to apply to the initializer word",
     )
     parser.add_argument(
         "--alias_tokens",
         type=str,
-        nargs='*',
+        nargs="*",
         default=[],
-        help="Tokens to create an alias for."
+        help="Tokens to create an alias for.",
     )
     parser.add_argument(
         "--inverted_initializer_tokens",
         type=str,
-        nargs='*',
-        help="A token to use as initializer word."
+        nargs="*",
+        help="A token to use as initializer word.",
     )
     parser.add_argument(
-        "--num_vectors",
-        type=int,
-        nargs='*',
-        help="Number of vectors per embedding."
+        "--num_vectors", type=int, nargs="*", help="Number of vectors per embedding."
     )
     parser.add_argument(
         "--exclude_collections",
         type=str,
-        nargs='*',
+        nargs="*",
         help="Exclude all items with a listed collection.",
     )
     parser.add_argument(
@@ -214,7 +215,7 @@ def parse_args():
         "--num_class_images",
         type=int,
         default=0,
-        help="How many class images to generate."
+        help="How many class images to generate.",
     )
     parser.add_argument(
         "--class_image_dir",
@@ -242,14 +243,11 @@ def parse_args():
     parser.add_argument(
         "--collection",
         type=str,
-        nargs='*',
+        nargs="*",
         help="A collection to filter the dataset.",
     )
     parser.add_argument(
-        "--seed",
-        type=int,
-        default=None,
-        help="A seed for reproducible training."
+        "--seed", type=int, default=None, help="A seed for reproducible training."
     )
     parser.add_argument(
         "--resolution",
@@ -270,18 +268,10 @@ def parse_args():
         "--input_pertubation",
         type=float,
         default=0,
-        help="The scale of input pretubation. Recommended 0.1."
-    )
-    parser.add_argument(
-        "--num_train_epochs",
-        type=int,
-        default=None
-    )
-    parser.add_argument(
-        "--num_train_steps",
-        type=int,
-        default=2000
+        help="The scale of input pretubation. Recommended 0.1.",
     )
+    parser.add_argument("--num_train_epochs", type=int, default=None)
+    parser.add_argument("--num_train_steps", type=int, default=2000)
     parser.add_argument(
         "--gradient_accumulation_steps",
         type=int,
@@ -289,22 +279,19 @@ def parse_args():
         help="Number of updates steps to accumulate before performing a backward/update pass.",
     )
     parser.add_argument(
-        "--lora_r",
-        type=int,
-        default=8,
-        help="Lora rank, only used if use_lora is True"
+        "--lora_r", type=int, default=8, help="Lora rank, only used if use_lora is True"
     )
     parser.add_argument(
         "--lora_alpha",
         type=int,
         default=32,
-        help="Lora alpha, only used if use_lora is True"
+        help="Lora alpha, only used if use_lora is True",
     )
     parser.add_argument(
         "--lora_dropout",
         type=float,
         default=0.0,
-        help="Lora dropout, only used if use_lora is True"
+        help="Lora dropout, only used if use_lora is True",
     )
     parser.add_argument(
         "--lora_bias",
@@ -344,7 +331,7 @@ def parse_args():
     parser.add_argument(
         "--train_text_encoder_cycles",
         default=999999,
-        help="Number of epochs the text encoder will be trained."
+        help="Number of epochs the text encoder will be trained.",
     )
     parser.add_argument(
         "--find_lr",
@@ -378,27 +365,31 @@ def parse_args():
         "--lr_scheduler",
         type=str,
         default="one_cycle",
-        choices=["linear", "cosine", "cosine_with_restarts", "polynomial",
-                 "constant", "constant_with_warmup", "one_cycle"],
-        help='The scheduler type to use.',
+        choices=[
+            "linear",
+            "cosine",
+            "cosine_with_restarts",
+            "polynomial",
+            "constant",
+            "constant_with_warmup",
+            "one_cycle",
+        ],
+        help="The scheduler type to use.",
     )
     parser.add_argument(
         "--lr_warmup_epochs",
         type=int,
         default=10,
-        help="Number of steps for the warmup in the lr scheduler."
+        help="Number of steps for the warmup in the lr scheduler.",
     )
     parser.add_argument(
-        "--lr_mid_point",
-        type=float,
-        default=0.3,
-        help="OneCycle schedule mid point."
+        "--lr_mid_point", type=float, default=0.3, help="OneCycle schedule mid point."
     )
     parser.add_argument(
         "--lr_cycles",
         type=int,
         default=None,
-        help="Number of restart cycles in the lr scheduler (if supported)."
+        help="Number of restart cycles in the lr scheduler (if supported).",
     )
     parser.add_argument(
         "--lr_warmup_func",
@@ -410,7 +401,7 @@ def parse_args():
         "--lr_warmup_exp",
         type=int,
         default=1,
-        help='If lr_warmup_func is "cos", exponent to modify the function'
+        help='If lr_warmup_func is "cos", exponent to modify the function',
     )
     parser.add_argument(
         "--lr_annealing_func",
@@ -422,69 +413,76 @@ def parse_args():
         "--lr_annealing_exp",
         type=int,
         default=3,
-        help='If lr_annealing_func is "half_cos" or "cos", exponent to modify the function'
+        help='If lr_annealing_func is "half_cos" or "cos", exponent to modify the function',
     )
     parser.add_argument(
         "--lr_min_lr",
         type=float,
         default=0.04,
-        help="Minimum learning rate in the lr scheduler."
-    )
-    parser.add_argument(
-        "--min_snr_gamma",
-        type=int,
-        default=5,
-        help="MinSNR gamma."
+        help="Minimum learning rate in the lr scheduler.",
     )
+    parser.add_argument("--min_snr_gamma", type=int, default=5, help="MinSNR gamma.")
     parser.add_argument(
         "--schedule_sampler",
         type=str,
         default="uniform",
         choices=["uniform", "loss-second-moment"],
-        help="Noise schedule sampler."
+        help="Noise schedule sampler.",
     )
     parser.add_argument(
         "--optimizer",
         type=str,
         default="adan",
-        choices=["adam", "adam8bit", "adan", "lion", "dadam", "dadan", "adafactor"],
-        help='Optimizer to use'
+        choices=[
+            "adam",
+            "adam8bit",
+            "adan",
+            "lion",
+            "dadam",
+            "dadan",
+            "dlion",
+            "adafactor",
+        ],
+        help="Optimizer to use",
     )
     parser.add_argument(
         "--dadaptation_d0",
         type=float,
         default=1e-6,
-        help="The d0 parameter for Dadaptation optimizers."
+        help="The d0 parameter for Dadaptation optimizers.",
+    )
+    parser.add_argument(
+        "--dadaptation_growth_rate",
+        type=float,
+        default=math.inf,
+        help="The growth_rate parameter for Dadaptation optimizers.",
     )
     parser.add_argument(
         "--adam_beta1",
         type=float,
         default=None,
-        help="The beta1 parameter for the Adam optimizer."
+        help="The beta1 parameter for the Adam optimizer.",
     )
     parser.add_argument(
         "--adam_beta2",
         type=float,
         default=None,
-        help="The beta2 parameter for the Adam optimizer."
+        help="The beta2 parameter for the Adam optimizer.",
     )
     parser.add_argument(
-        "--adam_weight_decay",
-        type=float,
-        default=2e-2,
-        help="Weight decay to use."
+        "--adam_weight_decay", type=float, default=2e-2, help="Weight decay to use."
     )
     parser.add_argument(
         "--adam_epsilon",
         type=float,
         default=1e-08,
-        help="Epsilon value for the Adam optimizer"
+        help="Epsilon value for the Adam optimizer",
     )
     parser.add_argument(
         "--adam_amsgrad",
         type=bool,
         default=False,
-        help="Amsgrad value for the Adam optimizer"
+        help="Amsgrad value for the Adam optimizer",
     )
     parser.add_argument(
         "--mixed_precision",
@@ -547,19 +545,19 @@ def parse_args():
         "--valid_set_size",
         type=int,
         default=None,
-        help="Number of images in the validation dataset."
+        help="Number of images in the validation dataset.",
     )
     parser.add_argument(
         "--valid_set_repeat",
         type=int,
         default=1,
-        help="Times the images in the validation dataset are repeated."
+        help="Times the images in the validation dataset are repeated.",
     )
     parser.add_argument(
         "--train_batch_size",
         type=int,
         default=1,
-        help="Batch size (per device) for the training dataloader."
+        help="Batch size (per device) for the training dataloader.",
     )
     parser.add_argument(
         "--sample_steps",
@@ -571,19 +569,10 @@ def parse_args():
         "--prior_loss_weight",
         type=float,
         default=1.0,
-        help="The weight of prior preservation loss."
-    )
-    parser.add_argument(
-        "--run_pti",
-        action="store_true",
-        help="Whether to run PTI."
-    )
-    parser.add_argument(
-        "--emb_alpha",
-        type=float,
-        default=1.0,
-        help="Embedding alpha"
+        help="The weight of prior preservation loss.",
     )
+    parser.add_argument("--run_pti", action="store_true", help="Whether to run PTI.")
+    parser.add_argument("--emb_alpha", type=float, default=1.0, help="Embedding alpha")
     parser.add_argument(
         "--emb_dropout",
         type=float,
@@ -591,27 +580,16 @@ def parse_args():
         help="Embedding dropout probability.",
     )
     parser.add_argument(
-        "--use_emb_decay",
-        action="store_true",
-        help="Whether to use embedding decay."
+        "--use_emb_decay", action="store_true", help="Whether to use embedding decay."
     )
     parser.add_argument(
-        "--emb_decay_target",
-        default=0.4,
-        type=float,
-        help="Embedding decay target."
+        "--emb_decay_target", default=0.4, type=float, help="Embedding decay target."
     )
     parser.add_argument(
-        "--emb_decay",
-        default=1e+2,
-        type=float,
-        help="Embedding decay factor."
+        "--emb_decay", default=1e2, type=float, help="Embedding decay factor."
     )
     parser.add_argument(
-        "--max_grad_norm",
-        default=1.0,
-        type=float,
-        help="Max gradient norm."
+        "--max_grad_norm", default=1.0, type=float, help="Max gradient norm."
     )
     parser.add_argument(
         "--noise_timesteps",
@@ -622,7 +600,7 @@ def parse_args():
         "--config",
         type=str,
         default=None,
-        help="Path to a JSON configuration file containing arguments for invoking this script."
+        help="Path to a JSON configuration file containing arguments for invoking this script.",
     )
 
     args = parser.parse_args()
@@ -649,29 +627,44 @@ def parse_args():
         args.placeholder_tokens = [args.placeholder_tokens]
 
     if isinstance(args.initializer_tokens, str):
-        args.initializer_tokens = [args.initializer_tokens] * len(args.placeholder_tokens)
+        args.initializer_tokens = [args.initializer_tokens] * len(
+            args.placeholder_tokens
+        )
 
     if len(args.placeholder_tokens) == 0:
-        args.placeholder_tokens = [f"<*{i}>" for i in range(len(args.initializer_tokens))]
+        args.placeholder_tokens = [
+            f"<*{i}>" for i in range(len(args.initializer_tokens))
+        ]
 
     if len(args.initializer_tokens) == 0:
         args.initializer_tokens = args.placeholder_tokens.copy()
 
     if len(args.placeholder_tokens) != len(args.initializer_tokens):
-        raise ValueError("--placeholder_tokens and --initializer_tokens must have the same number of items")
+        raise ValueError(
+            "--placeholder_tokens and --initializer_tokens must have the same number of items"
+        )
 
     if isinstance(args.inverted_initializer_tokens, str):
-        args.inverted_initializer_tokens = [args.inverted_initializer_tokens] * len(args.placeholder_tokens)
+        args.inverted_initializer_tokens = [args.inverted_initializer_tokens] * len(
+            args.placeholder_tokens
+        )
 
-    if isinstance(args.inverted_initializer_tokens, list) and len(args.inverted_initializer_tokens) != 0:
+    if (
+        isinstance(args.inverted_initializer_tokens, list)
+        and len(args.inverted_initializer_tokens) != 0
+    ):
         args.placeholder_tokens += [f"inv_{t}" for t in args.placeholder_tokens]
         args.initializer_tokens += args.inverted_initializer_tokens
 
     if isinstance(args.num_vectors, int):
         args.num_vectors = [args.num_vectors] * len(args.placeholder_tokens)
 
-    if isinstance(args.num_vectors, list) and len(args.placeholder_tokens) != len(args.num_vectors):
-        raise ValueError("--placeholder_tokens and --num_vectors must have the same number of items")
+    if isinstance(args.num_vectors, list) and len(args.placeholder_tokens) != len(
+        args.num_vectors
+    ):
+        raise ValueError(
+            "--placeholder_tokens and --num_vectors must have the same number of items"
+        )
 
     if args.alias_tokens is None:
         args.alias_tokens = []
@@ -695,15 +688,15 @@ def parse_args():
         raise ValueError("You must specify --output_dir")
 
     if args.adam_beta1 is None:
-        if args.optimizer in ('adam', 'adam8bit'):
+        if args.optimizer in ("adam", "adam8bit", "dadam"):
             args.adam_beta1 = 0.9
-        elif args.optimizer == 'lion':
+        elif args.optimizer in ("lion", "dlion"):
             args.adam_beta1 = 0.95
 
     if args.adam_beta2 is None:
-        if args.optimizer in ('adam', 'adam8bit'):
+        if args.optimizer in ("adam", "adam8bit", "dadam"):
             args.adam_beta2 = 0.999
-        elif args.optimizer == 'lion':
+        elif args.optimizer in ("lion", "dlion"):
             args.adam_beta2 = 0.98
 
     return args
@@ -719,7 +712,7 @@ def main():
     accelerator = Accelerator(
         log_with=LoggerType.TENSORBOARD,
         project_dir=f"{output_dir}",
-        mixed_precision=args.mixed_precision
+        mixed_precision=args.mixed_precision,
     )
 
     weight_dtype = torch.float32
@@ -728,6 +721,8 @@ def main():
     elif args.mixed_precision == "bf16":
         weight_dtype = torch.bfloat16
 
+    patch_xformers(weight_dtype)
+
     logging.basicConfig(filename=output_dir / "log.txt", level=logging.DEBUG)
 
     if args.seed is None:
@@ -737,12 +732,18 @@ def main():
 
     save_args(output_dir, args)
 
-    tokenizer, text_encoder, vae, unet, noise_scheduler, sample_scheduler = get_models(args.pretrained_model_name_or_path)
-    schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, noise_scheduler.config.num_train_timesteps)
-    
+    tokenizer, text_encoder, vae, unet, noise_scheduler, sample_scheduler = get_models(
+        args.pretrained_model_name_or_path
+    )
+    schedule_sampler = create_named_schedule_sampler(
+        args.schedule_sampler, noise_scheduler.config.num_train_timesteps
+    )
+
     def ensure_embeddings():
         if args.lora_text_encoder_emb:
-            raise ValueError("Can't use TI options when training token embeddings with LoRA")
+            raise ValueError(
+                "Can't use TI options when training token embeddings with LoRA"
+            )
         return patch_managed_embeddings(text_encoder, args.emb_alpha, args.emb_dropout)
 
     unet_config = LoraConfig(
@@ -757,7 +758,9 @@ def main():
     text_encoder_config = LoraConfig(
         r=args.lora_text_encoder_r,
         lora_alpha=args.lora_text_encoder_alpha,
-        target_modules=TEXT_ENCODER_TARGET_MODULES_WITH_EMBEDDING if args.lora_text_encoder_emb else TEXT_ENCODER_TARGET_MODULES,
+        target_modules=TEXT_ENCODER_TARGET_MODULES_WITH_EMBEDDING
+        if args.lora_text_encoder_emb
+        else TEXT_ENCODER_TARGET_MODULES,
         lora_dropout=args.lora_text_encoder_dropout,
         bias=args.lora_text_encoder_bias,
     )
@@ -787,7 +790,7 @@ def main():
 
     if len(args.alias_tokens) != 0:
         embeddings = ensure_embeddings()
-        
+
         alias_placeholder_tokens = args.alias_tokens[::2]
         alias_initializer_tokens = args.alias_tokens[1::2]
 
@@ -795,27 +798,33 @@ def main():
             tokenizer=tokenizer,
             embeddings=embeddings,
             placeholder_tokens=alias_placeholder_tokens,
-            initializer_tokens=alias_initializer_tokens
+            initializer_tokens=alias_initializer_tokens,
         )
         embeddings.persist()
-        print(f"Added {len(added_tokens)} aliases: {list(zip(alias_placeholder_tokens, added_tokens, alias_initializer_tokens, added_ids))}")
+        print(
+            f"Added {len(added_tokens)} aliases: {list(zip(alias_placeholder_tokens, added_tokens, alias_initializer_tokens, added_ids))}"
+        )
 
     placeholder_tokens = []
     placeholder_token_ids = []
 
     if args.embeddings_dir is not None:
         embeddings = ensure_embeddings()
-        
+
         embeddings_dir = Path(args.embeddings_dir)
         if not embeddings_dir.exists() or not embeddings_dir.is_dir():
             raise ValueError("--embeddings_dir must point to an existing directory")
 
-        added_tokens, added_ids = load_embeddings_from_dir(tokenizer, embeddings, embeddings_dir)
+        added_tokens, added_ids = load_embeddings_from_dir(
+            tokenizer, embeddings, embeddings_dir
+        )
 
         placeholder_tokens = added_tokens
         placeholder_token_ids = added_ids
 
-        print(f"Added {len(added_tokens)} tokens from embeddings dir: {list(zip(added_tokens, added_ids))}")
+        print(
+            f"Added {len(added_tokens)} tokens from embeddings dir: {list(zip(added_tokens, added_ids))}"
+        )
 
         if args.train_dir_embeddings:
             print("Training embeddings from embeddings dir")
@@ -824,7 +833,7 @@ def main():
 
     if len(args.placeholder_tokens) != 0 and not args.train_dir_embeddings:
         embeddings = ensure_embeddings()
-        
+
         placeholder_token_ids, initializer_token_ids = add_placeholder_tokens(
             tokenizer=tokenizer,
             embeddings=embeddings,
@@ -836,23 +845,34 @@ def main():
 
         placeholder_tokens = args.placeholder_tokens
 
-        stats = list(zip(
-            placeholder_tokens, placeholder_token_ids, args.initializer_tokens, initializer_token_ids
-        ))
+        stats = list(
+            zip(
+                placeholder_tokens,
+                placeholder_token_ids,
+                args.initializer_tokens,
+                initializer_token_ids,
+            )
+        )
         print(f"Training embeddings: {stats}")
 
     if args.scale_lr:
         args.learning_rate_unet = (
-            args.learning_rate_unet * args.gradient_accumulation_steps *
-            args.train_batch_size * accelerator.num_processes
+            args.learning_rate_unet
+            * args.gradient_accumulation_steps
+            * args.train_batch_size
+            * accelerator.num_processes
         )
         args.learning_rate_text = (
-            args.learning_rate_text * args.gradient_accumulation_steps *
-            args.train_batch_size * accelerator.num_processes
+            args.learning_rate_text
+            * args.gradient_accumulation_steps
+            * args.train_batch_size
+            * accelerator.num_processes
         )
         args.learning_rate_emb = (
-            args.learning_rate_emb * args.gradient_accumulation_steps *
-            args.train_batch_size * accelerator.num_processes
+            args.learning_rate_emb
+            * args.gradient_accumulation_steps
+            * args.train_batch_size
+            * accelerator.num_processes
         )
 
     if args.find_lr:
@@ -861,11 +881,13 @@ def main():
         args.learning_rate_emb = 1e-6
         args.lr_scheduler = "exponential_growth"
 
-    if args.optimizer == 'adam8bit':
+    if args.optimizer == "adam8bit":
         try:
             import bitsandbytes as bnb
         except ImportError:
-            raise ImportError("To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`.")
+            raise ImportError(
+                "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
+            )
 
         create_optimizer = partial(
             bnb.optim.AdamW8bit,
@@ -874,7 +896,7 @@ def main():
             eps=args.adam_epsilon,
             amsgrad=args.adam_amsgrad,
         )
-    elif args.optimizer == 'adam':
+    elif args.optimizer == "adam":
         create_optimizer = partial(
             torch.optim.AdamW,
             betas=(args.adam_beta1, args.adam_beta2),
@@ -882,11 +904,13 @@ def main():
             eps=args.adam_epsilon,
             amsgrad=args.adam_amsgrad,
         )
-    elif args.optimizer == 'adan':
+    elif args.optimizer == "adan":
         try:
             import timm.optim
         except ImportError:
-            raise ImportError("To use Adan, please install the PyTorch Image Models library: `pip install timm`.")
+            raise ImportError(
+                "To use Adan, please install the PyTorch Image Models library: `pip install timm`."
+            )
 
         create_optimizer = partial(
             timm.optim.Adan,
@@ -894,11 +918,13 @@ def main():
             eps=args.adam_epsilon,
             no_prox=True,
         )
-    elif args.optimizer == 'lion':
+    elif args.optimizer == "lion":
         try:
             import lion_pytorch
         except ImportError:
-            raise ImportError("To use Lion, please install the lion_pytorch library: `pip install lion-pytorch`.")
+            raise ImportError(
+                "To use Lion, please install the lion_pytorch library: `pip install lion-pytorch`."
+            )
 
         create_optimizer = partial(
             lion_pytorch.Lion,
@@ -906,7 +932,7 @@ def main():
             weight_decay=args.adam_weight_decay,
             use_triton=True,
         )
-    elif args.optimizer == 'adafactor':
+    elif args.optimizer == "adafactor":
         create_optimizer = partial(
             transformers.optimization.Adafactor,
             weight_decay=args.adam_weight_decay,
@@ -920,11 +946,13 @@ def main():
         args.learning_rate_unet = None
         args.learning_rate_text = None
         args.learning_rate_emb = None
-    elif args.optimizer == 'dadam':
+    elif args.optimizer == "dadam":
         try:
             import dadaptation
         except ImportError:
-            raise ImportError("To use DAdaptAdam, please install the dadaptation library: `pip install dadaptation`.")
+            raise ImportError(
+                "To use DAdaptAdam, please install the dadaptation library: `pip install dadaptation`."
+            )
 
         create_optimizer = partial(
             dadaptation.DAdaptAdam,
@@ -933,29 +961,35 @@ def main():
             eps=args.adam_epsilon,
             decouple=True,
             d0=args.dadaptation_d0,
+            growth_rate=args.dadaptation_growth_rate,
         )
 
         args.learning_rate_unet = 1.0
         args.learning_rate_text = 1.0
         args.learning_rate_emb = 1.0
-    elif args.optimizer == 'dadan':
+    elif args.optimizer == "dadan":
         try:
             import dadaptation
         except ImportError:
-            raise ImportError("To use DAdaptAdan, please install the dadaptation library: `pip install dadaptation`.")
+            raise ImportError(
+                "To use DAdaptAdan, please install the dadaptation library: `pip install dadaptation`."
+            )
 
         create_optimizer = partial(
             dadaptation.DAdaptAdan,
             weight_decay=args.adam_weight_decay,
             eps=args.adam_epsilon,
             d0=args.dadaptation_d0,
+            growth_rate=args.dadaptation_growth_rate,
         )
 
         args.learning_rate_unet = 1.0
         args.learning_rate_text = 1.0
         args.learning_rate_emb = 1.0
+    elif args.optimizer == "dlion":
+        raise ImportError("DLion has not been merged into dadaptation yet")
     else:
-        raise ValueError(f"Unknown --optimizer \"{args.optimizer}\"")
+        raise ValueError(f'Unknown --optimizer "{args.optimizer}"')
 
     trainer = partial(
         train,
@@ -1026,25 +1060,33 @@ def main():
 
     if args.run_pti and len(placeholder_tokens) != 0:
         embeddings = ensure_embeddings()
-        
-        filter_tokens = [token for token in args.filter_tokens if token in placeholder_tokens]
+
+        filter_tokens = [
+            token for token in args.filter_tokens if token in placeholder_tokens
+        ]
 
         pti_datamodule = create_datamodule(
             valid_set_size=0,
             batch_size=args.train_batch_size,
-            filter=partial(keyword_filter, filter_tokens, args.collection, args.exclude_collections),
+            filter=partial(
+                keyword_filter, filter_tokens, args.collection, args.exclude_collections
+            ),
         )
         pti_datamodule.setup()
 
         num_train_epochs = args.num_train_epochs
         pti_sample_frequency = args.sample_frequency
         if num_train_epochs is None:
-            num_train_epochs = math.ceil(
-                args.num_train_steps / len(pti_datamodule.train_dataset)
-            ) * args.gradient_accumulation_steps
-            pti_sample_frequency = math.ceil(num_train_epochs * (pti_sample_frequency / args.num_train_steps))
+            num_train_epochs = (
+                math.ceil(args.num_train_steps / len(pti_datamodule.train_dataset))
+                * args.gradient_accumulation_steps
+            )
+            pti_sample_frequency = math.ceil(
+                num_train_epochs * (pti_sample_frequency / args.num_train_steps)
+            )
         num_training_steps_per_epoch = math.ceil(
-            len(pti_datamodule.train_dataset) / args.gradient_accumulation_steps)
+            len(pti_datamodule.train_dataset) / args.gradient_accumulation_steps
+        )
         num_train_steps = num_training_steps_per_epoch * num_train_epochs
         if args.sample_num is not None:
             pti_sample_frequency = math.ceil(num_train_epochs / args.sample_num)
@@ -1060,11 +1102,15 @@ def main():
         print(f"============ PTI ============")
         print("")
 
-        pti_optimizer = create_optimizer([{
-            "params": text_encoder.text_model.embeddings.token_embedding.parameters(),
-            "lr": args.learning_rate_emb,
-            "weight_decay": 0,
-        }])
+        pti_optimizer = create_optimizer(
+            [
+                {
+                    "params": text_encoder.text_model.embeddings.token_embedding.parameters(),
+                    "lr": args.learning_rate_emb,
+                    "weight_decay": 0,
+                }
+            ]
+        )
 
         pti_lr_scheduler = create_lr_scheduler(
             "constant_with_warmup",
@@ -1113,11 +1159,16 @@ def main():
     num_train_epochs = args.num_train_epochs
     lora_sample_frequency = args.sample_frequency
     if num_train_epochs is None:
-        num_train_epochs = math.ceil(
-            args.num_train_steps / len(lora_datamodule.train_dataset)
-        ) * args.gradient_accumulation_steps
-        lora_sample_frequency = math.ceil(num_train_epochs * (lora_sample_frequency / args.num_train_steps))
-    num_training_steps_per_epoch = math.ceil(len(lora_datamodule.train_dataset) / args.gradient_accumulation_steps)
+        num_train_epochs = (
+            math.ceil(args.num_train_steps / len(lora_datamodule.train_dataset))
+            * args.gradient_accumulation_steps
+        )
+        lora_sample_frequency = math.ceil(
+            num_train_epochs * (lora_sample_frequency / args.num_train_steps)
+        )
+    num_training_steps_per_epoch = math.ceil(
+        len(lora_datamodule.train_dataset) / args.gradient_accumulation_steps
+    )
     num_train_steps = num_training_steps_per_epoch * num_train_epochs
     if args.sample_num is not None:
         lora_sample_frequency = math.ceil(num_train_epochs / args.sample_num)
@@ -1131,7 +1182,6 @@ def main():
 
     training_iter = 0
     auto_cycles = list(args.auto_cycles)
-    learning_rate_emb = args.learning_rate_emb
     learning_rate_unet = args.learning_rate_unet
     learning_rate_text = args.learning_rate_text
     lr_scheduler = args.lr_scheduler
@@ -1145,21 +1195,15 @@ def main():
 
     params_to_optimize = [
         {
-            "params": (
-                param
-                for param in unet.parameters()
-                if param.requires_grad
-            ),
+            "params": (param for param in unet.parameters() if param.requires_grad),
             "lr": learning_rate_unet,
         },
         {
             "params": (
-                param
-                for param in text_encoder.parameters()
-                if param.requires_grad
+                param for param in text_encoder.parameters() if param.requires_grad
             ),
             "lr": learning_rate_text,
-        }
+        },
     ]
     group_labels = ["unet", "text"]
 
@@ -1169,19 +1213,26 @@ def main():
         if len(auto_cycles) != 0:
             response = auto_cycles.pop(0)
         else:
-            response = input("\n### Choose action: [o] one_cycle, [w] warmup, [c] constant, [d] decay, [s] stop \n--> ")
+            response = input(
+                "\n### Choose action: [o] one_cycle, [w] warmup, [c] constant, [d] decay, [s] stop \n--> "
+            )
 
         if response.lower().strip() == "o":
-            if args.learning_rate_emb is not None:
-                learning_rate_emb = args.learning_rate_emb * 2
             if args.learning_rate_unet is not None:
-                learning_rate_unet = args.learning_rate_unet * 2
+                learning_rate_unet = (
+                    args.learning_rate_unet * 2 * (args.cycle_decay**training_iter)
+                )
             if args.learning_rate_text is not None:
-                learning_rate_text = args.learning_rate_text * 2
+                learning_rate_text = (
+                    args.learning_rate_text * 2 * (args.cycle_decay**training_iter)
+                )
         else:
-            learning_rate_emb = args.learning_rate_emb
-            learning_rate_unet = args.learning_rate_unet
-            learning_rate_text = args.learning_rate_text
+            learning_rate_unet = args.learning_rate_unet * (
+                args.cycle_decay**training_iter
+            )
+            learning_rate_text = args.learning_rate_text * (
+                args.cycle_decay**training_iter
+            )
 
         if response.lower().strip() == "o":
             lr_scheduler = "one_cycle"
@@ -1204,9 +1255,11 @@ def main():
         print("")
         print(f"============ LoRA cycle {training_iter + 1}: {response} ============")
         print("")
-        
-        for group, lr in zip(lora_optimizer.param_groups, [learning_rate_unet, learning_rate_text]):
-            group['lr'] = lr
+
+        for group, lr in zip(
+            lora_optimizer.param_groups, [learning_rate_unet, learning_rate_text]
+        ):
+            group["lr"] = lr
 
         lora_lr_scheduler = create_lr_scheduler(
             lr_scheduler,
@@ -1218,7 +1271,9 @@ def main():
             warmup_epochs=lr_warmup_epochs,
         )
 
-        lora_checkpoint_output_dir = output_dir / lora_project / f"model_{training_iter}"
+        lora_checkpoint_output_dir = (
+            output_dir / lora_project / f"model_{training_iter}"
+        )
 
         trainer(
             strategy=lora_strategy,
@@ -1246,12 +1301,6 @@ def main():
         )
 
         training_iter += 1
-        if learning_rate_emb is not None:
-            learning_rate_emb *= args.cycle_decay
-        if learning_rate_unet is not None:
-            learning_rate_unet *= args.cycle_decay
-        if learning_rate_text is not None:
-            learning_rate_text *= args.cycle_decay
 
     accelerator.end_training()