RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Added ability to use two seperate folders for datasets when doing image reference sliders

This commit is contained in:
Jaret Burkett 2023-08-18 11:44:33 -06:00
parent 8d09eb44ec
commit d51c4ca704
3 changed files with 134 additions and 59 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

118
extensions_built_in/image_reference_slider_trainer/ImageReferenceSliderTrainerProcess.py
View File

@ -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

70
toolkit/data_loader.py
View File

@ -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)

5
toolkit/stable_diffusion_model.py
View File

@ -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: