diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py
index 0016134..c1db6d6 100644
--- a/jobs/process/BaseSDTrainProcess.py
+++ b/jobs/process/BaseSDTrainProcess.py
@@ -2244,10 +2244,15 @@ class BaseSDTrainProcess(BaseTrainProcess):
ds.cache_tensors_to_disk = True
# load datasets if passed in the root process
if self.datasets is not None:
- self.data_loader = get_dataloader_from_datasets(self.datasets, self.train_config.batch_size, self.sd)
+ self.data_loader = get_dataloader_from_datasets(
+ self.datasets, self.train_config.batch_size, self.sd,
+ pool_datasets=self.train_config.pool_datasets,
+ )
if self.datasets_reg is not None:
- self.data_loader_reg = get_dataloader_from_datasets(self.datasets_reg, self.train_config.batch_size,
- self.sd)
+ self.data_loader_reg = get_dataloader_from_datasets(
+ self.datasets_reg, self.train_config.batch_size, self.sd,
+ pool_datasets=self.train_config.pool_datasets,
+ )
flush()
self.last_save_step = self.step_num
diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py
index 70dd08a..4f05d74 100644
--- a/toolkit/config_modules.py
+++ b/toolkit/config_modules.py
@@ -568,6 +568,9 @@ class TrainConfig:
self.unload_text_encoder = kwargs.get('unload_text_encoder', False)
# will toggle all datasets to cache text embeddings
self.cache_text_embeddings: bool = kwargs.get('cache_text_embeddings', False)
+ # merge all datasets into one file list so buckets/batches mix across folders;
+ # reg datasets pool separately from non-reg ones
+ self.pool_datasets: bool = kwargs.get('pool_datasets', False)
# for swapping which parameters are trained during training
self.do_paramiter_swapping = kwargs.get('do_paramiter_swapping', False)
# 0.1 is 10% of the parameters active at a time lower is less vram, higher is more
@@ -1109,10 +1112,28 @@ class DatasetConfig:
def preprocess_dataset_raw_config(raw_config: List[dict]) -> List[dict]:
"""
- This just splits up the datasets by resolutions so you dont have to do it manually
+ This just splits up the datasets by paths and resolutions so you dont have to do it manually
:param raw_config:
:return:
"""
+ # split up datasets by folder_path / dataset_path lists: one dataset per path,
+ # all other settings shared (pair with train.pool_datasets to mix them)
+ path_split_config = []
+ for dataset in raw_config:
+ split = False
+ for path_key in ['folder_path', 'dataset_path']:
+ paths = dataset.get(path_key, None)
+ if isinstance(paths, (list, tuple)):
+ for path in paths:
+ dataset_copy = dataset.copy()
+ dataset_copy[path_key] = path
+ path_split_config.append(dataset_copy)
+ split = True
+ break
+ if not split:
+ path_split_config.append(dataset)
+ raw_config = path_split_config
+
# split up datasets by resolutions
new_config = []
for dataset in raw_config:
diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py
index a39f586..eddb68a 100644
--- a/toolkit/data_loader.py
+++ b/toolkit/data_loader.py
@@ -390,11 +390,17 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
dataset_config: 'DatasetConfig',
batch_size=1,
sd: 'StableDiffusion' = None,
+ source_configs: List['DatasetConfig'] = None,
):
+ # a pooled dataset scans several configs into ONE file_list so buckets and
+ # batches mix across folders. Single-source is the same code with one config.
+ self.source_configs: List['DatasetConfig'] = list(source_configs) if source_configs else [dataset_config]
+ self.is_pooled = len(self.source_configs) > 1
self.dataset_config = dataset_config
# update bucket divisibility
- self.dataset_config.bucket_tolerance = sd.get_bucket_divisibility()
- self.is_video = dataset_config.num_frames > 1 or dataset_config.auto_frame_count
+ for config in self.source_configs:
+ config.bucket_tolerance = sd.get_bucket_divisibility()
+ self.is_video = any(c.num_frames > 1 or c.auto_frame_count for c in self.source_configs)
self.is_audio_model = hasattr(sd, 'is_audio_model') and sd.is_audio_model if sd is not None else False
# text-generating multimodal models take audio, image and video files in one dataset
self.is_multimodal_llm = getattr(sd, 'is_multimodal_llm', False) if sd is not None else False
@@ -403,12 +409,16 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
self.dataset_path = dataset_config.dataset_path
if self.dataset_path is None:
self.dataset_path = folder_path
+ if self.is_pooled:
+ self.dataset_path = f"pool of {len(self.source_configs)} datasets"
- self.is_caching_latents = dataset_config.cache_latents or dataset_config.cache_latents_to_disk
- self.is_caching_latents_to_memory = dataset_config.cache_latents
- self.is_caching_latents_to_disk = dataset_config.cache_latents_to_disk
- self.is_caching_clip_vision_to_disk = dataset_config.cache_clip_vision_to_disk
- self.is_generating_controls = len(dataset_config.controls) > 0
+ self.is_caching_latents = any(c.cache_latents or c.cache_latents_to_disk for c in self.source_configs)
+ self.is_caching_latents_to_memory = any(c.cache_latents for c in self.source_configs)
+ self.is_caching_latents_to_disk = any(c.cache_latents_to_disk for c in self.source_configs)
+ self.is_caching_clip_vision_to_disk = any(c.cache_clip_vision_to_disk for c in self.source_configs)
+ self.is_generating_controls = any(len(c.controls) > 0 for c in self.source_configs)
+ # mixin init read this off dataset_config only
+ self.is_caching_text_embeddings = any(c.cache_text_embeddings for c in self.source_configs)
self.epoch_num = 0
self.sd = sd
@@ -428,9 +438,45 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
self.resolution = dataset_config.resolution
self.caption_dict = None
self.file_list: List['FileItemDTO'] = []
+ self.transform = self._build_transform(dataset_config)
+
+ for config in self.source_configs:
+ self._scan_source(config)
+
+ if self.is_pooled:
+ print_acc(f"Pooled {len(self.source_configs)} datasets into {len(self.file_list)} items")
+
+ self.setup_epoch()
+
+ def _build_transform(self, dataset_config: 'DatasetConfig'):
+ if dataset_config.standardize_images:
+ if self.sd.is_xl or self.sd.is_vega or self.sd.is_ssd:
+ NormalizeMethod = NormalizeSDXLTransform
+ else:
+ NormalizeMethod = NormalizeSD15Transform
+
+ return transforms.Compose([
+ transforms.ToTensor(),
+ RescaleTransform(),
+ NormalizeMethod(),
+ ])
+ else:
+ return transforms.Compose([
+ transforms.ToTensor(),
+ RescaleTransform(),
+ ])
+
+ def _scan_source(self, dataset_config: 'DatasetConfig'):
+ # scans one config (folder or json) into FileItemDTOs and appends them to file_list.
+ # size database, captions and controls stay per source folder.
+ dataset_path = dataset_config.dataset_path
+ if dataset_path is None:
+ dataset_path = dataset_config.folder_path
+ transform = self.transform if dataset_config is self.dataset_config else self._build_transform(dataset_config)
+ caption_dict = None
# check if dataset_path is a folder or json
- if os.path.isdir(self.dataset_path):
+ if os.path.isdir(dataset_path):
extensions = image_extensions
if self.is_audio_model:
# only look for audio files
@@ -443,48 +489,34 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
extensions = video_extensions + image_extensions
# prune hidden dirs (.thumbs, .tmp) so their contents never train
file_list = []
- for root, dirs, files in os.walk(self.dataset_path):
+ for root, dirs, files in os.walk(dataset_path):
dirs[:] = [d for d in dirs if not d.startswith('.')]
file_list.extend(os.path.join(root, file) for file in files if file.lower().endswith(tuple(extensions)) and not file.startswith('.'))
else:
# assume json
- with open(self.dataset_path, 'r') as f:
- self.caption_dict = json.load(f)
+ with open(dataset_path, 'r') as f:
+ caption_dict = json.load(f)
# keys are file paths
- file_list = list(self.caption_dict.keys())
-
+ file_list = list(caption_dict.keys())
+ if self.caption_dict is None:
+ self.caption_dict = {}
+ self.caption_dict.update(caption_dict)
+
# remove items in the _controls_ folder
file_list = [x for x in file_list if not os.path.basename(os.path.dirname(x)) == "_controls"]
- if self.dataset_config.num_repeats > 1:
+ if dataset_config.num_repeats > 1:
# repeat the list
- file_list = file_list * self.dataset_config.num_repeats
-
- if self.dataset_config.standardize_images:
- if self.sd.is_xl or self.sd.is_vega or self.sd.is_ssd:
- NormalizeMethod = NormalizeSDXLTransform
- else:
- NormalizeMethod = NormalizeSD15Transform
-
- self.transform = transforms.Compose([
- transforms.ToTensor(),
- RescaleTransform(),
- NormalizeMethod(),
- ])
- else:
- self.transform = transforms.Compose([
- transforms.ToTensor(),
- RescaleTransform(),
- ])
+ file_list = file_list * dataset_config.num_repeats
# this might take a while
- print_acc(f"Dataset: {self.dataset_path}")
+ print_acc(f"Dataset: {dataset_path}")
if self.is_video:
print_acc(f" - Preprocessing video dimensions")
else:
print_acc(f" - Preprocessing image dimensions")
- dataset_folder = self.dataset_path
- if not os.path.isdir(self.dataset_path):
+ dataset_folder = dataset_path
+ if not os.path.isdir(dataset_path):
dataset_folder = os.path.dirname(dataset_folder)
dataset_size_file = os.path.join(dataset_folder, '.aitk_size.json')
@@ -492,20 +524,21 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
if os.path.exists(dataset_size_file):
try:
with open(dataset_size_file, 'r') as f:
- self.size_database = json.load(f)
+ size_database = json.load(f)
- if "__version__" not in self.size_database or self.size_database["__version__"] != dataloader_version:
+ if "__version__" not in size_database or size_database["__version__"] != dataloader_version:
print_acc("Upgrading size database to new version")
# old version, delete and recreate
- self.size_database = {}
+ size_database = {}
except Exception as e:
print_acc(f"Error loading size database: {dataset_size_file}")
print_acc(e)
- self.size_database = {}
+ size_database = {}
else:
- self.size_database = {}
+ size_database = {}
- self.size_database["__version__"] = dataloader_version
+ size_database["__version__"] = dataloader_version
+ self.size_database = size_database
# cache keys come from the model so a model can invalidate them on its own kwargs
latent_space_version = self.sd.get_latent_space_version() if self.sd is not None else "sd1"
@@ -518,6 +551,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
if hasattr(self.sd.unet, 'config') and hasattr(self.sd.unet.config, 'temporal_compression_ratio'):
temporal_compression = self.sd.unet.config.temporal_compression_ratio
+ source_items: List['FileItemDTO'] = []
bad_count = 0
for file in tqdm(file_list):
try:
@@ -526,8 +560,8 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
path=file,
is_audio_model=self.is_audio_model or (self.is_multimodal_llm and file.lower().endswith(tuple(audio_extensions))),
dataset_config=dataset_config,
- dataloader_transforms=self.transform,
- size_database=self.size_database,
+ dataloader_transforms=transform,
+ size_database=size_database,
dataset_root=dataset_folder,
encode_control_in_text_embeddings=self.sd.encode_control_in_text_embeddings if self.sd else False,
encode_first_frame_in_text_embeddings=getattr(self.sd, 'encode_first_frame_in_text_embeddings', False) if self.sd else False,
@@ -539,7 +573,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
temporal_compression=temporal_compression,
sample_rate=self.sd.sample_rate if (self.is_audio_model or self.is_multimodal_llm) and self.sd is not None else 48000,
)
- self.file_list.append(file_item)
+ source_items.append(file_item)
except Exception as e:
print_acc(traceback.format_exc())
if self.is_video:
@@ -551,47 +585,47 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
# save the size database
with open(dataset_size_file, 'w') as f:
- json.dump(self.size_database, f)
+ json.dump(size_database, f)
if self.is_video:
- num_videos = len([x for x in self.file_list if x.is_video])
- num_images = len(self.file_list) - num_videos
+ num_videos = len([x for x in source_items if x.is_video])
+ num_images = len(source_items) - num_videos
if num_images > 0:
print_acc(f" - Found {num_videos} videos and {num_images} images")
else:
print_acc(f" - Found {num_videos} videos")
- assert len(self.file_list) > 0, f"no videos found in {self.dataset_path}"
+ assert len(source_items) > 0, f"no videos found in {dataset_path}"
else:
- print_acc(f" - Found {len(self.file_list)} images")
- assert len(self.file_list) > 0, f"no images found in {self.dataset_path}"
+ print_acc(f" - Found {len(source_items)} images")
+ assert len(source_items) > 0, f"no images found in {dataset_path}"
# handle x axis flips
- if self.dataset_config.flip_x:
+ if dataset_config.flip_x:
print_acc(" - adding x axis flips")
- current_file_list = [x for x in self.file_list]
+ current_file_list = [x for x in source_items]
for file_item in current_file_list:
# create a copy that is flipped on the x axis
new_file_item = copy.deepcopy(file_item)
new_file_item.flip_x = True
- self.file_list.append(new_file_item)
+ source_items.append(new_file_item)
# handle y axis flips
- if self.dataset_config.flip_y:
+ if dataset_config.flip_y:
print_acc(" - adding y axis flips")
- current_file_list = [x for x in self.file_list]
+ current_file_list = [x for x in source_items]
for file_item in current_file_list:
# create a copy that is flipped on the y axis
new_file_item = copy.deepcopy(file_item)
new_file_item.flip_y = True
- self.file_list.append(new_file_item)
+ source_items.append(new_file_item)
- if self.dataset_config.flip_x or self.dataset_config.flip_y:
+ if dataset_config.flip_x or dataset_config.flip_y:
if self.is_video:
- print_acc(f" - Found {len(self.file_list)} videos after adding flips")
+ print_acc(f" - Found {len(source_items)} videos after adding flips")
else:
- print_acc(f" - Found {len(self.file_list)} images after adding flips")
+ print_acc(f" - Found {len(source_items)} images after adding flips")
- self.setup_epoch()
+ self.file_list.extend(source_items)
def setup_epoch(self):
if self.epoch_num == 0:
@@ -639,7 +673,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
def _get_single_item(self, index, _attempts=0) -> 'FileItemDTO':
file_item: 'FileItemDTO' = copy.deepcopy(self.file_list[index])
try:
- file_item.load_and_process_image(self.transform)
+ file_item.load_and_process_image(file_item.dataloader_transforms or self.transform)
except Exception as e:
print(f"Error loading image, skipping and loading a different one: {file_item.path} ({e})")
if _attempts >= 10:
@@ -681,6 +715,7 @@ def get_dataloader_from_datasets(
dataset_options,
batch_size=1,
sd: 'StableDiffusion' = None,
+ pool_datasets: bool = False,
) -> DataLoader:
if dataset_options is None or len(dataset_options) == 0:
return None
@@ -701,18 +736,34 @@ def get_dataloader_from_datasets(
dataset_config_list.append(DatasetConfig(**x))
for config in dataset_config_list:
+ if config.type != 'image':
+ raise ValueError(f"invalid dataset type: {config.type}")
+ if config.buckets:
+ has_buckets = True
+ if config.cache_latents or config.cache_latents_to_disk:
+ is_caching_latents = True
- if config.type == 'image':
+ if pool_datasets and len(dataset_config_list) > 1:
+ # one dataset over every config: buckets and batches mix across folders.
+ # callers already split reg and non-reg into separate loaders.
+ batch_sizes = {c.batch_size if c.batch_size is not None else batch_size for c in dataset_config_list}
+ if len(batch_sizes) > 1:
+ raise ValueError(
+ f"pool_datasets requires every dataset to use the same batch_size, got {sorted(batch_sizes)}"
+ )
+ if len({c.buckets for c in dataset_config_list}) > 1:
+ raise ValueError("pool_datasets requires buckets to be enabled or disabled on every dataset")
+ datasets.append(AiToolkitDataset(
+ dataset_config_list[0],
+ batch_size=batch_sizes.pop(),
+ sd=sd,
+ source_configs=dataset_config_list,
+ ))
+ else:
+ for config in dataset_config_list:
# dataset level batch_size overrides the train config batch_size when set
dataset_batch_size = config.batch_size if config.batch_size is not None else batch_size
- dataset = AiToolkitDataset(config, batch_size=dataset_batch_size, sd=sd)
- datasets.append(dataset)
- if config.buckets:
- has_buckets = True
- if config.cache_latents or config.cache_latents_to_disk:
- is_caching_latents = True
- else:
- raise ValueError(f"invalid dataset type: {config.type}")
+ datasets.append(AiToolkitDataset(config, batch_size=dataset_batch_size, sd=sd))
concatenated_dataset = ConcatDataset(datasets)
diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py
index b4cd1d0..e10b455 100644
--- a/toolkit/dataloader_mixins.py
+++ b/toolkit/dataloader_mixins.py
@@ -192,6 +192,43 @@ class Bucket:
self.file_list_idx: List[int] = []
+def get_batch_signature(config: 'DatasetConfig') -> tuple:
+ # every dataset-level value the batch DTO / trainer assumes is constant across a batch
+ control_path = config.control_path
+ if control_path is None:
+ num_control_paths = 0
+ elif isinstance(control_path, list):
+ num_control_paths = len(control_path)
+ else:
+ num_control_paths = 1
+ return (
+ num_control_paths,
+ tuple(config.controls),
+ config.control_from_same_folder,
+ config.num_controls_from_same_folder,
+ config.inpaint_path is not None,
+ config.mask_path is not None,
+ config.alpha_mask,
+ config.unconditional_path is not None,
+ config.clip_image_path is not None,
+ config.clip_image_from_same_folder,
+ config.cache_latents,
+ config.cache_latents_to_disk,
+ config.cache_tensors_to_disk,
+ config.cache_text_embeddings,
+ config.load_image_when_caching_latents,
+ config.standardize_images,
+ config.num_frames,
+ config.auto_frame_count,
+ config.fps,
+ config.do_i2v,
+ config.do_audio,
+ config.guidance_type,
+ config.diff_output_preservation,
+ len(config.extra_values),
+ )
+
+
class BucketsMixin:
def __init__(self):
self.buckets: Dict[str, Bucket] = {}
@@ -227,23 +264,33 @@ class BucketsMixin:
self.buckets = {} # clear it
config: 'DatasetConfig' = self.dataset_config
- resolution = config.resolution
bucket_tolerance = config.bucket_tolerance
file_list: List['FileItemDTO'] = self.file_list
+ # pooled datasets suffix the key so only batch-compatible sources share a bucket
+ is_pooled = getattr(self, 'is_pooled', False)
+ signature_ids: Dict[tuple, int] = {}
# for file_item in enumerate(file_list):
for idx, file_item in enumerate(file_list):
file_item: 'FileItemDTO' = file_item
+ item_config: 'DatasetConfig' = file_item.dataset_config
+ key_suffix = ''
+ if is_pooled:
+ sig = get_batch_signature(item_config)
+ if sig not in signature_ids:
+ signature_ids[sig] = len(signature_ids)
+ key_suffix = f'#s{signature_ids[sig]}'
if file_item.is_audio_model:
- bucket_key = f"{file_item.width}ms"
+ bucket_key = f"{file_item.width}ms{key_suffix}"
if bucket_key not in self.buckets:
self.buckets[bucket_key] = Bucket(file_item.width, 1)
self.buckets[bucket_key].file_list_idx.append(idx)
continue
- width = int(file_item.width * file_item.dataset_config.scale)
- height = int(file_item.height * file_item.dataset_config.scale)
+ resolution = item_config.resolution
+ width = int(file_item.width * item_config.scale)
+ height = int(file_item.height * item_config.scale)
- if self.dataset_config.square_crop:
+ if item_config.square_crop:
# we scale first so smallest size matches resolution
scale_factor_x = resolution / width
scale_factor_y = resolution / height
@@ -282,7 +329,7 @@ class BucketsMixin:
new_width = bucket_resolution["width"]
new_height = bucket_resolution["height"]
- if self.dataset_config.random_crop:
+ if item_config.random_crop:
# random crop
crop_x = random.randint(0, file_item.scale_to_width - new_width)
crop_y = random.randint(0, file_item.scale_to_height - new_height)
@@ -301,6 +348,7 @@ class BucketsMixin:
if self.is_video:
# images (1 frame) and videos must not mix in a batch
bucket_key += f'x{file_item.num_frames}f'
+ bucket_key += key_suffix
if bucket_key not in self.buckets:
self.buckets[bucket_key] = Bucket(file_item.crop_width, file_item.crop_height)
self.buckets[bucket_key].file_list_idx.append(idx)
@@ -311,7 +359,16 @@ class BucketsMixin:
if not quiet:
print_acc(f'Bucket sizes for {self.dataset_path}:')
for key, bucket in self.buckets.items():
- print_acc(f'{key}: {len(bucket.file_list_idx)} files')
+ if is_pooled:
+ # per-source breakdown shows that mixing actually happened
+ per_source: Dict[str, int] = {}
+ for i in bucket.file_list_idx:
+ src = file_list[i].dataset_config.folder_path or file_list[i].dataset_config.dataset_path
+ per_source[src] = per_source.get(src, 0) + 1
+ breakdown = ', '.join(f'{os.path.basename(os.path.normpath(k))}={v}' for k, v in per_source.items())
+ print_acc(f'{key}: {len(bucket.file_list_idx)} files ({breakdown})')
+ else:
+ print_acc(f'{key}: {len(bucket.file_list_idx)} files')
print_acc(f'{len(self.buckets)} buckets made')
@@ -1995,7 +2052,7 @@ class LatentCachingMixin:
cached_state_dict = load_file(prep_latent_path, device='cpu') if to_memory else None
return prep_item, prep_latent_path, cached_state_dict, False
# not saved to disk, load the image/video/audio
- prep_item.load_and_process_image(self.transform, only_load_latents=True)
+ prep_item.load_and_process_image(prep_item.dataloader_transforms or self.transform, only_load_latents=True)
except Exception as e:
print_acc(f"Error processing image: {prep_item.path}")
print_acc(f"Error: {str(e)}")
@@ -2114,7 +2171,7 @@ class LatentCachingMixin:
frames = None
# add batch dimension
cache_uint8 = getattr(self.sd, 'cache_latents_as_uint8', False)
- if self.dataset_config.cache_tensors_to_disk:
+ if file_item.dataset_config.cache_tensors_to_disk:
if not file_item.is_audio_model:
tensor_uint8 = _latent_to_uint8(file_item.tensor).cpu()
if to_disk:
@@ -2155,7 +2212,7 @@ class LatentCachingMixin:
raise e
# do first frame
is_video = file_item.is_video
- if is_video and self.dataset_config.do_i2v:
+ if is_video and file_item.dataset_config.do_i2v:
frames = file_item.tensor.unsqueeze(0).to(device, dtype=dtype)
if len(frames.shape) == 4:
first_frames = frames
@@ -2442,20 +2499,21 @@ class TextEmbeddingCachingMixin:
text_embedding_path = file_item.get_text_embedding_path(recalculate=True)
# (path, caption) pairs to encode for this item
encode_targets = [(text_embedding_path, file_item.caption)]
- if self.dataset_config.diff_output_preservation:
+ item_config = file_item.dataset_config
+ if item_config.diff_output_preservation:
dop_path = file_item.get_dop_text_embedding_path(recalculate=True)
if dop_path != text_embedding_path:
# trigger word was in the caption, cache the DOP version too
encode_targets.append((dop_path, file_item.caption_dop))
# dropout embeds are encoded as plain text (no control images)
dropout_target_paths = set()
- if self.dataset_config.caption_dropout_rate > 0:
+ if item_config.caption_dropout_rate > 0:
blank_path = file_item.get_blank_text_embedding_path(recalculate=True)
if blank_path != text_embedding_path:
# cache the dropout caption embedding (blank, or trigger word only)
encode_targets.append((blank_path, file_item.get_dropout_caption()))
dropout_target_paths.add(blank_path)
- if self.dataset_config.diff_output_preservation:
+ if item_config.diff_output_preservation:
# cache the DOP version of the dropout caption (class only)
dop_blank_path = file_item.get_dop_blank_text_embedding_path(recalculate=True)
if dop_blank_path not in [t[0] for t in encode_targets] + [text_embedding_path]:
@@ -2500,7 +2558,7 @@ class TextEmbeddingCachingMixin:
# exactly like its latent rows (frame count / trim)
ctrl_img_list.extend(control_video_paths)
if len(control_video_paths) > 0:
- self.sd._ref_video_dataset_config = self.dataset_config
+ self.sd._ref_video_dataset_config = item_config
if len(ctrl_img_list) == 0:
ctrl_img = None
@@ -2528,13 +2586,13 @@ class TextEmbeddingCachingMixin:
del prompt_embeds
elif (
getattr(self.sd, 'encode_first_frame_in_text_embeddings', False)
- and self.dataset_config.do_i2v
+ and item_config.do_i2v
and file_item.is_video
):
# video item: encode the clip's FIRST FRAME into the text embeddings
# as a vision reference, matching sampling (where the ctrl image goes
# into the embeds and is held as the clean first frames)
- file_item.load_and_process_image(self.transform, only_load_latents=True)
+ file_item.load_and_process_image(file_item.dataloader_transforms or self.transform, only_load_latents=True)
frames = file_item.tensor # (T, C, H, W) or (C, H, W), in [-1, 1]
first = frames[0] if frames.dim() == 4 else frames
ctrl_img = (
@@ -2576,7 +2634,7 @@ class TextEmbeddingCachingMixin:
file_item.get_dopsd_text_embedding_path(recalculate=True),
file_item.caption_dopsd,
)]
- if self.dataset_config.caption_dropout_rate > 0:
+ if item_config.caption_dropout_rate > 0:
dopsd_blank_path = file_item.get_dopsd_blank_text_embedding_path(recalculate=True)
if dopsd_blank_path != dopsd_targets[0][0]:
dopsd_targets.append((dopsd_blank_path, file_item.get_dopsd_dropout_caption()))
@@ -2588,10 +2646,10 @@ class TextEmbeddingCachingMixin:
if file_item.is_video:
# own path rides through the video-ref presentation
ctrl_img = [file_item.path]
- self.sd._ref_video_dataset_config = self.dataset_config
+ self.sd._ref_video_dataset_config = item_config
else:
# own bucketed pixels as the reference image
- file_item.load_and_process_image(self.transform, only_load_latents=True)
+ file_item.load_and_process_image(file_item.dataloader_transforms or self.transform, only_load_latents=True)
img = file_item.tensor # (C, H, W) in [-1, 1]
ctrl_img = [
((img + 1.0) / 2.0)
@@ -2670,7 +2728,11 @@ class CLIPCachingMixin:
# cache unconditionals
print_acc(f" - Caching {self.clip_vision_num_unconditional_cache} unconditional clip vision to disk")
- clip_vision_cache_path = os.path.join(self.dataset_config.clip_image_path, '_clip_vision_cache')
+ clip_image_path = next(
+ (c.clip_image_path for c in getattr(self, 'source_configs', [self.dataset_config]) if c.clip_image_path),
+ self.dataset_config.clip_image_path,
+ )
+ clip_vision_cache_path = os.path.join(clip_image_path, '_clip_vision_cache')
unconditional_paths = []
@@ -2826,7 +2888,7 @@ class ControlCachingMixin:
# use tqdm to show progress
for file_item in tqdm(self.file_list, desc=f'Generating Controls'):
- for control_type in self.dataset_config.controls:
+ for control_type in file_item.dataset_config.controls:
# generates the control if it is not already there
control_path = self.control_generator.get_control_path(file_item.path, control_type)
if control_path is not None:
diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx
index a03486f..3fd9f9e 100644
--- a/ui/src/app/jobs/new/SimpleJob.tsx
+++ b/ui/src/app/jobs/new/SimpleJob.tsx
@@ -1157,6 +1157,17 @@ export default function SimpleJob({
)}
+
+ {
+ setJobConfig(value ? true : undefined, 'config.process[0].train.pool_datasets');
+ }}
+ />
+
)}
diff --git a/ui/src/docs.tsx b/ui/src/docs.tsx
index d484dbd..7533d5d 100644
--- a/ui/src/docs.tsx
+++ b/ui/src/docs.tsx
@@ -288,6 +288,30 @@ const docs: { [key: string]: ConfigDoc } = {
),
},
+ 'train.pool_datasets': {
+ title: 'Pool Datasets',
+ description: (
+ <>
+ Off by default. Only useful when you have numerous datasets.
+
+
+ Normally every dataset is loaded as its own unit with its own buckets, so each batch is built from a single
+ dataset and each dataset keeps its own loader threads. With many datasets this means batches never mix images
+ from different folders and the number of loader threads grows with the dataset count.
+
+
+ Pooling merges all of your datasets into one shared set of buckets. Images from different folders that land in
+ the same bucket are mixed together in the same batch, and only one set of loader threads is used for the whole
+ pool. Regularization datasets are pooled separately from your training datasets, never together.
+
+
+ Each dataset keeps its own settings (captions, control images, repeats, flips, scale, resolution). Datasets whose
+ settings would not collate into one batch (for example a different number of control images, video vs image,
+ or different caching) are automatically kept in separate buckets within the pool. All pooled datasets must use
+ the same batch size.
+
+ ),
+ },
'train.do_differential_guidance': {
title: 'Differential Guidance',
description: (
diff --git a/ui/src/types.ts b/ui/src/types.ts
index eb6b1bf..7faa5f5 100644
--- a/ui/src/types.ts
+++ b/ui/src/types.ts
@@ -187,6 +187,7 @@ export interface TrainConfig {
loss_type: 'mse' | 'mae' | 'wavelet' | 'stepped';
do_differential_guidance?: boolean;
differential_guidance_scale?: number;
+ pool_datasets?: boolean;
audio_loss_multiplier?: number;
max_loss?: number | null;
validation_config?: ValidationConfig;