Code deduplication

1 files changed, 15 insertions, 56 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index ebcf802..da3a075 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -14,7 +14,6 @@ from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import LoggerType, set_seed
 from diffusers import AutoencoderKL, DDPMScheduler, DPMSolverMultistepScheduler, UNet2DConditionModel
-from diffusers.optimization import get_scheduler, get_cosine_with_hard_restarts_schedule_with_warmup
 import matplotlib.pyplot as plt
 from diffusers.training_utils import EMAModel
 from tqdm.auto import tqdm
@@ -24,8 +23,7 @@ from slugify import slugify
 from util import load_config, load_embeddings_from_dir
 from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 from data.csv import VlpnDataModule, VlpnDataItem
-from training.common import loss_step, generate_class_images
-from training.optimization import get_one_cycle_schedule
+from training.common import loss_step, generate_class_images, get_scheduler
 from training.lr import LRFinder
 from training.util import AverageMeter, CheckpointerBase, save_args
 from models.clip.embeddings import patch_managed_embeddings
@@ -750,35 +748,6 @@ def main():
         )
         return cond3 and cond4
 
-    def collate_fn(examples):
-        prompt_ids = [example["prompt_ids"] for example in examples]
-        nprompt_ids = [example["nprompt_ids"] for example in examples]
-
-        input_ids = [example["instance_prompt_ids"] for example in examples]
-        pixel_values = [example["instance_images"] for example in examples]
-
-        # concat class and instance examples for prior preservation
-        if args.num_class_images != 0 and "class_prompt_ids" in examples[0]:
-            input_ids += [example["class_prompt_ids"] for example in examples]
-            pixel_values += [example["class_images"] for example in examples]
-
-        pixel_values = torch.stack(pixel_values)
-        pixel_values = pixel_values.to(dtype=weight_dtype, memory_format=torch.contiguous_format)
-
-        prompts = prompt_processor.unify_input_ids(prompt_ids)
-        nprompts = prompt_processor.unify_input_ids(nprompt_ids)
-        inputs = prompt_processor.unify_input_ids(input_ids)
-
-        batch = {
-            "prompt_ids": prompts.input_ids,
-            "nprompt_ids": nprompts.input_ids,
-            "input_ids": inputs.input_ids,
-            "pixel_values": pixel_values,
-            "attention_mask": inputs.attention_mask,
-        }
-
-        return batch
-
     datamodule = VlpnDataModule(
         data_file=args.train_data_file,
         batch_size=args.train_batch_size,
@@ -798,7 +767,7 @@ def main():
         num_workers=args.dataloader_num_workers,
         seed=args.seed,
         filter=keyword_filter,
-        collate_fn=collate_fn
+        dtype=weight_dtype
     )
 
     datamodule.prepare_data()
@@ -829,33 +798,23 @@ def main():
         overrode_max_train_steps = True
     num_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
 
-    warmup_steps = args.lr_warmup_epochs * num_update_steps_per_epoch * args.gradient_accumulation_steps
-
-    if args.lr_scheduler == "one_cycle":
-        lr_min_lr = 0.04 if args.lr_min_lr is None else args.lr_min_lr / args.learning_rate
-        lr_scheduler = get_one_cycle_schedule(
-            optimizer=optimizer,
-            num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
-            warmup=args.lr_warmup_func,
-            annealing=args.lr_annealing_func,
-            warmup_exp=args.lr_warmup_exp,
-            annealing_exp=args.lr_annealing_exp,
-            min_lr=lr_min_lr,
-        )
-    elif args.lr_scheduler == "cosine_with_restarts":
-        lr_scheduler = get_cosine_with_hard_restarts_schedule_with_warmup(
-            optimizer=optimizer,
-            num_warmup_steps=warmup_steps,
-            num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
-            num_cycles=args.lr_cycles or math.ceil(math.sqrt(
-                ((args.max_train_steps - warmup_steps) / num_update_steps_per_epoch))),
-        )
+    if args.find_lr:
+        lr_scheduler = None
     else:
         lr_scheduler = get_scheduler(
             args.lr_scheduler,
             optimizer=optimizer,
-            num_warmup_steps=warmup_steps,
-            num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
+            min_lr=args.lr_min_lr,
+            lr=args.learning_rate,
+            warmup_func=args.lr_warmup_func,
+            annealing_func=args.lr_annealing_func,
+            warmup_exp=args.lr_warmup_exp,
+            annealing_exp=args.lr_annealing_exp,
+            cycles=args.lr_cycles,
+            warmup_epochs=args.lr_warmup_epochs,
+            max_train_steps=args.max_train_steps,
+            num_update_steps_per_epoch=num_update_steps_per_epoch,
+            gradient_accumulation_steps=args.gradient_accumulation_steps
         )
 
     unet, text_encoder, optimizer, train_dataloader, val_dataloader, lr_scheduler = accelerator.prepare(