Phase 2 - Input tooltips, autocorrect, synced run names, default names

modules/cloud/BaseCloud.py

@@ -26,9 +26,12 @@ def setup(self):
             self._make_tensorboard_tunnel()
 
     def download_output_model(self):
-        local=Path(self.config.local_output_model_destination)
-        remote=Path(self.config.output_model_destination)
-        self.file_sync.sync_down_file(local=local,remote=remote)
+        ext = self.config.output_model_format.file_extension()
+        name = f"{self.config.run_name}{ext}"
+        remote = Path(self.config.final_output_dir) / name
+        local_dir = getattr(self.config, "local_final_output_dir", self.config.final_output_dir)
+        local = Path(local_dir) / name
+        self.file_sync.sync_down_file(local=local, remote=remote)
         self.file_sync.sync_down_dir(local=local.with_suffix(local.suffix+"_embeddings"),
                            remote=remote.with_suffix(remote.suffix+"_embeddings"))
 

modules/trainer/BaseTrainer.py

@@ -12,6 +12,7 @@
 from modules.util.callbacks.TrainCallbacks import TrainCallbacks
 from modules.util.commands.TrainCommands import TrainCommands
 from modules.util.config.TrainConfig import TrainConfig
+from modules.util.time_util import generate_default_run_name
 from modules.util.TimedActionMixin import TimedActionMixin
 from modules.util.TrainProgress import TrainProgress
 
@@ -33,6 +34,9 @@ def __init__(self, config: TrainConfig, callbacks: TrainCallbacks, commands: Tra
         self.train_device = torch.device(self.config.train_device)
         self.temp_device = torch.device(self.config.temp_device)
 
+        if not self.config.run_name:
+            self.config.run_name = generate_default_run_name(self.config.training_method)
+
     @abstractmethod
     def start(self):
         pass

modules/trainer/CloudTrainer.py

@@ -162,7 +162,7 @@ def adjust(config, attribute: str, if_exists: bool=False):
         adjust(remote,"base_model_name", if_exists=True)
         adjust(remote.prior,"model_name", if_exists=True)
         adjust(remote.transformer,"model_name", if_exists=True)
-        adjust(remote,"output_model_destination")
+        adjust(remote,"final_output_dir")
         adjust(remote,"lora_model_name")
 
         adjust(remote.embedding,"model_name")

modules/trainer/GenericTrainer.py

@@ -70,7 +70,8 @@ def __init__(self, config: TrainConfig, callbacks: TrainCallbacks, commands: Tra
         if multi.is_master():
             tensorboard_log_dir = os.path.join(config.workspace_dir, "tensorboard")
             os.makedirs(Path(tensorboard_log_dir).absolute(), exist_ok=True)
-            self.tensorboard = SummaryWriter(os.path.join(tensorboard_log_dir, f"{config.save_filename_prefix}{get_string_timestamp()}"))
+            run_prefix = f"{config.run_name}-" if config.run_name else ""
+            self.tensorboard = SummaryWriter(os.path.join(tensorboard_log_dir, f"{run_prefix}{get_string_timestamp()}"))
             if config.tensorboard and not config.tensorboard_always_on:
                 super()._start_tensorboard()
 
@@ -164,7 +165,8 @@ def start(self):
     def __save_config_to_workspace(self):
         path = path_util.canonical_join(self.config.workspace_dir, "config")
         os.makedirs(Path(path).absolute(), exist_ok=True)
-        path = path_util.canonical_join(path, f"{self.config.save_filename_prefix}{get_string_timestamp()}.json")
+        run_prefix = f"{self.config.run_name}-" if self.config.run_name else ""
+        path = path_util.canonical_join(path, f"{run_prefix}{get_string_timestamp()}.json")
         with open(path, "w") as f:
             json.dump(self.config.to_pack_dict(secrets=False), f, indent=4)
 
@@ -183,17 +185,20 @@ def __prune_backups(self, backups_to_keep: int):
         backup_dirpath = os.path.join(self.config.workspace_dir, "backup")
         if os.path.exists(backup_dirpath):
             backup_directories = sorted(
-                [dirpath for dirpath in os.listdir(backup_dirpath) if
-                 os.path.isdir(os.path.join(backup_dirpath, dirpath))],
+                [name for name in os.listdir(backup_dirpath) if
+                 os.path.isdir(os.path.join(backup_dirpath, name))],
+                key=lambda n: TrainConfig._extract_backup_datetime(
+                    os.path.join(backup_dirpath, n), n
+                ),
                 reverse=True,
             )
 
-            for dirpath in backup_directories[backups_to_keep:]:
-                dirpath = os.path.join(backup_dirpath, dirpath)
+            for name in backup_directories[backups_to_keep:]:
+                full_path = os.path.join(backup_dirpath, name)
                 try:
-                    shutil.rmtree(dirpath)
+                    shutil.rmtree(full_path)
                 except Exception:
-                    print(f"Could not delete old rolling backup {dirpath}")
+                    print(f"Could not delete old rolling backup {full_path}")
 
         return
 
@@ -234,9 +239,10 @@ def __sample_loop(
                         f"{str(i)} - {safe_prompt}{folder_postfix}",
                     )
 
+                run_prefix = f"{self.config.run_name}-" if self.config.run_name else ""
                 sample_path = os.path.join(
                     sample_dir,
-                    f"{self.config.save_filename_prefix}{get_string_timestamp()}-training-sample-{train_progress.filename_string()}"
+                    f"{run_prefix}{get_string_timestamp()}-training-sample-{train_progress.filename_string()}"
                 )
 
                 def on_sample_default(sampler_output: ModelSamplerOutput):
@@ -433,7 +439,8 @@ def __backup(self, train_progress: TrainProgress, print_msg: bool = True, print_
 
         self.callbacks.on_update_status("Creating backup")
 
-        backup_name = f"{get_string_timestamp()}-backup-{train_progress.filename_string()}"
+        safe_prefix = path_util.safe_filename(f"{self.config.run_name}-", max_length=None) if self.config.run_name else ""
+        backup_name = f"{safe_prefix}{get_string_timestamp()}-backup-{train_progress.filename_string()}"
         backup_path = os.path.join(self.config.workspace_dir, "backup", backup_name)
 
         # Special case for schedule-free optimizers.
@@ -480,10 +487,11 @@ def __save(self, train_progress: TrainProgress, print_msg: bool = True, print_cb
 
         self.callbacks.on_update_status("Saving")
 
+        safe_prefix = path_util.safe_filename(f"{self.config.run_name}-", max_length=None) if self.config.run_name else ""
         save_path = os.path.join(
             self.config.workspace_dir,
             "save",
-            f"{self.config.save_filename_prefix}{get_string_timestamp()}-save-{train_progress.filename_string()}{self.config.output_model_format.file_extension()}"
+            f"{safe_prefix}{get_string_timestamp()}-save-{train_progress.filename_string()}{self.config.output_model_format.file_extension()}"
         )
         if print_msg:
             print_cb("Saving " + save_path)
@@ -853,13 +861,10 @@ def end(self):
 
                 if self.model.ema:
                     self.model.ema.copy_ema_to(self.parameters, store_temp=False)
-                if os.path.isdir(self.config.output_model_destination) and self.config.output_model_format.is_single_file():
-                    save_path = os.path.join(
-                        self.config.output_model_destination,
-                        f"{self.config.save_filename_prefix}{get_string_timestamp()}{self.config.output_model_format.file_extension()}"
-                    )
-                else:
-                    save_path = self.config.output_model_destination
+                save_path = os.path.join(
+                    self.config.final_output_dir,
+                    f"{self.config.run_name}{self.config.output_model_format.file_extension()}"
+                )
                 print("Saving " + save_path)
 
                 self.model_saver.save(

modules/ui/ConvertModelUI.py

@@ -113,7 +113,7 @@ def main_frame(self, master):
         components.path_entry(
             master, 5, 1, self.ui_state, "output_model_destination",
             mode="file",
-            io_type=PathIOType.MODEL,
+            io_type=PathIOType.OUTPUT,
         )
 
         self.button = components.button(master, 6, 1, "Convert", self.convert_model)

modules/ui/ModelTab.py

@@ -5,9 +5,11 @@
 from modules.util.enum.DataType import DataType
 from modules.util.enum.ModelFormat import ModelFormat
 from modules.util.enum.PathIOType import PathIOType
+from modules.util.enum.RunNameMode import RunNameMode
 from modules.util.enum.TrainingMethod import TrainingMethod
 from modules.util.ui import components
 from modules.util.ui.UIState import UIState
+from modules.util.ui.validation import RunNameValidator
 
 import customtkinter as ctk
 
@@ -613,14 +615,32 @@ def __create_output_components(
             allow_legacy_safetensors: bool = False,
             allow_comfy: bool = False,
     ) -> int:
-        # output model destination
-        components.label(frame, row, 0, "Model Output Destination",
-                         tooltip="Filename or directory where the output model is saved")
-        components.path_entry(
-            frame, row, 1, self.ui_state, "output_model_destination",
-            mode="file",
-            io_type=PathIOType.MODEL,
-        )
+        # run name mode
+        components.label(frame, row, 0, "Run Name",
+                         tooltip="String used as the filename and as a prefix for saves, backups, samples, and tensorboard. "
+                                 "The file extension is determined automatically by the Output Format")
+
+        run_name_frame = ctk.CTkFrame(frame, fg_color="transparent")
+        run_name_frame.grid(row=row, column=1, padx=0, pady=0, sticky="ew")
+        run_name_frame.grid_columnconfigure(1, weight=1)
+
+        run_name_entry = components.entry(run_name_frame, 0, 1, self.ui_state, "run_name",
+                         validator_factory=RunNameValidator, sticky="ew")
+
+        def _on_run_name_mode_change(_value=None):
+            mode_var = self.ui_state.get_var("run_name_mode")
+            mode_str = str(mode_var.get())
+            if mode_str == str(RunNameMode.CUSTOM):
+                run_name_entry.set_enabled()
+            else:
+                run_name_entry.set_disabled()
+
+        components.options_kv(run_name_frame, 0, 0, [
+            ("Default", RunNameMode.DEFAULT),
+            ("Friendly", RunNameMode.FRIENDLY),
+            ("Custom", RunNameMode.CUSTOM),
+        ], self.ui_state, "run_name_mode", command=_on_run_name_mode_change)
+        _on_run_name_mode_change()
 
         # output data type
         components.label(frame, row, 3, "Output Data Type",
@@ -635,6 +655,17 @@ def __create_output_components(
 
         row += 1
 
+        # final output directory
+        components.label(frame, row, 0, "Final Output Directory",
+                         tooltip="Directory where the final trained model is saved")
+        components.path_entry(
+            frame, row, 1, self.ui_state, "final_output_dir",
+            mode="dir",
+            io_type=PathIOType.OUTPUT,
+        )
+
+        row += 1
+
         # output format
         formats = []
         if allow_safetensors:

modules/ui/MuonAdamWindow.py

@@ -105,6 +105,6 @@ def create_adam_params_ui(self, master):
             components.label(master, row, col, title, tooltip=tooltip)
 
             if param_type != 'bool':
-                components.entry(master, row, col + 1, self.adam_ui_state, key)
+                components.entry(master, row, col + 1, self.adam_ui_state, key, allow_negative=True)
             else:
                 components.switch(master, row, col + 1, self.adam_ui_state, key)

modules/ui/OptimizerParamsWindow.py

@@ -237,7 +237,7 @@ def create_dynamic_ui(
                 self.toggle_muon_adam_button()
             elif type != 'bool':
                 components.entry(master, row, col + 1, self.optimizer_ui_state, key,
-                                 command=self.update_user_pref)
+                                 command=self.update_user_pref, allow_negative=True)
             else:
                 components.switch(master, row, col + 1, self.optimizer_ui_state, key,
                                   command=self.update_user_pref)