Implemented extended Dreambooth training

1 files changed, 4 insertions, 3 deletions

diff --git a/training/functional.py b/training/functional.py
index 5984ffb..f5c111e 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -494,10 +494,11 @@ def train(
     text_encoder: CLIPTextModel,
     vae: AutoencoderKL,
     noise_scheduler: DDPMScheduler,
-    train_dataloader: DataLoader,
-    val_dataloader: DataLoader,
     dtype: torch.dtype,
     seed: int,
+    project: str,
+    train_dataloader: DataLoader,
+    val_dataloader: DataLoader,
     optimizer: torch.optim.Optimizer,
     lr_scheduler: torch.optim.lr_scheduler._LRScheduler,
     callbacks_fn: Callable[..., TrainingCallbacks],
@@ -544,7 +545,7 @@ def train(
     )
 
     if accelerator.is_main_process:
-        accelerator.init_trackers("textual_inversion")
+        accelerator.init_trackers(project)
 
     train_loop(
         accelerator=accelerator,