mirror of
https://github.com/ostris/ai-toolkit.git
synced 2026-04-30 03:01:28 +00:00
Complete reqork of how slider training works and optimized it to hell. Can run entire algorythm in 1 batch now with less VRAM consumption than a quarter of it used to take
This commit is contained in:
Jaret Burkett
@@ -3,12 +3,14 @@
|
||||
import random
|
||||
from collections import OrderedDict
|
||||
import os
|
||||
from typing import Optional
|
||||
from typing import Optional, Union
|
||||
|
||||
from safetensors.torch import save_file, load_file
|
||||
import torch.utils.checkpoint as cp
|
||||
from tqdm import tqdm
|
||||
|
||||
from toolkit.config_modules import SliderConfig
|
||||
from toolkit.layers import CheckpointGradients
|
||||
from toolkit.paths import REPOS_ROOT
|
||||
import sys
|
||||
|
||||
@@ -16,88 +18,21 @@ from toolkit.stable_diffusion_model import PromptEmbeds
|
||||
from toolkit.train_tools import get_torch_dtype
|
||||
import gc
|
||||
from toolkit import train_tools
|
||||
from toolkit.prompt_utils import \
|
||||
EncodedPromptPair, ACTION_TYPES_SLIDER, \
|
||||
EncodedAnchor, concat_prompt_pairs, \
|
||||
concat_anchors, PromptEmbedsCache, encode_prompts_to_cache, build_prompt_pair_batch_from_cache, split_anchors, \
|
||||
split_prompt_pairs
|
||||
|
||||
import torch
|
||||
from .BaseSDTrainProcess import BaseSDTrainProcess
|
||||
|
||||
|
||||
class ACTION_TYPES_SLIDER:
|
||||
ERASE_NEGATIVE = 0
|
||||
ENHANCE_NEGATIVE = 1
|
||||
|
||||
|
||||
def flush():
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
|
||||
class EncodedPromptPair:
|
||||
def __init__(
|
||||
self,
|
||||
target_class,
|
||||
target_class_with_neutral,
|
||||
positive_target,
|
||||
positive_target_with_neutral,
|
||||
negative_target,
|
||||
negative_target_with_neutral,
|
||||
neutral,
|
||||
empty_prompt,
|
||||
both_targets,
|
||||
action=ACTION_TYPES_SLIDER.ERASE_NEGATIVE,
|
||||
multiplier=1.0,
|
||||
weight=1.0
|
||||
):
|
||||
self.target_class = target_class
|
||||
self.target_class_with_neutral = target_class_with_neutral
|
||||
self.positive_target = positive_target
|
||||
self.positive_target_with_neutral = positive_target_with_neutral
|
||||
self.negative_target = negative_target
|
||||
self.negative_target_with_neutral = negative_target_with_neutral
|
||||
self.neutral = neutral
|
||||
self.empty_prompt = empty_prompt
|
||||
self.both_targets = both_targets
|
||||
self.multiplier = multiplier
|
||||
self.action: int = action
|
||||
self.weight = weight
|
||||
|
||||
# simulate torch to for tensors
|
||||
def to(self, *args, **kwargs):
|
||||
self.target_class = self.target_class.to(*args, **kwargs)
|
||||
self.positive_target = self.positive_target.to(*args, **kwargs)
|
||||
self.positive_target_with_neutral = self.positive_target_with_neutral.to(*args, **kwargs)
|
||||
self.negative_target = self.negative_target.to(*args, **kwargs)
|
||||
self.negative_target_with_neutral = self.negative_target_with_neutral.to(*args, **kwargs)
|
||||
self.neutral = self.neutral.to(*args, **kwargs)
|
||||
self.empty_prompt = self.empty_prompt.to(*args, **kwargs)
|
||||
self.both_targets = self.both_targets.to(*args, **kwargs)
|
||||
return self
|
||||
|
||||
|
||||
class PromptEmbedsCache:
|
||||
prompts: dict[str, PromptEmbeds] = {}
|
||||
|
||||
def __setitem__(self, __name: str, __value: PromptEmbeds) -> None:
|
||||
self.prompts[__name] = __value
|
||||
|
||||
def __getitem__(self, __name: str) -> Optional[PromptEmbeds]:
|
||||
if __name in self.prompts:
|
||||
return self.prompts[__name]
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
class EncodedAnchor:
|
||||
def __init__(
|
||||
self,
|
||||
prompt,
|
||||
neg_prompt,
|
||||
multiplier=1.0
|
||||
):
|
||||
self.prompt = prompt
|
||||
self.neg_prompt = neg_prompt
|
||||
self.multiplier = multiplier
|
||||
|
||||
|
||||
class TrainSliderProcess(BaseSDTrainProcess):
|
||||
def __init__(self, process_id: int, job, config: OrderedDict):
|
||||
super().__init__(process_id, job, config)
|
||||
@@ -110,6 +45,8 @@ class TrainSliderProcess(BaseSDTrainProcess):
|
||||
self.prompt_cache = PromptEmbedsCache()
|
||||
self.prompt_pairs: list[EncodedPromptPair] = []
|
||||
self.anchor_pairs: list[EncodedAnchor] = []
|
||||
# keep track of prompt chunk size
|
||||
self.prompt_chunk_size = 1
|
||||
|
||||
def before_model_load(self):
|
||||
pass
|
||||
@@ -137,163 +74,57 @@ class TrainSliderProcess(BaseSDTrainProcess):
|
||||
|
||||
# get encoded latents for our prompts
|
||||
with torch.no_grad():
|
||||
if self.slider_config.prompt_tensors is not None:
|
||||
# check to see if it exists
|
||||
if os.path.exists(self.slider_config.prompt_tensors):
|
||||
# load it.
|
||||
self.print(f"Loading prompt tensors from {self.slider_config.prompt_tensors}")
|
||||
prompt_tensors = load_file(self.slider_config.prompt_tensors, device='cpu')
|
||||
# add them to the cache
|
||||
for prompt_txt, prompt_tensor in tqdm(prompt_tensors.items(), desc="Loading prompts", leave=False):
|
||||
if prompt_txt.startswith("te:"):
|
||||
prompt = prompt_txt[3:]
|
||||
# text_embeds
|
||||
text_embeds = prompt_tensor
|
||||
pooled_embeds = None
|
||||
# find pool embeds
|
||||
if f"pe:{prompt}" in prompt_tensors:
|
||||
pooled_embeds = prompt_tensors[f"pe:{prompt}"]
|
||||
# list of neutrals. Can come from file or be empty
|
||||
neutral_list = self.prompt_txt_list if self.prompt_txt_list is not None else [""]
|
||||
|
||||
# make it
|
||||
prompt_embeds = PromptEmbeds([text_embeds, pooled_embeds])
|
||||
cache[prompt] = prompt_embeds.to(device='cpu', dtype=torch.float32)
|
||||
# build the prompts to cache
|
||||
prompts_to_cache = []
|
||||
for neutral in neutral_list:
|
||||
for target in self.slider_config.targets:
|
||||
prompt_list = [
|
||||
f"{target.target_class}", # target_class
|
||||
f"{target.target_class} {neutral}", # target_class with neutral
|
||||
f"{target.positive}", # positive_target
|
||||
f"{target.positive} {neutral}", # positive_target with neutral
|
||||
f"{target.negative}", # negative_target
|
||||
f"{target.negative} {neutral}", # negative_target with neutral
|
||||
f"{neutral}", # neutral
|
||||
f"{target.positive} {target.negative}", # both targets
|
||||
f"{target.negative} {target.positive}", # both targets reverse
|
||||
]
|
||||
prompts_to_cache += prompt_list
|
||||
|
||||
if len(cache.prompts) == 0:
|
||||
print("Prompt tensors not found. Encoding prompts..")
|
||||
empty_prompt = ""
|
||||
# encode empty_prompt
|
||||
cache[empty_prompt] = self.sd.encode_prompt(empty_prompt)
|
||||
# remove duplicates
|
||||
prompts_to_cache = list(dict.fromkeys(prompts_to_cache))
|
||||
|
||||
neutral_list = self.prompt_txt_list if self.prompt_txt_list is not None else [""]
|
||||
|
||||
for neutral in tqdm(neutral_list, desc="Encoding prompts", leave=False):
|
||||
for target in self.slider_config.targets:
|
||||
prompt_list = [
|
||||
f"{target.target_class}", # target_class
|
||||
f"{target.target_class} {neutral}", # target_class with neutral
|
||||
f"{target.positive}", # positive_target
|
||||
f"{target.positive} {neutral}", # positive_target with neutral
|
||||
f"{target.negative}", # negative_target
|
||||
f"{target.negative} {neutral}", # negative_target with neutral
|
||||
f"{neutral}", # neutral
|
||||
f"{target.positive} {target.negative}", # both targets
|
||||
f"{target.negative} {target.positive}", # both targets
|
||||
]
|
||||
for p in prompt_list:
|
||||
# build the cache
|
||||
if cache[p] is None:
|
||||
cache[p] = self.sd.encode_prompt(p).to(device="cpu", dtype=torch.float32)
|
||||
|
||||
erase_negative = len(target.positive.strip()) == 0
|
||||
enhance_positive = len(target.negative.strip()) == 0
|
||||
|
||||
both = not erase_negative and not enhance_positive
|
||||
|
||||
if erase_negative and enhance_positive:
|
||||
raise ValueError("target must have at least one of positive or negative or both")
|
||||
# for slider we need to have an enhancer, an eraser, and then
|
||||
# an inverse with negative weights to balance the network
|
||||
# if we don't do this, we will get different contrast and focus.
|
||||
# we only perform actions of enhancing and erasing on the negative
|
||||
# todo work on way to do all of this in one shot
|
||||
if self.slider_config.prompt_tensors:
|
||||
print(f"Saving prompt tensors to {self.slider_config.prompt_tensors}")
|
||||
state_dict = {}
|
||||
for prompt_txt, prompt_embeds in cache.prompts.items():
|
||||
state_dict[f"te:{prompt_txt}"] = prompt_embeds.text_embeds.to("cpu",
|
||||
dtype=get_torch_dtype('fp16'))
|
||||
if prompt_embeds.pooled_embeds is not None:
|
||||
state_dict[f"pe:{prompt_txt}"] = prompt_embeds.pooled_embeds.to("cpu",
|
||||
dtype=get_torch_dtype(
|
||||
'fp16'))
|
||||
save_file(state_dict, self.slider_config.prompt_tensors)
|
||||
# encode them
|
||||
cache = encode_prompts_to_cache(
|
||||
prompt_list=prompts_to_cache,
|
||||
sd=self.sd,
|
||||
cache=cache,
|
||||
prompt_tensor_file=self.slider_config.prompt_tensors
|
||||
)
|
||||
|
||||
prompt_pairs = []
|
||||
for neutral in tqdm(neutral_list, desc="Encoding prompts", leave=False):
|
||||
prompt_batches = []
|
||||
for neutral in tqdm(neutral_list, desc="Building Prompt Pairs", leave=False):
|
||||
for target in self.slider_config.targets:
|
||||
erase_negative = len(target.positive.strip()) == 0
|
||||
enhance_positive = len(target.negative.strip()) == 0
|
||||
prompt_pair_batch = build_prompt_pair_batch_from_cache(
|
||||
cache=cache,
|
||||
target=target,
|
||||
neutral=neutral,
|
||||
|
||||
both = not erase_negative and not enhance_positive
|
||||
|
||||
if both or erase_negative:
|
||||
print("Encoding erase negative")
|
||||
prompt_pairs += [
|
||||
# erase standard
|
||||
EncodedPromptPair(
|
||||
target_class=cache[target.target_class],
|
||||
target_class_with_neutral=cache[f"{target.target_class} {neutral}"],
|
||||
positive_target=cache[f"{target.positive}"],
|
||||
positive_target_with_neutral=cache[f"{target.positive} {neutral}"],
|
||||
negative_target=cache[f"{target.negative}"],
|
||||
negative_target_with_neutral=cache[f"{target.negative} {neutral}"],
|
||||
neutral=cache[neutral],
|
||||
action=ACTION_TYPES_SLIDER.ERASE_NEGATIVE,
|
||||
multiplier=target.multiplier,
|
||||
both_targets=cache[f"{target.positive} {target.negative}"],
|
||||
empty_prompt=cache[""],
|
||||
weight=target.weight
|
||||
),
|
||||
]
|
||||
if both or enhance_positive:
|
||||
print("Encoding enhance positive")
|
||||
prompt_pairs += [
|
||||
# enhance standard, swap pos neg
|
||||
EncodedPromptPair(
|
||||
target_class=cache[target.target_class],
|
||||
target_class_with_neutral=cache[f"{target.target_class} {neutral}"],
|
||||
positive_target=cache[f"{target.negative}"],
|
||||
positive_target_with_neutral=cache[f"{target.negative} {neutral}"],
|
||||
negative_target=cache[f"{target.positive}"],
|
||||
negative_target_with_neutral=cache[f"{target.positive} {neutral}"],
|
||||
neutral=cache[neutral],
|
||||
action=ACTION_TYPES_SLIDER.ENHANCE_NEGATIVE,
|
||||
multiplier=target.multiplier,
|
||||
both_targets=cache[f"{target.positive} {target.negative}"],
|
||||
empty_prompt=cache[""],
|
||||
weight=target.weight
|
||||
),
|
||||
]
|
||||
# if both or enhance_positive:
|
||||
if both:
|
||||
print("Encoding erase positive (inverse)")
|
||||
prompt_pairs += [
|
||||
# erase inverted
|
||||
EncodedPromptPair(
|
||||
target_class=cache[target.target_class],
|
||||
target_class_with_neutral=cache[f"{target.target_class} {neutral}"],
|
||||
positive_target=cache[f"{target.negative}"],
|
||||
positive_target_with_neutral=cache[f"{target.negative} {neutral}"],
|
||||
negative_target=cache[f"{target.positive}"],
|
||||
negative_target_with_neutral=cache[f"{target.positive} {neutral}"],
|
||||
neutral=cache[neutral],
|
||||
action=ACTION_TYPES_SLIDER.ERASE_NEGATIVE,
|
||||
both_targets=cache[f"{target.positive} {target.negative}"],
|
||||
empty_prompt=cache[""],
|
||||
multiplier=target.multiplier * -1.0,
|
||||
weight=target.weight
|
||||
),
|
||||
]
|
||||
# if both or erase_negative:
|
||||
if both:
|
||||
print("Encoding enhance negative (inverse)")
|
||||
prompt_pairs += [
|
||||
# enhance inverted
|
||||
EncodedPromptPair(
|
||||
target_class=cache[target.target_class],
|
||||
target_class_with_neutral=cache[f"{target.target_class} {neutral}"],
|
||||
positive_target=cache[f"{target.positive}"],
|
||||
positive_target_with_neutral=cache[f"{target.positive} {neutral}"],
|
||||
negative_target=cache[f"{target.negative}"],
|
||||
negative_target_with_neutral=cache[f"{target.negative} {neutral}"],
|
||||
both_targets=cache[f"{target.positive} {target.negative}"],
|
||||
neutral=cache[neutral],
|
||||
action=ACTION_TYPES_SLIDER.ENHANCE_NEGATIVE,
|
||||
empty_prompt=cache[""],
|
||||
multiplier=target.multiplier * -1.0,
|
||||
weight=target.weight
|
||||
),
|
||||
]
|
||||
)
|
||||
if self.slider_config.batch_full_slide:
|
||||
# concat the prompt pairs
|
||||
# this allows us to run the entire 4 part process in one shot (for slider)
|
||||
self.prompt_chunk_size = 4
|
||||
concat_prompt_pair_batch = concat_prompt_pairs(prompt_pair_batch).to('cpu')
|
||||
prompt_pairs += [concat_prompt_pair_batch]
|
||||
else:
|
||||
self.prompt_chunk_size = 1
|
||||
# do them one at a time (probably not necessary after new optimizations)
|
||||
prompt_pairs += [x.to('cpu') for x in prompt_pair_batch]
|
||||
|
||||
# setup anchors
|
||||
anchor_pairs = []
|
||||
@@ -306,13 +137,26 @@ class TrainSliderProcess(BaseSDTrainProcess):
|
||||
if cache[prompt] == None:
|
||||
cache[prompt] = self.sd.encode_prompt(prompt)
|
||||
|
||||
anchor_batch = []
|
||||
# we get the prompt pair multiplier from first prompt pair
|
||||
# since they are all the same. We need to match their network polarity
|
||||
prompt_pair_multipliers = prompt_pairs[0].multiplier_list
|
||||
for prompt_multiplier in prompt_pair_multipliers:
|
||||
# match the network multiplier polarity
|
||||
anchor_scalar = 1.0 if prompt_multiplier > 0 else -1.0
|
||||
anchor_batch += [
|
||||
EncodedAnchor(
|
||||
prompt=cache[anchor.prompt],
|
||||
neg_prompt=cache[anchor.neg_prompt],
|
||||
multiplier=anchor.multiplier * anchor_scalar
|
||||
)
|
||||
]
|
||||
|
||||
anchor_pairs += [
|
||||
EncodedAnchor(
|
||||
prompt=cache[anchor.prompt],
|
||||
neg_prompt=cache[anchor.neg_prompt],
|
||||
multiplier=anchor.multiplier
|
||||
)
|
||||
concat_anchors(anchor_batch).to('cpu')
|
||||
]
|
||||
if len(anchor_pairs) > 0:
|
||||
self.anchor_pairs = anchor_pairs
|
||||
|
||||
# move to cpu to save vram
|
||||
# We don't need text encoder anymore, but keep it on cpu for sampling
|
||||
@@ -324,17 +168,13 @@ class TrainSliderProcess(BaseSDTrainProcess):
|
||||
self.sd.text_encoder.to("cpu")
|
||||
self.prompt_cache = cache
|
||||
self.prompt_pairs = prompt_pairs
|
||||
self.anchor_pairs = anchor_pairs
|
||||
# self.anchor_pairs = anchor_pairs
|
||||
flush()
|
||||
# end hook_before_train_loop
|
||||
|
||||
def hook_train_loop(self):
|
||||
dtype = get_torch_dtype(self.train_config.dtype)
|
||||
|
||||
# get random multiplier between 1 and 3
|
||||
rand_weight = 1
|
||||
# rand_weight = torch.rand((1,)).item() * 2 + 1
|
||||
|
||||
# get a random pair
|
||||
prompt_pair: EncodedPromptPair = self.prompt_pairs[
|
||||
torch.randint(0, len(self.prompt_pairs), (1,)).item()
|
||||
@@ -346,11 +186,10 @@ class TrainSliderProcess(BaseSDTrainProcess):
|
||||
height, width = self.slider_config.resolutions[
|
||||
torch.randint(0, len(self.slider_config.resolutions), (1,)).item()
|
||||
]
|
||||
if self.train_config.gradient_checkpointing:
|
||||
# may get disabled elsewhere
|
||||
self.sd.unet.enable_gradient_checkpointing()
|
||||
|
||||
weight = prompt_pair.weight
|
||||
multiplier = prompt_pair.multiplier
|
||||
|
||||
unet = self.sd.unet
|
||||
noise_scheduler = self.sd.noise_scheduler
|
||||
optimizer = self.optimizer
|
||||
lr_scheduler = self.lr_scheduler
|
||||
@@ -368,9 +207,6 @@ class TrainSliderProcess(BaseSDTrainProcess):
|
||||
guidance_scale=gs,
|
||||
)
|
||||
|
||||
# set network multiplier
|
||||
self.network.multiplier = multiplier * rand_weight
|
||||
|
||||
with torch.no_grad():
|
||||
self.sd.noise_scheduler.set_timesteps(
|
||||
self.train_config.max_denoising_steps, device=self.device_torch
|
||||
@@ -383,11 +219,14 @@ class TrainSliderProcess(BaseSDTrainProcess):
|
||||
1, self.train_config.max_denoising_steps, (1,)
|
||||
).item()
|
||||
|
||||
# for a complete slider, the batch size is 4 to begin with now
|
||||
true_batch_size = prompt_pair.target_class.text_embeds.shape[0] * self.train_config.batch_size
|
||||
|
||||
# get noise
|
||||
noise = self.sd.get_latent_noise(
|
||||
pixel_height=height,
|
||||
pixel_width=width,
|
||||
batch_size=self.train_config.batch_size,
|
||||
batch_size=true_batch_size,
|
||||
noise_offset=self.train_config.noise_offset,
|
||||
).to(self.device_torch, dtype=dtype)
|
||||
|
||||
@@ -397,7 +236,8 @@ class TrainSliderProcess(BaseSDTrainProcess):
|
||||
|
||||
with self.network:
|
||||
assert self.network.is_active
|
||||
self.network.multiplier = multiplier * rand_weight
|
||||
# pass the multiplier list to the network
|
||||
self.network.multiplier = prompt_pair.multiplier_list
|
||||
denoised_latents = self.sd.diffuse_some_steps(
|
||||
latents, # pass simple noise latents
|
||||
train_tools.concat_prompt_embeddings(
|
||||
@@ -410,19 +250,27 @@ class TrainSliderProcess(BaseSDTrainProcess):
|
||||
guidance_scale=3,
|
||||
)
|
||||
|
||||
# split the latents into out prompt pair chunks
|
||||
denoised_latent_chunks = torch.chunk(denoised_latents, self.prompt_chunk_size, dim=0)
|
||||
|
||||
noise_scheduler.set_timesteps(1000)
|
||||
|
||||
current_timestep = noise_scheduler.timesteps[
|
||||
int(timesteps_to * 1000 / self.train_config.max_denoising_steps)
|
||||
]
|
||||
|
||||
# flush() # 4.2GB to 3GB on 512x512
|
||||
|
||||
# 4.20 GB RAM for 512x512
|
||||
positive_latents = get_noise_pred(
|
||||
prompt_pair.positive_target, # negative prompt
|
||||
prompt_pair.negative_target, # positive prompt
|
||||
1,
|
||||
current_timestep,
|
||||
denoised_latents
|
||||
).to("cpu", dtype=torch.float32)
|
||||
)
|
||||
positive_latents.requires_grad = False
|
||||
positive_latents_chunks = torch.chunk(positive_latents, self.prompt_chunk_size, dim=0)
|
||||
|
||||
neutral_latents = get_noise_pred(
|
||||
prompt_pair.positive_target, # negative prompt
|
||||
@@ -430,7 +278,9 @@ class TrainSliderProcess(BaseSDTrainProcess):
|
||||
1,
|
||||
current_timestep,
|
||||
denoised_latents
|
||||
).to("cpu", dtype=torch.float32)
|
||||
)
|
||||
neutral_latents.requires_grad = False
|
||||
neutral_latents_chunks = torch.chunk(neutral_latents, self.prompt_chunk_size, dim=0)
|
||||
|
||||
unconditional_latents = get_noise_pred(
|
||||
prompt_pair.positive_target, # negative prompt
|
||||
@@ -438,87 +288,142 @@ class TrainSliderProcess(BaseSDTrainProcess):
|
||||
1,
|
||||
current_timestep,
|
||||
denoised_latents
|
||||
).to("cpu", dtype=torch.float32)
|
||||
|
||||
anchor_loss = None
|
||||
if len(self.anchor_pairs) > 0:
|
||||
# get a random anchor pair
|
||||
anchor: EncodedAnchor = self.anchor_pairs[
|
||||
torch.randint(0, len(self.anchor_pairs), (1,)).item()
|
||||
]
|
||||
with torch.no_grad():
|
||||
anchor_target_noise = get_noise_pred(
|
||||
anchor.prompt, anchor.neg_prompt, 1, current_timestep, denoised_latents
|
||||
).to("cpu", dtype=torch.float32)
|
||||
with self.network:
|
||||
# anchor whatever weight prompt pair is using
|
||||
pos_nem_mult = 1.0 if prompt_pair.multiplier > 0 else -1.0
|
||||
self.network.multiplier = anchor.multiplier * pos_nem_mult * rand_weight
|
||||
|
||||
anchor_pred_noise = get_noise_pred(
|
||||
anchor.prompt, anchor.neg_prompt, 1, current_timestep, denoised_latents
|
||||
).to("cpu", dtype=torch.float32)
|
||||
|
||||
self.network.multiplier = prompt_pair.multiplier * rand_weight
|
||||
|
||||
with self.network:
|
||||
self.network.multiplier = prompt_pair.multiplier * rand_weight
|
||||
target_latents = get_noise_pred(
|
||||
prompt_pair.positive_target,
|
||||
prompt_pair.target_class,
|
||||
1,
|
||||
current_timestep,
|
||||
denoised_latents
|
||||
).to("cpu", dtype=torch.float32)
|
||||
|
||||
# if self.logging_config.verbose:
|
||||
# self.print("target_latents:", target_latents[0, 0, :5, :5])
|
||||
|
||||
positive_latents.requires_grad = False
|
||||
neutral_latents.requires_grad = False
|
||||
unconditional_latents.requires_grad = False
|
||||
if len(self.anchor_pairs) > 0:
|
||||
anchor_target_noise.requires_grad = False
|
||||
anchor_loss = loss_function(
|
||||
anchor_target_noise,
|
||||
anchor_pred_noise,
|
||||
)
|
||||
erase = prompt_pair.action == ACTION_TYPES_SLIDER.ERASE_NEGATIVE
|
||||
guidance_scale = 1.0
|
||||
unconditional_latents.requires_grad = False
|
||||
unconditional_latents_chunks = torch.chunk(unconditional_latents, self.prompt_chunk_size, dim=0)
|
||||
|
||||
offset = guidance_scale * (positive_latents - unconditional_latents)
|
||||
flush() # 4.2GB to 3GB on 512x512
|
||||
|
||||
offset_neutral = neutral_latents
|
||||
if erase:
|
||||
offset_neutral -= offset
|
||||
else:
|
||||
# enhance
|
||||
offset_neutral += offset
|
||||
# 4.20 GB RAM for 512x512
|
||||
anchor_loss_float = None
|
||||
if len(self.anchor_pairs) > 0:
|
||||
with torch.no_grad():
|
||||
# get a random anchor pair
|
||||
anchor: EncodedAnchor = self.anchor_pairs[
|
||||
torch.randint(0, len(self.anchor_pairs), (1,)).item()
|
||||
]
|
||||
anchor.to(self.device_torch, dtype=dtype)
|
||||
|
||||
loss = loss_function(
|
||||
target_latents,
|
||||
offset_neutral,
|
||||
) * weight
|
||||
# first we get the target prediction without network active
|
||||
anchor_target_noise = get_noise_pred(
|
||||
anchor.neg_prompt, anchor.prompt, 1, current_timestep, denoised_latents
|
||||
# ).to("cpu", dtype=torch.float32)
|
||||
).requires_grad_(False)
|
||||
|
||||
loss_slide = loss.item()
|
||||
# to save vram, we will run these through separately while tracking grads
|
||||
# otherwise it consumes a ton of vram and this isn't our speed bottleneck
|
||||
anchor_chunks = split_anchors(anchor, self.prompt_chunk_size)
|
||||
anchor_target_noise_chunks = torch.chunk(anchor_target_noise, self.prompt_chunk_size, dim=0)
|
||||
assert len(anchor_chunks) == len(denoised_latent_chunks)
|
||||
|
||||
if anchor_loss is not None:
|
||||
loss += anchor_loss
|
||||
# 4.32 GB RAM for 512x512
|
||||
with self.network:
|
||||
assert self.network.is_active
|
||||
anchor_float_losses = []
|
||||
for anchor_chunk, denoised_latent_chunk, anchor_target_noise_chunk in zip(
|
||||
anchor_chunks, denoised_latent_chunks, anchor_target_noise_chunks
|
||||
):
|
||||
self.network.multiplier = anchor_chunk.multiplier_list
|
||||
|
||||
loss_float = loss.item()
|
||||
anchor_pred_noise = get_noise_pred(
|
||||
anchor_chunk.neg_prompt, anchor_chunk.prompt, 1, current_timestep, denoised_latent_chunk
|
||||
)
|
||||
# 9.42 GB RAM for 512x512 -> 4.20 GB RAM for 512x512 with new grad_checkpointing
|
||||
anchor_loss = loss_function(
|
||||
anchor_target_noise_chunk,
|
||||
anchor_pred_noise,
|
||||
)
|
||||
anchor_float_losses.append(anchor_loss.item())
|
||||
# compute anchor loss gradients
|
||||
# we will accumulate them later
|
||||
# this saves a ton of memory doing them separately
|
||||
anchor_loss.backward()
|
||||
del anchor_pred_noise
|
||||
del anchor_target_noise_chunk
|
||||
del anchor_loss
|
||||
flush()
|
||||
|
||||
loss = loss.to(self.device_torch)
|
||||
anchor_loss_float = sum(anchor_float_losses) / len(anchor_float_losses)
|
||||
del anchor_chunks
|
||||
del anchor_target_noise_chunks
|
||||
del anchor_target_noise
|
||||
# move anchor back to cpu
|
||||
anchor.to("cpu")
|
||||
flush()
|
||||
|
||||
prompt_pair_chunks = split_prompt_pairs(prompt_pair, self.prompt_chunk_size)
|
||||
assert len(prompt_pair_chunks) == len(denoised_latent_chunks)
|
||||
# 3.28 GB RAM for 512x512
|
||||
with self.network:
|
||||
assert self.network.is_active
|
||||
loss_list = []
|
||||
for prompt_pair_chunk, \
|
||||
denoised_latent_chunk, \
|
||||
positive_latents_chunk, \
|
||||
neutral_latents_chunk, \
|
||||
unconditional_latents_chunk \
|
||||
in zip(
|
||||
prompt_pair_chunks,
|
||||
denoised_latent_chunks,
|
||||
positive_latents_chunks,
|
||||
neutral_latents_chunks,
|
||||
unconditional_latents_chunks,
|
||||
):
|
||||
self.network.multiplier = prompt_pair_chunk.multiplier_list
|
||||
target_latents = get_noise_pred(
|
||||
prompt_pair_chunk.positive_target,
|
||||
prompt_pair_chunk.target_class,
|
||||
1,
|
||||
current_timestep,
|
||||
denoised_latent_chunk
|
||||
)
|
||||
|
||||
guidance_scale = 1.0
|
||||
|
||||
offset = guidance_scale * (positive_latents_chunk - unconditional_latents_chunk)
|
||||
|
||||
# make offset multiplier based on actions
|
||||
offset_multiplier_list = []
|
||||
for action in prompt_pair_chunk.action_list:
|
||||
if action == ACTION_TYPES_SLIDER.ERASE_NEGATIVE:
|
||||
offset_multiplier_list += [-1.0]
|
||||
elif action == ACTION_TYPES_SLIDER.ENHANCE_NEGATIVE:
|
||||
offset_multiplier_list += [1.0]
|
||||
|
||||
offset_multiplier = torch.tensor(offset_multiplier_list).to(offset.device, dtype=offset.dtype)
|
||||
# make offset multiplier match rank of offset
|
||||
offset_multiplier = offset_multiplier.view(offset.shape[0], 1, 1, 1)
|
||||
offset *= offset_multiplier
|
||||
|
||||
offset_neutral = neutral_latents_chunk
|
||||
# offsets are already adjusted on a per-batch basis
|
||||
offset_neutral += offset
|
||||
|
||||
# 16.15 GB RAM for 512x512 -> 4.20GB RAM for 512x512 with new grad_checkpointing
|
||||
loss = loss_function(
|
||||
target_latents,
|
||||
offset_neutral,
|
||||
) * prompt_pair_chunk.weight
|
||||
|
||||
loss.backward()
|
||||
loss_list.append(loss.item())
|
||||
del target_latents
|
||||
del offset_neutral
|
||||
del loss
|
||||
flush()
|
||||
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
lr_scheduler.step()
|
||||
|
||||
loss_float = sum(loss_list) / len(loss_list)
|
||||
if anchor_loss_float is not None:
|
||||
loss_float += anchor_loss_float
|
||||
|
||||
del (
|
||||
positive_latents,
|
||||
neutral_latents,
|
||||
unconditional_latents,
|
||||
target_latents,
|
||||
latents,
|
||||
latents
|
||||
)
|
||||
# move back to cpu
|
||||
prompt_pair.to("cpu")
|
||||
@@ -530,9 +435,9 @@ class TrainSliderProcess(BaseSDTrainProcess):
|
||||
loss_dict = OrderedDict(
|
||||
{'loss': loss_float},
|
||||
)
|
||||
if anchor_loss is not None:
|
||||
loss_dict['sl_l'] = loss_slide
|
||||
loss_dict['an_l'] = anchor_loss.item()
|
||||
if anchor_loss_float is not None:
|
||||
loss_dict['sl_l'] = loss_float
|
||||
loss_dict['an_l'] = anchor_loss_float
|
||||
|
||||
return loss_dict
|
||||
# end hook_train_loop
|
||||
|
||||
Reference in New Issue
Block a user