File size: 1,158 Bytes
acfbd40 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 | diff --git a/vendor/sd-scripts/train_network.py b/vendor/sd-scripts/train_network.py
index f9073c2..fae05b8 100644
--- a/vendor/sd-scripts/train_network.py
+++ b/vendor/sd-scripts/train_network.py
@@ -1524,6 +1524,9 @@ class NetworkTrainer:
logs = {"avr_loss": avr_loss} # , "lr": lr_scheduler.get_last_lr()[0]}
progress_bar.set_postfix(**{**max_mean_logs, **logs})
+ from realstyle_control import on_step as realstyle_on_step
+ realstyle_on_step(global_step, epoch + 1, len(train_dataloader), current_loss)
+
if is_tracking:
logs = self.generate_step_logs(
args,
@@ -1717,6 +1720,11 @@ class NetworkTrainer:
progress_bar.unpause()
optimizer_train_fn()
+ from realstyle_control import epoch_end as realstyle_epoch_end
+ if realstyle_epoch_end(epoch + 1, global_step, len(train_dataloader), loss_recorder.moving_average):
+ num_train_epochs = epoch + 1
+ break
+
# end of epoch
# metadata["ss_epoch"] = str(num_train_epochs)
|