From da2a79590f43dc4b47fb07226d0f5f6333dc10b6 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 1 Apr 2026 09:21:08 -0600 Subject: [PATCH] Add a merge network on save strength --- jobs/process/BaseSDTrainProcess.py | 2 +- toolkit/config_modules.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index 475176a..cf7c226 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -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 diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 09315ae..8e69d17 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -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)