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

Slider training functioning, time to perfect it

Browse Source
This commit is contained in:
Jaret Burkett
2023-07-21 22:06:49 -06:00
parent ddcd9069e1
commit 596e59dd6d
1 changed files with 130 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files

222
jobs/process/TrainSliderProcess.py
View File

@@ -3,6 +3,8 @@
import time import time
from collections import OrderedDict from collections import OrderedDict
import os import os
from typing import List
from toolkit.kohya_model_util import load_vae from toolkit.kohya_model_util import load_vae
from toolkit.lora_special import LoRASpecialNetwork from toolkit.lora_special import LoRASpecialNetwork
from toolkit.paths import REPOS_ROOT from toolkit.paths import REPOS_ROOT
@@ -99,6 +101,21 @@ class ModelConfig:
raise ValueError('name_or_path must be specified') raise ValueError('name_or_path must be specified')
class SliderTargetConfig:
def __init__(self, **kwargs):
self.target_class: str = kwargs.get('target_class', None)
self.positive: str = kwargs.get('positive', None)
self.negative: str = kwargs.get('negative', None)
class SliderConfig:
def __init__(self, **kwargs):
targets = kwargs.get('targets', [])
targets = [SliderTargetConfig(**target) for target in targets]
self.targets: List[SliderTargetConfig] = targets
self.resolutions: List[List[int]] = kwargs.get('resolutions', [[512, 512]])
class PromptSettingsOld: class PromptSettingsOld:
def __init__(self, **kwargs): def __init__(self, **kwargs):
self.target: str = kwargs.get('target', None) self.target: str = kwargs.get('target', None)
@@ -113,6 +130,24 @@ class PromptSettingsOld:
self.dynamic_crops: bool = kwargs.get('dynamic_crops', False) # default is False. only used when model is XL self.dynamic_crops: bool = kwargs.get('dynamic_crops', False) # default is False. only used when model is XL
class EncodedPromptPair:
def __init__(
self,
target_class,
positive,
negative,
neutral,
width=512,
height=512
):
self.target_class = target_class
self.positive = positive
self.negative = negative
self.neutral = neutral
self.width = width
self.height = height
class TrainSliderProcess(BaseTrainProcess): class TrainSliderProcess(BaseTrainProcess):
def __init__(self, process_id: int, job, config: OrderedDict): def __init__(self, process_id: int, job, config: OrderedDict):
super().__init__(process_id, job, config) super().__init__(process_id, job, config)
@@ -127,10 +162,9 @@ class TrainSliderProcess(BaseTrainProcess):
self.save_config = SaveConfig(**self.get_conf('save', {})) self.save_config = SaveConfig(**self.get_conf('save', {}))
self.sample_config = SampleConfig(**self.get_conf('sample', {})) self.sample_config = SampleConfig(**self.get_conf('sample', {}))
self.logging_config = LogingConfig(**self.get_conf('logging', {})) self.logging_config = LogingConfig(**self.get_conf('logging', {}))
self.slider_config = SliderConfig(**self.get_conf('slider', {}))
self.sd = None self.sd = None
self.prompt_settings = self.get_prompt_settings()
# added later # added later
self.network = None self.network = None
self.scheduler = None self.scheduler = None
@@ -142,14 +176,6 @@ class TrainSliderProcess(BaseTrainProcess):
param.data = -param.data param.data = -param.data
self.is_flipped = not self.is_flipped self.is_flipped = not self.is_flipped
def get_prompt_settings(self):
prompts = self.get_conf('prompts', required=True)
prompt_settings = [PromptSettingsOld(**prompt) for prompt in prompts]
# for i, prompt in enumerate(prompts):
# prompt_settings[i].fill_prompts(prompt)
return prompt_settings
def sample(self, step=None): def sample(self, step=None):
sample_folder = os.path.join(self.save_root, 'samples') sample_folder = os.path.join(self.save_root, 'samples')
if not os.path.exists(sample_folder): if not os.path.exists(sample_folder):
@@ -352,44 +378,38 @@ class TrainSliderProcess(BaseTrainProcess):
max_iterations=self.train_config.steps, max_iterations=self.train_config.steps,
lr_min=self.train_config.lr / 100, # not sure why leco did this, but ill do it to lr_min=self.train_config.lr / 100, # not sure why leco did this, but ill do it to
) )
criteria = torch.nn.MSELoss() loss_function = torch.nn.MSELoss()
if self.logging_config.verbose:
print("Prompts")
for settings in self.prompt_settings:
print(settings)
# debug
# debug_util.check_requires_grad(network)
# debug_util.check_training_mode(network)
cache = PromptEmbedsCache() cache = PromptEmbedsCache()
prompt_pairs: list[PromptEmbedsPair] = [] prompt_pairs: list[LatentPair] = []
# get encoded latents for our prompts
with torch.no_grad(): with torch.no_grad():
for settings in self.prompt_settings: neutral = ""
self.print(settings) for target in self.slider_config.targets:
for prompt in [ for resolution in self.slider_config.resolutions:
settings.target, width, height = resolution
settings.positive, for prompt in [
settings.neutral, target.target_class,
settings.unconditional, target.positive,
]: target.negative,
if cache[prompt] == None: neutral # empty neutral
cache[prompt] = train_util.encode_prompts( ]:
tokenizer, text_encoder, [prompt] if cache[prompt] == None:
) cache[prompt] = train_util.encode_prompts(
tokenizer, text_encoder, [prompt]
)
prompt_pairs.append( prompt_pairs.append(
PromptEmbedsPair( EncodedPromptPair(
criteria, target_class=cache[target.target_class],
cache[settings.target], positive=cache[target.positive],
cache[settings.positive], negative=cache[target.negative],
cache[settings.unconditional], neutral=cache[neutral],
cache[settings.neutral], width=width,
settings, height=height,
)
) )
)
# move to cpu to save vram # move to cpu to save vram
# tokenizer.to("cpu") # tokenizer.to("cpu")
@@ -400,7 +420,6 @@ class TrainSliderProcess(BaseTrainProcess):
self.print("Generating baseline samples before training") self.print("Generating baseline samples before training")
self.sample(0) self.sample(0)
self.progress_bar = tqdm(range(self.train_config.steps))
self.progress_bar = tqdm( self.progress_bar = tqdm(
total=self.train_config.steps, total=self.train_config.steps,
desc=self.job.name, desc=self.job.name,
@@ -408,6 +427,29 @@ class TrainSliderProcess(BaseTrainProcess):
) )
self.step_num = 0 self.step_num = 0
for step in range(self.train_config.steps): for step in range(self.train_config.steps):
# get a random pair
prompt_pair: EncodedPromptPair = prompt_pairs[
torch.randint(0, len(prompt_pairs), (1,)).item()
]
height = prompt_pair.height
width = prompt_pair.width
positive = prompt_pair.positive
target_class = prompt_pair.target_class
neutral = prompt_pair.neutral
negative = prompt_pair.negative
# swap every other step and invert lora to spread slider
do_swap = step % 2 == 0
if do_swap:
negative = prompt_pair.positive
positive = prompt_pair.negative
# set the network in a negative weight
self.network.multiplier = -1.0
with torch.no_grad(): with torch.no_grad():
noise_scheduler.set_timesteps( noise_scheduler.set_timesteps(
self.train_config.max_denoising_steps, device=self.device_torch self.train_config.max_denoising_steps, device=self.device_torch
@@ -415,34 +457,17 @@ class TrainSliderProcess(BaseTrainProcess):
optimizer.zero_grad() optimizer.zero_grad()
prompt_pair: PromptEmbedsPair = prompt_pairs[ # ger a random number of steps
torch.randint(0, len(prompt_pairs), (1,)).item()
]
# 1 ~ 49 random from 1 to 49
timesteps_to = torch.randint( timesteps_to = torch.randint(
1, self.train_config.max_denoising_steps, (1,) 1, self.train_config.max_denoising_steps, (1,)
).item() ).item()
height, width = (
prompt_pair.resolution,
prompt_pair.resolution,
)
if prompt_pair.dynamic_resolution:
height, width = train_util.get_random_resolution_in_bucket(
prompt_pair.resolution
)
if self.logging_config.verbose:
self.print("guidance_scale:", prompt_pair.guidance_scale)
self.print("resolution:", prompt_pair.resolution)
self.print("dynamic_resolution:", prompt_pair.dynamic_resolution)
if prompt_pair.dynamic_resolution:
self.print("bucketed resolution:", (height, width))
self.print("batch_size:", prompt_pair.batch_size)
latents = train_util.get_initial_latents( latents = train_util.get_initial_latents(
noise_scheduler, prompt_pair.batch_size, height, width, 1 noise_scheduler,
self.train_config.batch_size,
height,
width,
1
).to(self.device_torch, dtype=dtype) ).to(self.device_torch, dtype=dtype)
with self.network: with self.network:
@@ -453,9 +478,9 @@ class TrainSliderProcess(BaseTrainProcess):
noise_scheduler, noise_scheduler,
latents, # pass simple noise latents latents, # pass simple noise latents
train_util.concat_embeddings( train_util.concat_embeddings(
prompt_pair.unconditional, positive, # unconditional
prompt_pair.target, target_class, # target
prompt_pair.batch_size, self.train_config.batch_size,
), ),
start_timesteps=0, start_timesteps=0,
total_timesteps=timesteps_to, total_timesteps=timesteps_to,
@@ -468,16 +493,16 @@ class TrainSliderProcess(BaseTrainProcess):
int(timesteps_to * 1000 / self.train_config.max_denoising_steps) int(timesteps_to * 1000 / self.train_config.max_denoising_steps)
] ]
# with network: Only empty LoRA is enabled outside with network : # with network: 0 weight LoRA is enabled outside "with network:"
positive_latents = train_util.predict_noise( positive_latents = train_util.predict_noise( # positive_latents
unet, unet,
noise_scheduler, noise_scheduler,
current_timestep, current_timestep,
denoised_latents, denoised_latents,
train_util.concat_embeddings( train_util.concat_embeddings(
prompt_pair.unconditional, positive, # unconditional
prompt_pair.positive, negative, # positive
prompt_pair.batch_size, self.train_config.batch_size,
), ),
guidance_scale=1, guidance_scale=1,
).to("cpu", dtype=torch.float32) ).to("cpu", dtype=torch.float32)
@@ -487,9 +512,9 @@ class TrainSliderProcess(BaseTrainProcess):
current_timestep, current_timestep,
denoised_latents, denoised_latents,
train_util.concat_embeddings( train_util.concat_embeddings(
prompt_pair.unconditional, positive, # unconditional
prompt_pair.neutral, neutral, # neutral
prompt_pair.batch_size, self.train_config.batch_size,
), ),
guidance_scale=1, guidance_scale=1,
).to("cpu", dtype=torch.float32) ).to("cpu", dtype=torch.float32)
@@ -499,16 +524,12 @@ class TrainSliderProcess(BaseTrainProcess):
current_timestep, current_timestep,
denoised_latents, denoised_latents,
train_util.concat_embeddings( train_util.concat_embeddings(
prompt_pair.unconditional, positive, # unconditional
prompt_pair.unconditional, positive, # unconditional
prompt_pair.batch_size, self.train_config.batch_size,
), ),
guidance_scale=1, guidance_scale=1,
).to("cpu", dtype=torch.float32) ).to("cpu", dtype=torch.float32)
# if self.logging_config.verbose:
# self.print("positive_latents:", positive_latents[0, 0, :5, :5])
# self.print("neutral_latents:", neutral_latents[0, 0, :5, :5])
# self.print("unconditional_latents:", unconditional_latents[0, 0, :5, :5])
with self.network: with self.network:
target_latents = train_util.predict_noise( target_latents = train_util.predict_noise(
@@ -517,9 +538,9 @@ class TrainSliderProcess(BaseTrainProcess):
current_timestep, current_timestep,
denoised_latents, denoised_latents,
train_util.concat_embeddings( train_util.concat_embeddings(
prompt_pair.unconditional, positive, # unconditional
prompt_pair.target, target_class, # target
prompt_pair.batch_size, self.train_config.batch_size,
), ),
guidance_scale=1, guidance_scale=1,
).to("cpu", dtype=torch.float32) ).to("cpu", dtype=torch.float32)
@@ -531,12 +552,23 @@ class TrainSliderProcess(BaseTrainProcess):
neutral_latents.requires_grad = False neutral_latents.requires_grad = False
unconditional_latents.requires_grad = False unconditional_latents.requires_grad = False
loss = prompt_pair.loss( erase = True
target_latents=target_latents, guidance_scale = 1.0
positive_latents=positive_latents,
neutral_latents=neutral_latents, offset = guidance_scale * (positive_latents - unconditional_latents)
unconditional_latents=unconditional_latents,
offset_neutral = neutral_latents
if erase:
offset_neutral -= offset
else:
# enhance
offset_neutral += offset
loss = loss_function(
target_latents,
offset_neutral,
) )
loss_float = loss.item() loss_float = loss.item()
if self.train_config.optimizer.startswith('dadaptation'): if self.train_config.optimizer.startswith('dadaptation'):
learning_rate = ( learning_rate = (
@@ -561,6 +593,9 @@ class TrainSliderProcess(BaseTrainProcess):
) )
flush() flush()
# reset network
self.network.multiplier = 1.0
# don't do on first step # don't do on first step
if self.step_num != self.start_step: if self.step_num != self.start_step:
# pause progress bar # pause progress bar
@@ -594,8 +629,11 @@ class TrainSliderProcess(BaseTrainProcess):
# end of step # end of step
self.step_num = step self.step_num = step
print("")
self.save() self.save()
del ( del (
unet, unet,
noise_scheduler, noise_scheduler,