From edd860a4100f999c73fa6860a9f73a441d3d76c9 Mon Sep 17 00:00:00 2001
From: BitcrushedHeart <215311041+BitcrushedHeart@users.noreply.github.com>
Date: Sun, 24 May 2026 16:10:56 +0100
Subject: [PATCH] fix: don't drop epoch validation on mid-epoch resume

---
 .../modelLoader/mixin/InternalModelLoaderMixin.py    | 12 ++++++++++++
 modules/modelSaver/mixin/InternalModelSaverMixin.py  |  1 +
 modules/util/TimedActionMixin.py                     |  7 ++++++-
 modules/util/TrainProgress.py                        |  2 ++
 4 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/modules/modelLoader/mixin/InternalModelLoaderMixin.py b/modules/modelLoader/mixin/InternalModelLoaderMixin.py
index 1b9399973..21358714b 100644
--- a/modules/modelLoader/mixin/InternalModelLoaderMixin.py
+++ b/modules/modelLoader/mixin/InternalModelLoaderMixin.py
@@ -28,6 +28,18 @@ def _load_internal_data(
                     epoch_sample=meta['train_progress']['epoch_sample'],
                     global_step=meta['train_progress']['global_step'],
                 )
+                if 'last_action_epoch' in meta:
+                    train_progress.last_action_epoch = dict(meta['last_action_epoch'])
+                elif train_progress.epoch_step > 0:
+                    # Legacy backup taken mid-epoch: start-of-epoch actions
+                    # already fired in the pre-fix session. Pre-fill markers
+                    # to prevent a duplicate fire on the first resumed batch.
+                    train_progress.last_action_epoch = {
+                        'validate': train_progress.epoch,
+                        'sample': train_progress.epoch,
+                    }
+                else:
+                    train_progress.last_action_epoch = {}
 
             # optimizer
             with contextlib.suppress(FileNotFoundError):
diff --git a/modules/modelSaver/mixin/InternalModelSaverMixin.py b/modules/modelSaver/mixin/InternalModelSaverMixin.py
index 2f9513f49..99a79ad75 100644
--- a/modules/modelSaver/mixin/InternalModelSaverMixin.py
+++ b/modules/modelSaver/mixin/InternalModelSaverMixin.py
@@ -39,4 +39,5 @@ def _save_internal_data(
                     'epoch_sample': model.train_progress.epoch_sample,
                     'global_step': model.train_progress.global_step,
                 },
+                'last_action_epoch': dict(getattr(model.train_progress, 'last_action_epoch', {})),
             }, meta_file)
diff --git a/modules/util/TimedActionMixin.py b/modules/util/TimedActionMixin.py
index 370f2706b..5cf4c07b4 100644
--- a/modules/util/TimedActionMixin.py
+++ b/modules/util/TimedActionMixin.py
@@ -26,7 +26,12 @@ def repeating_action_needed(
                 if int(interval) == 0:
                     return False
                 if start_at_zero:
-                    return train_progress.epoch % int(interval) == 0 and train_progress.epoch_step == 0
+                    last = train_progress.last_action_epoch.get(name, -1)
+                    fire = train_progress.epoch > last \
+                        and train_progress.epoch % int(interval) == 0
+                    if fire:
+                        train_progress.last_action_epoch[name] = train_progress.epoch
+                    return fire
                 else:
                     # should actually be the last step of each epoch, but we don't know how many steps an epoch has
                     return train_progress.epoch % int(interval) == 0 and train_progress.epoch_step == 0 \
diff --git a/modules/util/TrainProgress.py b/modules/util/TrainProgress.py
index 7ffb79ffe..e9e5d72b3 100644
--- a/modules/util/TrainProgress.py
+++ b/modules/util/TrainProgress.py
@@ -5,11 +5,13 @@ def __init__(
             epoch_step: int = 0,
             epoch_sample: int = 0,
             global_step: int = 0,
+            last_action_epoch: dict | None = None,
     ):
         self.epoch = epoch
         self.epoch_step = epoch_step
         self.epoch_sample = epoch_sample
         self.global_step = global_step
+        self.last_action_epoch = last_action_epoch if last_action_epoch is not None else {}
 
     def next_step(self, batch_size: int):
         self.epoch_step += 1