Add support for save_last_n_epochs

bmaltais · Oct 30, 2024 · 021c6f5 · 021c6f5
1 parent 695b44a
commit 021c6f5
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 0 deletions.
diff --git a/kohya_gui/class_advanced_training.py b/kohya_gui/class_advanced_training.py
@@ -188,6 +188,12 @@ def list_vae_files(path):
                 precision=0,
                 info="(Optional) Save only the specified number of states (old models will be deleted)",
             )
+            self.save_last_n_epochs = gr.Number(
+                label="Save last N epochs",
+                value=self.config.get("advanced.save_last_n_epochs", 0),
+                precision=0,
+                info="(Optional) Save only the specified number of epochs (old epochs will be deleted)",
+            )
             self.save_last_n_epochs_state = gr.Number(
                 label="Save last N epochs state",
                 value=self.config.get("advanced.save_last_n_epochs_state", 0),

diff --git a/kohya_gui/dreambooth_gui.py b/kohya_gui/dreambooth_gui.py
@@ -165,6 +165,7 @@ def save_configuration(
     save_every_n_steps,
     save_last_n_steps,
     save_last_n_steps_state,
+    save_last_n_epochs,
     save_last_n_epochs_state,
     skip_cache_check,
     log_with,
@@ -372,6 +373,7 @@ def open_configuration(
     save_every_n_steps,
     save_last_n_steps,
     save_last_n_steps_state,
+    save_last_n_epochs,
     save_last_n_epochs_state,
     skip_cache_check,
     log_with,
@@ -574,6 +576,7 @@ def train_model(
     save_every_n_steps,
     save_last_n_steps,
     save_last_n_steps_state,
+    save_last_n_epochs,
     save_last_n_epochs_state,
     skip_cache_check,
     log_with,
@@ -990,6 +993,7 @@ def train_model(
         "save_last_n_steps_state": (
             save_last_n_steps_state if save_last_n_steps_state != 0 else None
         ),
+        "save_last_n_epochs": save_last_n_epochs if save_last_n_epochs != 0 else None,
         "save_last_n_epochs_state": (
             save_last_n_epochs_state if save_last_n_epochs_state != 0 else None
         ),
@@ -1343,6 +1347,7 @@ def dreambooth_tab(
             advanced_training.save_every_n_steps,
             advanced_training.save_last_n_steps,
             advanced_training.save_last_n_steps_state,
+            advanced_training.save_last_n_epochs,
             advanced_training.save_last_n_epochs_state,
             advanced_training.skip_cache_check,
             advanced_training.log_with,

diff --git a/kohya_gui/finetune_gui.py b/kohya_gui/finetune_gui.py
@@ -173,6 +173,7 @@ def save_configuration(
     save_every_n_steps,
     save_last_n_steps,
     save_last_n_steps_state,
+    save_last_n_epochs,
     save_last_n_epochs_state,
     skip_cache_check,
     log_with,
@@ -386,6 +387,7 @@ def open_configuration(
     save_every_n_steps,
     save_last_n_steps,
     save_last_n_steps_state,
+    save_last_n_epochs,
     save_last_n_epochs_state,
     skip_cache_check,
     log_with,
@@ -605,6 +607,7 @@ def train_model(
     save_every_n_steps,
     save_last_n_steps,
     save_last_n_steps_state,
+    save_last_n_epochs,
     save_last_n_epochs_state,
     skip_cache_check,
     log_with,
@@ -1051,6 +1054,7 @@ def train_model(
         "save_last_n_steps_state": (
             save_last_n_steps_state if save_last_n_steps_state != 0 else None
         ),
+        "save_last_n_epochs": save_last_n_epochs if save_last_n_epochs != 0 else None,
         "save_last_n_epochs_state": (
             save_last_n_epochs_state if save_last_n_epochs_state != 0 else None
         ),
@@ -1475,6 +1479,7 @@ def list_presets(path):
             advanced_training.save_every_n_steps,
             advanced_training.save_last_n_steps,
             advanced_training.save_last_n_steps_state,
+            advanced_training.save_last_n_epochs,
             advanced_training.save_last_n_epochs_state,
             advanced_training.skip_cache_check,
             advanced_training.log_with,

diff --git a/kohya_gui/lora_gui.py b/kohya_gui/lora_gui.py
@@ -185,6 +185,7 @@ def save_configuration(
     save_every_n_steps,
     save_last_n_steps,
     save_last_n_steps_state,
+    save_last_n_epochs,
     save_last_n_epochs_state,
     skip_cache_check,
     log_with,
@@ -450,6 +451,7 @@ def open_configuration(
     save_every_n_steps,
     save_last_n_steps,
     save_last_n_steps_state,
+    save_last_n_epochs,
     save_last_n_epochs_state,
     skip_cache_check,
     log_with,
@@ -749,6 +751,7 @@ def train_model(
     save_every_n_steps,
     save_last_n_steps,
     save_last_n_steps_state,
+    save_last_n_epochs,
     save_last_n_epochs_state,
     skip_cache_check,
     log_with,
@@ -1505,6 +1508,7 @@ def train_model(
         "save_last_n_steps_state": (
             save_last_n_steps_state if save_last_n_steps_state != 0 else None
         ),
+        "save_last_n_epochs": save_last_n_epochs if save_last_n_epochs != 0 else None,
         "save_last_n_epochs_state": (
             save_last_n_epochs_state if save_last_n_epochs_state != 0 else None
         ),
@@ -2648,6 +2652,7 @@ def update_LoRA_settings(
             advanced_training.save_every_n_steps,
             advanced_training.save_last_n_steps,
             advanced_training.save_last_n_steps_state,
+            advanced_training.save_last_n_epochs,
             advanced_training.save_last_n_epochs_state,
             advanced_training.skip_cache_check,
             advanced_training.log_with,

diff --git a/kohya_gui/textual_inversion_gui.py b/kohya_gui/textual_inversion_gui.py
@@ -158,6 +158,7 @@ def save_configuration(
     save_every_n_steps,
     save_last_n_steps,
     save_last_n_steps_state,
+    save_last_n_epochs,
     save_last_n_epochs_state,
     skip_cache_check,
     log_with,
@@ -323,6 +324,7 @@ def open_configuration(
     save_every_n_steps,
     save_last_n_steps,
     save_last_n_steps_state,
+    save_last_n_epochs,
     save_last_n_epochs_state,
     skip_cache_check,
     log_with,
@@ -481,6 +483,7 @@ def train_model(
     save_every_n_steps,
     save_last_n_steps,
     save_last_n_steps_state,
+    save_last_n_epochs,
     save_last_n_epochs_state,
     skip_cache_check,
     log_with,
@@ -845,6 +848,7 @@ def train_model(
         "save_last_n_steps_state": (
             save_last_n_steps_state if save_last_n_steps_state != 0 else None
         ),
+        "save_last_n_epochs": save_last_n_epochs if save_last_n_epochs != 0 else None,
         "save_last_n_epochs_state": (
             save_last_n_epochs_state if save_last_n_epochs_state != 0 else None
         ),
@@ -1226,6 +1230,7 @@ def list_embedding_files(path):
             advanced_training.save_every_n_steps,
             advanced_training.save_last_n_steps,
             advanced_training.save_last_n_steps_state,
+            advanced_training.save_last_n_epochs,
             advanced_training.save_last_n_epochs_state,
             advanced_training.skip_cache_check,
             advanced_training.log_with,