Merge pull request kohya-ss#129 from p1atdev/main

Add support for .jpeg images in glob
Prsaro · Jan 31, 2023 · 8c3a52e · 8c3a52e
2 parents 9abbee0 + 86f4e20
commit 8c3a52e
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 5 deletions.
diff --git 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
@@ -10,7 +10,7 @@
 
 
 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
@@ -9,15 +9,16 @@
 
 
 def main(args):
-
   image_paths = None
   train_data_dir_path = Path(args.train_data_dir)
   if args.recursive:
     image_paths = list(train_data_dir_path.rglob('*.jpg')) + \
+                  list(train_data_dir_path.rglob('*.jpeg')) + \
                   list(train_data_dir_path.rglob('*.png')) + \
                   list(train_data_dir_path.rglob('*.webp')) 
   else:
     image_paths = list(train_data_dir_path.glob('*.jpg')) + \
+                  list(train_data_dir_path.glob('*.jpeg')) + \
                   list(train_data_dir_path.glob('*.png')) + \
                   list(train_data_dir_path.glob('*.webp')) 
 

diff --git 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
@@ -37,6 +37,7 @@ def main(args):
 
   # 画像を読み込む
   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")) + \
       glob.glob(os.path.join(args.train_data_dir, "*.bmp"))