ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-01-27 00:49:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Bug fixes

Browse Source
This commit is contained in:
Jaret Burkett
2023-07-29 13:39:57 -06:00
parent 2305e55c82
commit 9cdf2dd6e4
2 changed files with 12 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
jobs/process/BaseSDTrainProcess.py
View File

@@ -282,12 +282,15 @@ class BaseSDTrainProcess(BaseTrainProcess):
# prepare meta
save_meta = get_meta_for_safetensors(self.meta, self.job.name)
if self.network is not None:
prev_multiplier = self.network.multiplier
self.network.multiplier = 1.0
# TODO handle dreambooth, fine tuning, etc
self.network.save_weights(
file_path,
dtype=get_torch_dtype(self.save_config.dtype),
metadata=save_meta
)
self.network.multiplier = prev_multiplier
else:
self.sd.save(
file_path,
@@ -639,7 +642,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
### HOOK ###
loss_dict = self.hook_train_loop()
if self.train_config.optimizer.startswith('dadaptation'):
if self.train_config.optimizer.lower().startswith('dadaptation'):
learning_rate = (
optimizer.param_groups[0]["d"] *
optimizer.param_groups[0]["lr"]

15
jobs/process/TrainSliderProcess.py
View File

@@ -115,12 +115,12 @@ class TrainSliderProcess(BaseSDTrainProcess):
cache[prompt] = self.sd.encode_prompt(prompt)
for resolution in self.slider_config.resolutions:
width, height = resolution
only_erase = len(target.positive.strip()) == 0
only_enhance = len(target.negative.strip()) == 0
erase_negative = len(target.positive.strip()) == 0
enhance_positive = len(target.negative.strip()) == 0
both = not only_erase and not only_enhance
both = not erase_negative and not enhance_positive
if only_erase and only_enhance:
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
@@ -128,7 +128,7 @@ class TrainSliderProcess(BaseSDTrainProcess):
# we only perform actions of enhancing and erasing on the negative
# todo work on way to do all of this in one shot
if both or only_erase:
if both or erase_negative:
prompt_pairs += [
# erase standard
EncodedPromptPair(
@@ -143,7 +143,7 @@ class TrainSliderProcess(BaseSDTrainProcess):
weight=target.weight
),
]
if both or only_enhance:
if both or enhance_positive:
prompt_pairs += [
# enhance standard, swap pos neg
EncodedPromptPair(
@@ -158,7 +158,7 @@ class TrainSliderProcess(BaseSDTrainProcess):
weight=target.weight
),
]
if both:
if both or enhance_positive:
prompt_pairs += [
# erase inverted
EncodedPromptPair(
@@ -173,6 +173,7 @@ class TrainSliderProcess(BaseSDTrainProcess):
weight=target.weight
),
]
if both or erase_negative:
prompt_pairs += [
# enhance inverted
EncodedPromptPair(