From 82aac2646945c8b483b9272d712eca16fe60b7aa Mon Sep 17 00:00:00 2001 From: rvhfxb <116002789+rvhfxb@users.noreply.github.com> Date: Wed, 8 Mar 2023 22:42:41 +0900 Subject: [PATCH] Update train_util.py --- library/train_util.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/library/train_util.py b/library/train_util.py index 75176e13..f1060cbb 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -906,10 +906,14 @@ class FineTuningDataset(BaseDataset): if os.path.exists(image_key): abs_path = image_key else: - # わりといい加減だがいい方法が思いつかん - abs_path = glob_images(subset.image_dir, image_key) - assert len(abs_path) >= 1, f"no image / 画像がありません: {image_key}" - abs_path = abs_path[0] + npz_path = os.path.join(glob.escape(train_data_dir), image_key + ".npz") + if os.path.exists(npz_path): + abs_path = npz_path + else: + # わりといい加減だがいい方法が思いつかん + abs_path = glob_images(subset.image_dir, image_key) + assert len(abs_path) >= 1, f"no image / 画像がありません: {image_key}" + abs_path = abs_path[0] caption = img_md.get('caption') tags = img_md.get('tags')