| 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 | |
| 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, | |
| 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) | |