Merge pull request #454 from kvzn/master

Fixed the bug that it did not support capital image extensions like .JPG, .JPEG
bmaltais · Mar 26, 2023 · ad018dd · ad018dd
2 parents b6332ce + 3f86e8f
commit ad018dd
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 17 deletions.
diff --git a/dreambooth_gui.py b/dreambooth_gui.py
@@ -358,11 +358,10 @@ def train_model(
         num_images = len(
             [
                 f
-                for f in os.listdir(os.path.join(train_data_dir, folder))
-                if f.endswith('.jpg')
-                or f.endswith('.jpeg')
-                or f.endswith('.png')
-                or f.endswith('.webp')
+                for f, lower_f in (
+                    (file, file.lower()) for file in os.listdir(os.path.join(train_data_dir, folder))
+                )
+                if lower_f.endswith(('.jpg', '.jpeg', '.png', '.webp'))
             ]
         )
 

diff --git a/finetune_gui.py b/finetune_gui.py
@@ -368,8 +368,10 @@ def train_model(
     image_num = len(
         [
             f
-            for f in os.listdir(image_folder)
-            if f.endswith('.jpg') or f.endswith('.png') or f.endswith('.webp')
+            for f, lower_f in (
+                (file, file.lower()) for file in os.listdir(image_folder)
+            )
+            if lower_f.endswith(('.jpg', '.jpeg', '.png', '.webp'))
         ]
     )
     print(f'image_num = {image_num}')

diff --git a/lora_gui.py b/lora_gui.py
@@ -419,14 +419,14 @@ def train_model(
         num_images = len(
             [
                 f
-                for f in os.listdir(os.path.join(train_data_dir, folder))
-                if f.endswith('.jpg')
-                or f.endswith('.jpeg')
-                or f.endswith('.png')
-                or f.endswith('.webp')
+                for f, lower_f in (
+                    (file, file.lower()) for file in os.listdir(os.path.join(train_data_dir, folder))
+                )
+                if lower_f.endswith(('.jpg', '.jpeg', '.png', '.webp'))
             ]
         )
 
+
         print(f'Folder {folder}: {num_images} images found')
 
         # Calculate the total number of steps for this folder

diff --git a/textual_inversion_gui.py b/textual_inversion_gui.py
@@ -375,11 +375,10 @@ def train_model(
         num_images = len(
             [
                 f
-                for f in os.listdir(os.path.join(train_data_dir, folder))
-                if f.endswith('.jpg')
-                or f.endswith('.jpeg')
-                or f.endswith('.png')
-                or f.endswith('.webp')
+                for f, lower_f in (
+                    (file, file.lower()) for file in os.listdir(os.path.join(train_data_dir, folder))
+                )
+                if lower_f.endswith(('.jpg', '.jpeg', '.png', '.webp'))
             ]
         )