github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Fixes to esrgan trainer. Moved logic for sd prompt embeddings out of diffusers pipeline so I can manipulate it

Browse Source
This commit is contained in:
Jaret Burkett
2023-09-16 17:41:07 -06:00
parent 27f343fc08
commit c698837241
11 changed files with 214 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
toolkit/config_modules.py
View File

@@ -3,6 +3,8 @@ import time
from typing import List, Optional, Literal, Union
import random
from toolkit.prompt_utils import PromptEmbeds
ImgExt = Literal['jpg', 'png', 'webp']
@@ -447,3 +449,11 @@ class GenerateImageConfig:
self.network_multiplier = float(content)
elif flag == 'gr':
self.guidance_rescale = float(content)
def post_process_embeddings(
self,
conditional_prompt_embeds: PromptEmbeds,
unconditional_prompt_embeds: Optional[PromptEmbeds] = None,
):
# this is called after prompt embeds are encoded. We can override them in the future here
pass

7
toolkit/data_loader.py
View File

@@ -53,6 +53,8 @@ class ImageDataset(Dataset, CaptionMixin):
else:
bad_count += 1
self.file_list = new_file_list
print(f" - Found {len(self.file_list)} images")
print(f" - Found {bad_count} images that are too small")
assert len(self.file_list) > 0, f"no images found in {self.path}"
@@ -90,7 +92,10 @@ class ImageDataset(Dataset, CaptionMixin):
scale_size = self.resolution
else:
scale_size = random.randint(self.resolution, int(min_img_size))
img = img.resize((scale_size, scale_size), Image.BICUBIC)
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)

6
toolkit/llvae.py
View File

@@ -7,7 +7,7 @@ import itertools
class LosslessLatentDecoder(nn.Module):
def __init__(self, in_channels, latent_depth, dtype=torch.float32):
super(LosslessLatentDecoder, self).__init__()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.latent_depth = latent_depth
self.in_channels = in_channels
self.out_channels = int(in_channels // (latent_depth * latent_depth))
@@ -46,7 +46,7 @@ class LosslessLatentDecoder(nn.Module):
class LosslessLatentEncoder(nn.Module):
def __init__(self, in_channels, latent_depth, dtype=torch.float32):
super(LosslessLatentEncoder, self).__init__()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.latent_depth = latent_depth
self.in_channels = in_channels
self.out_channels = int(in_channels * (latent_depth * latent_depth))
@@ -108,7 +108,7 @@ if __name__ == '__main__':
from PIL import Image
import torchvision.transforms as transforms
user_path = os.path.expanduser('~')
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.float32
input_path = os.path.join(user_path, "Pictures/sample_2_512.png")

16
toolkit/losses.py
View File

@@ -27,11 +27,17 @@ class ComparativeTotalVariation(torch.nn.Module):
# Gradient penalty
def get_gradient_penalty(critic, real, fake, device):
with torch.autocast(device_type='cuda'):
alpha = torch.rand(real.size(0), 1, 1, 1).to(device)
real = real.float()
fake = fake.float()
alpha = torch.rand(real.size(0), 1, 1, 1).to(device).float()
interpolates = (alpha * real + ((1 - alpha) * fake)).requires_grad_(True)
if torch.isnan(interpolates).any():
print('d_interpolates is nan')
d_interpolates = critic(interpolates)
fake = torch.ones(real.size(0), 1, device=device)
if torch.isnan(d_interpolates).any():
print('fake is nan')
gradients = torch.autograd.grad(
outputs=d_interpolates,
inputs=interpolates,
@@ -41,10 +47,14 @@ def get_gradient_penalty(critic, real, fake, device):
only_inputs=True,
)[0]
# see if any are nan
if torch.isnan(gradients).any():
print('gradients is nan')
gradients = gradients.view(gradients.size(0), -1)
gradient_norm = gradients.norm(2, dim=1)
gradient_penalty = ((gradient_norm - 1) ** 2).mean()
return gradient_penalty
return gradient_penalty.float()
class PatternLoss(torch.nn.Module):

6
toolkit/prompt_utils.py
View File

@@ -44,6 +44,12 @@ class PromptEmbeds:
self.pooled_embeds = self.pooled_embeds.detach()
return self
def clone(self):
if self.pooled_embeds is not None:
return PromptEmbeds([self.text_embeds.clone(), self.pooled_embeds.clone()])
else:
return PromptEmbeds(self.text_embeds.clone())
class EncodedPromptPair:
def __init__(

41
toolkit/stable_diffusion_model.py
View File

@@ -368,6 +368,19 @@ class StableDiffusion:
torch.manual_seed(gen_config.seed)
torch.cuda.manual_seed(gen_config.seed)
# encode the prompt ourselves so we can do fun stuff with embeddings
conditional_embeds = self.encode_prompt(gen_config.prompt, gen_config.prompt_2, force_all=True)
unconditional_embeds = self.encode_prompt(
gen_config.negative_prompt, gen_config.negative_prompt_2, force_all=True
)
# allow any manipulations to take place to embeddings
gen_config.post_process_embeddings(
conditional_embeds,
unconditional_embeds,
)
# todo do we disable text encoder here as well if disabled for model, or only do that for training?
if self.is_xl:
# fix guidance rescale for sdxl
@@ -382,10 +395,14 @@ class StableDiffusion:
extra['use_karras_sigmas'] = True
img = pipeline(
prompt=gen_config.prompt,
prompt_2=gen_config.prompt_2,
negative_prompt=gen_config.negative_prompt,
negative_prompt_2=gen_config.negative_prompt_2,
# prompt=gen_config.prompt,
# prompt_2=gen_config.prompt_2,
prompt_embeds=conditional_embeds.text_embeds,
pooled_prompt_embeds=conditional_embeds.pooled_embeds,
negative_prompt_embeds=unconditional_embeds.text_embeds,
negative_pooled_prompt_embeds=unconditional_embeds.pooled_embeds,
# negative_prompt=gen_config.negative_prompt,
# negative_prompt_2=gen_config.negative_prompt_2,
height=gen_config.height,
width=gen_config.width,
num_inference_steps=gen_config.num_inference_steps,
@@ -395,8 +412,10 @@ class StableDiffusion:
).images[0]
else:
img = pipeline(
prompt=gen_config.prompt,
negative_prompt=gen_config.negative_prompt,
# prompt=gen_config.prompt,
prompt_embeds=conditional_embeds.text_embeds,
negative_prompt_embeds=unconditional_embeds.text_embeds,
# negative_prompt=gen_config.negative_prompt,
height=gen_config.height,
width=gen_config.width,
num_inference_steps=gen_config.num_inference_steps,
@@ -625,21 +644,25 @@ class StableDiffusion:
# return latents_steps
return latents
def encode_prompt(self, prompt, num_images_per_prompt=1) -> PromptEmbeds:
def encode_prompt(self, prompt, prompt2=None, num_images_per_prompt=1, force_all=False) -> PromptEmbeds:
# sd1.5 embeddings are (bs, 77, 768)
prompt = prompt
# if it is not a list, make it one
if not isinstance(prompt, list):
prompt = [prompt]
if prompt2 is not None and not isinstance(prompt2, list):
prompt2 = [prompt2]
if self.is_xl:
return PromptEmbeds(
train_tools.encode_prompts_xl(
self.tokenizer,
self.text_encoder,
prompt,
prompt2,
num_images_per_prompt=num_images_per_prompt,
use_text_encoder_1=self.use_text_encoder_1,
use_text_encoder_2=self.use_text_encoder_2,
use_text_encoder_1=self.use_text_encoder_1 or force_all,
use_text_encoder_2=self.use_text_encoder_2 or force_all,
)
)
else:

13
toolkit/style.py
View File

@@ -33,12 +33,17 @@ class ContentLoss(nn.Module):
# Define the separate loss function
def separated_loss(y_pred, y_true):
y_pred = y_pred.float()
y_true = y_true.float()
diff = torch.abs(y_pred - y_true)
l2 = torch.sum(diff ** 2, dim=[1, 2, 3], keepdim=True) / 2.0
return 2. * l2 / content_size
# Calculate itemized loss
pred_itemized_loss = separated_loss(pred_layer, target_layer)
# check if is nan
if torch.isnan(pred_itemized_loss).any():
print('pred_itemized_loss is nan')
# Calculate the mean of itemized loss
loss = torch.mean(pred_itemized_loss, dim=(1, 2, 3), keepdim=True)
@@ -48,6 +53,7 @@ class ContentLoss(nn.Module):
def convert_to_gram_matrix(inputs):
inputs = inputs.float()
shape = inputs.size()
batch, filters, height, width = shape[0], shape[1], shape[2], shape[3]
size = height * width * filters
@@ -93,11 +99,14 @@ class StyleLoss(nn.Module):
target_grams = convert_to_gram_matrix(style_target)
pred_grams = convert_to_gram_matrix(preds)
itemized_loss = separated_loss(pred_grams, target_grams)
# check if is nan
if torch.isnan(itemized_loss).any():
print('itemized_loss is nan')
# reshape itemized loss to be (batch, 1, 1, 1)
itemized_loss = torch.unsqueeze(itemized_loss, dim=1)
# gram_size = (tf.shape(target_grams)[1] * tf.shape(target_grams)[2])
loss = torch.mean(itemized_loss, dim=(1, 2), keepdim=True)
self.loss = loss.to(input_dtype)
self.loss = loss.to(input_dtype).float()
return stacked_input.to(input_dtype)
@@ -149,7 +158,7 @@ def get_style_model_and_losses(
):
# content_layers = ['conv_4']
# style_layers = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']
content_layers = ['conv4_2']
content_layers = ['conv2_2', 'conv3_2', 'conv4_2', 'conv5_2']
style_layers = ['conv2_1', 'conv3_1', 'conv4_1']
cnn = models.vgg19(pretrained=True).features.to(device, dtype=dtype).eval()
# set all weights in the model to our dtype

6
toolkit/train_tools.py
View File

@@ -479,6 +479,7 @@ def encode_prompts_xl(
tokenizers: list['CLIPTokenizer'],
text_encoders: list[Union['CLIPTextModel', 'CLIPTextModelWithProjection']],
prompts: list[str],
prompts2: Union[list[str], None],
num_images_per_prompt: int = 1,
use_text_encoder_1: bool = True, # sdxl
use_text_encoder_2: bool = True # sdxl
@@ -486,11 +487,13 @@ def encode_prompts_xl(
# text_encoder and text_encoder_2's penuultimate layer's output
text_embeds_list = []
pooled_text_embeds = None # always text_encoder_2's pool
if prompts2 is None:
prompts2 = prompts
for idx, (tokenizer, text_encoder) in enumerate(zip(tokenizers, text_encoders)):
# todo, we are using a blank string to ignore that encoder for now.
# find a better way to do this (zeroing?, removing it from the unet?)
prompt_list_to_use = prompts
prompt_list_to_use = prompts if idx == 0 else prompts2
if idx == 0 and not use_text_encoder_1:
prompt_list_to_use = ["" for _ in prompts]
if idx == 1 and not use_text_encoder_2:
@@ -515,6 +518,7 @@ def text_encode(text_encoder: 'CLIPTextModel', tokens):
return text_encoder(tokens.to(text_encoder.device))[0]
def encode_prompts(
tokenizer: 'CLIPTokenizer',
text_encoder: 'CLIPTokenizer',