From 0c7baea88cfa98c5fab2898551c426f2d4fac4c6 Mon Sep 17 00:00:00 2001
From: feffy380 <114889020+feffy380@users.noreply.github.com>
Date: Sat, 23 Mar 2024 17:28:02 +0100
Subject: [PATCH] register reg images with correct subset
---
library/train_util.py | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/library/train_util.py b/library/train_util.py
index d076cf847..b69fb0950 100644
--- a/library/train_util.py
+++ b/library/train_util.py
@@ -1554,7 +1554,7 @@ def load_dreambooth_dir(subset: DreamBoothSubset):
for img_path, caption in zip(img_paths, captions):
info = ImageInfo(img_path, subset.num_repeats, caption, subset.is_reg, img_path)
if subset.is_reg:
- reg_infos.append(info)
+ reg_infos.append((info, subset))
else:
self.register_image(info, subset)
@@ -1575,7 +1575,7 @@ def load_dreambooth_dir(subset: DreamBoothSubset):
n = 0
first_loop = True
while n < num_train_images:
- for info in reg_infos:
+ for info, subset in reg_infos:
if first_loop:
self.register_image(info, subset)
n += info.num_repeats