4 files changed, 49 insertions, 23 deletions

diff --git a/data/csv.py b/data/csv.py
index 9770bec..c00ea07 100644
--- a/data/csv.py
+++ b/data/csv.py
@@ -143,7 +143,7 @@ class VlpnDataItem(NamedTuple):
 
 def keyword_filter(
     placeholder_tokens: Optional[list[str]],
-    collection: Optional[list[str]],
+    collections: Optional[list[str]],
     exclude_collections: Optional[list[str]],
     item: VlpnDataItem
 ):
@@ -152,11 +152,15 @@ def keyword_filter(
         for keyword in placeholder_tokens
         for part in item.prompt
     )
-    cond2 = collection is None or collection in item.collection
+    cond2 = collections is None or any(
+        collection in item.collection
+        for collection in collections
+    )
     cond3 = exclude_collections is None or not any(
         collection in item.collection
         for collection in exclude_collections
     )
+
     return cond1 and cond2 and cond3
 
 
diff --git a/train_dreambooth.py b/train_dreambooth.py
index 9345797..acb8287 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -4,6 +4,7 @@ import logging
 import itertools
 from pathlib import Path
 from functools import partial
+import math
 
 import torch
 import torch.utils.checkpoint
@@ -189,13 +190,12 @@ def parse_args():
     parser.add_argument(
         "--num_train_epochs",
         type=int,
-        default=100
+        default=None
     )
     parser.add_argument(
-        "--max_train_steps",
+        "--num_train_steps",
         type=int,
-        default=None,
-        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+        default=2000
     )
     parser.add_argument(
         "--gradient_accumulation_steps",
@@ -595,6 +595,11 @@ def main():
     )
     datamodule.setup()
 
+    num_train_epochs = args.num_train_epochs
+
+    if num_train_epochs is None:
+        num_train_epochs = math.ceil(len(datamodule.train_dataset) / args.num_train_steps)
+
     params_to_optimize = (unet.parameters(), )
     if args.train_text_encoder_epochs != 0:
         params_to_optimize += (
@@ -619,7 +624,7 @@ def main():
         annealing_exp=args.lr_annealing_exp,
         cycles=args.lr_cycles,
         end_lr=1e2,
-        train_epochs=args.num_train_epochs,
+        train_epochs=num_train_epochs,
         warmup_epochs=args.lr_warmup_epochs,
     )
 
@@ -631,7 +636,7 @@ def main():
         seed=args.seed,
         optimizer=optimizer,
         lr_scheduler=lr_scheduler,
-        num_train_epochs=args.num_train_epochs,
+        num_train_epochs=num_train_epochs,
         gradient_accumulation_steps=args.gradient_accumulation_steps,
         sample_frequency=args.sample_frequency,
         offset_noise_strength=args.offset_noise_strength,
diff --git a/train_lora.py b/train_lora.py
index 7ecddf0..a9c6e52 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -4,6 +4,7 @@ import logging
 import itertools
 from pathlib import Path
 from functools import partial
+import math
 
 import torch
 import torch.utils.checkpoint
@@ -178,13 +179,12 @@ def parse_args():
     parser.add_argument(
         "--num_train_epochs",
         type=int,
-        default=100
+        default=None
     )
     parser.add_argument(
-        "--max_train_steps",
+        "--num_train_steps",
         type=int,
-        default=None,
-        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+        default=2000
     )
     parser.add_argument(
         "--gradient_accumulation_steps",
@@ -627,6 +627,11 @@ def main():
     )
     datamodule.setup()
 
+    num_train_epochs = args.num_train_epochs
+
+    if num_train_epochs is None:
+        num_train_epochs = math.ceil(len(datamodule.train_dataset) / args.num_train_steps)
+
     optimizer = create_optimizer(
         itertools.chain(
             unet.parameters(),
@@ -647,7 +652,7 @@ def main():
         annealing_exp=args.lr_annealing_exp,
         cycles=args.lr_cycles,
         end_lr=1e2,
-        train_epochs=args.num_train_epochs,
+        train_epochs=num_train_epochs,
         warmup_epochs=args.lr_warmup_epochs,
     )
 
@@ -659,7 +664,7 @@ def main():
         seed=args.seed,
         optimizer=optimizer,
         lr_scheduler=lr_scheduler,
-        num_train_epochs=args.num_train_epochs,
+        num_train_epochs=num_train_epochs,
         gradient_accumulation_steps=args.gradient_accumulation_steps,
         sample_frequency=args.sample_frequency,
         offset_noise_strength=args.offset_noise_strength,
diff --git a/train_ti.py b/train_ti.py
index e4fd464..7bcc72f 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -3,6 +3,7 @@ import datetime
 import logging
 from functools import partial
 from pathlib import Path
+import math
 
 import torch
 import torch.utils.checkpoint
@@ -207,7 +208,12 @@ def parse_args():
     parser.add_argument(
         "--num_train_epochs",
         type=int,
-        default=100
+        default=None
+    )
+    parser.add_argument(
+        "--num_train_steps",
+        type=int,
+        default=2000
     )
     parser.add_argument(
         "--gradient_accumulation_steps",
@@ -513,13 +519,13 @@ def parse_args():
     if not isinstance(args.alias_tokens, list) or len(args.alias_tokens) % 2 != 0:
         raise ValueError("--alias_tokens must be a list with an even number of items")
 
-    args.alias_tokens += [
-        item
-        for pair in zip(args.placeholder_tokens, args.initializer_tokens)
-        for item in pair
-    ]
-
     if args.sequential:
+        args.alias_tokens += [
+            item
+            for pair in zip(args.placeholder_tokens, args.initializer_tokens)
+            for item in pair
+        ]
+
         if isinstance(args.train_data_template, str):
             args.train_data_template = [args.train_data_template] * len(args.placeholder_tokens)
 
@@ -607,6 +613,7 @@ def main():
             raise ValueError("--embeddings_dir must point to an existing directory")
 
         added_tokens, added_ids = load_embeddings_from_dir(tokenizer, embeddings, embeddings_dir)
+        embeddings.persist()
         print(f"Added {len(added_tokens)} tokens from embeddings dir: {list(zip(added_tokens, added_ids))}")
 
     if args.scale_lr:
@@ -682,7 +689,6 @@ def main():
         prior_loss_weight=args.prior_loss_weight if args.num_class_images != 0 else 0,
         no_val=args.valid_set_size == 0,
         strategy=textual_inversion_strategy,
-        num_train_epochs=args.num_train_epochs,
         gradient_accumulation_steps=args.gradient_accumulation_steps,
         sample_frequency=args.sample_frequency,
         checkpoint_frequency=args.checkpoint_frequency,
@@ -752,6 +758,11 @@ def main():
         )
         datamodule.setup()
 
+        num_train_epochs = args.num_train_epochs
+
+        if num_train_epochs is None:
+            num_train_epochs = math.ceil(len(datamodule.train_dataset) / args.num_train_steps)
+
         optimizer = create_optimizer(
             text_encoder.text_model.embeddings.temp_token_embedding.parameters(),
             lr=args.learning_rate,
@@ -769,7 +780,7 @@ def main():
             annealing_exp=args.lr_annealing_exp,
             cycles=args.lr_cycles,
             end_lr=1e3,
-            train_epochs=args.num_train_epochs,
+            train_epochs=num_train_epochs,
             warmup_epochs=args.lr_warmup_epochs,
         )
 
@@ -779,6 +790,7 @@ def main():
             val_dataloader=datamodule.val_dataloader,
             optimizer=optimizer,
             lr_scheduler=lr_scheduler,
+            num_train_epochs=num_train_epochs,
             # --
             sample_output_dir=sample_output_dir,
             placeholder_tokens=placeholder_tokens,