1 files changed, 37 insertions, 48 deletions

diff --git a/train_ti.py b/train_ti.py
index b4b602b..727b591 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -107,12 +107,6 @@ def parse_args():
         help="Exclude all items with a listed collection.",
     )
     parser.add_argument(
-        "--repeats",
-        type=int,
-        default=1,
-        help="How many times to repeat the training data."
-    )
-    parser.add_argument(
         "--output_dir",
         type=str,
         default="output/text-inversion",
@@ -722,7 +716,6 @@ def main():
         size=args.resolution,
         num_aspect_ratio_buckets=args.num_aspect_ratio_buckets,
         progressive_aspect_ratio_buckets=args.progressive_aspect_ratio_buckets,
-        repeats=args.repeats,
         dropout=args.tag_dropout,
         template_key=args.train_data_template,
         valid_set_size=args.valid_set_size,
@@ -733,7 +726,7 @@ def main():
     )
     datamodule.setup()
 
-    train_dataloaders = datamodule.train_dataloaders
+    train_dataloader = datamodule.train_dataloader
     val_dataloader = datamodule.val_dataloader
 
     if args.num_class_images != 0:
@@ -752,8 +745,7 @@ def main():
 
     # Scheduler and math around the number of training steps.
     overrode_max_train_steps = False
-    num_update_steps_per_dataloader = sum(len(dataloader) for dataloader in train_dataloaders)
-    num_update_steps_per_epoch = math.ceil(num_update_steps_per_dataloader / args.gradient_accumulation_steps)
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
     if args.max_train_steps is None:
         args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
         overrode_max_train_steps = True
@@ -790,10 +782,9 @@ def main():
             num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
         )
 
-    text_encoder, optimizer, val_dataloader, lr_scheduler = accelerator.prepare(
-        text_encoder, optimizer, val_dataloader, lr_scheduler
+    text_encoder, optimizer, train_dataloader, val_dataloader, lr_scheduler = accelerator.prepare(
+        text_encoder, optimizer, train_dataloader, val_dataloader, lr_scheduler
     )
-    train_dataloaders = accelerator.prepare(*train_dataloaders)
 
     # Move vae and unet to device
     vae.to(accelerator.device, dtype=weight_dtype)
@@ -811,8 +802,7 @@ def main():
         unet.eval()
 
     # We need to recalculate our total training steps as the size of the training dataloader may have changed.
-    num_update_steps_per_dataloader = sum(len(dataloader) for dataloader in train_dataloaders)
-    num_update_steps_per_epoch = math.ceil(num_update_steps_per_dataloader / args.gradient_accumulation_steps)
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
     if overrode_max_train_steps:
         args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
 
@@ -870,7 +860,7 @@ def main():
             accelerator,
             text_encoder,
             optimizer,
-            train_dataloaders[0],
+            train_dataloader,
             val_dataloader,
             loop,
             on_train=on_train,
@@ -949,48 +939,47 @@ def main():
             text_encoder.train()
 
             with on_train():
-                for train_dataloader in train_dataloaders:
-                    for step, batch in enumerate(train_dataloader):
-                        with accelerator.accumulate(text_encoder):
-                            loss, acc, bsz = loop(step, batch)
+                for step, batch in enumerate(train_dataloader):
+                    with accelerator.accumulate(text_encoder):
+                        loss, acc, bsz = loop(step, batch)
 
-                            accelerator.backward(loss)
+                        accelerator.backward(loss)
 
-                            optimizer.step()
-                            if not accelerator.optimizer_step_was_skipped:
-                                lr_scheduler.step()
-                            optimizer.zero_grad(set_to_none=True)
+                        optimizer.step()
+                        if not accelerator.optimizer_step_was_skipped:
+                            lr_scheduler.step()
+                        optimizer.zero_grad(set_to_none=True)
 
-                            avg_loss.update(loss.detach_(), bsz)
-                            avg_acc.update(acc.detach_(), bsz)
+                        avg_loss.update(loss.detach_(), bsz)
+                        avg_acc.update(acc.detach_(), bsz)
 
-                        # Checks if the accelerator has performed an optimization step behind the scenes
-                        if accelerator.sync_gradients:
-                            if args.use_ema:
-                                ema_embeddings.step(
-                                    text_encoder.text_model.embeddings.temp_token_embedding.parameters())
+                    # Checks if the accelerator has performed an optimization step behind the scenes
+                    if accelerator.sync_gradients:
+                        if args.use_ema:
+                            ema_embeddings.step(
+                                text_encoder.text_model.embeddings.temp_token_embedding.parameters())
 
-                            local_progress_bar.update(1)
-                            global_progress_bar.update(1)
+                        local_progress_bar.update(1)
+                        global_progress_bar.update(1)
 
-                            global_step += 1
+                        global_step += 1
 
-                        logs = {
-                            "train/loss": avg_loss.avg.item(),
-                            "train/acc": avg_acc.avg.item(),
-                            "train/cur_loss": loss.item(),
-                            "train/cur_acc": acc.item(),
-                            "lr": lr_scheduler.get_last_lr()[0],
-                        }
-                        if args.use_ema:
-                            logs["ema_decay"] = ema_embeddings.decay
+                    logs = {
+                        "train/loss": avg_loss.avg.item(),
+                        "train/acc": avg_acc.avg.item(),
+                        "train/cur_loss": loss.item(),
+                        "train/cur_acc": acc.item(),
+                        "lr": lr_scheduler.get_last_lr()[0],
+                    }
+                    if args.use_ema:
+                        logs["ema_decay"] = ema_embeddings.decay
 
-                        accelerator.log(logs, step=global_step)
+                    accelerator.log(logs, step=global_step)
 
-                        local_progress_bar.set_postfix(**logs)
+                    local_progress_bar.set_postfix(**logs)
 
-                        if global_step >= args.max_train_steps:
-                            break
+                    if global_step >= args.max_train_steps:
+                        break
 
             accelerator.wait_for_everyone()