From 3f86e8f7dff9e92e31c6f744ba4fb285049c2f14 Mon Sep 17 00:00:00 2001
From: KVZN <realkvz@gmail.com>
Date: Sun, 26 Mar 2023 17:58:41 +0800
Subject: [PATCH] Fixed the bug that it did not support capital image
 extensions like .JPG, .JPEG

---
 dreambooth_gui.py        |  9 ++++-----
 finetune_gui.py          |  6 ++++--
 lora_gui.py              | 10 +++++-----
 textual_inversion_gui.py |  9 ++++-----
 4 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/dreambooth_gui.py b/dreambooth_gui.py
index a72acff..fc85f1b 100644
--- 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
index f310521..2f12f7d 100644
--- 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
index 13223fe..91b4cc4 100644
--- 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
index 7b8c19c..d7ed087 100644
--- 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'))
             ]
         )