Merge pull request #454 from kvzn/master

Fixed the bug that it did not support capital image extensions like .JPG, .JPEG

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))
+                for f, lower_f in (
-                if f.endswith('.jpg')
+                    (file, file.lower()) for file in os.listdir(os.path.join(train_data_dir, folder))
-                or f.endswith('.jpeg')
+                )
-                or f.endswith('.png')
+                if lower_f.endswith(('.jpg', '.jpeg', '.png', '.webp'))
-                or f.endswith('.webp')
             ]
         )
         

finetune_gui.py

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

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))
+                for f, lower_f in (
-                if f.endswith('.jpg')
+                    (file, file.lower()) for file in os.listdir(os.path.join(train_data_dir, folder))
-                or f.endswith('.jpeg')
+                )
-                or f.endswith('.png')
+                if lower_f.endswith(('.jpg', '.jpeg', '.png', '.webp'))
-                or f.endswith('.webp')
             ]
         )
 
+
         print(f'Folder {folder}: {num_images} images found')
 
         # Calculate the total number of steps for this folder

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))
+                for f, lower_f in (
-                if f.endswith('.jpg')
+                    (file, file.lower()) for file in os.listdir(os.path.join(train_data_dir, folder))
-                or f.endswith('.jpeg')
+                )
-                or f.endswith('.png')
+                if lower_f.endswith(('.jpg', '.jpeg', '.png', '.webp'))
-                or f.endswith('.webp')
             ]
         )