6 files changed, 15 insertions, 9 deletions

diff --git a/data/csv.py b/data/csv.py
index c505230..a60733a 100644
--- a/data/csv.py
+++ b/data/csv.py
@@ -151,7 +151,7 @@ class CSVDataModule():
 
         num_images = len(items)
 
-        valid_set_size = self.valid_set_size if self.valid_set_size is not None else int(num_images * 0.2)
+        valid_set_size = self.valid_set_size if self.valid_set_size is not None else num_images // 10
         valid_set_size = max(valid_set_size, 1)
         train_set_size = num_images - valid_set_size
 
diff --git a/infer.py b/infer.py
index c4d1e0d..b29b136 100644
--- a/infer.py
+++ b/infer.py
@@ -295,11 +295,10 @@ def generate(output_dir, pipeline, args):
         for j, image in enumerate(images):
             image_dir = output_dir
             if use_subdirs:
-                idx = j % len(args.prompt)
-                image_dir = image_dir.joinpath(slugify(args.prompt[idx])[:100])
+                image_dir = image_dir.joinpath(slugify(args.prompt[j % len(args.prompt)])[:100])
             image_dir.mkdir(parents=True, exist_ok=True)
-            image.save(image_dir.joinpath(f"{seed}_{j}.png"))
-            image.save(image_dir.joinpath(f"{seed}_{j}.jpg"), quality=85)
+            image.save(image_dir.joinpath(f"{seed}_{j // len(args.prompt)}.png"))
+            image.save(image_dir.joinpath(f"{seed}_{j // len(args.prompt)}.jpg"), quality=85)
 
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
diff --git a/models/clip/embeddings.py b/models/clip/embeddings.py
index 9c3a56b..1280ebd 100644
--- a/models/clip/embeddings.py
+++ b/models/clip/embeddings.py
@@ -72,7 +72,9 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
 
         self.temp_token_ids = torch.cat([self.temp_token_ids, token_ids])
         self.temp_token_embedding.weight.data[token_ids] = initializer.to(
-            dtype=self.temp_token_embedding.weight.dtype)
+            device=self.temp_token_embedding.weight.device,
+            dtype=self.temp_token_embedding.weight.dtype,
+        )
 
     def load_embed(self, input_ids: list[int], filename: Path):
         with safe_open(filename, framework="pt", device="cpu") as file:
diff --git a/models/clip/tokenizer.py b/models/clip/tokenizer.py
index 11a3df0..4e97ab5 100644
--- a/models/clip/tokenizer.py
+++ b/models/clip/tokenizer.py
@@ -48,7 +48,7 @@ def shuffle_none(tokens: list[int]):
 
 
 def shuffle_auto(tokens: list[int]):
-    if len(tokens) >= 4:
+    if len(tokens) >= 5:
         return shuffle_between(tokens)
     if len(tokens) >= 3:
         return shuffle_trailing(tokens)
diff --git a/train_ti.py b/train_ti.py
index 1b60f64..8ada98c 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -889,7 +889,7 @@ def main():
             on_train=on_train,
             on_eval=on_eval,
         )
-        lr_finder.run(end_lr=1e3)
+        lr_finder.run(num_epochs=200, end_lr=1e3)
 
         plt.savefig(basepath.joinpath("lr.png"), dpi=300)
         plt.close()
diff --git a/training/lr.py b/training/lr.py
index c8dc040..3cdf994 100644
--- a/training/lr.py
+++ b/training/lr.py
@@ -12,7 +12,7 @@ from tqdm.auto import tqdm
 from training.util import AverageMeter
 
 
-def noop():
+def noop(*args, **kwards):
     pass
 
 
@@ -26,6 +26,7 @@ class LRFinder():
         val_dataloader,
         loss_fn: Union[Callable[[int, Any], Tuple[Any, Any, int]], Callable[[int, Any, bool], Tuple[Any, Any, int]]],
         on_train: Callable[[], None] = noop,
+        on_clip: Callable[[], None] = noop,
         on_eval: Callable[[], None] = noop
     ):
         self.accelerator = accelerator
@@ -35,6 +36,7 @@ class LRFinder():
         self.val_dataloader = val_dataloader
         self.loss_fn = loss_fn
         self.on_train = on_train
+        self.on_clip = on_clip
         self.on_eval = on_eval
 
         # self.model_state = copy.deepcopy(model.state_dict())
@@ -93,6 +95,9 @@ class LRFinder():
 
                     self.accelerator.backward(loss)
 
+                    if self.accelerator.sync_gradients:
+                        self.on_clip()
+
                     self.optimizer.step()
                     lr_scheduler.step()
                     self.optimizer.zero_grad(set_to_none=True)