diff --git a/ldm/data/data_loader.py b/ldm/data/data_loader.py index f0abe20..c11b413 100644 --- a/ldm/data/data_loader.py +++ b/ldm/data/data_loader.py @@ -116,19 +116,32 @@ class DataLoaderMultiAspect(): @staticmethod def __recurse_data_root(self, recurse_root): + multiply = 1 + multiply_path = os.path.join(recurse_root, "multiply.txt") + if os.path.exists(multiply_path): + try: + with open(multiply_path, encoding='utf-8', mode='r') as f: + multiply = int(float(f.read().strip())) + print(f" * DLMA multiply.txt in {recurse_root} set to {multiply}") + except: + print(f" *** Error reading multiply.txt in {recurse_root}, defaulting to 1") + pass + for f in os.listdir(recurse_root): current = os.path.join(recurse_root, f) if os.path.isfile(current): ext = os.path.splitext(f)[1] - if ext in ['.jpg', '.jpeg', '.png', '.bmp', '.webp']: - self.image_paths.append(current) + if ext in ['.jpg', '.jpeg', '.png', '.bmp', '.webp', '.jfif']: + # add image multiplyrepeats number of times + for _ in range(multiply): + self.image_paths.append(current) - sub_dirs = [] + sub_dirs = [] for d in os.listdir(recurse_root): current = os.path.join(recurse_root, d) - if os.path.isdir(current): + if os.path.isdir(current): sub_dirs.append(current) for dir in sub_dirs: