Input Validation, Tensorboard Tweaks + Drag-n-drop support

.gitignore

@@ -7,6 +7,9 @@
 debug*.ipynb
 debug*.py
 /debug*
+profile.json
+profile.html
+stacks.txt
 
 # user data
 /workspace*
@@ -16,7 +19,7 @@ debug*.py
 /training_user_settings
 /external
 /update.var
-config.json
+config*.json
 secrets.json
 *.zip
 

modules/trainer/BaseTrainer.py

@@ -1,6 +1,3 @@
-import os
-import subprocess
-import sys
 from abc import ABCMeta, abstractmethod
 
 from modules.model.BaseModel import BaseModel
@@ -23,8 +20,6 @@ class BaseTrainer(
     metaclass=ABCMeta,
 ):
 
-    tensorboard_subprocess: subprocess.Popen
-
     def __init__(self, config: TrainConfig, callbacks: TrainCallbacks, commands: TrainCommands):
         super().__init__()
         self.config = config
@@ -80,23 +75,3 @@ def create_model_sampler(self, model: BaseModel) -> BaseModelSampler:
             self.config.model_type,
             self.config.training_method
         )
-    def _start_tensorboard(self):
-        tensorboard_executable = os.path.join(os.path.dirname(sys.executable), "tensorboard")
-        tensorboard_log_dir = os.path.join(self.config.workspace_dir, "tensorboard")
-
-        tensorboard_args = [
-            tensorboard_executable,
-            "--logdir",
-            tensorboard_log_dir,
-            "--port",
-            str(self.config.tensorboard_port),
-            "--samples_per_plugin=images=100,scalars=10000",
-        ]
-
-        if self.config.tensorboard_expose:
-            tensorboard_args.append("--bind_all")
-
-        self.tensorboard_subprocess = subprocess.Popen(tensorboard_args)
-
-    def _stop_tensorboard(self):
-        self.tensorboard_subprocess.kill()

modules/trainer/CloudTrainer.py

@@ -13,6 +13,7 @@
 from modules.util.config.TrainConfig import TrainConfig
 from modules.util.enum.CloudAction import CloudAction
 from modules.util.enum.CloudType import CloudType
+from modules.util.tensorboard_util import get_tensorboard_manager
 
 
 class CloudTrainer(BaseTrainer):
@@ -29,8 +30,8 @@ def __init__(self, config: TrainConfig, callbacks: TrainCallbacks, commands: Tra
 
         tensorboard_log_dir = os.path.join(config.workspace_dir, "tensorboard")
         os.makedirs(Path(tensorboard_log_dir).absolute(), exist_ok=True)
-        if config.tensorboard and not config.cloud.tensorboard_tunnel and not config.tensorboard_always_on:
-            super()._start_tensorboard()
+
+        get_tensorboard_manager().handle_training_start(config)
 
         match config.cloud.type:
             case CloudType.RUNPOD:
@@ -113,8 +114,7 @@ def train(self):
 
     def end(self):
         try:
-            if self.config.tensorboard and not self.config.cloud.tensorboard_tunnel and not self.config.tensorboard_always_on:
-                super()._stop_tensorboard()
+            get_tensorboard_manager().handle_training_end(self.config)
 
             if self.config.cloud.delete_workspace and not self.error_caught and not self.commands.get_stop_command():
                 self.callbacks.on_update_status("Deleting remote workspace")

modules/trainer/GenericTrainer.py

@@ -30,6 +30,7 @@
 from modules.util.enum.TimeUnit import TimeUnit
 from modules.util.enum.TrainingMethod import TrainingMethod
 from modules.util.profiling_util import TorchMemoryRecorder, TorchProfiler
+from modules.util.tensorboard_util import get_tensorboard_manager
 from modules.util.time_util import get_string_timestamp
 from modules.util.torch_util import torch_gc
 from modules.util.TrainProgress import TrainProgress
@@ -70,9 +71,9 @@ 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()}"))
-            if config.tensorboard and not config.tensorboard_always_on:
-                super()._start_tensorboard()
+            self.tensorboard = SummaryWriter(os.path.join(tensorboard_log_dir, f"{config.get_run_name()}_{get_string_timestamp()}"))
+
+            get_tensorboard_manager().handle_training_start(config)
 
         self.model = None
         self.one_step_trained = False
@@ -164,7 +165,7 @@ 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")
+        path = path_util.canonical_join(path, f"{self.config.get_run_name()}_{get_string_timestamp()}.json")
         with open(path, "w") as f:
             json.dump(self.config.to_pack_dict(secrets=False), f, indent=4)
 
@@ -234,7 +235,7 @@ def __sample_loop(
 
                     sample_path = os.path.join(
                         sample_dir,
-                        f"{self.config.save_filename_prefix}{get_string_timestamp()}-training-sample-{train_progress.filename_string()}"
+                        f"{self.config.get_run_name()}_{get_string_timestamp()}-training-sample-{train_progress.filename_string()}"
                     )
 
                     def on_sample_default(sampler_output: ModelSamplerOutput):
@@ -431,7 +432,7 @@ 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()}"
+        backup_name = f"{get_string_timestamp()}-{self.config.get_run_name()}-backup-{train_progress.filename_string()}"
         backup_path = os.path.join(self.config.workspace_dir, "backup", backup_name)
 
         # Special case for schedule-free optimizers.
@@ -481,7 +482,7 @@ def __save(self, train_progress: TrainProgress, print_msg: bool = True, print_cb
         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"{self.config.save_filename_prefix}{get_string_timestamp()}-{self.config.get_run_name()}-save-{train_progress.filename_string()}{self.config.output_model_format.file_extension()}"
         )
         if print_msg:
             print_cb("Saving " + save_path)
@@ -851,7 +852,7 @@ def end(self):
                 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()}"
+                        f"{self.config.get_run_name()}_{get_string_timestamp()}{self.config.output_model_format.file_extension()}"
                     )
                 else:
                     save_path = self.config.output_model_destination
@@ -870,9 +871,7 @@ def end(self):
 
         if multi.is_master():
             self.tensorboard.close()
-
-            if self.config.tensorboard and not self.config.tensorboard_always_on:
-                super()._stop_tensorboard()
+            get_tensorboard_manager().handle_training_end(self.config)
 
         for handle in self.grad_hook_handles:
             handle.remove()

modules/ui/AdditionalEmbeddingsTab.py

@@ -89,7 +89,7 @@ def __init__(self, master, element, i, open_command, remove_command, clone_comma
         # embedding model names
         components.label(top_frame, 0, 2, "base embedding:",
                          tooltip="The base embedding to train on. Leave empty to create a new embedding")
-        components.file_entry(
+        components.path_entry(
             top_frame, 0, 3, self.ui_state, "model_name",
             path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x
         )

modules/ui/CloudTab.py

@@ -66,7 +66,7 @@ def __init__(self, master, train_config: TrainConfig, ui_state: UIState, parent)
 
         components.label(self.frame, 7, 0, "SSH keyfile path",
                  tooltip="Absolute path to the private key file used for SSH connections. Leave empty to rely on your system SSH configuration.")
-        components.file_entry(self.frame, 7, 1, self.ui_state, "secrets.cloud.key_file") #TODO Replace with path_entry in a future PR
+        components.path_entry(self.frame, 7, 1, self.ui_state, "secrets.cloud.key_file")
 
         components.label(self.frame, 8, 0, "SSH password",
                          tooltip="SSH password for password-based authentication. If you try to use native SCP requires sshpass to be installed. Leave empty to use key-based authentication.")

modules/ui/ConceptTab.py

@@ -179,13 +179,14 @@ def __init__(self, master, concept, i, open_command, remove_command, clone_comma
 
         self.grid_rowconfigure(1, weight=1)
 
-        # image
+        # image - start with placeholder, load actual image async
+        self._placeholder_image = self.__get_placeholder_image()
         self.image = ctk.CTkImage(
-            light_image=self.__get_preview_image(),
+            light_image=self._placeholder_image,
             size=(150, 150)
         )
-        image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=150, width=150)
-        image_label.grid(row=0, column=0)
+        self.image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=150, width=150)
+        self.image_label.grid(row=0, column=0)
 
         # name
         self.name_label = components.label(self, 1, 0, self.__get_display_name(), pad=5, wraplength=140)
@@ -224,15 +225,41 @@ def __init__(self, master, concept, i, open_command, remove_command, clone_comma
         )
         enabled_switch.place(x=110, y=0)
 
-        image_label.bind(
+        self.image_label.bind(
             "<Button-1>",
             lambda event: open_command(self.i, (self.ui_state, self.image_ui_state, self.text_ui_state))
         )
 
+        # defer actual image
+        self.after(50, self.__load_preview_image_deferred)
+
     def __randomize_seed(self, concept: ConceptConfig):
         concept.seed = ConceptConfig.default_values().seed
         return concept
 
+    _placeholder_image_cache = None
+
+    def __get_placeholder_image(self):
+        if ConceptWidget._placeholder_image_cache:
+            return ConceptWidget._placeholder_image_cache
+
+        try:
+            image = load_image("resources/icons/icon.png", convert_mode="RGBA")
+            size = min(image.width, image.height)
+            image = image.crop(((image.width - size) // 2, (image.height - size) // 2, (image.width + size) // 2, (image.height + size) // 2))
+        except (OSError):
+            image = Image.new("RGBA", (150, 150), (200, 200, 200, 255))
+
+        ConceptWidget._placeholder_image_cache = image.resize((150, 150), Image.Resampling.BILINEAR)
+        return ConceptWidget._placeholder_image_cache
+
+    def __load_preview_image_deferred(self):
+        try:
+            actual_image = self.__get_preview_image()
+            self.image.configure(light_image=actual_image)
+        except Exception:
+            pass  # Keep placeholder on error
+
     def __get_display_name(self):
         if self.concept.name:
             return self.concept.name

modules/ui/ConceptWindow.py

@@ -147,14 +147,14 @@ def __general_tab(self, master, concept: ConceptConfig):
         # path
         components.label(frame, 3, 0, "Path",
                          tooltip="Path where the training data is located")
-        components.dir_entry(frame, 3, 1, self.ui_state, "path")
+        components.path_entry(frame, 3, 1, self.ui_state, "path", path_type="directory")
         components.button(frame, 3, 2, text="download now", command=self.__download_dataset_threaded,
                           tooltip="Download dataset from Huggingface now, for the purpose of previewing and statistics. Otherwise, it will be downloaded when you start training. Path must be a Huggingface repository.")
 
         # prompt source
         components.label(frame, 4, 0, "Prompt Source",
                          tooltip="The source for prompts used during training. When selecting \"From single text file\", select a text file that contains a list of prompts")
-        prompt_path_entry = components.file_entry(frame, 4, 2, self.text_ui_state, "prompt_path")
+        prompt_path_entry = components.path_entry(frame, 4, 2, self.text_ui_state, "prompt_path")
 
         def set_prompt_path_entry_enabled(option: str):
             if option == 'concept':