jpeg

tengshaofeng · Feb 2, 2023 · 960ed6a · 960ed6a
1 parent 1a98d0d
commit 960ed6a
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/lora_diffusion/dataset.py b/lora_diffusion/dataset.py
@@ -139,9 +139,11 @@ def __init__(
             self.captions = open(f"{instance_data_root}/caption.txt").readlines()
 
         else:
-            possibily_src_images = glob.glob(
-                str(instance_data_root) + "/*.jpg"
-            ) + glob.glob(str(instance_data_root) + "/*.png")
+            possibily_src_images = (
+                glob.glob(str(instance_data_root) + "/*.jpg")
+                + glob.glob(str(instance_data_root) + "/*.png")
+                + glob.glob(str(instance_data_root) + "/*.jpeg")
+            )
             possibily_src_images = (
                 set(possibily_src_images)
                 - set(glob.glob(str(instance_data_root) + "/*mask.png"))