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)