Skip to content

Commit 0dba0fe

Browse files
authored
Merge pull request kohya-ss#1205 from feffy380/patch-1
register reg images with correct subset
2 parents bee5e84 + c9464b8 commit 0dba0fe

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

library/train_util.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1554,7 +1554,7 @@ def load_dreambooth_dir(subset: DreamBoothSubset):
15541554
for img_path, caption in zip(img_paths, captions):
15551555
info = ImageInfo(img_path, subset.num_repeats, caption, subset.is_reg, img_path)
15561556
if subset.is_reg:
1557-
reg_infos.append(info)
1557+
reg_infos.append((info, subset))
15581558
else:
15591559
self.register_image(info, subset)
15601560

@@ -1575,7 +1575,7 @@ def load_dreambooth_dir(subset: DreamBoothSubset):
15751575
n = 0
15761576
first_loop = True
15771577
while n < num_train_images:
1578-
for info in reg_infos:
1578+
for info, subset in reg_infos:
15791579
if first_loop:
15801580
self.register_image(info, subset)
15811581
n += info.num_repeats

0 commit comments

Comments
 (0)