Update

1 files changed, 2 insertions, 3 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index d722e68..48bdcf8 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -14,8 +14,7 @@ from slugify import slugify
 
 from util import load_config, load_embeddings_from_dir
 from data.csv import VlpnDataModule, keyword_filter
-from training.functional import train, generate_class_images, add_placeholder_tokens, get_models
-from training.strategy.ti import textual_inversion_strategy
+from training.functional import train, get_models
 from training.strategy.dreambooth import dreambooth_strategy
 from training.optimization import get_scheduler
 from training.util import save_args
@@ -610,7 +609,7 @@ def main():
     )
 
     trainer(
-        callbacks_fn=dreambooth_strategy,
+        strategy=dreambooth_strategy,
         project="dreambooth",
         train_dataloader=datamodule.train_dataloader,
         val_dataloader=datamodule.val_dataloader,