From d51c4ca7045af322fc9ee21c0213af61fc961ce9 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 18 Aug 2023 11:44:33 -0600 Subject: [PATCH] Added ability to use two seperate folders for datasets when doing image reference sliders --- .../ImageReferenceSliderTrainerProcess.py | 118 ++++++++++-------- toolkit/data_loader.py | 70 +++++++++-- toolkit/stable_diffusion_model.py | 5 + 3 files changed, 134 insertions(+), 59 deletions(-) diff --git a/extensions_built_in/image_reference_slider_trainer/ImageReferenceSliderTrainerProcess.py b/extensions_built_in/image_reference_slider_trainer/ImageReferenceSliderTrainerProcess.py index dba3a57f..568ee9ae 100644 --- a/extensions_built_in/image_reference_slider_trainer/ImageReferenceSliderTrainerProcess.py +++ b/extensions_built_in/image_reference_slider_trainer/ImageReferenceSliderTrainerProcess.py @@ -6,7 +6,7 @@ from contextlib import nullcontext from typing import Optional, Union, List from torch.utils.data import ConcatDataset, DataLoader from toolkit.data_loader import PairedImageDataset -from toolkit.prompt_utils import concat_prompt_embeds +from toolkit.prompt_utils import concat_prompt_embeds, split_prompt_embeds from toolkit.stable_diffusion_model import StableDiffusion, PromptEmbeds from toolkit.train_tools import get_torch_dtype import gc @@ -22,8 +22,18 @@ def flush(): class DatasetConfig: def __init__(self, **kwargs): + # can pass with a side by side pait or a folder with pos and neg folder self.pair_folder: str = kwargs.get('pair_folder', None) - self.network_weight: float = kwargs.get('network_weight', 1.0) + self.pos_folder: str = kwargs.get('pos_folder', None) + self.neg_folder: str = kwargs.get('neg_folder', None) + + self.network_weight: float = float(kwargs.get('network_weight', 1.0)) + self.pos_weight: float = float(kwargs.get('pos_weight', self.network_weight)) + self.neg_weight: float = float(kwargs.get('neg_weight', self.network_weight)) + # make sure they are all absolute values no negatives + self.pos_weight = abs(self.pos_weight) + self.neg_weight = abs(self.neg_weight) + self.target_class: str = kwargs.get('target_class', '') self.size: int = kwargs.get('size', 512) @@ -58,6 +68,10 @@ class ImageReferenceSliderTrainerProcess(BaseSDTrainProcess): 'size': dataset.size, 'default_prompt': dataset.target_class, 'network_weight': dataset.network_weight, + 'pos_weight': dataset.pos_weight, + 'neg_weight': dataset.neg_weight, + 'pos_folder': dataset.pos_folder, + 'neg_folder': dataset.neg_folder, } image_dataset = PairedImageDataset(config) datasets.append(image_dataset) @@ -81,10 +95,15 @@ class ImageReferenceSliderTrainerProcess(BaseSDTrainProcess): pass def hook_train_loop(self, batch): - do_mirror_loss = 'mirror' in self.slider_config.additional_losses - with torch.no_grad(): - imgs, prompts, base_network_weight = batch + imgs, prompts, network_weights = batch + network_pos_weight, network_neg_weight = network_weights + if isinstance(network_pos_weight, torch.Tensor): + network_pos_weight = network_pos_weight.item() + if isinstance(network_neg_weight, torch.Tensor): + network_neg_weight = network_neg_weight.item() + # if items in network_weight list are tensors, convert them to floats + dtype = get_torch_dtype(self.train_config.dtype) imgs: torch.Tensor = imgs.to(self.device_torch, dtype=dtype) # split batched images in half so left is negative and right is positive @@ -120,12 +139,7 @@ class ImageReferenceSliderTrainerProcess(BaseSDTrainProcess): noise_offset=self.train_config.noise_offset, ).to(self.device_torch, dtype=dtype) - if do_mirror_loss: - # mirror the noise - # torch shape is [batch, channels, height, width] - noise_negative = torch.flip(noise_positive.clone(), dims=[3]) - else: - noise_negative = noise_positive.clone() + noise_negative = noise_positive.clone() # Add noise to the latents according to the noise magnitude at each timestep # (this is the forward diffusion process) @@ -135,12 +149,11 @@ class ImageReferenceSliderTrainerProcess(BaseSDTrainProcess): noisy_latents = torch.cat([noisy_positive_latents, noisy_negative_latents], dim=0) noise = torch.cat([noise_positive, noise_negative], dim=0) timesteps = torch.cat([timesteps, timesteps], dim=0) - network_multiplier = [base_network_weight * 1.0, base_network_weight * -1.0] + network_multiplier = [network_pos_weight * 1.0, network_neg_weight * -1.0] flush() loss_float = None - loss_slide_float = None loss_mirror_float = None self.optimizer.zero_grad() @@ -157,48 +170,58 @@ class ImageReferenceSliderTrainerProcess(BaseSDTrainProcess): conditional_embeds = concat_prompt_embeds(embedding_list) conditional_embeds = concat_prompt_embeds([conditional_embeds, conditional_embeds]) - with self.network: - assert self.network.is_active + if self.model_config.is_xl: + # todo also allow for setting this for low ram in general, but sdxl spikes a ton on back prop + network_multiplier_list = network_multiplier + noisy_latent_list = torch.chunk(noisy_latents, 2, dim=0) + noise_list = torch.chunk(noise, 2, dim=0) + timesteps_list = torch.chunk(timesteps, 2, dim=0) + conditional_embeds_list = split_prompt_embeds(conditional_embeds) + else: + network_multiplier_list = [network_multiplier] + noisy_latent_list = [noisy_latents] + noise_list = [noise] + timesteps_list = [timesteps] + conditional_embeds_list = [conditional_embeds] - self.network.multiplier = network_multiplier + losses = [] + # allow to chunk it out to save vram + for network_multiplier, noisy_latents, noise, timesteps, conditional_embeds in zip( + network_multiplier_list, noisy_latent_list, noise_list, timesteps_list, conditional_embeds_list + ): + with self.network: + assert self.network.is_active - noise_pred = self.sd.predict_noise( - latents=noisy_latents, - conditional_embeddings=conditional_embeds, - timestep=timesteps, - ) + self.network.multiplier = network_multiplier - if self.sd.prediction_type == 'v_prediction': - # v-parameterization training - target = noise_scheduler.get_velocity(noisy_latents, noise, timesteps) - else: - target = noise + noise_pred = self.sd.predict_noise( + latents=noisy_latents.to(self.device_torch, dtype=dtype), + conditional_embeddings=conditional_embeds.to(self.device_torch, dtype=dtype), + timestep=timesteps, + ) + noise = noise.to(self.device_torch, dtype=dtype) - loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="none") - loss = loss.mean([1, 2, 3]) + if self.sd.prediction_type == 'v_prediction': + # v-parameterization training + target = noise_scheduler.get_velocity(noisy_latents, noise, timesteps) + else: + target = noise - # todo add snr gamma here + loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="none") + loss = loss.mean([1, 2, 3]) - loss = loss.mean() - loss_slide_float = loss.item() + # todo add snr gamma here - if do_mirror_loss: - noise_pred_pos, noise_pred_neg = torch.chunk(noise_pred, 2, dim=0) - # mirror the negative - noise_pred_neg = torch.flip(noise_pred_neg.clone(), dims=[3]) - loss_mirror = torch.nn.functional.mse_loss(noise_pred_pos.float(), noise_pred_neg.float(), - reduction="none") - loss_mirror = loss_mirror.mean([1, 2, 3]) - loss_mirror = loss_mirror.mean() - loss_mirror_float = loss_mirror.item() - loss += loss_mirror + loss = loss.mean() + loss_slide_float = loss.item() - loss_float = loss.item() + loss_float = loss.item() + losses.append(loss_float) - # back propagate loss to free ram - loss.backward() + # back propagate loss to free ram + loss.backward() + flush() - flush() # apply gradients optimizer.step() @@ -208,11 +231,8 @@ class ImageReferenceSliderTrainerProcess(BaseSDTrainProcess): self.network.multiplier = 1.0 loss_dict = OrderedDict( - {'loss': loss_float}, + {'loss': sum(losses) / len(losses) if len(losses) > 0 else 0.0} ) - if do_mirror_loss: - loss_dict['l/s'] = loss_slide_float - loss_dict['l/m'] = loss_mirror_float return loss_dict # end hook_train_loop diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index 2a58ab43..a6b2a7a0 100644 --- a/toolkit/data_loader.py +++ b/toolkit/data_loader.py @@ -147,12 +147,43 @@ class PairedImageDataset(Dataset): super().__init__() self.config = config self.size = self.get_config('size', 512) - self.path = self.get_config('path', required=True) + 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.file_list = [os.path.join(self.path, file) for file in os.listdir(self.path) if - file.lower().endswith(('.jpg', '.jpeg', '.png', '.webp'))] - print(f" - Found {len(self.file_list)} images") + 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(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(f" - Found {len(self.file_list)} images") self.transform = transforms.Compose([ transforms.ToTensor(), @@ -172,12 +203,31 @@ class PairedImageDataset(Dataset): return default def __getitem__(self, index): - img_path = self.file_list[index] - img = exif_transpose(Image.open(img_path)).convert('RGB') + 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') + # 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)) + + # 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 + img = exif_transpose(Image.open(img_path)).convert('RGB') + # see if prompt file exists + path_no_ext = os.path.splitext(img_path)[0] + prompt_path = path_no_ext + '.txt' - # 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() @@ -201,5 +251,5 @@ class PairedImageDataset(Dataset): img = img.resize((width, height), Image.BICUBIC) img = self.transform(img) - return img, prompt, self.network_weight + return img, prompt, (self.neg_weight, self.pos_weight) diff --git a/toolkit/stable_diffusion_model.py b/toolkit/stable_diffusion_model.py index dcdb2638..f0122a21 100644 --- a/toolkit/stable_diffusion_model.py +++ b/toolkit/stable_diffusion_model.py @@ -453,6 +453,8 @@ class StableDiffusion: if do_classifier_free_guidance: latent_model_input = torch.cat([latents] * 2) + else: + latent_model_input = latents latent_model_input = self.noise_scheduler.scale_model_input(latent_model_input, timestep) @@ -633,6 +635,9 @@ class StableDiffusion: key = prefix + k v = v.detach().clone() state_dict[key] = v.to("cpu", dtype=get_torch_dtype(save_dtype)) + # make sure there are not nan values + if torch.isnan(state_dict[key]).any(): + raise ValueError(f"NaN value in state dict: {key}") # todo see what logit scale is if self.is_xl: