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

Add a merge network on save strength

Browse Source
This commit is contained in:
Jaret Burkett
2026-04-01 09:21:08 -06:00
parent 853ffaf207
commit da2a79590f
2 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
jobs/process/BaseSDTrainProcess.py
View File

@@ -632,7 +632,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
print_acc("Merging network weights into full model for saving...")
self.network.merge_in(merge_weight=1.0)
self.network.merge_in(merge_weight=self.train_config.merge_network_on_save_weight)
# reset weights to zero
self.network.reset_weights()
self.network.is_merged_in = False

1
toolkit/config_modules.py
View File

@@ -390,6 +390,7 @@ class TrainConfig:
self.gradient_checkpointing = kwargs.get('gradient_checkpointing', True)
self.weight_jitter = kwargs.get('weight_jitter', 0.0)
self.merge_network_on_save = kwargs.get('merge_network_on_save', False)
self.merge_network_on_save_strength = kwargs.get('merge_network_on_save_strength', 1.0)
self.max_grad_norm = kwargs.get('max_grad_norm', 1.0)
self.start_step = kwargs.get('start_step', None)
self.free_u = kwargs.get('free_u', False)