From c20745b6e8344b0db73536c4f9b30382eb0fc263 Mon Sep 17 00:00:00 2001 From: unknown Date: Sun, 29 Jan 2023 22:30:45 +0900 Subject: [PATCH] fix: #53 --- finetune/make_captions.py | 2 +- finetune/merge_captions_to_metadata.py | 5 +++-- finetune/merge_dd_tags_to_metadata.py | 2 +- finetune/prepare_buckets_latents.py | 2 +- finetune/tag_images_by_wd14_tagger.py | 2 +- 5 files changed, 7 insertions(+), 6 deletions(-) diff --git a/finetune/make_captions.py b/finetune/make_captions.py index b02420bd..495450aa 100644 --- a/finetune/make_captions.py +++ b/finetune/make_captions.py @@ -31,7 +31,7 @@ def main(args): os.chdir('finetune') print(f"load images from {args.train_data_dir}") - image_paths = glob.glob(os.path.join(args.train_data_dir, "*.jpg")) + \ + image_paths = glob.glob(os.path.join(args.train_data_dir, "*.jpg")) + glob.glob(os.path.join(args.train_data_dir, "*.jpeg")) + \ glob.glob(os.path.join(args.train_data_dir, "*.png")) + glob.glob(os.path.join(args.train_data_dir, "*.webp")) print(f"found {len(image_paths)} images.") diff --git a/finetune/merge_captions_to_metadata.py b/finetune/merge_captions_to_metadata.py index 2da6356f..703f4f9d 100644 --- a/finetune/merge_captions_to_metadata.py +++ b/finetune/merge_captions_to_metadata.py @@ -10,7 +10,7 @@ from tqdm import tqdm def main(args): - image_paths = glob.glob(os.path.join(args.train_data_dir, "*.jpg")) + \ + image_paths = glob.glob(os.path.join(args.train_data_dir, "*.jpg")) + glob.glob(os.path.join(args.train_data_dir, "*.jpeg")) + \ glob.glob(os.path.join(args.train_data_dir, "*.png")) + glob.glob(os.path.join(args.train_data_dir, "*.webp")) print(f"found {len(image_paths)} images.") @@ -30,7 +30,8 @@ def main(args): for image_path in tqdm(image_paths): caption_path = os.path.splitext(image_path)[0] + args.caption_extension with open(caption_path, "rt", encoding='utf-8') as f: - caption = f.readlines()[0].strip() + lines = f.readlines() + caption = lines[0].strip() if len(lines) > 0 else "" image_key = image_path if args.full_path else os.path.splitext(os.path.basename(image_path))[0] if image_key not in metadata: diff --git a/finetune/merge_dd_tags_to_metadata.py b/finetune/merge_dd_tags_to_metadata.py index 8101ecd3..896c99ae 100644 --- a/finetune/merge_dd_tags_to_metadata.py +++ b/finetune/merge_dd_tags_to_metadata.py @@ -10,7 +10,7 @@ from tqdm import tqdm def main(args): - image_paths = glob.glob(os.path.join(args.train_data_dir, "*.jpg")) + \ + image_paths = glob.glob(os.path.join(args.train_data_dir, "*.jpg")) + glob.glob(os.path.join(args.train_data_dir, "*.jpeg")) + \ glob.glob(os.path.join(args.train_data_dir, "*.png")) + glob.glob(os.path.join(args.train_data_dir, "*.webp")) print(f"found {len(image_paths)} images.") diff --git a/finetune/prepare_buckets_latents.py b/finetune/prepare_buckets_latents.py index 00f847a1..87236c43 100644 --- a/finetune/prepare_buckets_latents.py +++ b/finetune/prepare_buckets_latents.py @@ -36,7 +36,7 @@ def get_latents(vae, images, weight_dtype): def main(args): - image_paths = glob.glob(os.path.join(args.train_data_dir, "*.jpg")) + \ + image_paths = glob.glob(os.path.join(args.train_data_dir, "*.jpg")) + glob.glob(os.path.join(args.train_data_dir, "*.jpeg")) + \ glob.glob(os.path.join(args.train_data_dir, "*.png")) + glob.glob(os.path.join(args.train_data_dir, "*.webp")) print(f"found {len(image_paths)} images.") diff --git a/finetune/tag_images_by_wd14_tagger.py b/finetune/tag_images_by_wd14_tagger.py index c5767894..60a9a890 100644 --- a/finetune/tag_images_by_wd14_tagger.py +++ b/finetune/tag_images_by_wd14_tagger.py @@ -36,7 +36,7 @@ def main(args): args.model_dir, SUB_DIR), force_download=True, force_filename=file) # 画像を読み込む - image_paths = glob.glob(os.path.join(args.train_data_dir, "*.jpg")) + \ + image_paths = glob.glob(os.path.join(args.train_data_dir, "*.jpg")) + glob.glob(os.path.join(args.train_data_dir, "*.jpeg")) + \ glob.glob(os.path.join(args.train_data_dir, "*.png")) + glob.glob(os.path.join(args.train_data_dir, "*.webp")) print(f"found {len(image_paths)} images.")