Spaces:
Running on Zero
Running on Zero
| import copy | |
| import json | |
| import os | |
| import random | |
| import traceback | |
| from concurrent.futures import ThreadPoolExecutor | |
| from functools import lru_cache | |
| from typing import List, TYPE_CHECKING | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from PIL.ImageOps import exif_transpose | |
| from torchvision import transforms | |
| from torch.utils.data import Dataset, DataLoader, ConcatDataset | |
| from tqdm import tqdm | |
| import albumentations as A | |
| from toolkit import image_utils | |
| from toolkit.buckets import get_bucket_for_image_size, BucketResolution | |
| from toolkit.config_modules import DatasetConfig, preprocess_dataset_raw_config | |
| from toolkit.dataloader_mixins import CaptionMixin, BucketsMixin, LatentCachingMixin, Augments, CLIPCachingMixin, ControlCachingMixin, TextEmbeddingCachingMixin | |
| from toolkit.data_transfer_object.data_loader import FileItemDTO, DataLoaderBatchDTO | |
| from toolkit.print import print_acc | |
| from toolkit.accelerator import get_accelerator | |
| import platform | |
| def is_native_windows(): | |
| return platform.system() == "Windows" and platform.release() != "2" | |
| def is_macos(): | |
| return platform.system() == "Darwin" | |
| if TYPE_CHECKING: | |
| from toolkit.stable_diffusion_model import StableDiffusion | |
| image_extensions = ['.jpg', '.jpeg', '.png', '.webp'] | |
| video_extensions = ['.mp4', '.avi', '.mov', '.webm', '.mkv', '.wmv', '.m4v', '.flv'] | |
| audio_extensions = ['.mp3', '.wav', '.flac', '.aac', '.ogg', '.m4a'] | |
| class RescaleTransform: | |
| """Transform to rescale images to the range [-1, 1].""" | |
| def __call__(self, image): | |
| return image * 2 - 1 | |
| class NormalizeSDXLTransform: | |
| """ | |
| Transforms the range from 0 to 1 to SDXL mean and std per channel based on avgs over thousands of images | |
| Mean: tensor([ 0.0002, -0.1034, -0.1879]) | |
| Standard Deviation: tensor([0.5436, 0.5116, 0.5033]) | |
| """ | |
| def __call__(self, image): | |
| return transforms.Normalize( | |
| mean=[0.0002, -0.1034, -0.1879], | |
| std=[0.5436, 0.5116, 0.5033], | |
| )(image) | |
| class NormalizeSD15Transform: | |
| """ | |
| Transforms the range from 0 to 1 to SDXL mean and std per channel based on avgs over thousands of images | |
| Mean: tensor([-0.1600, -0.2450, -0.3227]) | |
| Standard Deviation: tensor([0.5319, 0.4997, 0.5139]) | |
| """ | |
| def __call__(self, image): | |
| return transforms.Normalize( | |
| mean=[-0.1600, -0.2450, -0.3227], | |
| std=[0.5319, 0.4997, 0.5139], | |
| )(image) | |
| class ImageDataset(Dataset, CaptionMixin): | |
| def __init__(self, config): | |
| self.config = config | |
| self.name = self.get_config('name', 'dataset') | |
| self.path = self.get_config('path', required=True) | |
| self.scale = self.get_config('scale', 1) | |
| self.random_scale = self.get_config('random_scale', False) | |
| self.include_prompt = self.get_config('include_prompt', False) | |
| self.default_prompt = self.get_config('default_prompt', '') | |
| if self.include_prompt: | |
| self.caption_type = self.get_config('caption_ext', 'txt') | |
| else: | |
| self.caption_type = None | |
| # we always random crop if random scale is enabled | |
| self.random_crop = self.random_scale if self.random_scale else self.get_config('random_crop', False) | |
| self.resolution = self.get_config('resolution', 256) | |
| self.file_list = [os.path.join(self.path, file) for file in os.listdir(self.path) if | |
| file.lower().endswith(('.jpg', '.jpeg', '.png', '.webp'))] | |
| # this might take a while | |
| print_acc(f" - Preprocessing image dimensions") | |
| new_file_list = [] | |
| bad_count = 0 | |
| for file in tqdm(self.file_list): | |
| try: | |
| w, h = image_utils.get_image_size(file) | |
| except image_utils.UnknownImageFormat: | |
| img = exif_transpose(Image.open(file)) | |
| w, h = img.size | |
| # img = Image.open(file) | |
| if int(min([w, h]) * self.scale) >= self.resolution: | |
| new_file_list.append(file) | |
| else: | |
| bad_count += 1 | |
| self.file_list = new_file_list | |
| print_acc(f" - Found {len(self.file_list)} images") | |
| print_acc(f" - Found {bad_count} images that are too small") | |
| assert len(self.file_list) > 0, f"no images found in {self.path}" | |
| self.transform = transforms.Compose([ | |
| transforms.ToTensor(), | |
| RescaleTransform(), | |
| ]) | |
| def get_config(self, key, default=None, required=False): | |
| if key in self.config: | |
| value = self.config[key] | |
| return value | |
| elif required: | |
| raise ValueError(f'config file error. Missing "config.dataset.{key}" key') | |
| else: | |
| return default | |
| def __len__(self): | |
| return len(self.file_list) | |
| def __getitem__(self, index): | |
| img_path = self.file_list[index] | |
| try: | |
| img = exif_transpose(Image.open(img_path)).convert('RGB') | |
| except Exception as e: | |
| print_acc(f"Error opening image: {img_path}") | |
| print_acc(e) | |
| # make a noise image if we can't open it | |
| img = Image.fromarray(np.random.randint(0, 255, (1024, 1024, 3), dtype=np.uint8)) | |
| # Downscale the source image first | |
| img = img.resize((int(img.size[0] * self.scale), int(img.size[1] * self.scale)), Image.BICUBIC) | |
| min_img_size = min(img.size) | |
| if self.random_crop: | |
| if self.random_scale and min_img_size > self.resolution: | |
| if min_img_size < self.resolution: | |
| print_acc( | |
| f"Unexpected values: min_img_size={min_img_size}, self.resolution={self.resolution}, image file={img_path}") | |
| scale_size = self.resolution | |
| else: | |
| scale_size = random.randint(self.resolution, int(min_img_size)) | |
| scaler = scale_size / min_img_size | |
| scale_width = int((img.width + 5) * scaler) | |
| scale_height = int((img.height + 5) * scaler) | |
| img = img.resize((scale_width, scale_height), Image.BICUBIC) | |
| img = transforms.RandomCrop(self.resolution)(img) | |
| else: | |
| img = transforms.CenterCrop(min_img_size)(img) | |
| img = img.resize((self.resolution, self.resolution), Image.BICUBIC) | |
| img = self.transform(img) | |
| if self.include_prompt: | |
| prompt = self.get_caption_item(index) | |
| return img, prompt | |
| else: | |
| return img | |
| class AugmentedImageDataset(ImageDataset): | |
| def __init__(self, config): | |
| super().__init__(config) | |
| self.augmentations = self.get_config('augmentations', []) | |
| self.augmentations = [Augments(**aug) for aug in self.augmentations] | |
| augmentation_list = [] | |
| for aug in self.augmentations: | |
| # make sure method name is valid | |
| assert hasattr(A, aug.method_name), f"invalid augmentation method: {aug.method_name}" | |
| # get the method | |
| method = getattr(A, aug.method_name) | |
| # add the method to the list | |
| augmentation_list.append(method(**aug.params)) | |
| self.aug_transform = A.Compose(augmentation_list) | |
| self.original_transform = self.transform | |
| # replace transform so we get raw pil image | |
| self.transform = transforms.Compose([]) | |
| def __getitem__(self, index): | |
| # get the original image | |
| # image is a PIL image, convert to bgr | |
| pil_image = super().__getitem__(index) | |
| open_cv_image = np.array(pil_image) | |
| # Convert RGB to BGR | |
| open_cv_image = open_cv_image[:, :, ::-1].copy() | |
| # apply augmentations | |
| augmented = self.aug_transform(image=open_cv_image)["image"] | |
| # convert back to RGB tensor | |
| augmented = cv2.cvtColor(augmented, cv2.COLOR_BGR2RGB) | |
| # convert to PIL image | |
| augmented = Image.fromarray(augmented) | |
| # return both # return image as 0 - 1 tensor | |
| return transforms.ToTensor()(pil_image), transforms.ToTensor()(augmented) | |
| class PairedImageDataset(Dataset): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.config = config | |
| self.size = self.get_config('size', 512) | |
| self.path = self.get_config('path', None) | |
| self.pos_folder = self.get_config('pos_folder', None) | |
| self.neg_folder = self.get_config('neg_folder', None) | |
| self.default_prompt = self.get_config('default_prompt', '') | |
| self.network_weight = self.get_config('network_weight', 1.0) | |
| self.pos_weight = self.get_config('pos_weight', self.network_weight) | |
| self.neg_weight = self.get_config('neg_weight', self.network_weight) | |
| supported_exts = ('.jpg', '.jpeg', '.png', '.webp', '.JPEG', '.JPG', '.PNG', '.WEBP') | |
| if self.pos_folder is not None and self.neg_folder is not None: | |
| # find matching files | |
| self.pos_file_list = [os.path.join(self.pos_folder, file) for file in os.listdir(self.pos_folder) if | |
| file.lower().endswith(supported_exts)] | |
| self.neg_file_list = [os.path.join(self.neg_folder, file) for file in os.listdir(self.neg_folder) if | |
| file.lower().endswith(supported_exts)] | |
| matched_files = [] | |
| for pos_file in self.pos_file_list: | |
| pos_file_no_ext = os.path.splitext(pos_file)[0] | |
| for neg_file in self.neg_file_list: | |
| neg_file_no_ext = os.path.splitext(neg_file)[0] | |
| if os.path.basename(pos_file_no_ext) == os.path.basename(neg_file_no_ext): | |
| matched_files.append((neg_file, pos_file)) | |
| break | |
| # remove duplicates | |
| matched_files = [t for t in (set(tuple(i) for i in matched_files))] | |
| self.file_list = matched_files | |
| print_acc(f" - Found {len(self.file_list)} matching pairs") | |
| else: | |
| self.file_list = [os.path.join(self.path, file) for file in os.listdir(self.path) if | |
| file.lower().endswith(supported_exts)] | |
| print_acc(f" - Found {len(self.file_list)} images") | |
| self.transform = transforms.Compose([ | |
| transforms.ToTensor(), | |
| RescaleTransform(), | |
| ]) | |
| def get_all_prompts(self): | |
| prompts = [] | |
| for index in range(len(self.file_list)): | |
| prompts.append(self.get_prompt_item(index)) | |
| # remove duplicates | |
| prompts = list(set(prompts)) | |
| return prompts | |
| def __len__(self): | |
| return len(self.file_list) | |
| def get_config(self, key, default=None, required=False): | |
| if key in self.config: | |
| value = self.config[key] | |
| return value | |
| elif required: | |
| raise ValueError(f'config file error. Missing "config.dataset.{key}" key') | |
| else: | |
| return default | |
| def get_prompt_item(self, index): | |
| img_path_or_tuple = self.file_list[index] | |
| if isinstance(img_path_or_tuple, tuple): | |
| # check if either has a prompt file | |
| path_no_ext = os.path.splitext(img_path_or_tuple[0])[0] | |
| prompt_path = path_no_ext + '.txt' | |
| if not os.path.exists(prompt_path): | |
| path_no_ext = os.path.splitext(img_path_or_tuple[1])[0] | |
| prompt_path = path_no_ext + '.txt' | |
| else: | |
| img_path = img_path_or_tuple | |
| # see if prompt file exists | |
| path_no_ext = os.path.splitext(img_path)[0] | |
| prompt_path = path_no_ext + '.txt' | |
| if os.path.exists(prompt_path): | |
| with open(prompt_path, 'r', encoding='utf-8') as f: | |
| prompt = f.read() | |
| # remove any newlines | |
| prompt = prompt.replace('\n', ', ') | |
| # remove new lines for all operating systems | |
| prompt = prompt.replace('\r', ', ') | |
| prompt_split = prompt.split(',') | |
| # remove empty strings | |
| prompt_split = [p.strip() for p in prompt_split if p.strip()] | |
| # join back together | |
| prompt = ', '.join(prompt_split) | |
| else: | |
| prompt = self.default_prompt | |
| return prompt | |
| def __getitem__(self, index): | |
| img_path_or_tuple = self.file_list[index] | |
| if isinstance(img_path_or_tuple, tuple): | |
| # load both images | |
| img_path = img_path_or_tuple[0] | |
| img1 = exif_transpose(Image.open(img_path)).convert('RGB') | |
| img_path = img_path_or_tuple[1] | |
| img2 = exif_transpose(Image.open(img_path)).convert('RGB') | |
| # always use # 2 (pos) | |
| bucket_resolution = get_bucket_for_image_size( | |
| width=img2.width, | |
| height=img2.height, | |
| resolution=self.size, | |
| # divisibility=self. | |
| ) | |
| # images will be same base dimension, but may be trimmed. We need to shrink and then central crop | |
| if bucket_resolution['width'] > bucket_resolution['height']: | |
| img1_scale_to_height = bucket_resolution["height"] | |
| img1_scale_to_width = int(img1.width * (bucket_resolution["height"] / img1.height)) | |
| img2_scale_to_height = bucket_resolution["height"] | |
| img2_scale_to_width = int(img2.width * (bucket_resolution["height"] / img2.height)) | |
| else: | |
| img1_scale_to_width = bucket_resolution["width"] | |
| img1_scale_to_height = int(img1.height * (bucket_resolution["width"] / img1.width)) | |
| img2_scale_to_width = bucket_resolution["width"] | |
| img2_scale_to_height = int(img2.height * (bucket_resolution["width"] / img2.width)) | |
| img1_crop_height = bucket_resolution["height"] | |
| img1_crop_width = bucket_resolution["width"] | |
| img2_crop_height = bucket_resolution["height"] | |
| img2_crop_width = bucket_resolution["width"] | |
| # scale then center crop images | |
| img1 = img1.resize((img1_scale_to_width, img1_scale_to_height), Image.BICUBIC) | |
| img1 = transforms.CenterCrop((img1_crop_height, img1_crop_width))(img1) | |
| img2 = img2.resize((img2_scale_to_width, img2_scale_to_height), Image.BICUBIC) | |
| img2 = transforms.CenterCrop((img2_crop_height, img2_crop_width))(img2) | |
| # combine them side by side | |
| img = Image.new('RGB', (img1.width + img2.width, max(img1.height, img2.height))) | |
| img.paste(img1, (0, 0)) | |
| img.paste(img2, (img1.width, 0)) | |
| else: | |
| img_path = img_path_or_tuple | |
| img = exif_transpose(Image.open(img_path)).convert('RGB') | |
| height = self.size | |
| # determine width to keep aspect ratio | |
| width = int(img.size[0] * height / img.size[1]) | |
| # Downscale the source image first | |
| img = img.resize((width, height), Image.BICUBIC) | |
| prompt = self.get_prompt_item(index) | |
| img = self.transform(img) | |
| return img, prompt, (self.neg_weight, self.pos_weight) | |
| class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin, TextEmbeddingCachingMixin, BucketsMixin, CaptionMixin, Dataset): | |
| def __init__( | |
| self, | |
| 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 | |
| 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 | |
| super().__init__() | |
| folder_path = dataset_config.folder_path | |
| 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 = 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 | |
| if self.sd is None and self.is_caching_latents: | |
| raise ValueError(f"sd is required for caching latents") | |
| self.caption_type = dataset_config.caption_ext | |
| self.default_caption = dataset_config.default_caption | |
| self.random_scale = dataset_config.random_scale | |
| self.scale = dataset_config.scale | |
| self.batch_size = batch_size | |
| # built lazily inside the worker; an executor cannot be pickled to spawned workers | |
| self._batch_pool = None | |
| # we always random crop if random scale is enabled | |
| self.random_crop = self.random_scale if self.random_scale else dataset_config.random_crop | |
| 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(dataset_path): | |
| extensions = image_extensions | |
| if self.is_audio_model: | |
| # only look for audio files | |
| extensions = audio_extensions | |
| elif self.is_multimodal_llm: | |
| extensions = audio_extensions + image_extensions + (video_extensions if self.is_video else []) | |
| elif self.is_video: | |
| # look for videos and images. Video models can train on both; | |
| # images are bucketed separately as single-frame items | |
| extensions = video_extensions + image_extensions | |
| # prune hidden dirs (.thumbs, .tmp) so their contents never train | |
| file_list = [] | |
| 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(dataset_path, 'r') as f: | |
| caption_dict = json.load(f) | |
| # keys are file paths | |
| 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 dataset_config.num_repeats > 1: | |
| # repeat the list | |
| file_list = file_list * dataset_config.num_repeats | |
| # this might take a while | |
| 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 = 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') | |
| dataloader_version = "0.1.2" | |
| if os.path.exists(dataset_size_file): | |
| try: | |
| with open(dataset_size_file, 'r') as f: | |
| size_database = json.load(f) | |
| 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 | |
| size_database = {} | |
| except Exception as e: | |
| print_acc(f"Error loading size database: {dataset_size_file}") | |
| print_acc(e) | |
| size_database = {} | |
| else: | |
| size_database = {} | |
| 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" | |
| text_embedding_space_version = self.sd.get_text_embedding_space_version() if self.sd is not None else "sd1" | |
| temporal_compression = 8 | |
| if self.sd is not None: | |
| if hasattr(self.sd.vae, 'config') and hasattr(self.sd.vae.config, 'scale_factor_temporal'): | |
| temporal_compression = self.sd.vae.config.scale_factor_temporal | |
| 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: | |
| file_item = FileItemDTO( | |
| sd=self.sd, | |
| 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=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, | |
| dopsd_self_ref=getattr(self.sd, 'dopsd_self_ref', False) if self.sd else False, | |
| text_embedding_space_version=text_embedding_space_version, | |
| text_embedding_uses_target_size=getattr(self.sd, 'text_embedding_uses_target_size', False) if self.sd else False, | |
| te_padding_side=self.sd.te_padding_side if self.sd else "right", | |
| latent_space_version=latent_space_version, | |
| 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, | |
| ) | |
| source_items.append(file_item) | |
| except Exception as e: | |
| print_acc(traceback.format_exc()) | |
| if self.is_video: | |
| print_acc(f"Error processing video: {file}") | |
| else: | |
| print_acc(f"Error processing image: {file}") | |
| print_acc(e) | |
| bad_count += 1 | |
| # save the size database | |
| with open(dataset_size_file, 'w') as f: | |
| json.dump(size_database, f) | |
| if self.is_video: | |
| 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(source_items) > 0, f"no videos found in {dataset_path}" | |
| else: | |
| 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 dataset_config.flip_x: | |
| print_acc(" - adding x axis flips") | |
| 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 | |
| source_items.append(new_file_item) | |
| # handle y axis flips | |
| if dataset_config.flip_y: | |
| print_acc(" - adding y axis flips") | |
| 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 | |
| source_items.append(new_file_item) | |
| if dataset_config.flip_x or dataset_config.flip_y: | |
| if self.is_video: | |
| print_acc(f" - Found {len(source_items)} videos after adding flips") | |
| else: | |
| print_acc(f" - Found {len(source_items)} images after adding flips") | |
| self.file_list.extend(source_items) | |
| def setup_epoch(self): | |
| if self.epoch_num == 0: | |
| # initial setup | |
| # do not call for now | |
| if self.dataset_config.buckets: | |
| # setup buckets | |
| self.setup_buckets() | |
| if self.is_caching_latents: | |
| self.cache_latents_all_latents() | |
| if self.is_caching_clip_vision_to_disk: | |
| self.cache_clip_vision_to_disk() | |
| if self.is_caching_text_embeddings: | |
| self.cache_text_embeddings() | |
| if self.is_generating_controls: | |
| # always do this last | |
| self.setup_controls() | |
| self.epoch_num += 1 | |
| def __getstate__(self): | |
| # on Windows/macOS dataloader workers are spawned, which pickles the dataset. | |
| # sd (the model) is not picklable (weakrefs, cuda tensors) and is only needed | |
| # for caching, which runs in the main process before iteration starts. | |
| state = self.__dict__.copy() | |
| state['sd'] = None | |
| return state | |
| def __len__(self): | |
| if self.dataset_config.buckets: | |
| return len(self.batch_indices) | |
| return len(self.file_list) | |
| def _get_replacement_index(self, index) -> int: | |
| # when an image fails to load we have to swap in a different one. With buckets the | |
| # replacement must come from the same bucket so the collated shapes still match. | |
| if self.dataset_config.buckets: | |
| for bucket in self.buckets.values(): | |
| if index in bucket.file_list_idx: | |
| candidates = [i for i in bucket.file_list_idx if i != index] | |
| if candidates: | |
| return random.choice(candidates) | |
| break | |
| return random.randint(0, len(self.file_list) - 1) | |
| 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(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: | |
| # avoid infinite recursion if many files are corrupt | |
| raise | |
| new_index = self._get_replacement_index(index) | |
| return self._get_single_item(new_index, _attempts=_attempts + 1) | |
| file_item.load_caption(self.caption_dict) | |
| return file_item | |
| def __getitem__(self, item): | |
| if self.dataset_config.buckets: | |
| # for buckets we collate ourselves for now | |
| # todo allow a scheduler to dynamically make buckets | |
| # we collate ourselves | |
| if len(self.batch_indices) - 1 < item: | |
| # tried everything to solve this. No way to reset length when redoing things. Pick another index | |
| item = random.randint(0, len(self.batch_indices) - 1) | |
| idx_list = self.batch_indices[item] | |
| n_threads = min(self.dataset_config.batch_load_threads, len(idx_list)) | |
| if n_threads <= 1: | |
| return [self._get_single_item(idx) for idx in idx_list] | |
| if self._batch_pool is None or self._batch_pool._max_workers != n_threads: | |
| self._batch_pool = ThreadPoolExecutor(max_workers=n_threads) | |
| return list(self._batch_pool.map(self._get_single_item, idx_list)) | |
| else: | |
| # Dataloader is batching | |
| return self._get_single_item(item) | |
| def dto_collation(batch: List['FileItemDTO']): | |
| # must be a module level function so spawned dataloader workers can pickle it | |
| return DataLoaderBatchDTO( | |
| file_items=batch | |
| ) | |
| 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 | |
| datasets = [] | |
| has_buckets = False | |
| is_caching_latents = False | |
| dataset_config_list = [] | |
| # preprocess them all | |
| for dataset_option in dataset_options: | |
| if isinstance(dataset_option, DatasetConfig): | |
| dataset_config_list.append(dataset_option) | |
| else: | |
| # preprocess raw data | |
| split_configs = preprocess_dataset_raw_config([dataset_option]) | |
| for x in split_configs: | |
| 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 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 | |
| datasets.append(AiToolkitDataset(config, batch_size=dataset_batch_size, sd=sd)) | |
| concatenated_dataset = ConcatDataset(datasets) | |
| # todo build scheduler that can get buckets from all datasets that match | |
| # todo and evenly distribute reg images | |
| # check if is caching latents | |
| dataloader_kwargs = {} | |
| import multiprocessing as mp | |
| is_daemon = bool(getattr(mp.current_process(), 'daemon', False)) or bool(mp.current_process()._config.get('daemon', False)) | |
| dataloader_kwargs['num_workers'] = 0 if is_daemon else dataset_config_list[0].num_workers | |
| if dataloader_kwargs['num_workers'] > 0: | |
| dataloader_kwargs['prefetch_factor'] = dataset_config_list[0].prefetch_factor | |
| # keep workers alive across epochs. Without this, spawn platforms (Windows/macOS) | |
| # boot new worker processes every epoch, which can take longer than the epoch | |
| # itself on small datasets. The dataset is static after epoch 0 (setup_epoch only | |
| # does work on the first call) and per-epoch shuffling happens in the main process | |
| # sampler, so workers never hold stale state. | |
| dataloader_kwargs['persistent_workers'] = True | |
| # spawned workers re-import the full stack at boot and would repeat every | |
| # import-time warning the parent already printed. Children inherit these env | |
| # vars; the parent is unaffected since its imports already happened. | |
| os.environ.setdefault('PYTHONWARNINGS', 'ignore::FutureWarning') | |
| os.environ.setdefault('TORCH_LOGS', '-torch.utils._pytree') | |
| os.environ.setdefault('DIFFUSERS_VERBOSITY', 'error') | |
| os.environ.setdefault('NO_ALBUMENTATIONS_UPDATE', '1') | |
| # Enable pinned memory only when explicitly opted in via dataset config and | |
| # CUDA is available. pin_memory speeds up CPU->GPU transfer (helpful even at | |
| # num_workers=0), but page-locked RAM cannot be relocated by NVIDIA's | |
| # Windows driver shared-memory VRAM-overflow fallback, which can cause | |
| # severe PCIe thrashing for users at the VRAM ceiling. Off by default; opt | |
| # in via 'pin_memory: true' on the dataset config when VRAM headroom is | |
| # stable. | |
| if torch.cuda.is_available() and dataset_config_list[0].pin_memory: | |
| dataloader_kwargs['pin_memory'] = True | |
| if has_buckets: | |
| # make sure they all have buckets | |
| for dataset in datasets: | |
| assert dataset.dataset_config.buckets, f"buckets not found on dataset {dataset.dataset_config.folder_path}, you either need all buckets or none" | |
| data_loader = DataLoader( | |
| concatenated_dataset, | |
| batch_size=None, # we batch in the datasets for now | |
| drop_last=False, | |
| shuffle=True, | |
| collate_fn=dto_collation, # Use the custom collate function | |
| **dataloader_kwargs | |
| ) | |
| else: | |
| # without buckets the dataloader batches across all datasets at once, | |
| # so a dataset level batch_size cannot apply | |
| for config in dataset_config_list: | |
| if config.batch_size is not None: | |
| raise ValueError( | |
| f"Dataset level batch_size requires buckets to be enabled. Dataset {config.folder_path or config.dataset_path} has buckets disabled." | |
| ) | |
| data_loader = DataLoader( | |
| concatenated_dataset, | |
| batch_size=batch_size, | |
| shuffle=True, | |
| collate_fn=dto_collation, | |
| **dataloader_kwargs | |
| ) | |
| return data_loader | |
| def trigger_dataloader_setup_epoch(dataloader: DataLoader): | |
| # hacky but needed because of different types of datasets and dataloaders | |
| dataloader.len = None | |
| if isinstance(dataloader.dataset, list): | |
| for dataset in dataloader.dataset: | |
| if hasattr(dataset, 'datasets'): | |
| for sub_dataset in dataset.datasets: | |
| if hasattr(sub_dataset, 'setup_epoch'): | |
| sub_dataset.setup_epoch() | |
| sub_dataset.len = None | |
| elif hasattr(dataset, 'setup_epoch'): | |
| dataset.setup_epoch() | |
| dataset.len = None | |
| elif hasattr(dataloader.dataset, 'setup_epoch'): | |
| dataloader.dataset.setup_epoch() | |
| dataloader.dataset.len = None | |
| elif hasattr(dataloader.dataset, 'datasets'): | |
| dataloader.dataset.len = None | |
| for sub_dataset in dataloader.dataset.datasets: | |
| if hasattr(sub_dataset, 'setup_epoch'): | |
| sub_dataset.setup_epoch() | |
| sub_dataset.len = None | |
| def get_dataloader_datasets(dataloader: DataLoader): | |
| # hacky but needed because of different types of datasets and dataloaders | |
| if isinstance(dataloader.dataset, list): | |
| datasets = [] | |
| for dataset in dataloader.dataset: | |
| if hasattr(dataset, 'datasets'): | |
| for sub_dataset in dataset.datasets: | |
| datasets.append(sub_dataset) | |
| else: | |
| datasets.append(dataset) | |
| return datasets | |
| elif hasattr(dataloader.dataset, 'datasets'): | |
| return dataloader.dataset.datasets | |
| else: | |
| return [dataloader.dataset] | |