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

Save embeddings as their trigger to match auto and comfy style loading. Also, FINALLY found why gradients were wonkey and fixed it. The root problem is dropping out of network state before backward pass.

Browse Source
This commit is contained in:
Jaret Burkett
2023-09-09 12:02:07 -06:00
parent 408c50ead1
commit be804c9cf5
5 changed files with 64 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -69,29 +69,33 @@ class SDTrainer(BaseSDTrainProcess):
guidance_scale=1.0, guidance_scale=1.0,
) )
flush() flush()
# 9.18 gb # 9.18 gb
noise = noise.to(self.device_torch, dtype=dtype).detach() noise = noise.to(self.device_torch, dtype=dtype).detach()
if self.sd.prediction_type == 'v_prediction': if self.sd.prediction_type == 'v_prediction':
# v-parameterization training # v-parameterization training
target = self.sd.noise_scheduler.get_velocity(noisy_latents, noise, timesteps) target = self.sd.noise_scheduler.get_velocity(noisy_latents, noise, timesteps)
else: else:
target = noise target = noise
loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="none") loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="none")
loss = loss.mean([1, 2, 3]) loss = loss.mean([1, 2, 3])
if self.train_config.min_snr_gamma is not None and self.train_config.min_snr_gamma > 0.000001: if self.train_config.min_snr_gamma is not None and self.train_config.min_snr_gamma > 0.000001:
# add min_snr_gamma # add min_snr_gamma
loss = apply_snr_weight(loss, timesteps, self.sd.noise_scheduler, self.train_config.min_snr_gamma) loss = apply_snr_weight(loss, timesteps, self.sd.noise_scheduler, self.train_config.min_snr_gamma)
loss = loss.mean() loss = loss.mean()
# back propagate loss to free ram # IMPORTANT if gradient checkpointing do not leave with network when doing backward
loss.backward() # it will destroy the gradients. This is because the network is a context manager
torch.nn.utils.clip_grad_norm_(self.params, self.train_config.max_grad_norm) # and will change the multipliers back to 0.0 when exiting. They will be
flush() # 0.0 for the backward pass and the gradients will be 0.0
# I spent weeks on fighting this. DON'T DO IT
loss.backward()
torch.nn.utils.clip_grad_norm_(self.params, self.train_config.max_grad_norm)
flush()
# apply gradients # apply gradients
self.optimizer.step() self.optimizer.step()

72
jobs/process/BaseSDTrainProcess.py
View File

@@ -255,32 +255,37 @@ class BaseSDTrainProcess(BaseTrainProcess):
file_path = os.path.join(self.save_root, filename) file_path = os.path.join(self.save_root, filename)
# prepare meta # prepare meta
save_meta = get_meta_for_safetensors(self.meta, self.job.name) save_meta = get_meta_for_safetensors(self.meta, self.job.name)
if self.network is not None: if self.network is not None or self.embedding is not None:
prev_multiplier = self.network.multiplier if self.network is not None:
self.network.multiplier = 1.0 prev_multiplier = self.network.multiplier
if self.network_config.normalize: self.network.multiplier = 1.0
# apply the normalization if self.network_config.normalize:
self.network.apply_stored_normalizer() # apply the normalization
self.network.apply_stored_normalizer()
# if we are doing embedding training as well, add that # if we are doing embedding training as well, add that
embedding_dict = self.embedding.state_dict() if self.embedding else None embedding_dict = self.embedding.state_dict() if self.embedding else None
self.network.save_weights( self.network.save_weights(
file_path, file_path,
dtype=get_torch_dtype(self.save_config.dtype), dtype=get_torch_dtype(self.save_config.dtype),
metadata=save_meta, metadata=save_meta,
extra_state_dict=embedding_dict extra_state_dict=embedding_dict
) )
self.network.multiplier = prev_multiplier self.network.multiplier = prev_multiplier
# if we have an embedding as well, pair it with the network # if we have an embedding as well, pair it with the network
elif self.embedding is not None:
# for combo, above will get it # even if added to lora, still save the trigger version
# set current step if self.embedding is not None:
self.embedding.step = self.step_num emb_filename = f'{self.embed_config.trigger}{step_num}.safetensors'
# change filename to pt if that is set emb_file_path = os.path.join(self.save_root, emb_filename)
if self.embed_config.save_format == "pt": # for combo, above will get it
# replace extension # set current step
file_path = os.path.splitext(file_path)[0] + ".pt" self.embedding.step = self.step_num
self.embedding.save(file_path) # change filename to pt if that is set
if self.embed_config.save_format == "pt":
# replace extension
file_path = os.path.splitext(emb_file_path)[0] + ".pt"
self.embedding.save(file_path)
else: else:
self.sd.save( self.sd.save(
file_path, file_path,
@@ -316,17 +321,19 @@ class BaseSDTrainProcess(BaseTrainProcess):
# return loss # return loss
return 0.0 return 0.0
def get_latest_save_path(self): def get_latest_save_path(self, name=None):
if name == None:
name = self.job.name
# get latest saved step # get latest saved step
if os.path.exists(self.save_root): if os.path.exists(self.save_root):
latest_file = None latest_file = None
# pattern is {job_name}_{zero_filles_step}.safetensors or {job_name}.safetensors # pattern is {job_name}_{zero_filles_step}.safetensors or {job_name}.safetensors
pattern = f"{self.job.name}*.safetensors" pattern = f"{name}*.safetensors"
files = glob.glob(os.path.join(self.save_root, pattern)) files = glob.glob(os.path.join(self.save_root, pattern))
if len(files) > 0: if len(files) > 0:
latest_file = max(files, key=os.path.getctime) latest_file = max(files, key=os.path.getctime)
# try pt # try pt
pattern = f"{self.job.name}*.pt" pattern = f"{name}*.pt"
files = glob.glob(os.path.join(self.save_root, pattern)) files = glob.glob(os.path.join(self.save_root, pattern))
if len(files) > 0: if len(files) > 0:
latest_file = max(files, key=os.path.getctime) latest_file = max(files, key=os.path.getctime)
@@ -548,9 +555,12 @@ class BaseSDTrainProcess(BaseTrainProcess):
# we are doing embedding training as well # we are doing embedding training as well
self.embedding = Embedding( self.embedding = Embedding(
sd=self.sd, sd=self.sd,
embed_config=self.embed_config, embed_config=self.embed_config
state_dict=extra_weights
) )
latest_save_path = self.get_latest_save_path(self.embed_config.trigger)
# load last saved weights
if latest_save_path is not None:
self.embedding.load_embedding_from_file(latest_save_path, self.device_torch)
params.append({ params.append({
'params': self.embedding.get_trainable_params(), 'params': self.embedding.get_trainable_params(),
'lr': self.train_config.embedding_lr 'lr': self.train_config.embedding_lr
@@ -562,7 +572,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
sd=self.sd, sd=self.sd,
embed_config=self.embed_config embed_config=self.embed_config
) )
latest_save_path = self.get_latest_save_path() latest_save_path = self.get_latest_save_path(self.embed_config.trigger)
# load last saved weights # load last saved weights
if latest_save_path is not None: if latest_save_path is not None:
self.embedding.load_embedding_from_file(latest_save_path, self.device_torch) self.embedding.load_embedding_from_file(latest_save_path, self.device_torch)

5
toolkit/lora_special.py
View File

@@ -50,10 +50,7 @@ class LoRAModule(ToolkitModuleMixin, torch.nn.Module):
**kwargs **kwargs
): ):
"""if alpha == 0 or None, alpha is rank (no scaling).""" """if alpha == 0 or None, alpha is rank (no scaling)."""
super().__init__( super().__init__()
org_module=org_module,
parent=parent
)
self.lora_name = lora_name self.lora_name = lora_name
self.scalar = torch.tensor(1.0) self.scalar = torch.tensor(1.0)

7
toolkit/lycoris_special.py
View File

@@ -36,12 +36,7 @@ class LoConSpecialModule(ToolkitModuleMixin, LoConModule):
# call super of super # call super of super
torch.nn.Module.__init__(self) torch.nn.Module.__init__(self)
# call super of # call super of
super().__init__( super().__init__(call_super_init=False)
org_module=org_module,
call_super_init=False,
parent=parent,
**kwargs
)
self.lora_name = lora_name self.lora_name = lora_name
self.lora_dim = lora_dim self.lora_dim = lora_dim
self.cp = False self.cp = False

13
toolkit/network_mixins.py
View File

@@ -52,18 +52,9 @@ class ToolkitModuleMixin:
): ):
if call_super_init: if call_super_init:
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)
self.tk_orig_module: torch.nn.Module = kwargs.get('org_module', None)
self.tk_orig_parent = kwargs.get('parent', None)
self.is_checkpointing = False self.is_checkpointing = False
self.is_normalizing = False self.is_normalizing = False
self.normalize_scaler = 1.0 self.normalize_scaler = 1.0
# see if is conv or linear
self.is_conv = False
self.is_linear = False
if self.tk_orig_module.__class__.__name__ in LINEAR_MODULES:
self.is_linear = True
elif self.tk_orig_module.__class__.__name__ in CONV_MODULES:
self.is_conv = True
self._multiplier: Union[float, list, torch.Tensor] = 1.0 self._multiplier: Union[float, list, torch.Tensor] = 1.0
# this allows us to set different multipliers on a per item in a batch basis # this allows us to set different multipliers on a per item in a batch basis
@@ -140,10 +131,6 @@ class ToolkitModuleMixin:
lora_output_batch_size = lora_output.size(0) lora_output_batch_size = lora_output.size(0)
multiplier_batch_size = multiplier.size(0) multiplier_batch_size = multiplier.size(0)
if lora_output_batch_size != multiplier_batch_size: if lora_output_batch_size != multiplier_batch_size:
print(
f"Warning: lora_output_batch_size {lora_output_batch_size} != multiplier_batch_size {multiplier_batch_size}")
# doing cfg
# should be 1 for if total batch size was 1
num_interleaves = (lora_output_batch_size // 2) // multiplier_batch_size num_interleaves = (lora_output_batch_size // 2) // multiplier_batch_size
multiplier = multiplier.repeat_interleave(num_interleaves) multiplier = multiplier.repeat_interleave(num_interleaves)
# multiplier = 1.0 # multiplier = 1.0