ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-30 19:21:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Allow for a default caption file in the folder. Minor bug fixes.

Browse Source
This commit is contained in:
Jaret Burkett
2024-10-10 07:31:33 -06:00
parent 9452929300
commit ab22674980
4 changed files with 15 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -63,7 +63,7 @@ class SDTrainer(BaseSDTrainProcess):
self.is_bfloat = self.train_config.dtype == "bfloat16" or self.train_config.dtype == "bf16" self.is_bfloat = self.train_config.dtype == "bfloat16" or self.train_config.dtype == "bf16"
self.do_grad_scale = True self.do_grad_scale = True
if self.is_fine_tuning: if self.is_fine_tuning and self.is_bfloat:
self.do_grad_scale = False self.do_grad_scale = False
if self.adapter_config is not None: if self.adapter_config is not None:
if self.adapter_config.train: if self.adapter_config.train:

5
jobs/process/BaseSDTrainProcess.py
View File

@@ -1661,8 +1661,10 @@ class BaseSDTrainProcess(BaseTrainProcess):
batch_list = [] batch_list = []
for b in range(self.train_config.gradient_accumulation): for b in range(self.train_config.gradient_accumulation):
# keep track to alternate on an accumulation step for reg
batch_step = step
# don't do a reg step on sample or save steps as we dont want to normalize on those # don't do a reg step on sample or save steps as we dont want to normalize on those
if step % 2 == 0 and dataloader_reg is not None and not is_save_step and not is_sample_step: if batch_step % 2 == 0 and dataloader_reg is not None and not is_save_step and not is_sample_step:
try: try:
with self.timer('get_batch:reg'): with self.timer('get_batch:reg'):
batch = next(dataloader_iterator_reg) batch = next(dataloader_iterator_reg)
@@ -1698,6 +1700,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
else: else:
batch = None batch = None
batch_list.append(batch) batch_list.append(batch)
batch_step += 1
# setup accumulation # setup accumulation
if self.train_config.gradient_accumulation_steps == -1: if self.train_config.gradient_accumulation_steps == -1:

7
toolkit/dataloader_mixins.py
View File

@@ -121,6 +121,9 @@ class CaptionMixin:
prompt_path = path_no_ext + '.' + ext prompt_path = path_no_ext + '.' + ext
if os.path.exists(prompt_path): if os.path.exists(prompt_path):
break break
# allow folders to have a default prompt
default_prompt_path = os.path.join(os.path.dirname(img_path), 'default.txt')
if os.path.exists(prompt_path): if os.path.exists(prompt_path):
with open(prompt_path, 'r', encoding='utf-8') as f: with open(prompt_path, 'r', encoding='utf-8') as f:
@@ -131,6 +134,10 @@ class CaptionMixin:
if 'caption' in prompt: if 'caption' in prompt:
prompt = prompt['caption'] prompt = prompt['caption']
prompt = clean_caption(prompt)
elif os.path.exists(default_prompt_path):
with open(default_prompt_path, 'r', encoding='utf-8') as f:
prompt = f.read()
prompt = clean_caption(prompt) prompt = clean_caption(prompt)
else: else:
prompt = '' prompt = ''

6
toolkit/guidance.py
View File

@@ -481,9 +481,9 @@ def get_guided_loss_polarity(
loss = pred_loss + pred_neg_loss loss = pred_loss + pred_neg_loss
if sd.is_flow_matching: # if sd.is_flow_matching:
timestep_weight = sd.noise_scheduler.get_weights_for_timesteps(timesteps).to(loss.device, dtype=loss.dtype).detach() # timestep_weight = sd.noise_scheduler.get_weights_for_timesteps(timesteps).to(loss.device, dtype=loss.dtype).detach()
loss = loss * timestep_weight # loss = loss * timestep_weight
loss = loss.mean([1, 2, 3]) loss = loss.mean([1, 2, 3])