From fad971a32b106f1b2196c733f13b703822c788ce Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 4 Jul 2025 22:56:28 +1000 Subject: [PATCH 01/64] Move trainer cleanup under finally --- modules/ui/TrainUI.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index 7f66a7215..e0b4889be 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -623,16 +623,16 @@ def __training_thread_function(self): self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) error_caught = True traceback.print_exc() + finally: + trainer.end() - trainer.end() + # clear gpu memory + del trainer - # clear gpu memory - del trainer - - self.training_thread = None - self.training_commands = None - torch.clear_autocast_cache() - torch_gc() + self.training_thread = None + self.training_commands = None + torch.clear_autocast_cache() + torch_gc() if error_caught: self.on_update_status("error: check the console for more information") From 13efc33a701e42f1eeeca087ddc8c38b722aadee Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 4 Jul 2025 23:03:09 +1000 Subject: [PATCH 02/64] Tweak UI layout and labels --- modules/ui/TrainUI.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index e0b4889be..87a2d6892 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -211,20 +211,22 @@ def create_general_tab(self, master): components.dir_entry(frame, 5, 1, self.ui_state, "debug_dir") # tensorboard - components.label(frame, 6, 0, "Tensorboard", - tooltip="Starts the Tensorboard Web UI during training") + components.label(frame, 6, 0, "Tensorboard During Training", + tooltip="Starts the Tensorboard Web UI during training only") components.switch(frame, 6, 1, self.ui_state, "tensorboard") - components.label(frame, 7, 0, "Expose Tensorboard", - tooltip="Exposes Tensorboard Web UI to all network interfaces (makes it accessible from the network)") - components.switch(frame, 7, 1, self.ui_state, "tensorboard_expose") - components.label(frame, 7, 2, "Tensorboard Port", + components.label(frame, 7, 0, "Always-On Tensorboard", + tooltip="Keep Tensorboard accessible even when not training. Useful for monitoring completed training sessions.") + components.switch(frame, 7, 1, self.ui_state, "tensorboard_always_on", command=self._on_always_on_tensorboard_toggle) + + components.label(frame, 8, 0, "Tensorboard Port", tooltip="Port to use for Tensorboard link") - components.entry(frame, 7, 3, self.ui_state, "tensorboard_port") + components.entry(frame, 8, 1, self.ui_state, "tensorboard_port") + + components.label(frame, 8, 2, "Expose Tensorboard", + tooltip="Exposes Tensorboard Web UI to all network interfaces (makes it accessible from the network)") + components.switch(frame, 8, 3, self.ui_state, "tensorboard_expose") - components.label(frame, 8, 0, "Always-On Tensorboard", - tooltip="Keep Tensorboard accessible even when not training. Useful for monitoring completed training sessions.") - components.switch(frame, 8, 1, self.ui_state, "tensorboard_always_on", command=self._on_always_on_tensorboard_toggle) # validation components.label(frame, 9, 0, "Validation", From e5026c8e30b5b9729b7ce806dcecfd1976d10469 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Sat, 5 Jul 2025 00:29:51 +1000 Subject: [PATCH 03/64] Centralise tensorboard management + mutually exclusive toggles --- modules/ui/TrainUI.py | 116 ++++++++++++++++++++++++++---------------- 1 file changed, 72 insertions(+), 44 deletions(-) diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index 87a2d6892..7dde2bc0a 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -1,3 +1,4 @@ +import contextlib import json import os import subprocess @@ -94,7 +95,7 @@ def __init__(self): self.always_on_tensorboard_subprocess = None self.current_workspace_dir = self.train_config.workspace_dir - self._check_start_always_on_tensorboard() + self._initialize_tensorboard() self.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self._on_workspace_dir_change_trace) @@ -211,9 +212,9 @@ def create_general_tab(self, master): components.dir_entry(frame, 5, 1, self.ui_state, "debug_dir") # tensorboard - components.label(frame, 6, 0, "Tensorboard During Training", + components.label(frame, 6, 0, "Tensorboard (Train Only)", tooltip="Starts the Tensorboard Web UI during training only") - components.switch(frame, 6, 1, self.ui_state, "tensorboard") + components.switch(frame, 6, 1, self.ui_state, "tensorboard", command=self._on_train_only_tensorboard_toggle) components.label(frame, 7, 0, "Always-On Tensorboard", tooltip="Keep Tensorboard accessible even when not training. Useful for monitoring completed training sessions.") @@ -643,18 +644,12 @@ def __training_thread_function(self): self.training_button.configure(text="Start Training", state="normal") - if self.train_config.tensorboard_always_on and not self.always_on_tensorboard_subprocess: - self._start_always_on_tensorboard() - def start_training(self): if self.training_thread is None: self.save_default() self.training_button.configure(text="Stop Training", state="normal") - if self.train_config.tensorboard and not self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: - self._stop_always_on_tensorboard() - self.training_commands = TrainCommands() self.training_thread = threading.Thread(target=self.__training_thread_function) @@ -695,20 +690,19 @@ def save_now(self): if train_commands: train_commands.save() - def _check_start_always_on_tensorboard(self): - if self.train_config.tensorboard_always_on and not self.always_on_tensorboard_subprocess: + def _initialize_tensorboard(self): + """Initialize Tensorboard state on startup""" + if self.train_config.tensorboard_always_on: self._start_always_on_tensorboard() - def _start_always_on_tensorboard(self): - if self.always_on_tensorboard_subprocess: - self._stop_always_on_tensorboard() - + def _get_tensorboard_args(self): + """Build Tensorboard command arguments""" tensorboard_executable = os.path.join(os.path.dirname(sys.executable), "tensorboard") tensorboard_log_dir = os.path.join(self.train_config.workspace_dir, "tensorboard") os.makedirs(Path(tensorboard_log_dir).absolute(), exist_ok=True) - tensorboard_args = [ + args = [ tensorboard_executable, "--logdir", tensorboard_log_dir, @@ -718,44 +712,78 @@ def _start_always_on_tensorboard(self): ] if self.train_config.tensorboard_expose: - tensorboard_args.append("--bind_all") + args.append("--bind_all") + + return args + + def _start_tensorboard_subprocess(self): + """Start a new Tensorboard subprocess""" try: - self.always_on_tensorboard_subprocess = subprocess.Popen(tensorboard_args) + return subprocess.Popen(self._get_tensorboard_args()) except Exception: - self.always_on_tensorboard_subprocess = None + return None + + def _stop_tensorboard_subprocess(self, process): + """Gracefully stop a Tensorboard subprocess""" + if not process: + return + + try: + process.terminate() + process.wait(timeout=5) + except subprocess.TimeoutExpired: + with contextlib.suppress(Exception): + process.kill() + except Exception: + pass + + def _start_always_on_tensorboard(self): + """Start always-on Tensorboard if not already running""" + if not self.always_on_tensorboard_subprocess: + self.always_on_tensorboard_subprocess = self._start_tensorboard_subprocess() def _stop_always_on_tensorboard(self): + """Stop always-on Tensorboard""" if self.always_on_tensorboard_subprocess: - try: - self.always_on_tensorboard_subprocess.terminate() - self.always_on_tensorboard_subprocess.wait(timeout=5) - except subprocess.TimeoutExpired: - self.always_on_tensorboard_subprocess.kill() - except Exception: - pass - finally: - self.always_on_tensorboard_subprocess = None + self._stop_tensorboard_subprocess(self.always_on_tensorboard_subprocess) + self.always_on_tensorboard_subprocess = None + + def _restart_always_on_tensorboard(self): + """Restart always-on Tensorboard (for workspace directory changes)""" + if self.train_config.tensorboard_always_on: + self._stop_always_on_tensorboard() + self._start_always_on_tensorboard() + + def _on_train_only_tensorboard_toggle(self): + """Handle train-only Tensorboard toggle""" + if self.train_config.tensorboard: + # If train-only is enabled, disable always-on + if self.train_config.tensorboard_always_on: + self.ui_state.get_var("tensorboard_always_on").set(False) + self._stop_always_on_tensorboard() + + def _on_always_on_tensorboard_toggle(self): + """Handle always-on Tensorboard toggle""" + if self.train_config.tensorboard_always_on: + # If always-on is enabled, disable train-only + if self.train_config.tensorboard: + self.ui_state.get_var("tensorboard").set(False) + # Start always-on Tensorboard + self._start_always_on_tensorboard() + else: + # Stop always-on Tensorboard + self._stop_always_on_tensorboard() def _on_workspace_dir_change(self, new_workspace_dir: str): + """Handle direct workspace directory changes""" if new_workspace_dir != self.current_workspace_dir: self.current_workspace_dir = new_workspace_dir - - if self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: - self._start_always_on_tensorboard() + # Restart always-on Tensorboard with new workspace directory if it's enabled + if self.train_config.tensorboard_always_on: + self._restart_always_on_tensorboard() def _on_workspace_dir_change_trace(self, *args): + """Handle workspace directory changes via UI state trace""" new_workspace_dir = self.train_config.workspace_dir - if new_workspace_dir != self.current_workspace_dir: - self.current_workspace_dir = new_workspace_dir - - if self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: - self._start_always_on_tensorboard() - - def _on_always_on_tensorboard_toggle(self): - if self.train_config.tensorboard_always_on: - if not (self.training_thread and self.train_config.tensorboard): - self._start_always_on_tensorboard() - else: - if not (self.training_thread and self.train_config.tensorboard): - self._stop_always_on_tensorboard() + self._on_workspace_dir_change(new_workspace_dir) From 84c9f89c2b47a78d184d557cb8fee07233ad2586 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Sat, 5 Jul 2025 01:20:43 +1000 Subject: [PATCH 04/64] Create tensorboard_util that both BaseTrainer and TrainUI can import --- modules/trainer/BaseTrainer.py | 24 ++------- modules/ui/TrainUI.py | 58 ++-------------------- modules/util/tensorboard_util.py | 83 ++++++++++++++++++++++++++++++++ 3 files changed, 91 insertions(+), 74 deletions(-) create mode 100644 modules/util/tensorboard_util.py diff --git a/modules/trainer/BaseTrainer.py b/modules/trainer/BaseTrainer.py index 5cd913bce..dcd5e64d3 100644 --- a/modules/trainer/BaseTrainer.py +++ b/modules/trainer/BaseTrainer.py @@ -1,6 +1,4 @@ -import os import subprocess -import sys from abc import ABCMeta, abstractmethod from modules.model.BaseModel import BaseModel @@ -12,6 +10,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.tensorboard_util import start_filtered_tensorboard, stop_tensorboard from modules.util.TimedActionMixin import TimedActionMixin from modules.util.TrainProgress import TrainProgress @@ -84,23 +83,10 @@ 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 _start_tensorboard(self): + self.tensorboard_subprocess = start_filtered_tensorboard(self.config) def _stop_tensorboard(self): - self.tensorboard_subprocess.kill() + stop_tensorboard(self.tensorboard_subprocess) + self.tensorboard_subprocess = None diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index 7dde2bc0a..1bdefeba7 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -1,8 +1,4 @@ -import contextlib import json -import os -import subprocess -import sys import threading import traceback import webbrowser @@ -32,6 +28,7 @@ from modules.util.enum.ImageFormat import ImageFormat from modules.util.enum.ModelType import ModelType from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.tensorboard_util import start_filtered_tensorboard, stop_tensorboard from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress from modules.util.ui import components @@ -691,62 +688,16 @@ def save_now(self): train_commands.save() def _initialize_tensorboard(self): - """Initialize Tensorboard state on startup""" if self.train_config.tensorboard_always_on: self._start_always_on_tensorboard() - def _get_tensorboard_args(self): - """Build Tensorboard command arguments""" - tensorboard_executable = os.path.join(os.path.dirname(sys.executable), "tensorboard") - tensorboard_log_dir = os.path.join(self.train_config.workspace_dir, "tensorboard") - - os.makedirs(Path(tensorboard_log_dir).absolute(), exist_ok=True) - - args = [ - tensorboard_executable, - "--logdir", - tensorboard_log_dir, - "--port", - str(self.train_config.tensorboard_port), - "--samples_per_plugin=images=100,scalars=10000", - ] - - if self.train_config.tensorboard_expose: - args.append("--bind_all") - - return args - - - def _start_tensorboard_subprocess(self): - """Start a new Tensorboard subprocess""" - try: - return subprocess.Popen(self._get_tensorboard_args()) - except Exception: - return None - - def _stop_tensorboard_subprocess(self, process): - """Gracefully stop a Tensorboard subprocess""" - if not process: - return - - try: - process.terminate() - process.wait(timeout=5) - except subprocess.TimeoutExpired: - with contextlib.suppress(Exception): - process.kill() - except Exception: - pass - def _start_always_on_tensorboard(self): - """Start always-on Tensorboard if not already running""" if not self.always_on_tensorboard_subprocess: - self.always_on_tensorboard_subprocess = self._start_tensorboard_subprocess() + self.always_on_tensorboard_subprocess = start_filtered_tensorboard(self.train_config) def _stop_always_on_tensorboard(self): - """Stop always-on Tensorboard""" if self.always_on_tensorboard_subprocess: - self._stop_tensorboard_subprocess(self.always_on_tensorboard_subprocess) + stop_tensorboard(self.always_on_tensorboard_subprocess) self.always_on_tensorboard_subprocess = None def _restart_always_on_tensorboard(self): @@ -756,7 +707,6 @@ def _restart_always_on_tensorboard(self): self._start_always_on_tensorboard() def _on_train_only_tensorboard_toggle(self): - """Handle train-only Tensorboard toggle""" if self.train_config.tensorboard: # If train-only is enabled, disable always-on if self.train_config.tensorboard_always_on: @@ -764,7 +714,6 @@ def _on_train_only_tensorboard_toggle(self): self._stop_always_on_tensorboard() def _on_always_on_tensorboard_toggle(self): - """Handle always-on Tensorboard toggle""" if self.train_config.tensorboard_always_on: # If always-on is enabled, disable train-only if self.train_config.tensorboard: @@ -776,7 +725,6 @@ def _on_always_on_tensorboard_toggle(self): self._stop_always_on_tensorboard() def _on_workspace_dir_change(self, new_workspace_dir: str): - """Handle direct workspace directory changes""" if new_workspace_dir != self.current_workspace_dir: self.current_workspace_dir = new_workspace_dir # Restart always-on Tensorboard with new workspace directory if it's enabled diff --git a/modules/util/tensorboard_util.py b/modules/util/tensorboard_util.py new file mode 100644 index 000000000..c9b06cfe6 --- /dev/null +++ b/modules/util/tensorboard_util.py @@ -0,0 +1,83 @@ +import contextlib +import os +import subprocess +import sys +import threading +from pathlib import Path + +from modules.util.config.TrainConfig import TrainConfig + + +def get_tensorboard_args(config: TrainConfig) -> list[str]: + """Get Tensorboard command line arguments""" + tensorboard_executable = os.path.join(os.path.dirname(sys.executable), "tensorboard") + tensorboard_log_dir = os.path.join(config.workspace_dir, "tensorboard") + + os.makedirs(Path(tensorboard_log_dir).absolute(), exist_ok=True) + + args = [ + tensorboard_executable, + "--logdir", + tensorboard_log_dir, + "--port", + str(config.tensorboard_port), + "--samples_per_plugin=images=100,scalars=10000", + ] + + if config.tensorboard_expose: + args.append("--bind_all") + + return args + + +def start_filtered_tensorboard(config: TrainConfig) -> subprocess.Popen | None: + """Start Tensorboard with filtered output""" + try: + print("Starting Tensorboard please wait...") + env = os.environ.copy() + env['PYTHONWARNINGS'] = 'ignore::UserWarning:tensorboard.default' + + process = subprocess.Popen( + get_tensorboard_args(config), + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True + ) + + def filter_output(): + for line in iter(process.stdout.readline, ''): + # https://github.com/tensorflow/tensorboard/issues/7003 + line_lower = line.lower().strip() + if ( + "pkg_resources" in line_lower or + "installation not found" in line_lower + ): + continue # Skip these lines + + if line.strip(): + print(line, end='') + + threading.Thread(target=filter_output, daemon=True).start() + + return process + + except Exception as e: + print(f"Failed to start Tensorboard: {e}") + return None + + +def stop_tensorboard(process: subprocess.Popen | None): + """Gracefully stop the Tensorboard subprocess""" + if not process: + return + + try: + print("Stopping Tensorboard...") + process.terminate() + process.wait(timeout=5) + except subprocess.TimeoutExpired: + with contextlib.suppress(Exception): + process.kill() + except Exception: + pass From 6c7a7bc25dc6335a6b7ada07a19c4c9ca39b9554 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Sat, 5 Jul 2025 03:37:02 +1000 Subject: [PATCH 05/64] Move to option_kv to reduce clutter and make UI logic easier. --- modules/trainer/GenericTrainer.py | 5 +-- modules/ui/TrainUI.py | 51 ++++++++++++---------------- modules/util/config/BaseConfig.py | 17 +++++++--- modules/util/config/TrainConfig.py | 23 +++++++++++-- modules/util/enum/TensorboardMode.py | 7 ++++ modules/util/ui/UIState.py | 37 +++++++++++--------- 6 files changed, 84 insertions(+), 56 deletions(-) create mode 100644 modules/util/enum/TensorboardMode.py diff --git a/modules/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index 7eec514fc..f9bb2cf60 100644 --- a/modules/trainer/GenericTrainer.py +++ b/modules/trainer/GenericTrainer.py @@ -23,6 +23,7 @@ from modules.util.enum.ConceptType import ConceptType from modules.util.enum.FileType import FileType from modules.util.enum.ModelFormat import ModelFormat +from modules.util.enum.TensorboardMode import TensorboardMode from modules.util.enum.TimeUnit import TimeUnit from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.memory_util import TorchMemoryRecorder @@ -66,7 +67,7 @@ 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) 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: + if config.tensorboard_mode == TensorboardMode.TRAIN_ONLY: super()._start_tensorboard() self.model = None @@ -799,7 +800,7 @@ def end(self): self.tensorboard.close() - if self.config.tensorboard and not self.config.tensorboard_always_on: + if self.config.tensorboard_mode == TensorboardMode.TRAIN_ONLY: super()._stop_tensorboard() for handle in self.grad_hook_handles: diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index 1bdefeba7..3f081580c 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -27,6 +27,7 @@ from modules.util.enum.DataType import DataType from modules.util.enum.ImageFormat import ImageFormat from modules.util.enum.ModelType import ModelType +from modules.util.enum.TensorboardMode import TensorboardMode from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.tensorboard_util import start_filtered_tensorboard, stop_tensorboard from modules.util.torch_util import torch_gc @@ -95,6 +96,7 @@ def __init__(self): self._initialize_tensorboard() self.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self._on_workspace_dir_change_trace) + self.tensorboard_mode_trace_id = self.ui_state.add_var_trace("tensorboard_mode", self._on_tensorboard_mode_change) # Persistent profiling window. self.profiling_window = ProfilingWindow(self) @@ -106,6 +108,8 @@ def __close(self): self._stop_always_on_tensorboard() if hasattr(self, 'workspace_dir_trace_id'): self.ui_state.remove_var_trace("workspace_dir", self.workspace_dir_trace_id) + if hasattr(self, 'tensorboard_mode_trace_id'): + self.ui_state.remove_var_trace("tensorboard_mode", self.tensorboard_mode_trace_id) self.quit() def top_bar(self, master): @@ -209,22 +213,21 @@ def create_general_tab(self, master): components.dir_entry(frame, 5, 1, self.ui_state, "debug_dir") # tensorboard - components.label(frame, 6, 0, "Tensorboard (Train Only)", - tooltip="Starts the Tensorboard Web UI during training only") - components.switch(frame, 6, 1, self.ui_state, "tensorboard", command=self._on_train_only_tensorboard_toggle) - - components.label(frame, 7, 0, "Always-On Tensorboard", - tooltip="Keep Tensorboard accessible even when not training. Useful for monitoring completed training sessions.") - components.switch(frame, 7, 1, self.ui_state, "tensorboard_always_on", command=self._on_always_on_tensorboard_toggle) - - components.label(frame, 8, 0, "Tensorboard Port", + components.label(frame, 6, 0, "Tensorboard Mode", + tooltip="Off: Disabled.\nTrain Only: Active only during training.\nAlways On: Always available.") + components.options_kv(frame, 6, 1, [ + ("Off", TensorboardMode.OFF), + ("Train only", TensorboardMode.TRAIN_ONLY), + ("Always on", TensorboardMode.ALWAYS_ON), + ], self.ui_state, "tensorboard_mode") + + components.label(frame, 7, 0, "Tensorboard Port", tooltip="Port to use for Tensorboard link") - components.entry(frame, 8, 1, self.ui_state, "tensorboard_port") + components.entry(frame, 7, 1, self.ui_state, "tensorboard_port") - components.label(frame, 8, 2, "Expose Tensorboard", + components.label(frame, 7, 2, "Expose Tensorboard", tooltip="Exposes Tensorboard Web UI to all network interfaces (makes it accessible from the network)") - components.switch(frame, 8, 3, self.ui_state, "tensorboard_expose") - + components.switch(frame, 7, 3, self.ui_state, "tensorboard_expose") # validation components.label(frame, 9, 0, "Validation", @@ -688,7 +691,7 @@ def save_now(self): train_commands.save() def _initialize_tensorboard(self): - if self.train_config.tensorboard_always_on: + if self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON: self._start_always_on_tensorboard() def _start_always_on_tensorboard(self): @@ -702,33 +705,21 @@ def _stop_always_on_tensorboard(self): def _restart_always_on_tensorboard(self): """Restart always-on Tensorboard (for workspace directory changes)""" - if self.train_config.tensorboard_always_on: + if self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON: self._stop_always_on_tensorboard() self._start_always_on_tensorboard() - def _on_train_only_tensorboard_toggle(self): - if self.train_config.tensorboard: - # If train-only is enabled, disable always-on - if self.train_config.tensorboard_always_on: - self.ui_state.get_var("tensorboard_always_on").set(False) - self._stop_always_on_tensorboard() - - def _on_always_on_tensorboard_toggle(self): - if self.train_config.tensorboard_always_on: - # If always-on is enabled, disable train-only - if self.train_config.tensorboard: - self.ui_state.get_var("tensorboard").set(False) - # Start always-on Tensorboard + def _on_tensorboard_mode_change(self, *args): + if self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON: self._start_always_on_tensorboard() else: - # Stop always-on Tensorboard self._stop_always_on_tensorboard() def _on_workspace_dir_change(self, new_workspace_dir: str): if new_workspace_dir != self.current_workspace_dir: self.current_workspace_dir = new_workspace_dir # Restart always-on Tensorboard with new workspace directory if it's enabled - if self.train_config.tensorboard_always_on: + if self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON: self._restart_always_on_tensorboard() def _on_workspace_dir_change_trace(self, *args): diff --git a/modules/util/config/BaseConfig.py b/modules/util/config/BaseConfig.py index b38acd87d..8e8074fed 100644 --- a/modules/util/config/BaseConfig.py +++ b/modules/util/config/BaseConfig.py @@ -108,11 +108,18 @@ def from_dict(self, data: dict) -> 'BaseConfig': else: setattr(self, name, str(data[name])) elif issubclass_safe(self.types[name], Enum): - if isinstance(data[name], str): - if self.nullables[name]: - setattr(self, name, None if data[name] is None else self.types[name][data[name]]) - else: - setattr(self, name, self.types[name][data[name]]) + if data[name] is None and self.nullables[name]: + setattr(self, name, None) + elif isinstance(data[name], str): + try: + # Try to set by value first (for StrEnum) + setattr(self, name, self.types[name](data[name])) + except ValueError: + try: + # Fallback to setting by name for other Enum types + setattr(self, name, self.types[name][data[name]]) + except KeyError: + print(f"Could not set {name} as {str(data[name])}") else: setattr(self, name, data[name]) elif self.types[name] is bool: diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index ffd0a0ef2..c13d3f5cf 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -22,6 +22,7 @@ from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.ModelType import ModelType, PeftType from modules.util.enum.Optimizer import Optimizer +from modules.util.enum.TensorboardMode import TensorboardMode from modules.util.enum.TimestepDistribution import TimestepDistribution from modules.util.enum.TimeUnit import TimeUnit from modules.util.enum.TrainingMethod import TrainingMethod @@ -269,7 +270,7 @@ class TrainConfig(BaseConfig): debug_dir: str workspace_dir: str cache_dir: str - tensorboard: bool + tensorboard_mode: TensorboardMode tensorboard_expose: bool tensorboard_always_on: bool tensorboard_port: str @@ -453,7 +454,7 @@ class TrainConfig(BaseConfig): def __init__(self, data: list[(str, Any, type, bool)]): super().__init__( data, - config_version=6, + config_version=7, config_migrations={ 0: self.__migration_0, 1: self.__migration_1, @@ -461,6 +462,7 @@ def __init__(self, data: list[(str, Any, type, bool)]): 3: self.__migration_3, 4: self.__migration_4, 5: self.__migration_5, + 6: self.__migration_6, } ) @@ -626,6 +628,21 @@ def __migration_5(self, data: dict) -> dict: return migrated_data + def __migration_6(self, data: dict) -> dict: + migrated_data = data.copy() + + tensorboard = migrated_data.pop("tensorboard", True) + tensorboard_always_on = migrated_data.pop("tensorboard_always_on", False) + + if tensorboard_always_on: + migrated_data["tensorboard_mode"] = TensorboardMode.ALWAYS_ON + elif tensorboard: + migrated_data["tensorboard_mode"] = TensorboardMode.TRAIN_ONLY + else: + migrated_data["tensorboard_mode"] = TensorboardMode.OFF + + return migrated_data + def weight_dtypes(self) -> ModelWeightDtypes: return ModelWeightDtypes( self.train_dtype, @@ -769,7 +786,7 @@ def default_values() -> 'TrainConfig': data.append(("debug_dir", "debug", str, False)) data.append(("workspace_dir", "workspace/run", str, False)) data.append(("cache_dir", "workspace-cache/run", str, False)) - data.append(("tensorboard", True, bool, False)) + data.append(("tensorboard_mode", TensorboardMode.TRAIN_ONLY, TensorboardMode, False)) data.append(("tensorboard_expose", False, bool, False)) data.append(("tensorboard_always_on", False, bool, False)) data.append(("tensorboard_port", 6006, int, False)) diff --git a/modules/util/enum/TensorboardMode.py b/modules/util/enum/TensorboardMode.py new file mode 100644 index 000000000..aa3274f59 --- /dev/null +++ b/modules/util/enum/TensorboardMode.py @@ -0,0 +1,7 @@ +from enum import StrEnum + + +class TensorboardMode(StrEnum): + OFF = "off" + TRAIN_ONLY = "train_only" + ALWAYS_ON = "always_on" diff --git a/modules/util/ui/UIState.py b/modules/util/ui/UIState.py index 258bba40c..943675c7c 100644 --- a/modules/util/ui/UIState.py +++ b/modules/util/ui/UIState.py @@ -68,22 +68,27 @@ def update(_0, _1, _2): return update def __set_enum_var(self, obj, is_dict, name, var, var_type, nullable): - if is_dict: - def update(_0, _1, _2): - string_var = var.get() - if (string_var == "" or string_var == "None") and nullable: - obj[name] = None - else: - obj[name] = var_type[string_var] - self.__call_var_traces(name) - else: - def update(_0, _1, _2): - string_var = var.get() - if (string_var == "" or string_var == "None") and nullable: - setattr(obj, name, None) - else: - setattr(obj, name, var_type[string_var]) - self.__call_var_traces(name) + def update(_0, _1, _2): + string_var = var.get() + value_to_set = None + + if not ((string_var == "" or string_var == "None") and nullable): + try: + # Try to set by value first (for StrEnum) + value_to_set = var_type(string_var) + except ValueError: + try: + # Fallback to setting by name for other Enum types + value_to_set = var_type[string_var] + except KeyError: + print(f"Could not set {name} as {string_var}") + + if is_dict: + obj[name] = value_to_set + else: + setattr(obj, name, value_to_set) + + self.__call_var_traces(name) return update From aef07addfb93235a12d2cdfcdc5f1e662c66b552 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Sat, 5 Jul 2025 04:24:09 +1000 Subject: [PATCH 06/64] Made the default value always on --- modules/util/config/TrainConfig.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index c13d3f5cf..aad54672b 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -634,12 +634,10 @@ def __migration_6(self, data: dict) -> dict: tensorboard = migrated_data.pop("tensorboard", True) tensorboard_always_on = migrated_data.pop("tensorboard_always_on", False) - if tensorboard_always_on: + if tensorboard_always_on or tensorboard: migrated_data["tensorboard_mode"] = TensorboardMode.ALWAYS_ON - elif tensorboard: - migrated_data["tensorboard_mode"] = TensorboardMode.TRAIN_ONLY else: - migrated_data["tensorboard_mode"] = TensorboardMode.OFF + migrated_data["tensorboard_mode"] = TensorboardMode.ALWAYS_ON return migrated_data @@ -786,7 +784,7 @@ def default_values() -> 'TrainConfig': data.append(("debug_dir", "debug", str, False)) data.append(("workspace_dir", "workspace/run", str, False)) data.append(("cache_dir", "workspace-cache/run", str, False)) - data.append(("tensorboard_mode", TensorboardMode.TRAIN_ONLY, TensorboardMode, False)) + data.append(("tensorboard_mode", TensorboardMode.ALWAYS_ON, TensorboardMode, False)) data.append(("tensorboard_expose", False, bool, False)) data.append(("tensorboard_always_on", False, bool, False)) data.append(("tensorboard_port", 6006, int, False)) From 7f7796cfad86c0b97372e4eea193286e8a95fed0 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Sat, 12 Jul 2025 22:50:26 +1000 Subject: [PATCH 07/64] Potential fix for cloud always-on conflict. --- modules/trainer/BaseTrainer.py | 16 ++++++++++++---- modules/trainer/CloudTrainer.py | 5 +++-- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/modules/trainer/BaseTrainer.py b/modules/trainer/BaseTrainer.py index dcd5e64d3..d4fa67d80 100644 --- a/modules/trainer/BaseTrainer.py +++ b/modules/trainer/BaseTrainer.py @@ -22,7 +22,7 @@ class BaseTrainer( metaclass=ABCMeta, ): - tensorboard_subprocess: subprocess.Popen + tensorboard_subprocess: subprocess.Popen = None def __init__(self, config: TrainConfig, callbacks: TrainCallbacks, commands: TrainCommands): super().__init__() @@ -85,8 +85,16 @@ def create_model_sampler(self, model: BaseModel) -> BaseModelSampler: ) def _start_tensorboard(self): - self.tensorboard_subprocess = start_filtered_tensorboard(self.config) + is_remote_with_tunnel = (self.config.cloud and self.config.cloud.tensorboard_tunnel) + if is_remote_with_tunnel: + print("Remote training with TensorBoard TCP tunnel enabled.") + self.tensorboard_subprocess = None + else: + print("Starting local TensorBoard instance.") + self.tensorboard_subprocess = start_filtered_tensorboard(self.config) def _stop_tensorboard(self): - stop_tensorboard(self.tensorboard_subprocess) - self.tensorboard_subprocess = None + if self.tensorboard_subprocess: + print("Stopping local TensorBoard instance.") + stop_tensorboard(self.tensorboard_subprocess) + self.tensorboard_subprocess = None diff --git a/modules/trainer/CloudTrainer.py b/modules/trainer/CloudTrainer.py index bcab32428..7e9484d56 100644 --- a/modules/trainer/CloudTrainer.py +++ b/modules/trainer/CloudTrainer.py @@ -30,7 +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: + + if config.tensorboard and not config.tensorboard_always_on: super()._start_tensorboard() match config.cloud.type: @@ -114,7 +115,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: + if self.config.tensorboard and not self.config.tensorboard_always_on: super()._stop_tensorboard() if self.config.cloud.delete_workspace and not self.error_caught and not self.commands.get_stop_command(): From 15256d8e772280bf6e2a29178e4269e53d48802c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 2 Sep 2025 14:40:00 +0000 Subject: [PATCH 08/64] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- modules/ui/TrainUI.py | 3 --- modules/util/config/TrainConfig.py | 2 +- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index 0e9ba4d28..9f0a7b2e0 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -1,9 +1,6 @@ import ctypes import json -import os import platform -import subprocess -import sys import threading import traceback import webbrowser diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 98bf7ea3f..000dbe79c 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -643,7 +643,7 @@ def __migration_6(self, data: dict) -> dict: migrated_data["lora_layer_preset"] = "full" return migrated_data - + def __migration_7(self, data: dict) -> dict: migrated_data = data.copy() From 272e8d47343ebe94ee8ea01d5fd4acead682c78d Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Mon, 13 Oct 2025 03:14:07 +1100 Subject: [PATCH 09/64] Adds initial drag-n-drop fuctionality --- modules/ui/ModelTab.py | 32 +++++++++- modules/ui/TrainUI.py | 4 +- modules/util/ui/components.py | 107 +++++++++++++++++++++++++++++++++- 3 files changed, 138 insertions(+), 5 deletions(-) diff --git a/modules/ui/ModelTab.py b/modules/ui/ModelTab.py index d27118bb3..d313a82d2 100644 --- a/modules/ui/ModelTab.py +++ b/modules/ui/ModelTab.py @@ -480,10 +480,35 @@ def __create_output_components( allow_diffusers: bool = False, allow_legacy_safetensors: bool = False, ) -> int: + def get_valid_extensions() -> list[str]: + format_str = self.ui_state.get_var("output_model_format").get() + + # Map format to extensions + format_extensions = { + "SAFETENSORS": [".safetensors"], + "DIFFUSERS": [".json"], + "LEGACY_SAFETENSORS": [".safetensors"], + } + + return format_extensions.get(format_str, []) + + def update_output_validation(_=None): + + output_var = self.ui_state.get_var("output_model_destination") + current_value = output_var.get() + + if current_value: + output_var.set(current_value) + # output model destination components.label(self.scroll_frame, row, 0, "Model Output Destination", tooltip="Filename or directory where the output model is saved") - components.file_entry(self.scroll_frame, row, 1, self.ui_state, "output_model_destination", is_output=True) + + components.file_entry( + self.scroll_frame, row, 1, self.ui_state, "output_model_destination", + is_output=True, + valid_extensions=get_valid_extensions() + ) # output data type components.label(self.scroll_frame, row, 3, "Output Data Type", @@ -509,7 +534,10 @@ def __create_output_components( components.label(self.scroll_frame, row, 0, "Output Format", tooltip="Format to use when saving the output model") - components.options_kv(self.scroll_frame, row, 1, formats, self.ui_state, "output_model_format") + components.options_kv( + self.scroll_frame, row, 1, formats, self.ui_state, "output_model_format", + command=update_output_validation # Trigger validation when format changes + ) # include config components.label(self.scroll_frame, row, 3, "Include Config", diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index d3da16162..68e6dcd4f 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -47,6 +47,7 @@ import customtkinter as ctk from customtkinter import AppearanceModeTracker +from tkinterdnd2 import TkinterDnD # chunk for forcing Windows to ignore DPI scaling when moving between monitors # fixes the long standing transparency bug https://github.com/Nerogar/OneTrainer/issues/90 @@ -55,7 +56,7 @@ # https://learn.microsoft.com/en-us/windows/win32/hidpi/setting-the-default-dpi-awareness-for-a-process#setting-default-awareness-programmatically ctypes.windll.shcore.SetProcessDpiAwareness(1) # PROCESS_SYSTEM_DPI_AWARE -class TrainUI(ctk.CTk): +class TrainUI(ctk.CTk, TkinterDnD.DnDWrapper): set_step_progress: Callable[[int, int], None] set_epoch_progress: Callable[[int, int], None] @@ -92,6 +93,7 @@ class TrainUI(ctk.CTk): def __init__(self): super().__init__() + self.TkdndVersion = TkinterDnD._require(self) self.title("OneTrainer") self.geometry("1100x740") diff --git a/modules/util/ui/components.py b/modules/util/ui/components.py index 3043f32fc..f4d7a0a71 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -1,4 +1,5 @@ import contextlib +import os from collections.abc import Callable from tkinter import filedialog from typing import Any @@ -11,10 +12,48 @@ import customtkinter as ctk from customtkinter.windows.widgets.scaling import CTkScalingBaseClass from PIL import Image +from tkinterdnd2 import DND_FILES PAD = 10 +def _drop_enter(event): + event.widget.focus_force() + return event.action + + +def _drop_leave(event): + return event.action + + +def _create_drop_handler(entry_widget, ui_state, var_name, command=None): + def drop(event): + if event.data: + file_path = event.data.strip('{}') + + ui_state.get_var(var_name).set(file_path) + + entry_widget.focus_force() + entry_widget.event_generate('') + entry_widget.event_generate('') # mark as touched + entry_widget.event_generate('') + + if command: + command(file_path) + return event.action + return drop + + +def _register_drop_target(entry_widget, ui_state, var_name, command=None): + try: + entry_widget.drop_target_register(DND_FILES) + entry_widget.dnd_bind('<>', _drop_enter) + entry_widget.dnd_bind('<>', _drop_leave) + entry_widget.dnd_bind('<>', _create_drop_handler(entry_widget, ui_state, var_name, command)) + except Exception: + pass + + def app_title(master, row, column): frame = ctk.CTkFrame(master) frame.grid(row=row, column=column, padx=5, pady=5, sticky="nsew") @@ -49,6 +88,7 @@ def entry( wide_tooltip: bool = False, width: int = 140, sticky: str = "new", + custom_validator: Callable[[str], tuple[bool, str]] = None, ): var = ui_state.get_var(var_name) trace_id = None @@ -122,6 +162,12 @@ def fail(_reason: str): elif declared_type is bool: if value.lower() not in ("true", "false", "0", "1"): return fail("Invalid bool") + + if custom_validator: + is_valid, error_msg = custom_validator(value) + if not is_valid: + return fail(error_msg) + return success() except ValueError: return fail("Invalid value") @@ -206,13 +252,49 @@ def file_entry( allow_model_files: bool = True, allow_image_files: bool = False, command: Callable[[str], None] = None, + valid_extensions: list[str] = None, ): frame = ctk.CTkFrame(master, fg_color="transparent") frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") frame.grid_columnconfigure(0, weight=1) - entry(frame,row=0, column=0, ui_state=ui_state, var_name=var_name) + def validate_file_path(value: str) -> tuple[bool, str]: + if not value: + return True, "" + + if os.path.isdir(value): + return False, "Must be a file, not a directory" + + # For output files, validate extension if specified + if is_output: + # Check if user has entered a file extension (contains a period after last path separator) + if '.' in os.path.basename(value): + if valid_extensions: + file_ext = os.path.splitext(value)[1].lower() + if file_ext not in valid_extensions: + valid_exts_str = ', '.join(valid_extensions) + return False, f"Invalid extension. Expected: {valid_exts_str}" + + # Check that parent directory exists + parent_dir = os.path.dirname(value) + if parent_dir and not os.path.isdir(parent_dir): + return False, "Parent directory does not exist" + return True, "" + + # For input files, verify the file exists + if not os.path.isfile(value): + return False, "File does not exist" + + return True, "" + + entry_widget = entry( + frame, row=0, column=0, ui_state=ui_state, var_name=var_name, + custom_validator=validate_file_path, command=command + ) + + # Register drag-and-drop support + _register_drop_target(entry_widget, ui_state, var_name, command) def __open_dialog(): filetypes = [ @@ -256,7 +338,28 @@ def dir_entry(master, row, column, ui_state: UIState, var_name: str, command: Ca frame.grid_columnconfigure(0, weight=1) - entry(frame, row=0, column=0, ui_state=ui_state, var_name=var_name) + def validate_and_normalize_directory(value: str) -> tuple[bool, str]: + if not value: + return True, "" + + # convert filepaths to dir + if os.path.isfile(value): + normalized = os.path.dirname(value) + ui_state.get_var(var_name).set(normalized) + return True, "" + + # check if valid dir + if os.path.isdir(value): + return True, "" + + return True, "" + + entry_widget = entry( + frame, row=0, column=0, ui_state=ui_state, var_name=var_name, + custom_validator=validate_and_normalize_directory, command=command + ) + + _register_drop_target(entry_widget, ui_state, var_name, command) def __open_dialog(): dir_path = filedialog.askdirectory() From 15bad8b95652839159c15662c250d1e0b6a017e0 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Mon, 13 Oct 2025 18:37:14 +1100 Subject: [PATCH 10/64] Add dnd to conceptsTab and small cleanup --- modules/ui/ConfigList.py | 66 ++++++++++++++++++++++++----------- modules/util/ui/components.py | 20 +++++++++++ 2 files changed, 65 insertions(+), 21 deletions(-) diff --git a/modules/ui/ConfigList.py b/modules/ui/ConfigList.py index 6fa6a8562..bb2ddb310 100644 --- a/modules/ui/ConfigList.py +++ b/modules/ui/ConfigList.py @@ -128,6 +128,14 @@ def __create_configs_dropdown(self): ) self._update_toggle_button_text() + def _create_placeholder(self): + placeholder = ctk.CTkLabel( + self.element_list, + text="Either click the 'Add Concept' button or drag and drop a directory here to add a new concept", + text_color="gray50" + ) + placeholder.grid(row=0, column=0, padx=20, pady=20) + def _create_element_list(self): if not self.from_external_file: self.current_config = getattr(self.train_config, self.attr_name) @@ -142,6 +150,9 @@ def _create_element_list(self): if self.is_full_width: self.element_list.grid_columnconfigure(0, weight=1) + if len(self.current_config) == 0: + self._create_placeholder() + for i, element in enumerate(self.current_config): widget = self.create_widget( self.element_list, element, i, @@ -154,6 +165,21 @@ def _create_element_list(self): widget.place_in_list() + components.register_dir_drop_target(self.master, self.__handle_dir_drop) + + def __handle_dir_drop(self, dir_path: str): + if any(hasattr(el, 'path') and el.path == dir_path for el in self.current_config): + return + + new_element = self.create_new_element() + + if hasattr(new_element, 'path'): + new_element.path = dir_path + + self.current_config.append(new_element) + self._create_element_list() + self.save_current_config() + def __load_available_config_names(self): if os.path.isdir(self.config_dir): for path in os.listdir(self.config_dir): @@ -161,6 +187,7 @@ def __load_available_config_names(self): if path.endswith(".json") and os.path.isfile(path): name = os.path.basename(path) name = os.path.splitext(name)[0] + name = path_util.safe_filename(name) self.configs.append((name, path)) if len(self.configs) == 0: @@ -177,11 +204,16 @@ def __create_config(self, name: str): def __add_config(self): dialogs.StringInputDialog(self.master, "name", "Name", self.__create_config) - def __add_element(self): + def __add_element_to_list(self, new_element): i = len(self.current_config) - new_element = self.create_new_element() + + # placeholder removal + if i == 0: + for child in self.element_list.winfo_children(): + child.destroy() self.current_config.append(new_element) + widget = self.create_widget( self.element_list, new_element, i, self.__open_element_window, @@ -190,39 +222,31 @@ def __add_element(self): self.save_current_config ) self.widgets.append(widget) - widget.place_in_list() - self.save_current_config() + def __add_element(self): + new_element = self.create_new_element() + self.__add_element_to_list(new_element) + def __clone_element(self, clone_i, modify_element_fun=None): - i = len(self.current_config) new_element = copy.deepcopy(self.current_config[clone_i]) if modify_element_fun is not None: new_element = modify_element_fun(new_element) - self.current_config.append(new_element) - widget = self.create_widget( - self.element_list, new_element, i, - self.__open_element_window, - self.__remove_element, - self.__clone_element, - self.save_current_config - ) - self.widgets.append(widget) - - widget.place_in_list() - - self.save_current_config() + self.__add_element_to_list(new_element) def __remove_element(self, remove_i): self.current_config.pop(remove_i) self.widgets.pop(remove_i).destroy() - for i, widget in enumerate(self.widgets): - widget.i = i - widget.place_in_list() + if len(self.current_config) == 0: + self._create_placeholder() + else: + for i, widget in enumerate(self.widgets): + widget.i = i + widget.place_in_list() self.save_current_config() diff --git a/modules/util/ui/components.py b/modules/util/ui/components.py index f4d7a0a71..efd57112d 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -244,6 +244,26 @@ def new_destroy(): return component +def register_dir_drop_target(widget, drop_callback: Callable[[str], None]): + def drop_handler(event): + if event.data: + path = event.data.strip('{}') + + if os.path.isfile(path): + path = os.path.dirname(path) + + if os.path.isdir(path): + drop_callback(path) + + return event.action + + try: + widget.drop_target_register(DND_FILES) + widget.dnd_bind('<>', _drop_enter) + widget.dnd_bind('<>', _drop_leave) + widget.dnd_bind('<>', drop_handler) + except Exception: + pass def file_entry( master, row, column, ui_state: UIState, var_name: str, From 3ca797365bcf597d0de7b775097e61ff69f34d4f Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Mon, 13 Oct 2025 18:53:51 +1100 Subject: [PATCH 11/64] Add tkinterdnd2 to requirements --- requirements-global.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements-global.txt b/requirements-global.txt index 85ee87118..ae2ef6b49 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -51,6 +51,7 @@ scalene==1.5.51 # ui customtkinter==5.2.2 +tkinterdnd2==0.4.3 # cloud runpod==1.7.10 From edd700112f7442bdf0815168dc1f5b675a4ad992 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Tue, 14 Oct 2025 00:39:40 +1100 Subject: [PATCH 12/64] Allow multi dnd for concepts and cleanup validation further + move to utils --- modules/ui/ConfigList.py | 2 +- modules/ui/ModelTab.py | 3 +- modules/util/path_util.py | 63 ++++++++++++++++++ modules/util/ui/components.py | 120 ++++------------------------------ modules/util/ui/ui_utils.py | 91 ++++++++++++++++++++++++++ 5 files changed, 171 insertions(+), 108 deletions(-) diff --git a/modules/ui/ConfigList.py b/modules/ui/ConfigList.py index bb2ddb310..86630c945 100644 --- a/modules/ui/ConfigList.py +++ b/modules/ui/ConfigList.py @@ -165,7 +165,7 @@ def _create_element_list(self): widget.place_in_list() - components.register_dir_drop_target(self.master, self.__handle_dir_drop) + components.register_concept_drop_target(self.master, self.__handle_dir_drop) def __handle_dir_drop(self, dir_path: str): if any(hasattr(el, 'path') and el.path == dir_path for el in self.current_config): diff --git a/modules/ui/ModelTab.py b/modules/ui/ModelTab.py index d313a82d2..bf4ac9479 100644 --- a/modules/ui/ModelTab.py +++ b/modules/ui/ModelTab.py @@ -507,7 +507,8 @@ def update_output_validation(_=None): components.file_entry( self.scroll_frame, row, 1, self.ui_state, "output_model_destination", is_output=True, - valid_extensions=get_valid_extensions() + valid_extensions=get_valid_extensions(), + path_type="file" # Allow both file and directory paths ) # output data type diff --git a/modules/util/path_util.py b/modules/util/path_util.py index 21932b9c2..1e0d4affd 100644 --- a/modules/util/path_util.py +++ b/modules/util/path_util.py @@ -52,3 +52,66 @@ def supported_video_extensions() -> set[str]: def is_supported_video_extension(extension: str) -> bool: return extension.lower() in SUPPORTED_VIDEO_EXTENSIONS + + +def validate_file_path( + value: str, + is_output: bool = False, + valid_extensions: list[str] | None = None, + path_type: str = "file", # "file" or "directory" +) -> tuple[bool, str]: + if not value: + return True, "" + + # windows invalid characters + if is_output and os.name == 'nt': + invalid_chars = ['<', '>', '|', '\0', '"'] + path_parts = [os.path.basename(value)] + + dirname = os.path.dirname(value) + if dirname: + path_parts.extend( + part for part in dirname.replace('\\', '/').split('/') + if part and not (len(part) == 2 and part[1] == ':') + ) + + for part in path_parts: + for char in invalid_chars: + if char in part: + location = "filename" if part == path_parts[0] else "path" + return False, f"Invalid character in {location}: '{char}'" + + if is_output: + if path_type == "directory" and os.path.isfile(value): + return False, "Must be a directory, not a file" + + if path_type == "file": + if os.path.isdir(value): + return False, "Must be a file, not a directory" + + # Validate extension if applicable + if valid_extensions and "." in os.path.basename(value): + file_ext = os.path.splitext(value)[1].lower() + if file_ext not in valid_extensions: + return False, f"Invalid extension. Expected: {', '.join(valid_extensions)}" + + # Validate parent directory exists + parent_dir = os.path.dirname(value) + if parent_dir and not os.path.isdir(parent_dir): + return False, "Parent directory does not exist" + + return True, "" + + if path_type == "directory": + if os.path.isfile(value) or (not os.path.isdir(value) and os.path.splitext(value)[1]): + value = os.path.dirname(value) or '.' + + if not os.path.isdir(value): + return False, "Directory does not exist" + else: # file + if os.path.isdir(value): + return False, "Must be a file, not a directory" + if not os.path.isfile(value): + return False, "File does not exist" + + return True, "" diff --git a/modules/util/ui/components.py b/modules/util/ui/components.py index efd57112d..0a3afdef3 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -1,59 +1,21 @@ import contextlib -import os from collections.abc import Callable from tkinter import filedialog from typing import Any from modules.util.enum.TimeUnit import TimeUnit -from modules.util.path_util import supported_image_extensions +from modules.util.path_util import supported_image_extensions, validate_file_path from modules.util.ui.ToolTip import ToolTip +from modules.util.ui.ui_utils import _register_drop_target from modules.util.ui.UIState import UIState import customtkinter as ctk from customtkinter.windows.widgets.scaling import CTkScalingBaseClass from PIL import Image -from tkinterdnd2 import DND_FILES PAD = 10 -def _drop_enter(event): - event.widget.focus_force() - return event.action - - -def _drop_leave(event): - return event.action - - -def _create_drop_handler(entry_widget, ui_state, var_name, command=None): - def drop(event): - if event.data: - file_path = event.data.strip('{}') - - ui_state.get_var(var_name).set(file_path) - - entry_widget.focus_force() - entry_widget.event_generate('') - entry_widget.event_generate('') # mark as touched - entry_widget.event_generate('') - - if command: - command(file_path) - return event.action - return drop - - -def _register_drop_target(entry_widget, ui_state, var_name, command=None): - try: - entry_widget.drop_target_register(DND_FILES) - entry_widget.dnd_bind('<>', _drop_enter) - entry_widget.dnd_bind('<>', _drop_leave) - entry_widget.dnd_bind('<>', _create_drop_handler(entry_widget, ui_state, var_name, command)) - except Exception: - pass - - def app_title(master, row, column): frame = ctk.CTkFrame(master) frame.grid(row=row, column=column, padx=5, pady=5, sticky="nsew") @@ -244,76 +206,35 @@ def new_destroy(): return component -def register_dir_drop_target(widget, drop_callback: Callable[[str], None]): - def drop_handler(event): - if event.data: - path = event.data.strip('{}') - - if os.path.isfile(path): - path = os.path.dirname(path) - - if os.path.isdir(path): - drop_callback(path) - - return event.action - - try: - widget.drop_target_register(DND_FILES) - widget.dnd_bind('<>', _drop_enter) - widget.dnd_bind('<>', _drop_leave) - widget.dnd_bind('<>', drop_handler) - except Exception: - pass def file_entry( master, row, column, ui_state: UIState, var_name: str, - is_output: bool = False, path_modifier: Callable[[str], str] = None, + is_output: bool = False, allow_model_files: bool = True, allow_image_files: bool = False, command: Callable[[str], None] = None, valid_extensions: list[str] = None, + path_type: str = "file", # "file" or "directory" ): frame = ctk.CTkFrame(master, fg_color="transparent") frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") frame.grid_columnconfigure(0, weight=1) - def validate_file_path(value: str) -> tuple[bool, str]: - if not value: - return True, "" - - if os.path.isdir(value): - return False, "Must be a file, not a directory" + # Only use metadata as fallback if is_output wasn't explicitly set + if not is_output: + meta = ui_state.get_field_metadata(var_name) + is_output = getattr(meta, 'is_output', False) - # For output files, validate extension if specified - if is_output: - # Check if user has entered a file extension (contains a period after last path separator) - if '.' in os.path.basename(value): - if valid_extensions: - file_ext = os.path.splitext(value)[1].lower() - if file_ext not in valid_extensions: - valid_exts_str = ', '.join(valid_extensions) - return False, f"Invalid extension. Expected: {valid_exts_str}" - - # Check that parent directory exists - parent_dir = os.path.dirname(value) - if parent_dir and not os.path.isdir(parent_dir): - return False, "Parent directory does not exist" - return True, "" - - # For input files, verify the file exists - if not os.path.isfile(value): - return False, "File does not exist" - - return True, "" + def validate_file_path_wrapper(value: str) -> tuple[bool, str]: + return validate_file_path(value, is_output, valid_extensions, path_type) entry_widget = entry( frame, row=0, column=0, ui_state=ui_state, var_name=var_name, - custom_validator=validate_file_path, command=command + custom_validator=validate_file_path_wrapper, command=command ) - # Register drag-and-drop support _register_drop_target(entry_widget, ui_state, var_name, command) def __open_dialog(): @@ -358,25 +279,12 @@ def dir_entry(master, row, column, ui_state: UIState, var_name: str, command: Ca frame.grid_columnconfigure(0, weight=1) - def validate_and_normalize_directory(value: str) -> tuple[bool, str]: - if not value: - return True, "" - - # convert filepaths to dir - if os.path.isfile(value): - normalized = os.path.dirname(value) - ui_state.get_var(var_name).set(normalized) - return True, "" - - # check if valid dir - if os.path.isdir(value): - return True, "" - - return True, "" + def validate_dir_path_wrapper(value: str) -> tuple[bool, str]: + return validate_file_path(value, is_output=False, valid_extensions=None, path_type="directory") entry_widget = entry( frame, row=0, column=0, ui_state=ui_state, var_name=var_name, - custom_validator=validate_and_normalize_directory, command=command + custom_validator=validate_dir_path_wrapper, command=command ) _register_drop_target(entry_widget, ui_state, var_name, command) diff --git a/modules/util/ui/ui_utils.py b/modules/util/ui/ui_utils.py index c1c9d0a68..a9a79e578 100644 --- a/modules/util/ui/ui_utils.py +++ b/modules/util/ui/ui_utils.py @@ -1,3 +1,4 @@ +import os import platform import sys import tkinter as tk @@ -7,6 +8,7 @@ from typing import Any from customtkinter import CTk, CTkToplevel +from tkinterdnd2 import DND_FILES def bind_mousewheel( @@ -99,3 +101,92 @@ def set_icon(): except Exception as e: print(f"Failed to set window icon: {e}") + +def _drop_enter(event): + event.widget.focus_force() + return event.action + + +def _drop_leave(event): + return event.action + + +def _create_drop_handler(entry_widget, ui_state, var_name, command=None): + def drop(event): + if event.data: + paths = _parse_dropped_paths(event.data) + if paths: + file_path = paths[0] # only take the first path for single-value entries + + ui_state.get_var(var_name).set(file_path) + + entry_widget.focus_force() + entry_widget.event_generate('') + entry_widget.event_generate('') # mark as touched + entry_widget.event_generate('') + + if command: + command(file_path) + return event.action + return drop + +def _parse_dropped_paths(event_data: str) -> list[str]: + paths = [] + current_path = "" + in_braces = False + + for char in event_data: + if char == '{': + in_braces = True + elif char == '}': + in_braces = False + if current_path: + paths.append(current_path) + current_path = "" + elif char == ' ' and not in_braces: + if current_path: + paths.append(current_path) + current_path = "" + else: + current_path += char + + if current_path: + paths.append(current_path) + + return [p.strip() for p in paths if p.strip()] + +def _register_drop_target(entry_widget, ui_state, var_name, command=None): + try: + entry_widget.drop_target_register(DND_FILES) + entry_widget.dnd_bind('<>', _drop_enter) + entry_widget.dnd_bind('<>', _drop_leave) + entry_widget.dnd_bind('<>', _create_drop_handler(entry_widget, ui_state, var_name, command)) + except Exception: + pass + +def register_concept_drop_target(widget, drop_callback: Callable[[str], None], allow_multiple: bool = True): + def drop_handler(event): + if not event.data: + return event.action + + paths = _parse_dropped_paths(event.data) + + for path in paths: + if os.path.isfile(path): + path = os.path.dirname(path) + + if os.path.isdir(path): + drop_callback(path) + + if not allow_multiple: + break + + return event.action + + try: + widget.drop_target_register(DND_FILES) + widget.dnd_bind('<>', _drop_enter) + widget.dnd_bind('<>', _drop_leave) + widget.dnd_bind('<>', drop_handler) + except Exception: + pass From c216f19912b0a54f0f9424a74f5e632673e2295c Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Tue, 14 Oct 2025 01:31:47 +1100 Subject: [PATCH 13/64] Adds XL Lora output correction to saving + bugfix --- modules/modelSaver/BaseModelSaver.py | 13 +++++++++++++ .../modelSaver/StableDiffusionXLLoRAModelSaver.py | 6 ++++++ modules/ui/ConfigList.py | 4 ++-- 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/modules/modelSaver/BaseModelSaver.py b/modules/modelSaver/BaseModelSaver.py index 4ad07a49a..0928225b0 100644 --- a/modules/modelSaver/BaseModelSaver.py +++ b/modules/modelSaver/BaseModelSaver.py @@ -11,6 +11,19 @@ class BaseModelSaver(metaclass=ABCMeta): def __init__(self): super().__init__() + def _ensure_correct_extension( + self, + output_model_format: ModelFormat, + output_model_destination: str, + ) -> str: + + if output_model_format.is_single_file(): + expected_extension = output_model_format.file_extension() + if not output_model_destination.endswith(expected_extension): + return output_model_destination + expected_extension + + return output_model_destination + @abstractmethod def save( self, diff --git a/modules/modelSaver/StableDiffusionXLLoRAModelSaver.py b/modules/modelSaver/StableDiffusionXLLoRAModelSaver.py index fe7ebf679..972a44aa3 100644 --- a/modules/modelSaver/StableDiffusionXLLoRAModelSaver.py +++ b/modules/modelSaver/StableDiffusionXLLoRAModelSaver.py @@ -24,6 +24,12 @@ def save( output_model_destination: str, dtype: torch.dtype | None, ): + # Ensure the output destination has the correct file extension + output_model_destination = self._ensure_correct_extension( + output_model_format, + output_model_destination + ) + lora_model_saver = StableDiffusionXLLoRASaver() embedding_model_saver = StableDiffusionXLEmbeddingSaver() diff --git a/modules/ui/ConfigList.py b/modules/ui/ConfigList.py index 86630c945..01b5b0274 100644 --- a/modules/ui/ConfigList.py +++ b/modules/ui/ConfigList.py @@ -8,7 +8,7 @@ from modules.util.config.BaseConfig import BaseConfig from modules.util.config.TrainConfig import TrainConfig from modules.util.path_util import write_json_atomic -from modules.util.ui import components, dialogs +from modules.util.ui import components, dialogs, ui_utils from modules.util.ui.UIState import UIState import customtkinter as ctk @@ -165,7 +165,7 @@ def _create_element_list(self): widget.place_in_list() - components.register_concept_drop_target(self.master, self.__handle_dir_drop) + ui_utils.register_concept_drop_target(self.master, self.__handle_dir_drop) def __handle_dir_drop(self, dir_path: str): if any(hasattr(el, 'path') and el.path == dir_path for el in self.current_config): From f759a8f8e665709af8fa0a87bf3997b0403a0d8b Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Tue, 14 Oct 2025 21:52:54 +1100 Subject: [PATCH 14/64] Move the migration to migration 8 --- modules/util/config/TrainConfig.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index be975938f..181bf8d07 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -506,6 +506,7 @@ def __init__(self, data: list[(str, Any, type, bool)]): 5: self.__migration_5, 6: self.__migration_6, 7: self.__migration_7, + 8: self.__migration_8, } ) @@ -686,6 +687,18 @@ def __migration_6(self, data: dict) -> dict: def __migration_7(self, data: dict) -> dict: migrated_data = data.copy() + # Migrate lora_layers to layer_filter + if "lora_layers" in migrated_data: + migrated_data["layer_filter"] = migrated_data.pop("lora_layers") + if "lora_layer_preset" in migrated_data: + migrated_data["layer_filter_preset"] = migrated_data.pop("lora_layer_preset") + if "lora_layers_regex" in migrated_data: + migrated_data["layer_filter_regex"] = migrated_data.pop("lora_layers_regex") + + return migrated_data + + def __migration_8(self, data: dict) -> dict: + migrated_data = data.copy() # Migrate old tensorboard booleans to enum tensorboard = migrated_data.pop("tensorboard", True) tensorboard_always_on = migrated_data.pop("tensorboard_always_on", False) @@ -697,16 +710,6 @@ def __migration_7(self, data: dict) -> dict: else: migrated_data["tensorboard_mode"] = TensorboardMode.OFF - # Migrate lora_layers to layer_filter - if "lora_layers" in migrated_data: - migrated_data["layer_filter"] = migrated_data.pop("lora_layers") - if "lora_layer_preset" in migrated_data: - migrated_data["layer_filter_preset"] = migrated_data.pop("lora_layer_preset") - if "lora_layers_regex" in migrated_data: - migrated_data["layer_filter_regex"] = migrated_data.pop("lora_layers_regex") - - return migrated_data - def weight_dtypes(self) -> ModelWeightDtypes: return ModelWeightDtypes( self.train_dtype, From 784fb3d2c7268d144f4a7034bf975a6369275a67 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Tue, 14 Oct 2025 21:57:22 +1100 Subject: [PATCH 15/64] Attempt at fixing debounced workspace bug + cleanup --- modules/trainer/BaseTrainer.py | 16 ----- modules/trainer/GenericTrainer.py | 5 -- modules/ui/TrainUI.py | 109 ++++++++++++++++++++---------- modules/util/tensorboard_util.py | 29 ++++++-- 4 files changed, 97 insertions(+), 62 deletions(-) diff --git a/modules/trainer/BaseTrainer.py b/modules/trainer/BaseTrainer.py index 53bbd1fbe..3d9f9e34b 100644 --- a/modules/trainer/BaseTrainer.py +++ b/modules/trainer/BaseTrainer.py @@ -10,7 +10,6 @@ from modules.util.callbacks.TrainCallbacks import TrainCallbacks from modules.util.commands.TrainCommands import TrainCommands from modules.util.config.TrainConfig import TrainConfig -from modules.util.tensorboard_util import start_filtered_tensorboard, stop_tensorboard from modules.util.TimedActionMixin import TimedActionMixin from modules.util.TrainProgress import TrainProgress @@ -79,18 +78,3 @@ def create_model_sampler(self, model: BaseModel) -> BaseModelSampler: self.config.model_type, self.config.training_method ) - - def _start_tensorboard(self): - is_remote_with_tunnel = (self.config.cloud and self.config.cloud.tensorboard_tunnel) - if is_remote_with_tunnel: - print("Remote training with TensorBoard TCP tunnel enabled.") - self.tensorboard_subprocess = None - else: - print("Starting local TensorBoard instance.") - self.tensorboard_subprocess = start_filtered_tensorboard(self.config) - - def _stop_tensorboard(self): - if self.tensorboard_subprocess: - print("Stopping local TensorBoard instance.") - stop_tensorboard(self.tensorboard_subprocess) - self.tensorboard_subprocess = None diff --git a/modules/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index 362f5d770..c57d32c8a 100644 --- a/modules/trainer/GenericTrainer.py +++ b/modules/trainer/GenericTrainer.py @@ -26,7 +26,6 @@ from modules.util.enum.EMAMode import EMAMode from modules.util.enum.FileType import FileType from modules.util.enum.ModelFormat import ModelFormat -from modules.util.enum.TensorboardMode import TensorboardMode from modules.util.enum.TimeUnit import TimeUnit from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.memory_util import TorchMemoryRecorder @@ -71,8 +70,6 @@ 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) self.tensorboard = SummaryWriter(os.path.join(tensorboard_log_dir, f"{config.save_filename_prefix}{get_string_timestamp()}")) - if config.tensorboard_mode == TensorboardMode.TRAIN_ONLY: - super()._start_tensorboard() self.model = None self.one_step_trained = False @@ -859,8 +856,6 @@ def end(self): if multi.is_master(): self.tensorboard.close() - if self.config.tensorboard_mode == TensorboardMode.TRAIN_ONLY: - super()._stop_tensorboard() for handle in self.grad_hook_handles: handle.remove() diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index 53522fcd2..f4a94aa9d 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -1,6 +1,8 @@ +import contextlib import ctypes import datetime import json +import logging import platform import threading import time @@ -47,6 +49,8 @@ import customtkinter as ctk from customtkinter import AppearanceModeTracker +logger = logging.getLogger(__name__) + # chunk for forcing Windows to ignore DPI scaling when moving between monitors # fixes the long standing transparency bug https://github.com/Nerogar/OneTrainer/issues/90 if platform.system() == "Windows": @@ -131,10 +135,23 @@ def __init__(self): self.always_on_tensorboard_subprocess = None self.current_workspace_dir = self.train_config.workspace_dir - self._initialize_tensorboard() + self._previous_tensorboard_mode = self.train_config.tensorboard_mode + + # Initialize tensorboard if needed + if self._previous_tensorboard_mode == TensorboardMode.ALWAYS_ON: + self._start_always_on_tensorboard() + + # Add debounced trace for workspace_dir changes + self.workspace_dir_debounce_id = None + self.workspace_dir_trace_id = self.ui_state.add_var_trace( + "workspace_dir", + lambda: self._on_workspace_dir_change_debounced() + ) - self.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self._on_workspace_dir_change_trace) - self.tensorboard_mode_trace_id = self.ui_state.add_var_trace("tensorboard_mode", self._on_tensorboard_mode_change) + self.tensorboard_mode_trace_id = self.ui_state.add_var_trace( + "tensorboard_mode", + lambda: self._on_tensorboard_mode_change() + ) # Persistent profiling window. self.profiling_window = ProfilingWindow(self) @@ -241,7 +258,8 @@ def create_general_tab(self, master): # workspace dir components.label(frame, 0, 0, "Workspace Directory", tooltip="The directory where all files of this training run are saved") - components.dir_entry(frame, 0, 1, self.ui_state, "workspace_dir", command=self._on_workspace_dir_change) + # Remove command parameter - we'll use the trace instead for debouncing + components.dir_entry(frame, 0, 1, self.ui_state, "workspace_dir") # cache dir components.label(frame, 0, 2, "Cache Directory", @@ -746,6 +764,10 @@ def open_sample_ui(self): def __training_thread_function(self): error_caught = False + # Start Tensorboard if in TRAIN_ONLY mode + if self.train_config.tensorboard_mode == TensorboardMode.TRAIN_ONLY: + self.always_on_tensorboard_subprocess = start_filtered_tensorboard(self.train_config) + self.training_callbacks = TrainCallbacks( on_update_train_progress=self.on_update_train_progress, on_update_status=self.on_update_status, @@ -780,12 +802,16 @@ def __training_thread_function(self): self.on_update_status("Stopped") self.delete_eta_label() - # queue UI update on Tk main thread; _set_training_button_idle applies shared styles, avoid potential race/crash - self.after(0, self._set_training_button_idle) + # Stop Tensorboard if it was in TRAIN_ONLY mode + if self.train_config.tensorboard_mode == TensorboardMode.TRAIN_ONLY: + self._stop_always_on_tensorboard() - if self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON and not self.always_on_tensorboard_subprocess: + # Restart if ALWAYS_ON mode (in case it was stopped during training) + if self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON and self.always_on_tensorboard_subprocess is None: self.after(0, self._start_always_on_tensorboard) + self.after(0, self._set_training_button_idle) + def start_training(self): if self.training_thread is None: self.save_default() @@ -831,46 +857,61 @@ def save_now(self): if train_commands: train_commands.save() - def _initialize_tensorboard(self): - if self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON: - self._start_always_on_tensorboard() - def _start_always_on_tensorboard(self): - if not self.always_on_tensorboard_subprocess: + if self.always_on_tensorboard_subprocess is None: self.always_on_tensorboard_subprocess = start_filtered_tensorboard(self.train_config) def _stop_always_on_tensorboard(self): - if self.always_on_tensorboard_subprocess: + if self.always_on_tensorboard_subprocess is not None: stop_tensorboard(self.always_on_tensorboard_subprocess) self.always_on_tensorboard_subprocess = None - def _restart_always_on_tensorboard(self): - """Restart always-on Tensorboard (for workspace directory changes)""" + def _on_workspace_dir_change_debounced(self, *args): + """Debounce workspace directory changes to avoid rapid restarts""" + if self.workspace_dir_debounce_id is not None: + with contextlib.suppress(Exception): + self.after_cancel(self.workspace_dir_debounce_id) + self.workspace_dir_debounce_id = self.after(1500, self._on_workspace_dir_change_actual) + + def _on_workspace_dir_change_actual(self): + if (new := self.train_config.workspace_dir) == self.current_workspace_dir: + return + self.current_workspace_dir = new if self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON: + logger.info("Restarting Tensorboard due to workspace change") self._stop_always_on_tensorboard() self._start_always_on_tensorboard() - def _on_tensorboard_mode_change(self, *args): - if self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON: - self._start_always_on_tensorboard() - else: + def _on_tensorboard_mode_change(self): + new_mode = self.train_config.tensorboard_mode + old_mode = self._previous_tensorboard_mode + + if new_mode == old_mode: + return + + logger.debug(f"Tensorboard mode changed: {old_mode} -> {new_mode}") + + # If training is active, handle mode changes immediately + if self.training_thread is not None: + if new_mode == TensorboardMode.OFF: + # Stop tensorboard regardless of previous mode + self._stop_always_on_tensorboard() + elif old_mode == TensorboardMode.OFF: + # Turning it back on during training (works for both TRAIN_ONLY and ALWAYS_ON) + self.always_on_tensorboard_subprocess = start_filtered_tensorboard(self.train_config) + # If switching between TRAIN_ONLY and ALWAYS_ON during training, keep running + + self._previous_tensorboard_mode = new_mode + return + + # Not training - manage normally + if old_mode == TensorboardMode.ALWAYS_ON: self._stop_always_on_tensorboard() - def _on_workspace_dir_change(self, new_workspace_dir: str): - if new_workspace_dir != self.current_workspace_dir: - self.current_workspace_dir = new_workspace_dir - # Restart always-on Tensorboard with new workspace directory if it's enabled - if self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON: - self._restart_always_on_tensorboard() - - def _on_workspace_dir_change_trace(self, *args): - """Handle workspace directory changes via UI state trace""" - new_workspace_dir = self.train_config.workspace_dir - if new_workspace_dir != self.current_workspace_dir: - self.current_workspace_dir = new_workspace_dir - - if self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON and self.always_on_tensorboard_subprocess: - self._start_always_on_tensorboard() + if new_mode == TensorboardMode.ALWAYS_ON: + self._start_always_on_tensorboard() + + self._previous_tensorboard_mode = new_mode def _set_training_button_style(self, mode: str): if not self.training_button: diff --git a/modules/util/tensorboard_util.py b/modules/util/tensorboard_util.py index c9b06cfe6..f4f376155 100644 --- a/modules/util/tensorboard_util.py +++ b/modules/util/tensorboard_util.py @@ -9,7 +9,9 @@ def get_tensorboard_args(config: TrainConfig) -> list[str]: - """Get Tensorboard command line arguments""" + """ + Generates args required to start Tensorboard + """ tensorboard_executable = os.path.join(os.path.dirname(sys.executable), "tensorboard") tensorboard_log_dir = os.path.join(config.workspace_dir, "tensorboard") @@ -31,7 +33,7 @@ def get_tensorboard_args(config: TrainConfig) -> list[str]: def start_filtered_tensorboard(config: TrainConfig) -> subprocess.Popen | None: - """Start Tensorboard with filtered output""" + """Start Tensorboard filtering the annoying warnings Google doesnt plan to fix""" try: print("Starting Tensorboard please wait...") env = os.environ.copy() @@ -75,9 +77,22 @@ def stop_tensorboard(process: subprocess.Popen | None): try: print("Stopping Tensorboard...") process.terminate() - process.wait(timeout=5) - except subprocess.TimeoutExpired: + + try: + process.wait(timeout=5) + print("Tensorboard stopped.") + except subprocess.TimeoutExpired: + with contextlib.suppress(Exception): + process.kill() + try: + process.wait(timeout=5) + print("Tensorboard killed.") + except subprocess.TimeoutExpired: + print("Failed to kill Tensorboard within timeout.") + + except Exception as e: + print(f"Error stopping Tensorboard: {e}") + finally: with contextlib.suppress(Exception): - process.kill() - except Exception: - pass + if process.stdout: + process.stdout.close() From 1a442e59efdffaca2d38db36e0931408bf3003a9 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Tue, 14 Oct 2025 22:19:02 +1100 Subject: [PATCH 16/64] Fix returning the migrated data --- modules/util/config/TrainConfig.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 181bf8d07..790a93f44 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -710,6 +710,8 @@ def __migration_8(self, data: dict) -> dict: else: migrated_data["tensorboard_mode"] = TensorboardMode.OFF + return migrated_data + def weight_dtypes(self) -> ModelWeightDtypes: return ModelWeightDtypes( self.train_dtype, From 10caaf3e291857c3f3faeb8081d2d26b0fd0ee16 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Tue, 14 Oct 2025 23:15:57 +1100 Subject: [PATCH 17/64] Remove debug logging and add is_output for dir_entry as well. --- modules/ui/TrainUI.py | 21 +++++++-------------- modules/util/ui/components.py | 9 ++++++--- 2 files changed, 13 insertions(+), 17 deletions(-) diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index dcce4cfb5..b672743e6 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -2,7 +2,6 @@ import ctypes import datetime import json -import logging import platform import threading import time @@ -50,8 +49,6 @@ from customtkinter import AppearanceModeTracker from tkinterdnd2 import TkinterDnD -logger = logging.getLogger(__name__) - # chunk for forcing Windows to ignore DPI scaling when moving between monitors # fixes the long standing transparency bug https://github.com/Nerogar/OneTrainer/issues/90 if platform.system() == "Windows": @@ -261,12 +258,12 @@ def create_general_tab(self, master): components.label(frame, 0, 0, "Workspace Directory", tooltip="The directory where all files of this training run are saved") # Remove command parameter - we'll use the trace instead for debouncing - components.dir_entry(frame, 0, 1, self.ui_state, "workspace_dir") + components.dir_entry(frame, 0, 1, self.ui_state, "workspace_dir", is_output=True) # cache dir components.label(frame, 0, 2, "Cache Directory", tooltip="The directory where cached data is saved") - components.dir_entry(frame, 0, 3, self.ui_state, "cache_dir") + components.dir_entry(frame, 0, 3, self.ui_state, "cache_dir", is_output=True) # continue from previous backup components.label(frame, 2, 0, "Continue from last backup", @@ -285,7 +282,7 @@ def create_general_tab(self, master): components.label(frame, 4, 2, "Debug Directory", tooltip="The directory where debug data is saved") - components.dir_entry(frame, 4, 3, self.ui_state, "debug_dir") + components.dir_entry(frame, 4, 3, self.ui_state, "debug_dir", is_output=True) # tensorboard components.label(frame, 6, 0, "Tensorboard Mode", @@ -880,7 +877,7 @@ def _on_workspace_dir_change_actual(self): return self.current_workspace_dir = new if self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON: - logger.info("Restarting Tensorboard due to workspace change") + print("Restarting Tensorboard due to workspace change") self._stop_always_on_tensorboard() self._start_always_on_tensorboard() @@ -891,22 +888,18 @@ def _on_tensorboard_mode_change(self): if new_mode == old_mode: return - logger.debug(f"Tensorboard mode changed: {old_mode} -> {new_mode}") - - # If training is active, handle mode changes immediately + # training if self.training_thread is not None: if new_mode == TensorboardMode.OFF: - # Stop tensorboard regardless of previous mode self._stop_always_on_tensorboard() elif old_mode == TensorboardMode.OFF: - # Turning it back on during training (works for both TRAIN_ONLY and ALWAYS_ON) + # if user enabled during training, start it self.always_on_tensorboard_subprocess = start_filtered_tensorboard(self.train_config) - # If switching between TRAIN_ONLY and ALWAYS_ON during training, keep running self._previous_tensorboard_mode = new_mode return - # Not training - manage normally + # not training if old_mode == TensorboardMode.ALWAYS_ON: self._stop_always_on_tensorboard() diff --git a/modules/util/ui/components.py b/modules/util/ui/components.py index 0a3afdef3..2c015bd03 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -222,7 +222,6 @@ def file_entry( frame.grid_columnconfigure(0, weight=1) - # Only use metadata as fallback if is_output wasn't explicitly set if not is_output: meta = ui_state.get_field_metadata(var_name) is_output = getattr(meta, 'is_output', False) @@ -273,14 +272,18 @@ def __open_dialog(): return frame -def dir_entry(master, row, column, ui_state: UIState, var_name: str, command: Callable[[str], None] = None): +def dir_entry(master, row, column, ui_state: UIState, var_name: str, command: Callable[[str], None] = None, is_output: bool = False): frame = ctk.CTkFrame(master, fg_color="transparent") frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") frame.grid_columnconfigure(0, weight=1) + if not is_output: + meta = ui_state.get_field_metadata(var_name) + is_output = getattr(meta, 'is_output', False) + def validate_dir_path_wrapper(value: str) -> tuple[bool, str]: - return validate_file_path(value, is_output=False, valid_extensions=None, path_type="directory") + return validate_file_path(value, is_output=is_output, valid_extensions=None, path_type="directory") entry_widget = entry( frame, row=0, column=0, ui_state=ui_state, var_name=var_name, From a65350ae6a4a9677a1e655e3a2008c7256238431 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Wed, 15 Oct 2025 01:11:32 +1100 Subject: [PATCH 18/64] Account for cloud tunnel and add whitespace trim for dir_entry on last item only --- modules/ui/TrainUI.py | 20 ++++++++++++++------ modules/util/ui/components.py | 14 ++++++++++++++ 2 files changed, 28 insertions(+), 6 deletions(-) diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index b672743e6..67d185b96 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -763,8 +763,9 @@ def open_sample_ui(self): def __training_thread_function(self): error_caught = False - # Start Tensorboard if in TRAIN_ONLY mode - if self.train_config.tensorboard_mode == TensorboardMode.TRAIN_ONLY: + # Start Tensorboard if TRAIN_ONLY mode and not cloud tunnel + if (self.train_config.tensorboard_mode == TensorboardMode.TRAIN_ONLY and + not (self.train_config.cloud.enabled and self.train_config.cloud.tensorboard_tunnel)): self.always_on_tensorboard_subprocess = start_filtered_tensorboard(self.train_config) self.training_callbacks = TrainCallbacks( @@ -801,12 +802,15 @@ def __training_thread_function(self): self.on_update_status("Stopped") self.delete_eta_label() - # Stop Tensorboard if it was in TRAIN_ONLY mode - if self.train_config.tensorboard_mode == TensorboardMode.TRAIN_ONLY: + # Stop Tensorboard if TRAIN_ONLY mode + if (self.train_config.tensorboard_mode == TensorboardMode.TRAIN_ONLY and + not (self.train_config.cloud.enabled and self.train_config.cloud.tensorboard_tunnel)): self._stop_always_on_tensorboard() - # Restart if ALWAYS_ON mode (in case it was stopped during training) - if self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON and self.always_on_tensorboard_subprocess is None: + # Restart if ALWAYS_ON mode + if (self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON and + self.always_on_tensorboard_subprocess is None + and not (self.train_config.cloud.enabled and self.train_config.cloud.tensorboard_tunnel)): self.after(0, self._start_always_on_tensorboard) self.after(0, self._set_training_button_idle) @@ -857,6 +861,10 @@ def save_now(self): train_commands.save() def _start_always_on_tensorboard(self): + # stop if cloud tunnel + if self.train_config.cloud.enabled and self.train_config.cloud.tensorboard_tunnel: + return + if self.always_on_tensorboard_subprocess is None: self.always_on_tensorboard_subprocess = start_filtered_tensorboard(self.train_config) diff --git a/modules/util/ui/components.py b/modules/util/ui/components.py index 2c015bd03..59b16459c 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -1,5 +1,6 @@ import contextlib from collections.abc import Callable +from pathlib import Path from tkinter import filedialog from typing import Any @@ -283,6 +284,19 @@ def dir_entry(master, row, column, ui_state: UIState, var_name: str, command: Ca is_output = getattr(meta, 'is_output', False) def validate_dir_path_wrapper(value: str) -> tuple[bool, str]: + # trim whitespaces from the last path component, nothing else + if not value: + return validate_file_path(value, is_output=is_output, valid_extensions=None, path_type="directory") + + path = Path(value) + trimmed_name = path.name.strip() + + if trimmed_name != path.name: + sep = '\\' if '\\' in value else '/' + trimmed_path = f"{path.parent}{sep}{trimmed_name}" if str(path.parent) != '.' else trimmed_name + ui_state.get_var(var_name).set(trimmed_path) + value = trimmed_path + return validate_file_path(value, is_output=is_output, valid_extensions=None, path_type="directory") entry_widget = entry( From 9c69bbdde20d79f32a854aa1074bc45ee381d1c7 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Thu, 16 Oct 2025 23:51:50 +1100 Subject: [PATCH 19/64] Fix(?) cloud regression with tensorboard --- modules/trainer/CloudTrainer.py | 7 +- modules/ui/TrainUI.py | 77 +++------------ modules/util/tensorboard_util.py | 156 +++++++++++++++++++++++++++++++ 3 files changed, 174 insertions(+), 66 deletions(-) diff --git a/modules/trainer/CloudTrainer.py b/modules/trainer/CloudTrainer.py index 3592d0a19..592eb7348 100644 --- a/modules/trainer/CloudTrainer.py +++ b/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 TensorboardManager class CloudTrainer(BaseTrainer): @@ -30,8 +31,7 @@ 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.tensorboard_always_on: - super()._start_tensorboard() + TensorboardManager().handle_training_start(config) match config.cloud.type: case CloudType.RUNPOD: @@ -114,8 +114,7 @@ def train(self): def end(self): try: - if self.config.tensorboard and not self.config.tensorboard_always_on: - super()._stop_tensorboard() + TensorboardManager().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") diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index 67d185b96..fdccc4e30 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -36,7 +36,7 @@ from modules.util.enum.ModelType import ModelType from modules.util.enum.TensorboardMode import TensorboardMode from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.tensorboard_util import start_filtered_tensorboard, stop_tensorboard +from modules.util.tensorboard_util import TensorboardManager from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress from modules.util.ui import components @@ -132,13 +132,12 @@ def __init__(self): self.training_callbacks = None self.training_commands = None - self.always_on_tensorboard_subprocess = None self.current_workspace_dir = self.train_config.workspace_dir - self._previous_tensorboard_mode = self.train_config.tensorboard_mode + self.tensorboard_manager = TensorboardManager() # Initialize tensorboard if needed - if self._previous_tensorboard_mode == TensorboardMode.ALWAYS_ON: - self._start_always_on_tensorboard() + if self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON: + self.tensorboard_manager.start(self.train_config) # Add debounced trace for workspace_dir changes self.workspace_dir_debounce_id = None @@ -159,7 +158,7 @@ def __init__(self): def __close(self): self.top_bar_component.save_default() - self._stop_always_on_tensorboard() + self.tensorboard_manager.stop() if hasattr(self, 'workspace_dir_trace_id'): self.ui_state.remove_var_trace("workspace_dir", self.workspace_dir_trace_id) if hasattr(self, 'tensorboard_mode_trace_id'): @@ -763,10 +762,9 @@ def open_sample_ui(self): def __training_thread_function(self): error_caught = False - # Start Tensorboard if TRAIN_ONLY mode and not cloud tunnel - if (self.train_config.tensorboard_mode == TensorboardMode.TRAIN_ONLY and - not (self.train_config.cloud.enabled and self.train_config.cloud.tensorboard_tunnel)): - self.always_on_tensorboard_subprocess = start_filtered_tensorboard(self.train_config) + # Use centralized manager + self.tensorboard_manager.handle_training_start(self.train_config) + self.training_callbacks = TrainCallbacks( on_update_train_progress=self.on_update_train_progress, @@ -802,16 +800,8 @@ def __training_thread_function(self): self.on_update_status("Stopped") self.delete_eta_label() - # Stop Tensorboard if TRAIN_ONLY mode - if (self.train_config.tensorboard_mode == TensorboardMode.TRAIN_ONLY and - not (self.train_config.cloud.enabled and self.train_config.cloud.tensorboard_tunnel)): - self._stop_always_on_tensorboard() - - # Restart if ALWAYS_ON mode - if (self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON and - self.always_on_tensorboard_subprocess is None - and not (self.train_config.cloud.enabled and self.train_config.cloud.tensorboard_tunnel)): - self.after(0, self._start_always_on_tensorboard) + # Use centralized manager + self.tensorboard_manager.handle_training_end(self.train_config) self.after(0, self._set_training_button_idle) @@ -860,19 +850,6 @@ def save_now(self): if train_commands: train_commands.save() - def _start_always_on_tensorboard(self): - # stop if cloud tunnel - if self.train_config.cloud.enabled and self.train_config.cloud.tensorboard_tunnel: - return - - if self.always_on_tensorboard_subprocess is None: - self.always_on_tensorboard_subprocess = start_filtered_tensorboard(self.train_config) - - def _stop_always_on_tensorboard(self): - if self.always_on_tensorboard_subprocess is not None: - stop_tensorboard(self.always_on_tensorboard_subprocess) - self.always_on_tensorboard_subprocess = None - def _on_workspace_dir_change_debounced(self, *args): """Debounce workspace directory changes to avoid rapid restarts""" if self.workspace_dir_debounce_id is not None: @@ -884,37 +861,13 @@ def _on_workspace_dir_change_actual(self): if (new := self.train_config.workspace_dir) == self.current_workspace_dir: return self.current_workspace_dir = new - if self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON: - print("Restarting Tensorboard due to workspace change") - self._stop_always_on_tensorboard() - self._start_always_on_tensorboard() + self.tensorboard_manager.handle_workspace_change(self.train_config) def _on_tensorboard_mode_change(self): - new_mode = self.train_config.tensorboard_mode - old_mode = self._previous_tensorboard_mode - - if new_mode == old_mode: - return - - # training - if self.training_thread is not None: - if new_mode == TensorboardMode.OFF: - self._stop_always_on_tensorboard() - elif old_mode == TensorboardMode.OFF: - # if user enabled during training, start it - self.always_on_tensorboard_subprocess = start_filtered_tensorboard(self.train_config) - - self._previous_tensorboard_mode = new_mode - return - - # not training - if old_mode == TensorboardMode.ALWAYS_ON: - self._stop_always_on_tensorboard() - - if new_mode == TensorboardMode.ALWAYS_ON: - self._start_always_on_tensorboard() - - self._previous_tensorboard_mode = new_mode + self.tensorboard_manager.handle_mode_change( + self.train_config, + is_training=self.training_thread is not None + ) def _set_training_button_style(self, mode: str): if not self.training_button: diff --git a/modules/util/tensorboard_util.py b/modules/util/tensorboard_util.py index f4f376155..c47f61c1d 100644 --- a/modules/util/tensorboard_util.py +++ b/modules/util/tensorboard_util.py @@ -6,6 +6,162 @@ from pathlib import Path from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.TensorboardMode import TensorboardMode + + +class TensorboardManager: + """Singleton manager for Tensorboard processes across GUI, CLI, and Cloud training""" + + _instance = None + _lock = threading.Lock() + + def __new__(cls): + if cls._instance is None: + with cls._lock: + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance._initialized = False + return cls._instance + + def __init__(self): + if self._initialized: + return + self._process = None + self._config = None + self._mode = None + self._previous_mode = None + self._initialized = True + + def should_start(self, config: TrainConfig, is_training: bool = False) -> bool: + """Determine if tensorboard should start based on mode and context""" + # Skip if cloud tunnel is enabled + if config.cloud.enabled and config.cloud.tensorboard_tunnel: + return False + + mode = config.tensorboard_mode + + if mode == TensorboardMode.OFF: + return False + elif mode == TensorboardMode.ALWAYS_ON: + return not is_training # Start only if not already running during training + elif mode == TensorboardMode.TRAIN_ONLY: + return is_training + + return False + + def start(self, config: TrainConfig, mode: TensorboardMode | None = None) -> bool: + """Start tensorboard if not already running""" + if self._process is not None: + return False # Already running + + # Use provided mode or config's mode + effective_mode = mode or config.tensorboard_mode + + # Check if we should start based on mode + if effective_mode == TensorboardMode.OFF: + return False + + # Skip if cloud tunnel is enabled + if config.cloud.enabled and config.cloud.tensorboard_tunnel: + return False + + self._process = start_filtered_tensorboard(config) + self._config = config + self._mode = effective_mode + self._previous_mode = effective_mode + return self._process is not None + + def stop(self) -> bool: + """Stop tensorboard if running""" + if self._process is None: + return False + + stop_tensorboard(self._process) + self._process = None + self._config = None + self._mode = None + return True + + def restart(self, config: TrainConfig) -> bool: + """Restart tensorboard with new config""" + self.stop() + return self.start(config) + + def is_running(self) -> bool: + """Check if tensorboard is currently running""" + return self._process is not None + + def get_mode(self) -> TensorboardMode | None: + """Get current tensorboard mode""" + return self._mode + + def handle_training_start(self, config: TrainConfig): + """Handle tensorboard when training starts""" + # If TRAIN_ONLY mode, start it + if config.tensorboard_mode == TensorboardMode.TRAIN_ONLY: + if not (config.cloud.enabled and config.cloud.tensorboard_tunnel): + self.start(config) + + def handle_training_end(self, config: TrainConfig): + """Handle tensorboard when training ends""" + # If TRAIN_ONLY mode, stop it + if config.tensorboard_mode == TensorboardMode.TRAIN_ONLY: + self.stop() + + # If ALWAYS_ON mode, restart it (in case it was stopped) + elif config.tensorboard_mode == TensorboardMode.ALWAYS_ON: + if not (config.cloud.enabled and config.cloud.tensorboard_tunnel): + if not self.is_running(): + self.start(config) + + def handle_mode_change(self, config: TrainConfig, is_training: bool = False): + """Handle tensorboard mode changes""" + new_mode = config.tensorboard_mode + + # Initialize previous mode on first call + if self._previous_mode is None: + self._previous_mode = new_mode + # Start if ALWAYS_ON and not training + if new_mode == TensorboardMode.ALWAYS_ON and not is_training: + if not (config.cloud.enabled and config.cloud.tensorboard_tunnel): + self.start(config) + return + + old_mode = self._previous_mode + + if new_mode == old_mode: + return + + # Skip if cloud tunnel + if config.cloud.enabled and config.cloud.tensorboard_tunnel: + self.stop() + self._previous_mode = new_mode + return + + # During training + if is_training: + if new_mode == TensorboardMode.OFF: + self.stop() + elif old_mode == TensorboardMode.OFF and new_mode in (TensorboardMode.TRAIN_ONLY, TensorboardMode.ALWAYS_ON): + self.start(config) + self._previous_mode = new_mode + return + + # Not training + if old_mode == TensorboardMode.ALWAYS_ON: + self.stop() + + if new_mode == TensorboardMode.ALWAYS_ON: + self.start(config) + + self._previous_mode = new_mode + + def handle_workspace_change(self, config: TrainConfig): + """Handle workspace directory changes""" + if config.tensorboard_mode == TensorboardMode.ALWAYS_ON: + if not (config.cloud.enabled and config.cloud.tensorboard_tunnel): + print("Restarting Tensorboard due to workspace change") + self.restart(config) def get_tensorboard_args(config: TrainConfig) -> list[str]: From e8c30d1f7154c467c92e15880913136e4c534566 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Sun, 9 Nov 2025 05:14:34 +1100 Subject: [PATCH 20/64] Input validation logic --- modules/ui/ModelTab.py | 37 +- modules/ui/TrainUI.py | 53 +- modules/util/config/TrainConfig.py | 10 + modules/util/path_util.py | 63 --- modules/util/ui/ToolTip.py | 151 +++++- modules/util/ui/components.py | 799 +++++++++++++++++++++-------- modules/util/ui/validation.py | 415 +++++++++++++++ 7 files changed, 1186 insertions(+), 342 deletions(-) create mode 100644 modules/util/ui/validation.py diff --git a/modules/ui/ModelTab.py b/modules/ui/ModelTab.py index bf4ac9479..76f4abcfb 100644 --- a/modules/ui/ModelTab.py +++ b/modules/ui/ModelTab.py @@ -398,7 +398,7 @@ def __create_base_components( if has_vae: # base model components.label(self.scroll_frame, row, 0, "VAE Override", - tooltip="Directory or Hugging Face repository of a VAE model in diffusers format. Can be used to override the VAE included in the base model. Using a safetensor VAE file will cause an error that the model cannot be loaded.") + tooltip="Directory or Hugging Face repository of a VAE model in diffusers format. Can be used to override the VAE included in the base model. \n\n Using a safetensor VAE file will cause an error that the model cannot be loaded.") components.file_entry( self.scroll_frame, row, 1, self.ui_state, "vae.model_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x @@ -480,35 +480,15 @@ def __create_output_components( allow_diffusers: bool = False, allow_legacy_safetensors: bool = False, ) -> int: - def get_valid_extensions() -> list[str]: - format_str = self.ui_state.get_var("output_model_format").get() - - # Map format to extensions - format_extensions = { - "SAFETENSORS": [".safetensors"], - "DIFFUSERS": [".json"], - "LEGACY_SAFETENSORS": [".safetensors"], - } - - return format_extensions.get(format_str, []) - - def update_output_validation(_=None): - - output_var = self.ui_state.get_var("output_model_destination") - current_value = output_var.get() - - if current_value: - output_var.set(current_value) - # output model destination components.label(self.scroll_frame, row, 0, "Model Output Destination", - tooltip="Filename or directory where the output model is saved") + tooltip="Filename or directory where the *final* epoch output model is saved") - components.file_entry( - self.scroll_frame, row, 1, self.ui_state, "output_model_destination", - is_output=True, - valid_extensions=get_valid_extensions(), - path_type="file" # Allow both file and directory paths + components.model_output_entry( + self.scroll_frame, row, 1, self.ui_state, + var_name="output_model_destination", + format_var_name="output_model_format", + method_var_name="training_method", ) # output data type @@ -536,8 +516,7 @@ def update_output_validation(_=None): components.label(self.scroll_frame, row, 0, "Output Format", tooltip="Format to use when saving the output model") components.options_kv( - self.scroll_frame, row, 1, formats, self.ui_state, "output_model_format", - command=update_output_validation # Trigger validation when format changes + self.scroll_frame, row, 1, formats, self.ui_state, "output_model_format" ) # include config diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index fdccc4e30..f5b8767ec 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -309,49 +309,62 @@ def create_general_tab(self, master): tooltip="The interval used when validate training") components.time_entry(frame, 8, 3, self.ui_state, "validate_after", "validate_after_unit") + # input validation settings + components.label(frame, 9, 0, "Auto-correct Input", + tooltip="Automatically correct file paths and extensions where possible") + components.switch(frame, 9, 1, self.ui_state, "validation_auto_correct") + + components.label(frame, 9, 2, "Show Validation Tooltips", + tooltip="Show tooltips with validation messages for errors and warnings") + components.switch(frame, 9, 3, self.ui_state, "validation_show_tooltips") + + components.label(frame, 10, 0, "Use Friendly Names", + tooltip="Use friendly names instead of timestamps for auto-corrected/generated filenames") + components.switch(frame, 10, 1, self.ui_state, "use_friendly_names") + # device - components.label(frame, 10, 0, "Dataloader Threads", + components.label(frame, 11, 0, "Dataloader Threads", tooltip="Number of threads used for the data loader. Increase if your GPU has room during caching, decrease if it's going out of memory during caching.") - components.entry(frame, 10, 1, self.ui_state, "dataloader_threads") + components.entry(frame, 11, 1, self.ui_state, "dataloader_threads") - components.label(frame, 11, 0, "Train Device", + components.label(frame, 12, 0, "Train Device", tooltip="The device used for training. Can be \"cuda\", \"cuda:0\", \"cuda:1\" etc. Default:\"cuda\". Must be \"cuda\" for multi-GPU training.") - components.entry(frame, 11, 1, self.ui_state, "train_device") + components.entry(frame, 12, 1, self.ui_state, "train_device") - components.label(frame, 12, 0, "Multi-GPU", + components.label(frame, 13, 0, "Multi-GPU", tooltip="Enable multi-GPU training") - components.switch(frame, 12, 1, self.ui_state, "multi_gpu") - components.label(frame, 12, 2, "Device Indexes", + components.switch(frame, 13, 1, self.ui_state, "multi_gpu") + components.label(frame, 13, 2, "Device Indexes", tooltip="Multi-GPU: A comma-separated list of device indexes. If empty, all your GPUs are used. With a list such as \"0,1,3,4\" you can omit a GPU, for example an on-board graphics GPU.") - components.entry(frame, 12, 3, self.ui_state, "device_indexes") + components.entry(frame, 13, 3, self.ui_state, "device_indexes") - components.label(frame, 13, 0, "Sequential model setup", + components.label(frame, 14, 0, "Sequential model setup", tooltip="Multi-GPU: If enabled, loading and setting up the model is done for each GPU one after the other. This is slower, but can reduce peak RAM usage.") - components.switch(frame, 13, 1, self.ui_state, "sequential_model_setup") + components.switch(frame, 14, 1, self.ui_state, "sequential_model_setup") - components.label(frame, 14, 0, "Gradient Reduce Precision", + components.label(frame, 15, 0, "Gradient Reduce Precision", tooltip="WEIGHT_DTYPE: Reduce gradients between GPUs in your weight data type; can be imprecise, but more efficient than float32\n" "WEIGHT_DTYPE_STOCHASTIC: Sum up the gradients in your weight data type, but average them in float32 and stochastically round if your weight data type is bfloat16\n" "FLOAT_32: Reduce gradients in float32\n" "FLOAT_32_STOCHASTIC: Reduce gradients in float32; use stochastic rounding to bfloat16 if your weight data type is bfloat16", wide_tooltip=True) - components.options(frame, 14, 1, [str(x) for x in list(GradientReducePrecision)], self.ui_state, + components.options(frame, 15, 1, [str(x) for x in list(GradientReducePrecision)], self.ui_state, "gradient_reduce_precision") - components.label(frame, 14, 2, "Fused Gradient Reduce", + components.label(frame, 15, 2, "Fused Gradient Reduce", tooltip="Multi-GPU: Gradient synchronisation during the backward pass. Can be more efficient, especially with Async Gradient Reduce") - components.switch(frame, 14, 3, self.ui_state, "fused_gradient_reduce") + components.switch(frame, 15, 3, self.ui_state, "fused_gradient_reduce") - components.label(frame, 15, 0, "Async Gradient Reduce", + components.label(frame, 16, 0, "Async Gradient Reduce", tooltip="Multi-GPU: Asynchroniously start the gradient reduce operations during the backward pass. Can be more efficient, but requires some VRAM.") - components.switch(frame, 15, 1, self.ui_state, "async_gradient_reduce") - components.label(frame, 15, 2, "Buffer size (MB)", + components.switch(frame, 16, 1, self.ui_state, "async_gradient_reduce") + components.label(frame, 16, 2, "Buffer size (MB)", tooltip="Multi-GPU: Maximum VRAM for \"Async Gradient Reduce\", in megabytes. A multiple of this value can be needed if combined with \"Fused Back Pass\" and/or \"Layer offload fraction\"") - components.entry(frame, 15, 3, self.ui_state, "async_gradient_reduce_buffer") + components.entry(frame, 16, 3, self.ui_state, "async_gradient_reduce_buffer") - components.label(frame, 16, 0, "Temp Device", + components.label(frame, 17, 0, "Temp Device", tooltip="The device used to temporarily offload models while they are not used. Default:\"cpu\"") - components.entry(frame, 16, 1, self.ui_state, "temp_device") + components.entry(frame, 17, 1, self.ui_state, "temp_device") frame.pack(fill="both", expand=1) return frame diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 49510c178..6c81fbb6a 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -313,6 +313,11 @@ class TrainConfig(BaseConfig): continue_last_backup: bool include_train_config: ConfigPart + # validation settings + validation_auto_correct: bool + validation_show_tooltips: bool + use_friendly_names: bool + # multi-GPU multi_gpu: bool device_indexes: str @@ -869,6 +874,11 @@ def default_values() -> 'TrainConfig': data.append(("continue_last_backup", False, bool, False)) data.append(("include_train_config", ConfigPart.NONE, ConfigPart, False)) + # validation settings + data.append(("validation_auto_correct", True, bool, False)) + data.append(("validation_show_tooltips", True, bool, False)) + data.append(("use_friendly_names", False, bool, False)) + #multi-GPU data.append(("multi_gpu", False, bool, False)) data.append(("device_indexes", "", str, False)) diff --git a/modules/util/path_util.py b/modules/util/path_util.py index 1e0d4affd..21932b9c2 100644 --- a/modules/util/path_util.py +++ b/modules/util/path_util.py @@ -52,66 +52,3 @@ def supported_video_extensions() -> set[str]: def is_supported_video_extension(extension: str) -> bool: return extension.lower() in SUPPORTED_VIDEO_EXTENSIONS - - -def validate_file_path( - value: str, - is_output: bool = False, - valid_extensions: list[str] | None = None, - path_type: str = "file", # "file" or "directory" -) -> tuple[bool, str]: - if not value: - return True, "" - - # windows invalid characters - if is_output and os.name == 'nt': - invalid_chars = ['<', '>', '|', '\0', '"'] - path_parts = [os.path.basename(value)] - - dirname = os.path.dirname(value) - if dirname: - path_parts.extend( - part for part in dirname.replace('\\', '/').split('/') - if part and not (len(part) == 2 and part[1] == ':') - ) - - for part in path_parts: - for char in invalid_chars: - if char in part: - location = "filename" if part == path_parts[0] else "path" - return False, f"Invalid character in {location}: '{char}'" - - if is_output: - if path_type == "directory" and os.path.isfile(value): - return False, "Must be a directory, not a file" - - if path_type == "file": - if os.path.isdir(value): - return False, "Must be a file, not a directory" - - # Validate extension if applicable - if valid_extensions and "." in os.path.basename(value): - file_ext = os.path.splitext(value)[1].lower() - if file_ext not in valid_extensions: - return False, f"Invalid extension. Expected: {', '.join(valid_extensions)}" - - # Validate parent directory exists - parent_dir = os.path.dirname(value) - if parent_dir and not os.path.isdir(parent_dir): - return False, "Parent directory does not exist" - - return True, "" - - if path_type == "directory": - if os.path.isfile(value) or (not os.path.isdir(value) and os.path.splitext(value)[1]): - value = os.path.dirname(value) or '.' - - if not os.path.isdir(value): - return False, "Directory does not exist" - else: # file - if os.path.isdir(value): - return False, "Must be a file, not a directory" - if not os.path.isfile(value): - return False, "File does not exist" - - return True, "" diff --git a/modules/util/ui/ToolTip.py b/modules/util/ui/ToolTip.py index be80e2765..9de7264da 100644 --- a/modules/util/ui/ToolTip.py +++ b/modules/util/ui/ToolTip.py @@ -1,56 +1,153 @@ +import contextlib + import customtkinter as ctk class ToolTip: - """ - create a tooltip for a given widget - """ + MAX_TOOLTIP_LENGTH = 350 + EDGE_MARGIN = 10 - def __init__(self, widget, text='widget info', x_position=20, wide=False): + def __init__(self, widget, text='widget info', x_position=20, wide=False, + hover_only=True, track_movement=False): self.widget = widget self.text = text self.x_position = x_position + self.wraplength = 350 if wide else 180 + self.hover_only = hover_only + self.track_movement = track_movement - self.waittime = 500 # miliseconds - self.wraplength = 180 if not wide else 350 # pixels - self.widget.bind("", self.enter) - self.widget.bind("", self.leave) - self.widget.bind("", self.leave) + self.waittime = 500 self.id = None self.tw = None + self._after_id = None + + if hover_only: + self.widget.bind("", self.enter) + self.widget.bind("", self.leave) + self.widget.bind("", self.leave) + + if track_movement: + self._root = widget.winfo_toplevel() + self._root.bind("", self._on_move, add="+") + self._root.bind("", self.hide, add="+") + self._root.bind("", self.hide, add="+") def enter(self, event=None): self.schedule() def leave(self, event=None): self.unschedule() - self.hidetip() + self.hide() def schedule(self): self.unschedule() self.id = self.widget.after(self.waittime, self.showtip) def unschedule(self): - id = self.id - self.id = None - if id: - self.widget.after_cancel(id) + if self.id: + self.widget.after_cancel(self.id) + self.id = None + + def show(self, text=None, color=None, duration_ms=None, fg_color="#2b2b2b"): + if text: + self.text = text[:self.MAX_TOOLTIP_LENGTH] + ("..." if len(text) > self.MAX_TOOLTIP_LENGTH else "") + + if self._after_id: + with contextlib.suppress(Exception): + self.widget.after_cancel(self._after_id) + self._after_id = None + + self.hide() - def showtip(self, event=None): - x = y = 0 - x, y, cx, cy = self.widget.bbox("insert") - x += self.widget.winfo_rootx() + 25 - y += self.widget.winfo_rooty() + self.x_position - # creates a toplevel window self.tw = ctk.CTkToplevel(self.widget) - # Leaves only the label and removes the app window self.tw.wm_overrideredirect(True) + if self.track_movement: + self.tw.attributes("-topmost", True) + + frame = ctk.CTkFrame(self.tw, fg_color=fg_color, corner_radius=6) + frame.pack(padx=0, pady=0) + + label = ctk.CTkLabel( + frame, + text=self.text, + justify='left', + wraplength=self.wraplength, + text_color=color or "white", + fg_color="transparent", + padx=8, + pady=6 if self.track_movement else 8 + ) + label.pack() + + self._position() + + if duration_ms: + self._after_id = self.widget.after(duration_ms, self.hide) + + def showtip(self, event=None): + self.show() + + def _position(self): + if not self.tw or not self.tw.winfo_exists(): + return + + x = self.widget.winfo_rootx() + 25 + y = self.widget.winfo_rooty() + self.x_position + + if self.track_movement: + y = self.widget.winfo_rooty() + self.widget.winfo_height() + 4 + + self.tw.update_idletasks() + tip_width = self.tw.winfo_reqwidth() + tip_height = self.tw.winfo_reqheight() + screen_width = self.widget.winfo_screenwidth() + screen_height = self.widget.winfo_screenheight() + + x = max(10, min(self.widget.winfo_rootx(), screen_width - tip_width - 10)) + + if y + tip_height > screen_height - 10: + y = self.widget.winfo_rooty() - tip_height - 4 + + y = max(10, y) + self.tw.wm_geometry(f"+{x}+{y}") - label = ctk.CTkLabel(self.tw, text=self.text, justify='left', wraplength=self.wraplength) - label.pack(padx=8, pady=8) - def hidetip(self): - tw = self.tw + def _on_move(self, _evt=None): + if self.tw and self.tw.winfo_exists(): + self._position() + + def hide(self, _evt=None): + if self._after_id: + with contextlib.suppress(Exception): + self.widget.after_cancel(self._after_id) + self._after_id = None + + if self.tw and self.tw.winfo_exists(): + self.tw.destroy() self.tw = None - if tw: - tw.destroy() + + hidetip = hide + + def show_error(self, message: str, duration_ms: int | None = 7000): + self.show(message, color="#ff6b6b", duration_ms=duration_ms) + + def show_warning(self, message: str, duration_ms: int | None = 7000): + self.show(message, color="#ff9500", duration_ms=duration_ms) + + def show_info(self, message: str, duration_ms: int | None = 5000): + self.show(message, color="#5bc0de", duration_ms=duration_ms) + + def is_visible(self) -> bool: + return self.tw is not None and self.tw.winfo_exists() + + def destroy(self): + self.hide() + if self.track_movement: + with contextlib.suppress(Exception): + self._root.unbind("", self._on_move) + self._root.unbind("", self.hide) + self._root.unbind("", self.hide) + + def __del__(self): + with contextlib.suppress(Exception): + self.destroy() diff --git a/modules/util/ui/components.py b/modules/util/ui/components.py index 59b16459c..1db41be8a 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -1,14 +1,17 @@ import contextlib +import tkinter as tk from collections.abc import Callable +from dataclasses import dataclass from pathlib import Path from tkinter import filedialog -from typing import Any +from typing import Any, Literal from modules.util.enum.TimeUnit import TimeUnit -from modules.util.path_util import supported_image_extensions, validate_file_path +from modules.util.path_util import supported_image_extensions from modules.util.ui.ToolTip import ToolTip from modules.util.ui.ui_utils import _register_drop_target from modules.util.ui.UIState import UIState +from modules.util.ui.validation import ValidationResult, validate_basic_type, validate_file_path import customtkinter as ctk from customtkinter.windows.widgets.scaling import CTkScalingBaseClass @@ -16,183 +19,430 @@ PAD = 10 +@dataclass(frozen=True) +class ComponentValidationSettings: + debounce_stop_typing_ms: int = 1700 + debounced_invalid_revert_ms: int = 1000 + focusout_invalid_revert_ms: int = 1200 -def app_title(master, row, column): - frame = ctk.CTkFrame(master) - frame.grid(row=row, column=column, padx=5, pady=5, sticky="nsew") - image_component = ctk.CTkImage( - Image.open("resources/icons/icon.png").resize((40, 40), Image.Resampling.LANCZOS), - size=(40, 40) - ) - image_label_component = ctk.CTkLabel(frame, image=image_component, text="") - image_label_component.grid(row=0, column=0, padx=PAD, pady=PAD) +COMPONENT_VALIDATION_SETTINGS = ComponentValidationSettings() - label_component = ctk.CTkLabel(frame, text="OneTrainer", font=ctk.CTkFont(size=20, weight="bold")) - label_component.grid(row=0, column=1, padx=(0, PAD), pady=PAD) +@dataclass +class ValidationState: + """Tracks validation state for UI components.""" + status: Literal['error', 'warning'] | None = None + message: str = "" + def clear(self): + self.status = None + self.message = "" -def label(master, row, column, text, pad=PAD, tooltip=None, wide_tooltip=False, wraplength=0): - component = ctk.CTkLabel(master, text=text, wraplength=wraplength) - component.grid(row=row, column=column, padx=pad, pady=pad, sticky="nw") - if tooltip: - ToolTip(component, tooltip, wide=wide_tooltip) - return component - + def set_error(self, message: str): + self.status = 'error' + self.message = message -def entry( - master, - row, - column, - ui_state: UIState, - var_name: str, - command: Callable[[], None] = None, - tooltip: str = "", - wide_tooltip: bool = False, - width: int = 140, - sticky: str = "new", - custom_validator: Callable[[str], tuple[bool, str]] = None, -): - var = ui_state.get_var(var_name) - trace_id = None - if command: - trace_id = ui_state.add_var_trace(var_name, command) + def set_warning(self, message: str): + self.status = 'warning' + self.message = message - component = ctk.CTkEntry(master, textvariable=var, width=width) - component.grid(row=row, column=column, padx=PAD, pady=PAD, sticky=sticky) - try: - original_border_color = component.cget("border_color") - except Exception: - original_border_color = "gray50" +class EntryValidationHandler: + """Handles UI coordination for entry widget validation with debouncing and visual feedback.""" - error_border_color = "#dc3545" + def __init__( + self, + component: ctk.CTkEntry, + var: tk.Variable, + var_name: str, + ui_state: UIState, + custom_validator: Callable[[str], ValidationResult] | None = None, + validation_state: ValidationState | None = None, + ): + self.component = component + self.var = var + self.var_name = var_name + self.ui_state = ui_state + self.custom_validator = custom_validator + self.validation_state = validation_state # Optional external validation state + + # Get border colors + try: + self.original_border_color = component.cget("border_color") + except Exception: + self.original_border_color = "gray50" - validation_after_id = None - revert_after_id = None - touched = False + self.error_border_color = "#dc3545" + self.warning_border_color = "#ff9500" - DEBOUNCE_STOP_TYPING_MS = 1500 - DEBOUNCED_INVALID_REVERT_MS = 1000 - FOCUSOUT_INVALID_REVERT_MS = 1200 + # Validation state + self.validation_after_id = None + self.revert_after_id = None + self.touched = False + self.last_valid_value = var.get() - last_valid_value = var.get() + # Create tooltip (hover_only=False to prevent default hover behavior, wide=True for validation messages) + self.validation_tooltip = ToolTip(component, text="", hover_only=False, track_movement=True, wide=True) + component._validation_tooltip = self.validation_tooltip - def validate_value(value: str, revert_delay_ms: int | None) -> bool: - nonlocal revert_after_id, last_valid_value - meta = ui_state.get_field_metadata(var_name) - declared_type = meta.type - nullable = meta.nullable - default_val = meta.default + def should_show_tooltip(self) -> bool: + """Check if validation tooltips are enabled in settings.""" + try: + show_tooltips_var = self.ui_state.get_var("validation_show_tooltips") + return bool(show_tooltips_var.get()) + except Exception: + return True - if revert_after_id: - with contextlib.suppress(Exception): - component.after_cancel(revert_after_id) - revert_after_id = None + def validate_value(self, value: str, revert_delay_ms: int | None) -> bool: + """ + Validate the current value and update UI accordingly. - def success(): - nonlocal last_valid_value - component.configure(border_color=original_border_color) - last_valid_value = value - return True + Returns: + True if validation passed, False otherwise. + """ - def do_revert(): - var.set(last_valid_value) - component.configure(border_color=original_border_color) + # Cancel any pending revert + if self.revert_after_id: + with contextlib.suppress(Exception): + self.component.after_cancel(self.revert_after_id) + self.revert_after_id = None - def fail(_reason: str): - nonlocal revert_after_id - component.configure(border_color=error_border_color) - if revert_delay_ms is not None: - revert_after_id = component.after(revert_delay_ms, do_revert) - else: - do_revert() - return False + meta = self.ui_state.get_field_metadata(self.var_name) - if value == "": - if nullable: - return success() - if declared_type is str: - if default_val == "": - return success() - return fail("Value required") + # Perform basic type validation + basic_result = validate_basic_type( + value=value, + declared_type=meta.type, + nullable=meta.nullable, + default_val=meta.default + ) - try: - if declared_type is int: - int(value) - elif declared_type is float: - float(value) - elif declared_type is bool: - if value.lower() not in ("true", "false", "0", "1"): - return fail("Invalid bool") - - if custom_validator: - is_valid, error_msg = custom_validator(value) - if not is_valid: - return fail(error_msg) - - return success() - except ValueError: - return fail("Invalid value") - - def debounced_validate(*_): - nonlocal validation_after_id, revert_after_id - if not touched: - if validation_after_id: + # If basic validation fails, handle failure + if not basic_result.ok: + return self._fail(basic_result.message, revert_delay_ms) + + # Custom validation + if self.custom_validator: + result = self.custom_validator(value) + + # Handle based on status + if result.status == 'warning': + return self._warning(result.message, value) + elif not result.ok: + return self._fail(result.message, revert_delay_ms) + + return self._success(value) + + def _success(self, value: str) -> bool: + """Handle successful validation.""" + self.component.configure(border_color=self.original_border_color) + self.validation_tooltip.hide() + self.last_valid_value = value + return True + + def _warning(self, reason: str, value: str) -> bool: + """Handle warning state with orange border.""" + self.component.configure(border_color=self.warning_border_color) + if self.should_show_tooltip(): + self.validation_tooltip.show_warning(reason) + self.last_valid_value = value + + # Update external validation state if provided + if self.validation_state: + self.validation_state.set_warning(reason) + + return True + + def _fail(self, reason: str, revert_delay_ms: int | None) -> bool: + """Handle validation failure.""" + self.component.configure(border_color=self.error_border_color) + if self.should_show_tooltip(): + self.validation_tooltip.show_error(reason) + + # Update external validation state if provided + if self.validation_state: + self.validation_state.set_error(reason) + + if revert_delay_ms is not None: + self.revert_after_id = self.component.after(revert_delay_ms, self._do_revert) + else: + self._do_revert() + return False + + def _do_revert(self): + """Revert to last valid value.""" + self.var.set(self.last_valid_value) + self.component.configure(border_color=self.original_border_color) + self.validation_tooltip.hide() + + def debounced_validate(self, *_): + """Validate with debouncing - called on variable changes.""" + if not self.touched: + if self.validation_after_id: with contextlib.suppress(Exception): - component.after_cancel(validation_after_id) - validation_after_id = None + self.component.after_cancel(self.validation_after_id) + self.validation_after_id = None return - if revert_after_id: + + # Clear validation UI immediately when field is modified + self.validation_tooltip.hide() + self.component.configure(border_color=self.original_border_color) + + # Cancel pending operations + if self.revert_after_id: with contextlib.suppress(Exception): - component.after_cancel(revert_after_id) - revert_after_id = None - if validation_after_id: + self.component.after_cancel(self.revert_after_id) + self.revert_after_id = None + if self.validation_after_id: with contextlib.suppress(Exception): - component.after_cancel(validation_after_id) - validation_after_id = component.after( - DEBOUNCE_STOP_TYPING_MS, - lambda: validate_value(var.get(), DEBOUNCED_INVALID_REVERT_MS) + self.component.after_cancel(self.validation_after_id) + + # Schedule validation + self.validation_after_id = self.component.after( + COMPONENT_VALIDATION_SETTINGS.debounce_stop_typing_ms, + lambda: self.validate_value( + self.var.get(), + COMPONENT_VALIDATION_SETTINGS.debounced_invalid_revert_ms + ) ) - validation_trace_name = var.trace_add("write", debounced_validate) + def on_focus_in(self, _e=None): + """Handle focus in event.""" + self.touched = False + # Re-show tooltip if there's an existing validation state + if self.validation_state and self.should_show_tooltip(): + if self.validation_state.status == 'error': + self.validation_tooltip.show_error(self.validation_state.message, duration_ms=None) + elif self.validation_state.status == 'warning': + self.validation_tooltip.show_warning(self.validation_state.message, duration_ms=None) + + def on_user_input(self, _e=None): + """Handle user input event.""" + self.touched = True + + def on_focus_out(self, _e=None): + """Handle focus out event - validate if user interacted.""" + # Always hide tooltip on focus out + self.validation_tooltip.hide() + + if self.touched: + self.validate_value( + self.var.get(), + COMPONENT_VALIDATION_SETTINGS.focusout_invalid_revert_ms + ) + + def cleanup(self): + """Cleanup resources.""" + if self.validation_after_id: + with contextlib.suppress(Exception): + self.component.after_cancel(self.validation_after_id) + if self.revert_after_id: + with contextlib.suppress(Exception): + self.component.after_cancel(self.revert_after_id) + - def on_focus_in(_e=None): - nonlocal touched - touched = False +class ModelOutputValidator: + """ + Smart validator for model output paths with auto-correction. + Delegates validation logic to validation.py. + """ + + def __init__( + self, + var: tk.Variable, + ui_state: UIState, + format_var_name: str = "output_model_format", + method_var_name: str = "training_method", + prefix_var_name: str = "save_filename_prefix", + ): + self.var = var + self.ui_state = ui_state + self.format_var = ui_state.get_var(format_var_name) + self.method_var = ui_state.get_var(method_var_name) + self.prefix_var = ui_state.get_var(prefix_var_name) + self.autocorrect_var = ui_state.get_var("validation_auto_correct") + self.friendly_names_var = ui_state.get_var("use_friendly_names") + + self.state = ValidationState() + self._trace_ids: dict[str, str] = {} + + def validate(self, value: str) -> ValidationResult: + """ + Validate model output path using validation.py logic. + + Returns: + ValidationResult with status (success/warning/error) + """ + from modules.util.enum.ModelFormat import ModelFormat + from modules.util.enum.TrainingMethod import TrainingMethod + from modules.util.ui.validation import validate_destination + + value = value.strip() + if not value: + self.state.clear() + return ValidationResult(ok=True, corrected=None, message="", status='success') + + try: + output_format = ModelFormat[self.format_var.get()] + except (KeyError, ValueError): + output_format = ModelFormat.SAFETENSORS + + try: + training_method = TrainingMethod[self.method_var.get()] + except (KeyError, ValueError): + training_method = TrainingMethod.FINE_TUNE + + try: + autocorrect = bool(self.autocorrect_var.get()) + except Exception: + autocorrect = True - def on_user_input(_e=None): - nonlocal touched - touched = True + prefix = self.prefix_var.get() if self.prefix_var else "" + + try: + use_friendly_names = bool(self.friendly_names_var.get()) + except Exception: + use_friendly_names = False + + # Use validation logic from validation.py + result = validate_destination( + value, + output_format, + training_method, + autocorrect=autocorrect, + prefix=prefix, + use_friendly_names=use_friendly_names, + is_output=True + ) + + # Update internal state for hover behavior + if result.status == 'error': + self.state.set_error(result.message) + elif result.status == 'warning': + self.state.set_warning(result.message) + else: + self.state.clear() + + # Apply correction if available + if result.corrected and result.corrected != value: + self.var.set(result.corrected) + + return result + + def trigger_validation(self, *_): + self.validate(self.var.get()) + + def on_value_change(self, *_): + self.state.clear() + + def setup_traces(self): + self._trace_ids['format'] = self.format_var.trace_add("write", self.trigger_validation) + self._trace_ids['method'] = self.method_var.trace_add("write", self.trigger_validation) + self._trace_ids['prefix'] = self.prefix_var.trace_add("write", self.trigger_validation) + self._trace_ids['friendly_names'] = self.friendly_names_var.trace_add("write", self.trigger_validation) + + def cleanup_traces(self): + for key, trace_id in self._trace_ids.items(): + with contextlib.suppress(Exception): + var = getattr(self, f"{key}_var", None) + if var: + var.trace_remove("write", trace_id) + self._trace_ids.clear() - def on_focus_out(_e=None): - # only validate on focus-out if the user interacted with the field. - if touched: - validate_value(var.get(), FOCUSOUT_INVALID_REVERT_MS) - component.bind("", on_focus_in) - component.bind("", on_user_input) - component.bind("<>", on_user_input) - component.bind("<>", on_user_input) - component.bind("", on_focus_out) +# UI components +def entry( + master, + row, + column, + ui_state: UIState, + var_name: str, + command: Callable[[], None] = None, + tooltip: str = "", + wide_tooltip: bool = False, + width: int = 140, + sticky: str = "new", + custom_validator: Callable[[str], ValidationResult] | None = None, + validation_state: ValidationState = None, + enable_hover_validation: bool = False, +): + """ + Create an entry widget with optional validation. + + Args: + validation_state: Optional ValidationState object for stateful validators. + If provided, enables hover behavior to re-show validation tooltips. + enable_hover_validation: If True, enables hover behavior to show validation state tooltips. + """ + var = ui_state.get_var(var_name) + trace_id = None + if command: + trace_id = ui_state.add_var_trace(var_name, command) + + component = ctk.CTkEntry(master, textvariable=var, width=width) + component.grid(row=row, column=column, padx=PAD, pady=PAD, sticky=sticky) + + # Create validation handler + validation_handler = EntryValidationHandler( + component=component, + var=var, + var_name=var_name, + ui_state=ui_state, + custom_validator=custom_validator, + validation_state=validation_state, + ) + + # Setup variable trace for validation + validation_trace_name = var.trace_add("write", validation_handler.debounced_validate) + + # Bind focus and input events + component.bind("", validation_handler.on_focus_in) + component.bind("", validation_handler.on_user_input) + component.bind("<>", validation_handler.on_user_input) + component.bind("<>", validation_handler.on_user_input) + component.bind("", validation_handler.on_focus_out) + + # Optional hover validation behavior + if enable_hover_validation and validation_state: + show_tooltips_var = ui_state.get_var("validation_show_tooltips") + validation_tooltip = validation_handler.validation_tooltip + + def on_hover_enter(_e=None): + """Re-show tooltip on hover if validation error/warning exists.""" + try: + show_tooltips = bool(show_tooltips_var.get()) + except Exception: + show_tooltips = True + + if not show_tooltips: + return + + # Show tooltip without auto-hide when hovering + if validation_state.status == 'error': + validation_tooltip.show_error(validation_state.message, duration_ms=None) + elif validation_state.status == 'warning': + validation_tooltip.show_warning(validation_state.message, duration_ms=None) + + def on_hover_leave(_e=None): + """Hide tooltip when mouse leaves the field.""" + validation_tooltip.hide() + + component.bind("", on_hover_enter, add="+") + component.bind("", on_hover_leave, add="+") + + # Override destroy to cleanup resources original_destroy = component.destroy def new_destroy(): # 'temporary' fix until https://github.com/TomSchimansky/CustomTkinter/pull/2077 is merged # unfortunately Tom has admitted to forgetting about how to maintain CTK so this likely will never be merged - nonlocal validation_after_id, revert_after_id if component._textvariable_callback_name: component._textvariable.trace_remove("write", component._textvariable_callback_name) component._textvariable_callback_name = "" - if validation_after_id: - with contextlib.suppress(Exception): - component.after_cancel(validation_after_id) - if revert_after_id: - with contextlib.suppress(Exception): - component.after_cancel(revert_after_id) - + # Cleanup validation + validation_handler.cleanup() var.trace_remove("write", validation_trace_name) if command is not None and trace_id is not None: @@ -208,119 +458,262 @@ def new_destroy(): return component -def file_entry( - master, row, column, ui_state: UIState, var_name: str, - path_modifier: Callable[[str], str] = None, - is_output: bool = False, - allow_model_files: bool = True, - allow_image_files: bool = False, +def model_output_entry( + master, + row, + column, + ui_state: UIState, + var_name: str = "output_model_destination", + format_var_name: str = "output_model_format", + method_var_name: str = "training_method", + prefix_var_name: str = "save_filename_prefix", command: Callable[[str], None] = None, - valid_extensions: list[str] = None, - path_type: str = "file", # "file" or "directory" ): + """ + Specialized entry for model output paths with smart validation. + Automatically updates validation when format, method, or prefix changes. + """ frame = ctk.CTkFrame(master, fg_color="transparent") frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") - frame.grid_columnconfigure(0, weight=1) - if not is_output: - meta = ui_state.get_field_metadata(var_name) - is_output = getattr(meta, 'is_output', False) + var = ui_state.get_var(var_name) - def validate_file_path_wrapper(value: str) -> tuple[bool, str]: - return validate_file_path(value, is_output, valid_extensions, path_type) + validator = ModelOutputValidator( + var=var, + ui_state=ui_state, + format_var_name=format_var_name, + method_var_name=method_var_name, + prefix_var_name=prefix_var_name, + ) + # Create entry with smart validation and hover behavior entry_widget = entry( frame, row=0, column=0, ui_state=ui_state, var_name=var_name, - custom_validator=validate_file_path_wrapper, command=command + custom_validator=validator.validate, + validation_state=validator.state, + enable_hover_validation=True, + command=command ) - _register_drop_target(entry_widget, ui_state, var_name, command) + # Setup reactive traces to re-validate when related fields change + validator.setup_traces() - def __open_dialog(): - filetypes = [ - ("All Files", "*.*"), - ] + # Clear validation state when value changes (before validation runs) + validation_trace = var.trace_add("write", validator.on_value_change) - if allow_model_files: - filetypes.extend([ - ("Diffusers", "model_index.json"), - ("Checkpoint", "*.ckpt *.pt *.bin"), - ("Safetensors", "*.safetensors"), - ]) - if allow_image_files: - filetypes.extend([ - ("Image", ' '.join([f"*.{x}" for x in supported_image_extensions()])), - ]) - - if is_output: - file_path = filedialog.asksaveasfilename(filetypes=filetypes) - else: - file_path = filedialog.askopenfilename(filetypes=filetypes) + # Cleanup on destroy + original_destroy = entry_widget.destroy - if file_path: - if path_modifier: - file_path = path_modifier(file_path) + def new_destroy(): + with contextlib.suppress(Exception): + var.trace_remove("write", validation_trace) + validator.cleanup_traces() + original_destroy() - ui_state.get_var(var_name).set(file_path) + entry_widget.destroy = new_destroy + # Enable drag-and-drop + _register_drop_target(entry_widget, ui_state, var_name, command) + + # Browse button + def open_dialog(): + filetypes = [ + ("All Files", "*.*"), + ("Diffusers", "model_index.json"), + ("Checkpoint", "*.ckpt *.pt *.bin"), + ("Safetensors", "*.safetensors"), + ] + selected_path = filedialog.asksaveasfilename(filetypes=filetypes) + if selected_path: + var.set(selected_path) if command: - command(file_path) + command(selected_path) - button_component = ctk.CTkButton(frame, text="...", width=40, command=__open_dialog) + button_component = ctk.CTkButton(frame, text="...", width=40, command=open_dialog) button_component.grid(row=0, column=1, padx=(0, PAD), pady=PAD, sticky="nsew") return frame -def dir_entry(master, row, column, ui_state: UIState, var_name: str, command: Callable[[str], None] = None, is_output: bool = False): +def path_entry( + master, + row, + column, + ui_state: UIState, + var_name: str, + path_modifier: Callable[[str], str] = None, + is_output: bool = False, + path_type: str = "file", # "file" or "directory" + command: Callable[[str], None] = None, + allow_model_files: bool = True, + allow_image_files: bool = False, + valid_extensions: list[str] = None, +): + """ + General-purpose path entry component for files and directories. + + For model output destination with smart validation, use model_output_entry() instead. + + Args: + master: Parent widget + row, column: Grid position + ui_state: UI state manager + var_name: Variable name in ui_state + path_modifier: Optional function to modify selected paths + is_output: Whether this is an output path + path_type: "file" or "directory" + command: Optional callback when path changes + allow_model_files: Show model file types in dialog (file mode only) + allow_image_files: Show image file types in dialog (file mode only) + valid_extensions: List of valid extensions (file mode only) + + Returns: + Frame containing entry and browse button + """ frame = ctk.CTkFrame(master, fg_color="transparent") frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") - frame.grid_columnconfigure(0, weight=1) + # Determine if this is an output path if not is_output: meta = ui_state.get_field_metadata(var_name) is_output = getattr(meta, 'is_output', False) - def validate_dir_path_wrapper(value: str) -> tuple[bool, str]: - # trim whitespaces from the last path component, nothing else - if not value: - return validate_file_path(value, is_output=is_output, valid_extensions=None, path_type="directory") - - path = Path(value) - trimmed_name = path.name.strip() + var = ui_state.get_var(var_name) - if trimmed_name != path.name: - sep = '\\' if '\\' in value else '/' - trimmed_path = f"{path.parent}{sep}{trimmed_name}" if str(path.parent) != '.' else trimmed_name - ui_state.get_var(var_name).set(trimmed_path) - value = trimmed_path + # Simple path validation with auto-trimming for directories + def simple_path_validator(value: str) -> ValidationResult: + if not value: + return validate_file_path(value, is_output, valid_extensions, path_type) + + # Auto-trim whitespace from directory names + if path_type == "directory": + path = Path(value) + trimmed_name = path.name.strip() + if trimmed_name != path.name: + sep = '\\' if '\\' in value else '/' + trimmed_path = f"{path.parent}{sep}{trimmed_name}" if str(path.parent) != '.' else trimmed_name + var.set(trimmed_path) + value = trimmed_path - return validate_file_path(value, is_output=is_output, valid_extensions=None, path_type="directory") + return validate_file_path(value, is_output, valid_extensions, path_type) + # Create entry with validation entry_widget = entry( frame, row=0, column=0, ui_state=ui_state, var_name=var_name, - custom_validator=validate_dir_path_wrapper, command=command + custom_validator=simple_path_validator, command=command ) + # Enable drag-and-drop _register_drop_target(entry_widget, ui_state, var_name, command) - def __open_dialog(): - dir_path = filedialog.askdirectory() + # Browse button + def open_dialog(): + if path_type == "directory": + selected_path = filedialog.askdirectory() + else: + # Build file type filters + filetypes = [("All Files", "*.*")] + + if allow_model_files: + filetypes.extend([ + ("Diffusers", "model_index.json"), + ("Checkpoint", "*.ckpt *.pt *.bin"), + ("Safetensors", "*.safetensors"), + ]) + if allow_image_files: + filetypes.extend([ + ("Image", ' '.join([f"*.{x}" for x in supported_image_extensions()])), + ]) + + if is_output: + selected_path = filedialog.asksaveasfilename(filetypes=filetypes) + else: + selected_path = filedialog.askopenfilename(filetypes=filetypes) + + if selected_path: + if path_modifier: + selected_path = path_modifier(selected_path) - if dir_path: - ui_state.get_var(var_name).set(dir_path) + var.set(selected_path) if command: - command(dir_path) + command(selected_path) - button_component = ctk.CTkButton(frame, text="...", width=40, command=__open_dialog) + button_component = ctk.CTkButton(frame, text="...", width=40, command=open_dialog) button_component.grid(row=0, column=1, padx=(0, PAD), pady=PAD, sticky="nsew") return frame +def file_entry( + master, row, column, ui_state: UIState, var_name: str, + path_modifier: Callable[[str], str] = None, + is_output: bool = False, + allow_model_files: bool = True, + allow_image_files: bool = False, + command: Callable[[str], None] = None, + valid_extensions: list[str] = None, + path_type: str = "file", +): + return path_entry( + master=master, + row=row, + column=column, + ui_state=ui_state, + var_name=var_name, + path_modifier=path_modifier, + is_output=is_output, + path_type=path_type, + command=command, + allow_model_files=allow_model_files, + allow_image_files=allow_image_files, + valid_extensions=valid_extensions, + ) + + +def dir_entry( + master, row, column, ui_state: UIState, var_name: str, + command: Callable[[str], None] = None, + is_output: bool = False +): + return path_entry( + master=master, + row=row, + column=column, + ui_state=ui_state, + var_name=var_name, + is_output=is_output, + path_type="directory", + command=command, + ) + + +def app_title(master, row, column): + frame = ctk.CTkFrame(master) + frame.grid(row=row, column=column, padx=5, pady=5, sticky="nsew") + + image_component = ctk.CTkImage( + Image.open("resources/icons/icon.png").resize((40, 40), Image.Resampling.LANCZOS), + size=(40, 40) + ) + image_label_component = ctk.CTkLabel(frame, image=image_component, text="") + image_label_component.grid(row=0, column=0, padx=PAD, pady=PAD) + + label_component = ctk.CTkLabel(frame, text="OneTrainer", font=ctk.CTkFont(size=20, weight="bold")) + label_component.grid(row=0, column=1, padx=(0, PAD), pady=PAD) + + +def label(master, row, column, text, pad=PAD, tooltip=None, wide_tooltip=False, wraplength=0): + component = ctk.CTkLabel(master, text=text, wraplength=wraplength) + component.grid(row=row, column=column, padx=pad, pady=pad, sticky="nw") + if tooltip: + ToolTip(component, tooltip, wide=wide_tooltip) + return component + + + def time_entry(master, row, column, ui_state: UIState, var_name: str, unit_var_name, supports_time_units: bool = True): frame = ctk.CTkFrame(master, fg_color="transparent") frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") diff --git a/modules/util/ui/validation.py b/modules/util/ui/validation.py new file mode 100644 index 000000000..9bd8a4456 --- /dev/null +++ b/modules/util/ui/validation.py @@ -0,0 +1,415 @@ +""" +Input validation utilities for OneTrainer UI components. + +Provides validation for file paths, directories, and model output destinations +with support for auto-correction and detailed error messaging. +""" + +from __future__ import annotations + +import os +import re +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Any, Literal + +from modules.util.enum.ModelFormat import ModelFormat +from modules.util.enum.TrainingMethod import TrainingMethod + +from friendlywords import generate as generate_friendly + + +@dataclass(frozen=True) +class ValidationSettings: + datetime_format: str = "%Y%m%d_%H%M%S" + +SETTINGS = ValidationSettings() + +# Regex patterns +ONLY_WHITESPACE = re.compile(r"^\s+$") +TRAILING_SLASH_RE = re.compile(r"[\\/]$") +ENDS_WITH_EXT = re.compile(r"\.[A-Za-z0-9]+$") + +# Invalid characters/names +INVALID_CHARS_WIN = set('<>:"|?*') +INVALID_NAMES = {"", "."} + +ValidationStatus = Literal['success', 'warning', 'error'] + +@dataclass +class ValidationResult: + """Result of a validation operation.""" + ok: bool + corrected: str | None + message: str = "" + status: ValidationStatus | None = None + + def __post_init__(self): + """Auto-determine status from ok field if not explicitly set.""" + if self.status is None: + self.status = 'success' if self.ok else 'error' + + @property + def warning(self) -> bool: + """Backward compatibility: whether this is a warning.""" + return self.status == 'warning' + +# Helper to create common results +def _result(ok: bool, corrected: str | None = None, message: str = "", status: ValidationStatus | None = None) -> ValidationResult: + return ValidationResult(ok=ok, corrected=corrected, message=message, status=status) + + +def _is_windows() -> bool: + return os.name == "nt" + +def _has_invalid_chars(path: str) -> bool: + """Check if path contains invalid characters.""" + if _is_windows(): + has_drive = len(path) >= 2 and path[1] == ':' and path[0].isalpha() + start_idx = 2 if has_drive else 0 + return any(c in INVALID_CHARS_WIN for c in path[start_idx:]) + return "\x00" in path + +def _safe_exists(path: Path, check_dir: bool = True) -> bool: + """Safely check if a path exists.""" + try: + return path.is_dir() if check_dir else path.parent.exists() + except OSError: + return False + +def _folder_exists(path: Path) -> bool: + return _safe_exists(path, check_dir=True) + +def _parent_exists(path: Path) -> bool: + return _safe_exists(path, check_dir=False) + +def _format_path(path: Path, separator: str) -> str: + """Format path with preferred separator.""" + return str(path).replace('\\', separator) + +def _has_extension(path: str | Path, extension: str) -> bool: + """Check if path has the given extension (case-insensitive).""" + return str(path).lower().endswith(extension.lower()) + +def _get_extension(path: str | Path) -> str | None: + """Extract extension from path, or None if no extension.""" + match = ENDS_WITH_EXT.search(str(path)) + return match.group(0).lower() if match else None + +def _is_partial_match(current: str, required: str) -> bool: + """Check if current extension is a partial match of required.""" + chars_after_dot = len(current) - 1 + return chars_after_dot >= 3 and required.startswith(current) + +def generate_default_filename( + training_method: TrainingMethod, + prefix: str = "", + extension: str = ".safetensors", + use_friendly_names: bool = False, +) -> str: + """Generate a default filename for model output.""" + if use_friendly_names: + friendly = generate_friendly(1 if prefix else 2, separator="" if prefix else "-") + name = f"{prefix}_{friendly}" if prefix else friendly + else: + method_str = str(training_method).lower().replace('_', '-') + timestamp = datetime.now().strftime(SETTINGS.datetime_format) + name = f"{prefix}_{timestamp}" if prefix else f"{method_str}_{timestamp}" + return f"{name}{extension}" + +def get_allowed_formats_for_method(training_method: TrainingMethod) -> set[ModelFormat]: + """Get the allowed output formats for a given training method.""" + if training_method in (TrainingMethod.EMBEDDING, TrainingMethod.LORA): + return {ModelFormat.SAFETENSORS, ModelFormat.CKPT} + return {ModelFormat.SAFETENSORS, ModelFormat.CKPT, ModelFormat.DIFFUSERS} + +@dataclass +class ValidationContext: + """Context for validation operations.""" + raw: str + user_input: str + output_format: ModelFormat + training_method: TrainingMethod + autocorrect: bool + prefix: str + use_friendly_names: bool + is_output: bool + separator: str + + +def _validate_basic_input(ctx: ValidationContext) -> ValidationResult | None: + """Perform basic validation checks on input. Returns None to continue.""" + + # Check for whitespace-only input + if ONLY_WHITESPACE.match(ctx.raw): + return _result(False, "", "Input cleared: empty or invalid.") + + # Check for empty or invalid names + if not ctx.user_input or ctx.user_input in INVALID_NAMES: + return _result(False, "", "Input cleared: empty or invalid.") + + # Check for invalid characters + if _has_invalid_chars(ctx.user_input): + return _result(False, None, "Input contains invalid path characters.") + + # Check format/method compatibility + if ctx.output_format not in get_allowed_formats_for_method(ctx.training_method): + return _result(False, None, f"{ctx.training_method} cannot output {ctx.output_format} format.") + + # For input paths, verify existence + if not ctx.is_output: + if not Path(ctx.user_input).exists(): + return _result(False, None, "Input path does not exist.") + return _result(True, None, "") + + return None + +def _validate_diffusers_path(ctx: ValidationContext) -> ValidationResult: + """Validate DIFFUSERS format path (directory-based).""" + ext_match = ENDS_WITH_EXT.search(ctx.user_input) + + if ext_match and not ctx.autocorrect: + return _result(False, None, "DIFFUSERS expects a directory, not a file. Add a trailing slash.") + + # Remove extension if present + sanitized = ctx.user_input[:ext_match.start()] if ext_match else ctx.user_input + check_path = Path(sanitized) + + # Check parent directory exists + if check_path.parent != Path('.') and not _parent_exists(check_path): + return _result(False, None, "Parent folder does not exist.") + + # Warn if directory already exists + if not ctx.autocorrect and _folder_exists(check_path): + msg = "Directory already exists. Stripped extension for DIFFUSERS." if ext_match else "Directory already exists." + return _result(True, sanitized if ext_match else None, msg, 'warning') + + # Extension was stripped + if ext_match: + return _result(True, sanitized, f"Stripped extension for {ctx.output_format}.") + + return _result(True, None, "") + +def _validate_single_file_path(ctx: ValidationContext) -> ValidationResult: + """Validate single-file format path (SAFETENSORS, CKPT).""" + required_ext = ctx.output_format.file_extension() + + # Check if input looks like a directory + is_dir_like = bool(TRAILING_SLASH_RE.search(ctx.user_input)) + path = Path(ctx.user_input) + + # If no slash but exists as folder, treat as directory + if not is_dir_like and '/' not in ctx.user_input and '\\' not in ctx.user_input and _folder_exists(path): + is_dir_like = True + + # Handle directory-like input + if is_dir_like: + if ctx.autocorrect: + if not _folder_exists(path): + return _result(False, None, "Directory does not exist.") + default_name = generate_default_filename(ctx.training_method, ctx.prefix, required_ext, ctx.use_friendly_names) + corrected = ctx.user_input + default_name + return _result(True, corrected, f"Appended default filename: {default_name}") + return _result(False, None, "Path points to a folder; a filename is required.") + + # Check parent directory exists + if path.parent != Path('.') and not _parent_exists(path): + return _result(False, None, "Parent folder does not exist.") + + base = path.name + + # Handle trailing periods + if base.endswith('.'): + if ctx.autocorrect: + new_base = base.rstrip('.') + new_path = _format_path(path.with_name(new_base), ctx.separator) + return _result(True, new_path, "Removed trailing period(s).") + return _result(False, None, "Filename cannot end with a period.") + + # Check if already has correct extension + if _has_extension(base, required_ext): + if path.exists(): + return _result(True, None, "WARNING: File already exists and will be overwritten.", 'warning') + return _result(True, None, "") + + # Check for partial or incorrect extension + ext_match = ENDS_WITH_EXT.search(base) + if ext_match: + current_ext = ext_match.group(0).lower() + chars_after_dot = len(current_ext) - 1 + start_idx = ext_match.start() + + # Extension has at least 3 characters after dot + if chars_after_dot >= 3: + # Check if it's a partial match (e.g., ".safe" for ".safetensors") + if _is_partial_match(current_ext, required_ext): + if ctx.autocorrect: + new_base = base[:start_idx] + required_ext + new_path = _format_path(path.with_name(new_base), ctx.separator) + return _result(True, new_path, f"Completed {current_ext} -> {required_ext}") + return _result(False, None, f"Extension is incomplete; expected {required_ext}") + else: + # Different extension entirely + if ctx.autocorrect: + new_base = base[:start_idx] + required_ext + new_path = _format_path(path.with_name(new_base), ctx.separator) + return _result(True, new_path, f"Replaced {current_ext} with {required_ext}") + return _result(False, None, f"Wrong extension; expected {required_ext}") + + # Extension too short for auto-correction + if 1 <= chars_after_dot < 3 and required_ext.startswith(current_ext): + return _result(False, None, f"Extension {current_ext} is too short; Type at least 3 characters for auto-correction to occur.") + + # Unrecognized short extension - append required + if ctx.autocorrect: + corrected = _format_path(Path(path.parent, base + required_ext), ctx.separator) + return _result(True, corrected, f"Appended {required_ext} extension.") + return _result(False, None, f"Filename must end with {required_ext}") + + # No extension - append required extension + if ctx.autocorrect: + corrected = _format_path(Path(path.parent, base + required_ext), ctx.separator) + return _result(True, corrected, f"Appended {required_ext} extension.") + return _result(False, None, f"Filename must end with {required_ext}") + +def validate_destination( + raw: str, + output_format: ModelFormat, + training_method: TrainingMethod, + autocorrect: bool = True, + prefix: str = "", + use_friendly_names: bool = False, + is_output: bool = True, +) -> ValidationResult: + """ + Validate and optionally auto-correct a model output destination path. + + This function handles: + - Empty/invalid input + - Format/method compatibility + - Directory vs file paths + - File extension validation and correction + - Parent directory existence + - Existing file warnings + + Args: + raw: Raw user input string + output_format: Target output format (DIFFUSERS, SAFETENSORS, CKPT) + training_method: Training method being used + autocorrect: Whether to auto-correct invalid input + prefix: Optional prefix for auto-generated filenames + use_friendly_names: Use friendly names instead of timestamps + is_output: Whether this is an output path (True) or input path (False) + + Returns: + ValidationResult with ok status, optional corrected value, and message + """ + ctx = ValidationContext( + raw=raw, + user_input=raw.strip(), + output_format=output_format, + training_method=training_method, + autocorrect=autocorrect, + prefix=prefix, + use_friendly_names=use_friendly_names, + is_output=is_output, + separator='/' if '/' in raw else '\\' + ) + + # Basic validation checks + if basic_check := _validate_basic_input(ctx): + return basic_check + + # Format-specific validation + if output_format == ModelFormat.DIFFUSERS: + return _validate_diffusers_path(ctx) + else: + return _validate_single_file_path(ctx) + +def validate_basic_type( + value: str, + declared_type: type, + nullable: bool, + default_val: Any, +) -> ValidationResult: + """ + Validate basic Python types (int, float, bool, str). + + Args: + value: The string value to validate + declared_type: The expected Python type (int, float, bool, str) + nullable: Whether None/empty is allowed + default_val: The default value for this field + + Returns: + ValidationResult with validation outcome + """ + # Handle empty values + if value == "": + if nullable or (declared_type is str and default_val == ""): + return _result(True, None, "") + return _result(False, None, "Value required") + + # Type validation + try: + if declared_type is int: + int(value) + elif declared_type is float: + float(value) + elif declared_type is bool: + if value.lower() not in ("true", "false", "0", "1"): + return _result(False, None, "Invalid bool") + return _result(True, None, "") + except ValueError: + type_name = declared_type.__name__ if hasattr(declared_type, '__name__') else str(declared_type) + return _result(False, None, f"Invalid {type_name}") + +def validate_file_path( + value: str, + is_output: bool = False, + valid_extensions: list[str] | None = None, + path_type: str = "file", +) -> ValidationResult: + """ + Validate a file or directory path. + + Args: + value: Path string to validate + is_output: Whether this is an output path + valid_extensions: List of valid extensions (with dots), or None for any + path_type: Either "file" or "directory" + + Returns: + ValidationResult with validation outcome + """ + if not value: + return _result(True, None, "") + + # Check for invalid characters + if _has_invalid_chars(value): + return _result(False, None, "Input contains invalid path characters.") + + path = Path(value) + + if path_type == "directory": + if is_output: + if path.parent != Path('.') and not _parent_exists(path): + return _result(False, None, "Parent folder does not exist.") + else: + if not _folder_exists(path): + return _result(False, None, "Directory does not exist.") + else: + # For files + if valid_extensions: + if not any(_has_extension(value, ext) for ext in valid_extensions): + return _result(False, None, f"File must have one of these extensions: {', '.join(valid_extensions)}") + + if is_output: + if path.parent != Path('.') and not _parent_exists(path): + return _result(False, None, "Parent folder does not exist.") + else: + if not path.exists(): + return _result(False, None, "File does not exist") + + return _result(True, None, "") From b6609438df24f9e26e00c7273362745aa17441d5 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Mon, 10 Nov 2025 02:00:10 +1100 Subject: [PATCH 21/64] Trimming the fat from components.py and remove legacy file_entry/dir_entry --- modules/ui/AdditionalEmbeddingsTab.py | 2 +- modules/ui/ConceptWindow.py | 4 +- modules/ui/ConvertModelUI.py | 4 +- modules/ui/LoraTab.py | 2 +- modules/ui/ModelTab.py | 14 +- modules/ui/SampleFrame.py | 4 +- modules/ui/TrainUI.py | 10 +- modules/util/ui/components.py | 408 +++++++------------------- 8 files changed, 119 insertions(+), 329 deletions(-) diff --git a/modules/ui/AdditionalEmbeddingsTab.py b/modules/ui/AdditionalEmbeddingsTab.py index 465f3ac16..0ce9daed2 100644 --- a/modules/ui/AdditionalEmbeddingsTab.py +++ b/modules/ui/AdditionalEmbeddingsTab.py @@ -85,7 +85,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 ) diff --git a/modules/ui/ConceptWindow.py b/modules/ui/ConceptWindow.py index 60b195c75..5c0c35d56 100644 --- a/modules/ui/ConceptWindow.py +++ b/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': diff --git a/modules/ui/ConvertModelUI.py b/modules/ui/ConvertModelUI.py index 685aeaa85..8000103e2 100644 --- a/modules/ui/ConvertModelUI.py +++ b/modules/ui/ConvertModelUI.py @@ -82,7 +82,7 @@ def main_frame(self, master): # input name components.label(master, 2, 0, "Input name", tooltip="Filename, directory or hugging face repository of the base model") - components.file_entry( + components.path_entry( master, 2, 1, self.ui_state, "input_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x ) @@ -107,7 +107,7 @@ def main_frame(self, master): # output model destination components.label(master, 5, 0, "Model Output Destination", tooltip="Filename or directory where the output model is saved") - components.file_entry(master, 5, 1, self.ui_state, "output_model_destination", is_output=True) + components.path_entry(master, 5, 1, self.ui_state, "output_model_destination", is_output=True) self.button = components.button(master, 6, 1, "Convert", self.convert_model) diff --git a/modules/ui/LoraTab.py b/modules/ui/LoraTab.py index 9b852bcd5..0bcb0d616 100644 --- a/modules/ui/LoraTab.py +++ b/modules/ui/LoraTab.py @@ -58,7 +58,7 @@ def setup_lora(self, peft_type: PeftType): # lora model name components.label(master, 0, 0, f"{name} base model", tooltip=f"The base {name} to train on. Leave empty to create a new {name}") - entry = components.file_entry( + entry = components.path_entry( master, 0, 1, self.ui_state, "lora_model_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x ) diff --git a/modules/ui/ModelTab.py b/modules/ui/ModelTab.py index 94608f4b1..1a8c259ad 100644 --- a/modules/ui/ModelTab.py +++ b/modules/ui/ModelTab.py @@ -288,7 +288,7 @@ def __create_base_dtype_components(self, row: int) -> int: # base model components.label(self.scroll_frame, row, 0, "Base Model", tooltip="Filename, directory or Hugging Face repository of the base model") - components.file_entry( + components.path_entry( self.scroll_frame, row, 1, self.ui_state, "base_model_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x ) @@ -333,7 +333,7 @@ def __create_base_components( # prior model components.label(self.scroll_frame, row, 0, "Prior Model", tooltip="Filename, directory or Hugging Face repository of the prior model") - components.file_entry( + components.path_entry( self.scroll_frame, row, 1, self.ui_state, "prior.model_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x ) @@ -351,7 +351,7 @@ def __create_base_components( # transformer model components.label(self.scroll_frame, row, 0, "Override Transformer / GGUF", tooltip="Can be used to override the transformer in the base model. Safetensors and GGUF files are supported, local and on Huggingface. If a GGUF file is used, the DataType must also be set to GGUF") - components.file_entry( + components.path_entry( self.scroll_frame, row, 1, self.ui_state, "transformer.model_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x ) @@ -405,7 +405,7 @@ def __create_base_components( # text encoder 4 weight dtype components.label(self.scroll_frame, row, 0, "Text Encoder 4 Override", tooltip="Filename, directory or Hugging Face repository of the text encoder 4 model") - components.file_entry( + components.path_entry( self.scroll_frame, row, 1, self.ui_state, "text_encoder_4.model_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x ) @@ -422,7 +422,7 @@ def __create_base_components( # base model components.label(self.scroll_frame, row, 0, "VAE Override", tooltip="Directory or Hugging Face repository of a VAE model in diffusers format. Can be used to override the VAE included in the base model. \n\n Using a safetensor VAE file will cause an error that the model cannot be loaded.") - components.file_entry( + components.path_entry( self.scroll_frame, row, 1, self.ui_state, "vae.model_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x ) @@ -441,7 +441,7 @@ def __create_effnet_encoder_components(self, row: int): # effnet encoder model components.label(self.scroll_frame, row, 0, "Effnet Encoder Model", tooltip="Filename, directory or Hugging Face repository of the effnet encoder model") - components.file_entry( + components.path_entry( self.scroll_frame, row, 1, self.ui_state, "effnet_encoder.model_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x ) @@ -464,7 +464,7 @@ def __create_decoder_components( # decoder model components.label(self.scroll_frame, row, 0, "Decoder Model", tooltip="Filename, directory or Hugging Face repository of the decoder model") - components.file_entry( + components.path_entry( self.scroll_frame, row, 1, self.ui_state, "decoder.model_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x ) diff --git a/modules/ui/SampleFrame.py b/modules/ui/SampleFrame.py index 89c056f8c..0dcb9ecd5 100644 --- a/modules/ui/SampleFrame.py +++ b/modules/ui/SampleFrame.py @@ -108,7 +108,7 @@ def __init__( # base image path components.label(bottom_frame, 6, 0, "base image path:", tooltip="The base image used when inpainting.") - components.file_entry(bottom_frame, 6, 1, self.ui_state, "base_image_path", + components.path_entry(bottom_frame, 6, 1, self.ui_state, "base_image_path", allow_model_files=False, allow_image_files=True, ) @@ -116,7 +116,7 @@ def __init__( # mask image path components.label(bottom_frame, 7, 2, "mask image path:", tooltip="The mask used when inpainting.") - components.file_entry(bottom_frame, 7, 3, self.ui_state, "mask_image_path", + components.path_entry(bottom_frame, 7, 3, self.ui_state, "mask_image_path", allow_model_files=False, allow_image_files=True, ) diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index f5b8767ec..20fff9d3c 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -257,12 +257,12 @@ def create_general_tab(self, master): components.label(frame, 0, 0, "Workspace Directory", tooltip="The directory where all files of this training run are saved") # Remove command parameter - we'll use the trace instead for debouncing - components.dir_entry(frame, 0, 1, self.ui_state, "workspace_dir", is_output=True) + components.path_entry(frame, 0, 1, self.ui_state, "workspace_dir", is_output=True, path_type="directory") # cache dir components.label(frame, 0, 2, "Cache Directory", tooltip="The directory where cached data is saved") - components.dir_entry(frame, 0, 3, self.ui_state, "cache_dir", is_output=True) + components.path_entry(frame, 0, 3, self.ui_state, "cache_dir", is_output=True, path_type="directory") # continue from previous backup components.label(frame, 2, 0, "Continue from last backup", @@ -281,7 +281,7 @@ def create_general_tab(self, master): components.label(frame, 4, 2, "Debug Directory", tooltip="The directory where debug data is saved") - components.dir_entry(frame, 4, 3, self.ui_state, "debug_dir", is_output=True) + components.path_entry(frame, 4, 3, self.ui_state, "debug_dir", is_output=True, path_type="directory") # tensorboard components.label(frame, 6, 0, "Tensorboard Mode", @@ -514,7 +514,7 @@ def lora_tab(self, master): # lora model name components.label(frame, 0, 0, "LoRA base model", tooltip="The base LoRA to train on. Leave empty to create a new LoRA") - components.file_entry( + components.path_entry( frame, 0, 1, self.ui_state, "lora_model_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x ) @@ -561,7 +561,7 @@ def embedding_tab(self, master): # embedding model name components.label(frame, 0, 0, "Base embedding", tooltip="The base embedding to train on. Leave empty to create a new embedding") - components.file_entry( + components.path_entry( frame, 0, 1, self.ui_state, "embedding.model_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x ) diff --git a/modules/util/ui/components.py b/modules/util/ui/components.py index 1db41be8a..4a45e5d5f 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -6,12 +6,14 @@ from tkinter import filedialog from typing import Any, Literal +from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.TimeUnit import TimeUnit +from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.path_util import supported_image_extensions from modules.util.ui.ToolTip import ToolTip from modules.util.ui.ui_utils import _register_drop_target from modules.util.ui.UIState import UIState -from modules.util.ui.validation import ValidationResult, validate_basic_type, validate_file_path +from modules.util.ui.validation import ValidationResult, validate_basic_type, validate_destination, validate_file_path import customtkinter as ctk from customtkinter.windows.widgets.scaling import CTkScalingBaseClass @@ -19,6 +21,28 @@ PAD = 10 +# Common filetypes for dialogs +MODEL_FILETYPES = [ + ("All Files", "*.*"), + ("Diffusers", "model_index.json"), + ("Checkpoint", "*.ckpt *.pt *.bin"), + ("Safetensors", "*.safetensors"), +] + + +def _safe_bool(var, default=True): + try: + return bool(var.get()) + except Exception: + return default + + +def _wrap_dropdown_destroy(component): + # temporary fix until https://github.com/TomSchimansky/CustomTkinter/pull/2246 is merged + orig_destroy = component._dropdown_menu.destroy + component._dropdown_menu.destroy = lambda: (orig_destroy(), CTkScalingBaseClass.destroy(component._dropdown_menu)) + + @dataclass(frozen=True) class ComponentValidationSettings: debounce_stop_typing_ms: int = 1700 @@ -30,25 +54,17 @@ class ComponentValidationSettings: @dataclass class ValidationState: - """Tracks validation state for UI components.""" status: Literal['error', 'warning'] | None = None message: str = "" def clear(self): - self.status = None - self.message = "" - - def set_error(self, message: str): - self.status = 'error' - self.message = message + self.status, self.message = None, "" - def set_warning(self, message: str): - self.status = 'warning' - self.message = message + def set_status(self, status: Literal['error', 'warning'], message: str): + self.status, self.message = status, message class EntryValidationHandler: - """Handles UI coordination for entry widget validation with debouncing and visual feedback.""" def __init__( self, @@ -64,7 +80,7 @@ def __init__( self.var_name = var_name self.ui_state = ui_state self.custom_validator = custom_validator - self.validation_state = validation_state # Optional external validation state + self.validation_state = validation_state # Get border colors try: @@ -75,37 +91,34 @@ def __init__( self.error_border_color = "#dc3545" self.warning_border_color = "#ff9500" - # Validation state self.validation_after_id = None self.revert_after_id = None self.touched = False self.last_valid_value = var.get() - # Create tooltip (hover_only=False to prevent default hover behavior, wide=True for validation messages) self.validation_tooltip = ToolTip(component, text="", hover_only=False, track_movement=True, wide=True) component._validation_tooltip = self.validation_tooltip - def should_show_tooltip(self) -> bool: - """Check if validation tooltips are enabled in settings.""" - try: - show_tooltips_var = self.ui_state.get_var("validation_show_tooltips") - return bool(show_tooltips_var.get()) - except Exception: - return True + def _cancel_after(self, after_id): + with contextlib.suppress(Exception): + self.component.after_cancel(after_id) - def validate_value(self, value: str, revert_delay_ms: int | None) -> bool: - """ - Validate the current value and update UI accordingly. + def _reset_border(self): + self.component.configure(border_color=self.original_border_color) + + def should_show_tooltip(self) -> bool: + return _safe_bool(self.ui_state.get_var("validation_show_tooltips")) - Returns: - True if validation passed, False otherwise. - """ + def _show_validation_tooltip(self): + if not self.should_show_tooltip() or not self.validation_state: + return + if self.validation_state.status == 'error': + self.validation_tooltip.show_error(self.validation_state.message, duration_ms=None) + elif self.validation_state.status == 'warning': + self.validation_tooltip.show_warning(self.validation_state.message, duration_ms=None) - # Cancel any pending revert - if self.revert_after_id: - with contextlib.suppress(Exception): - self.component.after_cancel(self.revert_after_id) - self.revert_after_id = None + def validate_value(self, value: str, revert_delay_ms: int | None) -> bool: + self._cancel_after(self.revert_after_id) meta = self.ui_state.get_field_metadata(self.var_name) @@ -134,34 +147,29 @@ def validate_value(self, value: str, revert_delay_ms: int | None) -> bool: return self._success(value) def _success(self, value: str) -> bool: - """Handle successful validation.""" - self.component.configure(border_color=self.original_border_color) + self._reset_border() self.validation_tooltip.hide() self.last_valid_value = value return True def _warning(self, reason: str, value: str) -> bool: - """Handle warning state with orange border.""" self.component.configure(border_color=self.warning_border_color) if self.should_show_tooltip(): self.validation_tooltip.show_warning(reason) self.last_valid_value = value - # Update external validation state if provided if self.validation_state: - self.validation_state.set_warning(reason) + self.validation_state.set_status('warning', reason) return True def _fail(self, reason: str, revert_delay_ms: int | None) -> bool: - """Handle validation failure.""" self.component.configure(border_color=self.error_border_color) if self.should_show_tooltip(): self.validation_tooltip.show_error(reason) - # Update external validation state if provided if self.validation_state: - self.validation_state.set_error(reason) + self.validation_state.set_status('error', reason) if revert_delay_ms is not None: self.revert_after_id = self.component.after(revert_delay_ms, self._do_revert) @@ -170,34 +178,21 @@ def _fail(self, reason: str, revert_delay_ms: int | None) -> bool: return False def _do_revert(self): - """Revert to last valid value.""" self.var.set(self.last_valid_value) - self.component.configure(border_color=self.original_border_color) + self._reset_border() self.validation_tooltip.hide() def debounced_validate(self, *_): - """Validate with debouncing - called on variable changes.""" if not self.touched: - if self.validation_after_id: - with contextlib.suppress(Exception): - self.component.after_cancel(self.validation_after_id) - self.validation_after_id = None + self._cancel_after(self.validation_after_id) return - # Clear validation UI immediately when field is modified self.validation_tooltip.hide() - self.component.configure(border_color=self.original_border_color) + self._reset_border() - # Cancel pending operations - if self.revert_after_id: - with contextlib.suppress(Exception): - self.component.after_cancel(self.revert_after_id) - self.revert_after_id = None - if self.validation_after_id: - with contextlib.suppress(Exception): - self.component.after_cancel(self.validation_after_id) + self._cancel_after(self.revert_after_id) + self._cancel_after(self.validation_after_id) - # Schedule validation self.validation_after_id = self.component.after( COMPONENT_VALIDATION_SETTINGS.debounce_stop_typing_ms, lambda: self.validate_value( @@ -207,22 +202,13 @@ def debounced_validate(self, *_): ) def on_focus_in(self, _e=None): - """Handle focus in event.""" self.touched = False - # Re-show tooltip if there's an existing validation state - if self.validation_state and self.should_show_tooltip(): - if self.validation_state.status == 'error': - self.validation_tooltip.show_error(self.validation_state.message, duration_ms=None) - elif self.validation_state.status == 'warning': - self.validation_tooltip.show_warning(self.validation_state.message, duration_ms=None) + self._show_validation_tooltip() def on_user_input(self, _e=None): - """Handle user input event.""" self.touched = True def on_focus_out(self, _e=None): - """Handle focus out event - validate if user interacted.""" - # Always hide tooltip on focus out self.validation_tooltip.hide() if self.touched: @@ -232,20 +218,11 @@ def on_focus_out(self, _e=None): ) def cleanup(self): - """Cleanup resources.""" - if self.validation_after_id: - with contextlib.suppress(Exception): - self.component.after_cancel(self.validation_after_id) - if self.revert_after_id: - with contextlib.suppress(Exception): - self.component.after_cancel(self.revert_after_id) + for after_id in (self.validation_after_id, self.revert_after_id): + self._cancel_after(after_id) class ModelOutputValidator: - """ - Smart validator for model output paths with auto-correction. - Delegates validation logic to validation.py. - """ def __init__( self, @@ -266,43 +243,24 @@ def __init__( self.state = ValidationState() self._trace_ids: dict[str, str] = {} - def validate(self, value: str) -> ValidationResult: - """ - Validate model output path using validation.py logic. + def _get_enum_value(self, var, default_enum): + with contextlib.suppress(KeyError, ValueError): + return type(default_enum)[var.get()] + return default_enum - Returns: - ValidationResult with status (success/warning/error) - """ - from modules.util.enum.ModelFormat import ModelFormat - from modules.util.enum.TrainingMethod import TrainingMethod - from modules.util.ui.validation import validate_destination + def validate(self, value: str) -> ValidationResult: value = value.strip() if not value: self.state.clear() return ValidationResult(ok=True, corrected=None, message="", status='success') - try: - output_format = ModelFormat[self.format_var.get()] - except (KeyError, ValueError): - output_format = ModelFormat.SAFETENSORS - - try: - training_method = TrainingMethod[self.method_var.get()] - except (KeyError, ValueError): - training_method = TrainingMethod.FINE_TUNE - - try: - autocorrect = bool(self.autocorrect_var.get()) - except Exception: - autocorrect = True + output_format = self._get_enum_value(self.format_var, ModelFormat.SAFETENSORS) + training_method = self._get_enum_value(self.method_var, TrainingMethod.FINE_TUNE) + autocorrect = _safe_bool(self.autocorrect_var) prefix = self.prefix_var.get() if self.prefix_var else "" - - try: - use_friendly_names = bool(self.friendly_names_var.get()) - except Exception: - use_friendly_names = False + use_friendly_names = _safe_bool(self.friendly_names_var, default=False) # Use validation logic from validation.py result = validate_destination( @@ -315,37 +273,31 @@ def validate(self, value: str) -> ValidationResult: is_output=True ) - # Update internal state for hover behavior - if result.status == 'error': - self.state.set_error(result.message) - elif result.status == 'warning': - self.state.set_warning(result.message) + if result.status: + self.state.set_status(result.status, result.message) else: self.state.clear() - # Apply correction if available if result.corrected and result.corrected != value: self.var.set(result.corrected) return result - def trigger_validation(self, *_): - self.validate(self.var.get()) - - def on_value_change(self, *_): - self.state.clear() - def setup_traces(self): - self._trace_ids['format'] = self.format_var.trace_add("write", self.trigger_validation) - self._trace_ids['method'] = self.method_var.trace_add("write", self.trigger_validation) - self._trace_ids['prefix'] = self.prefix_var.trace_add("write", self.trigger_validation) - self._trace_ids['friendly_names'] = self.friendly_names_var.trace_add("write", self.trigger_validation) + trace_vars = [ + ('format', self.format_var), + ('method', self.method_var), + ('prefix', self.prefix_var), + ('friendly_names', self.friendly_names_var), + ] + for key, var in trace_vars: + self._trace_ids[key] = var.trace_add("write", lambda *_: self.validate(self.var.get())) def cleanup_traces(self): for key, trace_id in self._trace_ids.items(): - with contextlib.suppress(Exception): - var = getattr(self, f"{key}_var", None) - if var: + var = getattr(self, f"{key}_var", None) + if var: + with contextlib.suppress(Exception): var.trace_remove("write", trace_id) self._trace_ids.clear() @@ -367,14 +319,7 @@ def entry( validation_state: ValidationState = None, enable_hover_validation: bool = False, ): - """ - Create an entry widget with optional validation. - - Args: - validation_state: Optional ValidationState object for stateful validators. - If provided, enables hover behavior to re-show validation tooltips. - enable_hover_validation: If True, enables hover behavior to show validation state tooltips. - """ + var = ui_state.get_var(var_name) trace_id = None if command: @@ -383,7 +328,6 @@ def entry( component = ctk.CTkEntry(master, textvariable=var, width=width) component.grid(row=row, column=column, padx=PAD, pady=PAD, sticky=sticky) - # Create validation handler validation_handler = EntryValidationHandler( component=component, var=var, @@ -392,56 +336,40 @@ def entry( custom_validator=custom_validator, validation_state=validation_state, ) - - # Setup variable trace for validation validation_trace_name = var.trace_add("write", validation_handler.debounced_validate) # Bind focus and input events component.bind("", validation_handler.on_focus_in) - component.bind("", validation_handler.on_user_input) - component.bind("<>", validation_handler.on_user_input) - component.bind("<>", validation_handler.on_user_input) component.bind("", validation_handler.on_focus_out) - - # Optional hover validation behavior + for event in ("", "<>", "<>"): + component.bind(event, validation_handler.on_user_input) if enable_hover_validation and validation_state: show_tooltips_var = ui_state.get_var("validation_show_tooltips") validation_tooltip = validation_handler.validation_tooltip def on_hover_enter(_e=None): - """Re-show tooltip on hover if validation error/warning exists.""" - try: - show_tooltips = bool(show_tooltips_var.get()) - except Exception: - show_tooltips = True - - if not show_tooltips: + if not _safe_bool(show_tooltips_var): return - # Show tooltip without auto-hide when hovering if validation_state.status == 'error': validation_tooltip.show_error(validation_state.message, duration_ms=None) elif validation_state.status == 'warning': validation_tooltip.show_warning(validation_state.message, duration_ms=None) def on_hover_leave(_e=None): - """Hide tooltip when mouse leaves the field.""" validation_tooltip.hide() component.bind("", on_hover_enter, add="+") component.bind("", on_hover_leave, add="+") - # Override destroy to cleanup resources original_destroy = component.destroy def new_destroy(): # 'temporary' fix until https://github.com/TomSchimansky/CustomTkinter/pull/2077 is merged - # unfortunately Tom has admitted to forgetting about how to maintain CTK so this likely will never be merged if component._textvariable_callback_name: component._textvariable.trace_remove("write", component._textvariable_callback_name) component._textvariable_callback_name = "" - # Cleanup validation validation_handler.cleanup() var.trace_remove("write", validation_trace_name) @@ -469,10 +397,6 @@ def model_output_entry( prefix_var_name: str = "save_filename_prefix", command: Callable[[str], None] = None, ): - """ - Specialized entry for model output paths with smart validation. - Automatically updates validation when format, method, or prefix changes. - """ frame = ctk.CTkFrame(master, fg_color="transparent") frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") frame.grid_columnconfigure(0, weight=1) @@ -487,7 +411,6 @@ def model_output_entry( prefix_var_name=prefix_var_name, ) - # Create entry with smart validation and hover behavior entry_widget = entry( frame, row=0, column=0, ui_state=ui_state, var_name=var_name, custom_validator=validator.validate, @@ -495,36 +418,19 @@ def model_output_entry( enable_hover_validation=True, command=command ) - - # Setup reactive traces to re-validate when related fields change validator.setup_traces() - - # Clear validation state when value changes (before validation runs) - validation_trace = var.trace_add("write", validator.on_value_change) - - # Cleanup on destroy original_destroy = entry_widget.destroy def new_destroy(): - with contextlib.suppress(Exception): - var.trace_remove("write", validation_trace) validator.cleanup_traces() original_destroy() entry_widget.destroy = new_destroy - # Enable drag-and-drop _register_drop_target(entry_widget, ui_state, var_name, command) - # Browse button def open_dialog(): - filetypes = [ - ("All Files", "*.*"), - ("Diffusers", "model_index.json"), - ("Checkpoint", "*.ckpt *.pt *.bin"), - ("Safetensors", "*.safetensors"), - ] - selected_path = filedialog.asksaveasfilename(filetypes=filetypes) + selected_path = filedialog.asksaveasfilename(filetypes=MODEL_FILETYPES) if selected_path: var.set(selected_path) if command: @@ -550,27 +456,6 @@ def path_entry( allow_image_files: bool = False, valid_extensions: list[str] = None, ): - """ - General-purpose path entry component for files and directories. - - For model output destination with smart validation, use model_output_entry() instead. - - Args: - master: Parent widget - row, column: Grid position - ui_state: UI state manager - var_name: Variable name in ui_state - path_modifier: Optional function to modify selected paths - is_output: Whether this is an output path - path_type: "file" or "directory" - command: Optional callback when path changes - allow_model_files: Show model file types in dialog (file mode only) - allow_image_files: Show image file types in dialog (file mode only) - valid_extensions: List of valid extensions (file mode only) - - Returns: - Frame containing entry and browse button - """ frame = ctk.CTkFrame(master, fg_color="transparent") frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") frame.grid_columnconfigure(0, weight=1) @@ -584,11 +469,8 @@ def path_entry( # Simple path validation with auto-trimming for directories def simple_path_validator(value: str) -> ValidationResult: - if not value: - return validate_file_path(value, is_output, valid_extensions, path_type) - # Auto-trim whitespace from directory names - if path_type == "directory": + if value and path_type == "directory": path = Path(value) trimmed_name = path.name.strip() if trimmed_name != path.name: @@ -613,19 +495,12 @@ def open_dialog(): if path_type == "directory": selected_path = filedialog.askdirectory() else: - # Build file type filters filetypes = [("All Files", "*.*")] if allow_model_files: - filetypes.extend([ - ("Diffusers", "model_index.json"), - ("Checkpoint", "*.ckpt *.pt *.bin"), - ("Safetensors", "*.safetensors"), - ]) + filetypes.extend(MODEL_FILETYPES[1:]) # Skip "All Files" since we already added it if allow_image_files: - filetypes.extend([ - ("Image", ' '.join([f"*.{x}" for x in supported_image_extensions()])), - ]) + filetypes.append(("Image", ' '.join(f"*.{x}" for x in supported_image_extensions()))) if is_output: selected_path = filedialog.asksaveasfilename(filetypes=filetypes) @@ -647,49 +522,6 @@ def open_dialog(): return frame -def file_entry( - master, row, column, ui_state: UIState, var_name: str, - path_modifier: Callable[[str], str] = None, - is_output: bool = False, - allow_model_files: bool = True, - allow_image_files: bool = False, - command: Callable[[str], None] = None, - valid_extensions: list[str] = None, - path_type: str = "file", -): - return path_entry( - master=master, - row=row, - column=column, - ui_state=ui_state, - var_name=var_name, - path_modifier=path_modifier, - is_output=is_output, - path_type=path_type, - command=command, - allow_model_files=allow_model_files, - allow_image_files=allow_image_files, - valid_extensions=valid_extensions, - ) - - -def dir_entry( - master, row, column, ui_state: UIState, var_name: str, - command: Callable[[str], None] = None, - is_output: bool = False -): - return path_entry( - master=master, - row=row, - column=column, - ui_state=ui_state, - var_name=var_name, - is_output=is_output, - path_type="directory", - command=command, - ) - - def app_title(master, row, column): frame = ctk.CTkFrame(master) frame.grid(row=row, column=column, padx=5, pady=5, sticky="nsew") @@ -698,8 +530,7 @@ def app_title(master, row, column): Image.open("resources/icons/icon.png").resize((40, 40), Image.Resampling.LANCZOS), size=(40, 40) ) - image_label_component = ctk.CTkLabel(frame, image=image_component, text="") - image_label_component.grid(row=0, column=0, padx=PAD, pady=PAD) + ctk.CTkLabel(frame, image=image_component, text="").grid(row=0, column=0, padx=PAD, pady=PAD) label_component = ctk.CTkLabel(frame, text="OneTrainer", font=ctk.CTkFont(size=20, weight="bold")) label_component.grid(row=0, column=1, padx=(0, PAD), pady=PAD) @@ -713,7 +544,6 @@ def label(master, row, column, text, pad=PAD, tooltip=None, wide_tooltip=False, return component - def time_entry(master, row, column, ui_state: UIState, var_name: str, unit_var_name, supports_time_units: bool = True): frame = ctk.CTkFrame(master, fg_color="transparent") frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") @@ -760,18 +590,7 @@ def options(master, row, column, values, ui_state: UIState, var_name: str, comma component = ctk.CTkOptionMenu(master, values=values, variable=ui_state.get_var(var_name), command=command) component.grid(row=row, column=column, padx=PAD, pady=(PAD, PAD), sticky="new") - # temporary fix until https://github.com/TomSchimansky/CustomTkinter/pull/2246 is merged - def create_destroy(component): - orig_destroy = component.destroy - - def destroy(self): - orig_destroy() - CTkScalingBaseClass.destroy(self) - - return destroy - - destroy = create_destroy(component._dropdown_menu) - component._dropdown_menu.destroy = lambda: destroy(component._dropdown_menu) + _wrap_dropdown_destroy(component) return component @@ -792,18 +611,7 @@ def options_adv(master, row, column, values, ui_state: UIState, var_name: str, if command: command(ui_state.get_var(var_name).get()) # call command once to set the initial value - # temporary fix until https://github.com/TomSchimansky/CustomTkinter/pull/2246 is merged - def create_destroy(component): - orig_destroy = component.destroy - - def destroy(self): - orig_destroy() - CTkScalingBaseClass.destroy(self) - - return destroy - - destroy = create_destroy(component._dropdown_menu) - component._dropdown_menu.destroy = lambda: destroy(component._dropdown_menu) + _wrap_dropdown_destroy(component) return frame, {'component': component, 'button_component': button_component} @@ -811,10 +619,8 @@ def destroy(self): def options_kv(master, row, column, values: list[tuple[str, Any]], ui_state: UIState, var_name: str, command: Callable[[Any], None] = None): var = ui_state.get_var(var_name) - keys = [key for key, value in values] - # if the current value is not valid, select the first option - if var.get() not in [str(value) for key, value in values] and len(keys) > 0: + if var.get() not in [str(value) for key, value in values] and values: var.set(values[0][1]) deactivate_update_var = False @@ -830,7 +636,7 @@ def update_component(text): deactivate_update_var = False break - component = ctk.CTkOptionMenu(master, values=keys, command=update_component) + component = ctk.CTkOptionMenu(master, values=[key for key, _ in values], command=update_component) component.grid(row=row, column=column, padx=PAD, pady=(PAD, PAD), sticky="new") def update_var(): @@ -846,18 +652,7 @@ def update_var(): var.trace_add("write", lambda _0, _1, _2: update_var()) update_var() # call update_var once to set the initial value - # temporary fix until https://github.com/TomSchimansky/CustomTkinter/pull/2246 is merged - def create_destroy(component): - orig_destroy = component.destroy - - def destroy(self): - orig_destroy() - CTkScalingBaseClass.destroy(self) - - return destroy - - destroy = create_destroy(component._dropdown_menu) - component._dropdown_menu.destroy = lambda: destroy(component._dropdown_menu) + _wrap_dropdown_destroy(component) return component @@ -878,19 +673,14 @@ def switch( component = ctk.CTkSwitch(master, variable=var, text=text, command=command) component.grid(row=row, column=column, padx=PAD, pady=(PAD, PAD), sticky="new") - def create_destroy(component): - orig_destroy = component.destroy - - def destroy(self): - if command is not None: - ui_state.remove_var_trace(var_name, trace_id) - - orig_destroy() + original_destroy = component.destroy - return destroy + def new_destroy(): + if command is not None: + ui_state.remove_var_trace(var_name, trace_id) + original_destroy() - destroy = create_destroy(component) - component.destroy = lambda: destroy(component) + component.destroy = new_destroy return component From 60407ad3dfc3280be3555922e3639e921fe178e7 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Mon, 10 Nov 2025 02:13:17 +1100 Subject: [PATCH 22/64] Further trim fat from components.py and remove model_output_entry --- modules/ui/ModelTab.py | 4 +- modules/util/ui/components.py | 128 +++++++++++++--------------------- 2 files changed, 53 insertions(+), 79 deletions(-) diff --git a/modules/ui/ModelTab.py b/modules/ui/ModelTab.py index 1a8c259ad..ea59dc7e7 100644 --- a/modules/ui/ModelTab.py +++ b/modules/ui/ModelTab.py @@ -507,9 +507,11 @@ def __create_output_components( components.label(self.scroll_frame, row, 0, "Model Output Destination", tooltip="Filename or directory where the *final* epoch output model is saved") - components.model_output_entry( + components.path_entry( self.scroll_frame, row, 1, self.ui_state, var_name="output_model_destination", + is_output=True, + use_model_validator=True, format_var_name="output_model_format", method_var_name="training_method", ) diff --git a/modules/util/ui/components.py b/modules/util/ui/components.py index 4a45e5d5f..0dab246a2 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -386,62 +386,6 @@ def new_destroy(): return component -def model_output_entry( - master, - row, - column, - ui_state: UIState, - var_name: str = "output_model_destination", - format_var_name: str = "output_model_format", - method_var_name: str = "training_method", - prefix_var_name: str = "save_filename_prefix", - command: Callable[[str], None] = None, -): - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") - frame.grid_columnconfigure(0, weight=1) - - var = ui_state.get_var(var_name) - - validator = ModelOutputValidator( - var=var, - ui_state=ui_state, - format_var_name=format_var_name, - method_var_name=method_var_name, - prefix_var_name=prefix_var_name, - ) - - entry_widget = entry( - frame, row=0, column=0, ui_state=ui_state, var_name=var_name, - custom_validator=validator.validate, - validation_state=validator.state, - enable_hover_validation=True, - command=command - ) - validator.setup_traces() - original_destroy = entry_widget.destroy - - def new_destroy(): - validator.cleanup_traces() - original_destroy() - - entry_widget.destroy = new_destroy - - _register_drop_target(entry_widget, ui_state, var_name, command) - - def open_dialog(): - selected_path = filedialog.asksaveasfilename(filetypes=MODEL_FILETYPES) - if selected_path: - var.set(selected_path) - if command: - command(selected_path) - - button_component = ctk.CTkButton(frame, text="...", width=40, command=open_dialog) - button_component.grid(row=0, column=1, padx=(0, PAD), pady=PAD, sticky="nsew") - - return frame - - def path_entry( master, row, @@ -455,6 +399,10 @@ def path_entry( allow_model_files: bool = True, allow_image_files: bool = False, valid_extensions: list[str] = None, + use_model_validator: bool = False, + format_var_name: str = "output_model_format", + method_var_name: str = "training_method", + prefix_var_name: str = "save_filename_prefix", ): frame = ctk.CTkFrame(master, fg_color="transparent") frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") @@ -467,26 +415,55 @@ def path_entry( var = ui_state.get_var(var_name) - # Simple path validation with auto-trimming for directories - def simple_path_validator(value: str) -> ValidationResult: - # Auto-trim whitespace from directory names - if value and path_type == "directory": - path = Path(value) - trimmed_name = path.name.strip() - if trimmed_name != path.name: - sep = '\\' if '\\' in value else '/' - trimmed_path = f"{path.parent}{sep}{trimmed_name}" if str(path.parent) != '.' else trimmed_name - var.set(trimmed_path) - value = trimmed_path - - return validate_file_path(value, is_output, valid_extensions, path_type) + # Set up validator and validation state + if use_model_validator: + validator = ModelOutputValidator( + var=var, ui_state=ui_state, + format_var_name=format_var_name, + method_var_name=method_var_name, + prefix_var_name=prefix_var_name, + ) + validation_state = validator.state + custom_validator = validator.validate + else: + validator = None + validation_state = None + + def simple_path_validator(value: str) -> ValidationResult: + if value and path_type == "directory": + path = Path(value) + trimmed_name = path.name.strip() + if trimmed_name != path.name: + sep = '\\' if '\\' in value else '/' + parent_str = str(path.parent) + trimmed_path = f"{parent_str}{sep}{trimmed_name}" if parent_str != '.' else trimmed_name + var.set(trimmed_path) + value = trimmed_path + return validate_file_path(value, is_output, valid_extensions, path_type) + + custom_validator = simple_path_validator # Create entry with validation entry_widget = entry( frame, row=0, column=0, ui_state=ui_state, var_name=var_name, - custom_validator=simple_path_validator, command=command + custom_validator=custom_validator, + validation_state=validation_state, + enable_hover_validation=bool(validator), + command=command ) + # Set up validator traces and cleanup + if validator: + validator.setup_traces() + original_destroy = entry_widget.destroy + + def new_destroy(): + if validator: + validator.cleanup_traces() + original_destroy() + + entry_widget.destroy = new_destroy + # Enable drag-and-drop _register_drop_target(entry_widget, ui_state, var_name, command) @@ -496,23 +473,18 @@ def open_dialog(): selected_path = filedialog.askdirectory() else: filetypes = [("All Files", "*.*")] - if allow_model_files: - filetypes.extend(MODEL_FILETYPES[1:]) # Skip "All Files" since we already added it + filetypes.extend(MODEL_FILETYPES[1:]) if allow_image_files: filetypes.append(("Image", ' '.join(f"*.{x}" for x in supported_image_extensions()))) - if is_output: - selected_path = filedialog.asksaveasfilename(filetypes=filetypes) - else: - selected_path = filedialog.askopenfilename(filetypes=filetypes) + selected_path = (filedialog.asksaveasfilename(filetypes=filetypes) if is_output + else filedialog.askopenfilename(filetypes=filetypes)) if selected_path: if path_modifier: selected_path = path_modifier(selected_path) - var.set(selected_path) - if command: command(selected_path) From 079d1a5d80c51d6ffe77aa3d8c6e87973196261c Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Mon, 10 Nov 2025 02:23:29 +1100 Subject: [PATCH 23/64] More trimming, testing as I go along --- modules/util/ui/ToolTip.py | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/modules/util/ui/ToolTip.py b/modules/util/ui/ToolTip.py index 9de7264da..a2f4294ef 100644 --- a/modules/util/ui/ToolTip.py +++ b/modules/util/ui/ToolTip.py @@ -5,7 +5,6 @@ class ToolTip: MAX_TOOLTIP_LENGTH = 350 - EDGE_MARGIN = 10 def __init__(self, widget, text='widget info', x_position=20, wide=False, hover_only=True, track_movement=False): @@ -48,15 +47,17 @@ def unschedule(self): self.widget.after_cancel(self.id) self.id = None - def show(self, text=None, color=None, duration_ms=None, fg_color="#2b2b2b"): - if text: - self.text = text[:self.MAX_TOOLTIP_LENGTH] + ("..." if len(text) > self.MAX_TOOLTIP_LENGTH else "") - + def _cancel_scheduled_hide(self): if self._after_id: with contextlib.suppress(Exception): self.widget.after_cancel(self._after_id) self._after_id = None + def show(self, text=None, color=None, duration_ms=None, fg_color="#2b2b2b"): + if text: + self.text = text[:self.MAX_TOOLTIP_LENGTH] + ("..." if len(text) > self.MAX_TOOLTIP_LENGTH else "") + + self._cancel_scheduled_hide() self.hide() self.tw = ctk.CTkToplevel(self.widget) @@ -65,7 +66,7 @@ def show(self, text=None, color=None, duration_ms=None, fg_color="#2b2b2b"): self.tw.attributes("-topmost", True) frame = ctk.CTkFrame(self.tw, fg_color=fg_color, corner_radius=6) - frame.pack(padx=0, pady=0) + frame.pack() label = ctk.CTkLabel( frame, @@ -117,17 +118,12 @@ def _on_move(self, _evt=None): self._position() def hide(self, _evt=None): - if self._after_id: - with contextlib.suppress(Exception): - self.widget.after_cancel(self._after_id) - self._after_id = None + self._cancel_scheduled_hide() if self.tw and self.tw.winfo_exists(): self.tw.destroy() self.tw = None - hidetip = hide - def show_error(self, message: str, duration_ms: int | None = 7000): self.show(message, color="#ff6b6b", duration_ms=duration_ms) From 04963d5d6c121344633974211c8f5fc2db6c1e7d Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Mon, 10 Nov 2025 02:42:15 +1100 Subject: [PATCH 24/64] Even more fat trimming. --- modules/util/tensorboard_util.py | 71 +++++++++----------------------- modules/util/ui/components.py | 4 +- modules/util/ui/ui_utils.py | 28 ++++++------- 3 files changed, 34 insertions(+), 69 deletions(-) diff --git a/modules/util/tensorboard_util.py b/modules/util/tensorboard_util.py index c47f61c1d..0ade3ecda 100644 --- a/modules/util/tensorboard_util.py +++ b/modules/util/tensorboard_util.py @@ -10,7 +10,6 @@ class TensorboardManager: - """Singleton manager for Tensorboard processes across GUI, CLI, and Cloud training""" _instance = None _lock = threading.Lock() @@ -32,37 +31,26 @@ def __init__(self): self._previous_mode = None self._initialized = True + def _is_cloud_tunnel_enabled(self, config: TrainConfig) -> bool: + return config.cloud.enabled and config.cloud.tensorboard_tunnel + def should_start(self, config: TrainConfig, is_training: bool = False) -> bool: - """Determine if tensorboard should start based on mode and context""" - # Skip if cloud tunnel is enabled - if config.cloud.enabled and config.cloud.tensorboard_tunnel: + if self._is_cloud_tunnel_enabled(config): return False mode = config.tensorboard_mode - if mode == TensorboardMode.OFF: return False - elif mode == TensorboardMode.ALWAYS_ON: - return not is_training # Start only if not already running during training - elif mode == TensorboardMode.TRAIN_ONLY: - return is_training - - return False + if mode == TensorboardMode.ALWAYS_ON: + return not is_training + return mode == TensorboardMode.TRAIN_ONLY and is_training def start(self, config: TrainConfig, mode: TensorboardMode | None = None) -> bool: - """Start tensorboard if not already running""" if self._process is not None: - return False # Already running - - # Use provided mode or config's mode - effective_mode = mode or config.tensorboard_mode - - # Check if we should start based on mode - if effective_mode == TensorboardMode.OFF: return False - # Skip if cloud tunnel is enabled - if config.cloud.enabled and config.cloud.tensorboard_tunnel: + effective_mode = mode or config.tensorboard_mode + if effective_mode == TensorboardMode.OFF or self._is_cloud_tunnel_enabled(config): return False self._process = start_filtered_tensorboard(config) @@ -72,7 +60,6 @@ def start(self, config: TrainConfig, mode: TensorboardMode | None = None) -> boo return self._process is not None def stop(self) -> bool: - """Stop tensorboard if running""" if self._process is None: return False @@ -83,48 +70,35 @@ def stop(self) -> bool: return True def restart(self, config: TrainConfig) -> bool: - """Restart tensorboard with new config""" self.stop() return self.start(config) def is_running(self) -> bool: - """Check if tensorboard is currently running""" return self._process is not None def get_mode(self) -> TensorboardMode | None: - """Get current tensorboard mode""" return self._mode def handle_training_start(self, config: TrainConfig): - """Handle tensorboard when training starts""" - # If TRAIN_ONLY mode, start it if config.tensorboard_mode == TensorboardMode.TRAIN_ONLY: - if not (config.cloud.enabled and config.cloud.tensorboard_tunnel): + if not self._is_cloud_tunnel_enabled(config): self.start(config) def handle_training_end(self, config: TrainConfig): - """Handle tensorboard when training ends""" - # If TRAIN_ONLY mode, stop it if config.tensorboard_mode == TensorboardMode.TRAIN_ONLY: self.stop() - - # If ALWAYS_ON mode, restart it (in case it was stopped) elif config.tensorboard_mode == TensorboardMode.ALWAYS_ON: - if not (config.cloud.enabled and config.cloud.tensorboard_tunnel): - if not self.is_running(): - self.start(config) + if not self._is_cloud_tunnel_enabled(config) and not self.is_running(): + self.start(config) def handle_mode_change(self, config: TrainConfig, is_training: bool = False): - """Handle tensorboard mode changes""" new_mode = config.tensorboard_mode - # Initialize previous mode on first call if self._previous_mode is None: self._previous_mode = new_mode - # Start if ALWAYS_ON and not training - if new_mode == TensorboardMode.ALWAYS_ON and not is_training: - if not (config.cloud.enabled and config.cloud.tensorboard_tunnel): - self.start(config) + if (new_mode == TensorboardMode.ALWAYS_ON and not is_training + and not self._is_cloud_tunnel_enabled(config)): + self.start(config) return old_mode = self._previous_mode @@ -132,13 +106,11 @@ def handle_mode_change(self, config: TrainConfig, is_training: bool = False): if new_mode == old_mode: return - # Skip if cloud tunnel - if config.cloud.enabled and config.cloud.tensorboard_tunnel: + if self._is_cloud_tunnel_enabled(config): self.stop() self._previous_mode = new_mode return - # During training if is_training: if new_mode == TensorboardMode.OFF: self.stop() @@ -157,17 +129,13 @@ def handle_mode_change(self, config: TrainConfig, is_training: bool = False): self._previous_mode = new_mode def handle_workspace_change(self, config: TrainConfig): - """Handle workspace directory changes""" if config.tensorboard_mode == TensorboardMode.ALWAYS_ON: - if not (config.cloud.enabled and config.cloud.tensorboard_tunnel): + if not self._is_cloud_tunnel_enabled(config): print("Restarting Tensorboard due to workspace change") self.restart(config) def get_tensorboard_args(config: TrainConfig) -> list[str]: - """ - Generates args required to start Tensorboard - """ tensorboard_executable = os.path.join(os.path.dirname(sys.executable), "tensorboard") tensorboard_log_dir = os.path.join(config.workspace_dir, "tensorboard") @@ -189,7 +157,7 @@ def get_tensorboard_args(config: TrainConfig) -> list[str]: def start_filtered_tensorboard(config: TrainConfig) -> subprocess.Popen | None: - """Start Tensorboard filtering the annoying warnings Google doesnt plan to fix""" + """Start Tensorboard, filtering the annoying warnings Google doesnt plan to fix""" try: print("Starting Tensorboard please wait...") env = os.environ.copy() @@ -211,7 +179,7 @@ def filter_output(): "pkg_resources" in line_lower or "installation not found" in line_lower ): - continue # Skip these lines + continue if line.strip(): print(line, end='') @@ -226,7 +194,6 @@ def filter_output(): def stop_tensorboard(process: subprocess.Popen | None): - """Gracefully stop the Tensorboard subprocess""" if not process: return diff --git a/modules/util/ui/components.py b/modules/util/ui/components.py index 0dab246a2..0fd953f36 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -11,7 +11,7 @@ from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.path_util import supported_image_extensions from modules.util.ui.ToolTip import ToolTip -from modules.util.ui.ui_utils import _register_drop_target +from modules.util.ui.ui_utils import register_drop_target from modules.util.ui.UIState import UIState from modules.util.ui.validation import ValidationResult, validate_basic_type, validate_destination, validate_file_path @@ -465,7 +465,7 @@ def new_destroy(): entry_widget.destroy = new_destroy # Enable drag-and-drop - _register_drop_target(entry_widget, ui_state, var_name, command) + register_drop_target(entry_widget, ui_state, var_name, command) # Browse button def open_dialog(): diff --git a/modules/util/ui/ui_utils.py b/modules/util/ui/ui_utils.py index 99f9bff53..506d1ae58 100644 --- a/modules/util/ui/ui_utils.py +++ b/modules/util/ui/ui_utils.py @@ -107,7 +107,6 @@ def _drop_enter(event): event.widget.focus_force() return event.action - def _drop_leave(event): return event.action @@ -141,27 +140,28 @@ def _parse_dropped_paths(event_data: str) -> list[str]: in_braces = True elif char == '}': in_braces = False - if current_path: - paths.append(current_path) - current_path = "" elif char == ' ' and not in_braces: - if current_path: - paths.append(current_path) - current_path = "" + pass # Will append below else: current_path += char + continue + + # Common path appending logic for '}' and ' ' + if current_path: + paths.append(current_path) + current_path = "" if current_path: paths.append(current_path) return [p.strip() for p in paths if p.strip()] -def _register_drop_target(entry_widget, ui_state, var_name, command=None): +def register_drop_target(entry_widget, ui_state, var_name, command=None): try: entry_widget.drop_target_register(DND_FILES) - entry_widget.dnd_bind('<>', _drop_enter) - entry_widget.dnd_bind('<>', _drop_leave) - entry_widget.dnd_bind('<>', _create_drop_handler(entry_widget, ui_state, var_name, command)) + for event, handler in [('<>', _drop_enter), ('<>', _drop_leave), + ('<>', _create_drop_handler(entry_widget, ui_state, var_name, command))]: + entry_widget.dnd_bind(event, handler) except Exception: pass @@ -173,12 +173,9 @@ def drop_handler(event): paths = _parse_dropped_paths(event.data) for path in paths: - if os.path.isfile(path): - path = os.path.dirname(path) - + path = os.path.dirname(path) if os.path.isfile(path) else path if os.path.isdir(path): drop_callback(path) - if not allow_multiple: break @@ -191,6 +188,7 @@ def drop_handler(event): widget.dnd_bind('<>', drop_handler) except Exception: pass + class DebounceTimer: def __init__(self, widget, delay_ms: int, callback: Callable[..., Any]): self.widget = widget From 57a5933e1315a3257ebe7faf6bb55bf5fcfc45ca Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Mon, 10 Nov 2025 02:54:58 +1100 Subject: [PATCH 25/64] Trimming the fat from validation.py --- modules/util/ui/validation.py | 100 +++++++--------------------------- 1 file changed, 20 insertions(+), 80 deletions(-) diff --git a/modules/util/ui/validation.py b/modules/util/ui/validation.py index 9bd8a4456..24d67d465 100644 --- a/modules/util/ui/validation.py +++ b/modules/util/ui/validation.py @@ -1,9 +1,4 @@ -""" -Input validation utilities for OneTrainer UI components. - -Provides validation for file paths, directories, and model output destinations -with support for auto-correction and detailed error messaging. -""" +"""Input validation utilities for UI components.""" from __future__ import annotations @@ -26,12 +21,10 @@ class ValidationSettings: SETTINGS = ValidationSettings() -# Regex patterns ONLY_WHITESPACE = re.compile(r"^\s+$") TRAILING_SLASH_RE = re.compile(r"[\\/]$") ENDS_WITH_EXT = re.compile(r"\.[A-Za-z0-9]+$") -# Invalid characters/names INVALID_CHARS_WIN = set('<>:"|?*') INVALID_NAMES = {"", "."} @@ -39,7 +32,6 @@ class ValidationSettings: @dataclass class ValidationResult: - """Result of a validation operation.""" ok: bool corrected: str | None message: str = "" @@ -55,7 +47,6 @@ def warning(self) -> bool: """Backward compatibility: whether this is a warning.""" return self.status == 'warning' -# Helper to create common results def _result(ok: bool, corrected: str | None = None, message: str = "", status: ValidationStatus | None = None) -> ValidationResult: return ValidationResult(ok=ok, corrected=corrected, message=message, status=status) @@ -64,7 +55,6 @@ def _is_windows() -> bool: return os.name == "nt" def _has_invalid_chars(path: str) -> bool: - """Check if path contains invalid characters.""" if _is_windows(): has_drive = len(path) >= 2 and path[1] == ':' and path[0].isalpha() start_idx = 2 if has_drive else 0 @@ -72,7 +62,6 @@ def _has_invalid_chars(path: str) -> bool: return "\x00" in path def _safe_exists(path: Path, check_dir: bool = True) -> bool: - """Safely check if a path exists.""" try: return path.is_dir() if check_dir else path.parent.exists() except OSError: @@ -85,22 +74,22 @@ def _parent_exists(path: Path) -> bool: return _safe_exists(path, check_dir=False) def _format_path(path: Path, separator: str) -> str: - """Format path with preferred separator.""" return str(path).replace('\\', separator) +def _check_parent_exists(path: Path) -> ValidationResult | None: + if path.parent != Path('.') and not _parent_exists(path): + return _result(False, None, "Parent folder does not exist.") + return None + def _has_extension(path: str | Path, extension: str) -> bool: - """Check if path has the given extension (case-insensitive).""" return str(path).lower().endswith(extension.lower()) def _get_extension(path: str | Path) -> str | None: - """Extract extension from path, or None if no extension.""" match = ENDS_WITH_EXT.search(str(path)) return match.group(0).lower() if match else None def _is_partial_match(current: str, required: str) -> bool: - """Check if current extension is a partial match of required.""" - chars_after_dot = len(current) - 1 - return chars_after_dot >= 3 and required.startswith(current) + return len(current) > 3 and required.startswith(current) def generate_default_filename( training_method: TrainingMethod, @@ -108,7 +97,6 @@ def generate_default_filename( extension: str = ".safetensors", use_friendly_names: bool = False, ) -> str: - """Generate a default filename for model output.""" if use_friendly_names: friendly = generate_friendly(1 if prefix else 2, separator="" if prefix else "-") name = f"{prefix}_{friendly}" if prefix else friendly @@ -119,14 +107,12 @@ def generate_default_filename( return f"{name}{extension}" def get_allowed_formats_for_method(training_method: TrainingMethod) -> set[ModelFormat]: - """Get the allowed output formats for a given training method.""" if training_method in (TrainingMethod.EMBEDDING, TrainingMethod.LORA): return {ModelFormat.SAFETENSORS, ModelFormat.CKPT} return {ModelFormat.SAFETENSORS, ModelFormat.CKPT, ModelFormat.DIFFUSERS} @dataclass class ValidationContext: - """Context for validation operations.""" raw: str user_input: str output_format: ModelFormat @@ -139,25 +125,19 @@ class ValidationContext: def _validate_basic_input(ctx: ValidationContext) -> ValidationResult | None: - """Perform basic validation checks on input. Returns None to continue.""" - - # Check for whitespace-only input + """Perform basic input checks; return a ValidationResult to stop or None to continue.""" if ONLY_WHITESPACE.match(ctx.raw): return _result(False, "", "Input cleared: empty or invalid.") - # Check for empty or invalid names if not ctx.user_input or ctx.user_input in INVALID_NAMES: return _result(False, "", "Input cleared: empty or invalid.") - # Check for invalid characters if _has_invalid_chars(ctx.user_input): return _result(False, None, "Input contains invalid path characters.") - # Check format/method compatibility if ctx.output_format not in get_allowed_formats_for_method(ctx.training_method): return _result(False, None, f"{ctx.training_method} cannot output {ctx.output_format} format.") - # For input paths, verify existence if not ctx.is_output: if not Path(ctx.user_input).exists(): return _result(False, None, "Input path does not exist.") @@ -166,44 +146,37 @@ def _validate_basic_input(ctx: ValidationContext) -> ValidationResult | None: return None def _validate_diffusers_path(ctx: ValidationContext) -> ValidationResult: - """Validate DIFFUSERS format path (directory-based).""" + """Validate DIFFUSERS (directory-based) paths.""" ext_match = ENDS_WITH_EXT.search(ctx.user_input) if ext_match and not ctx.autocorrect: return _result(False, None, "DIFFUSERS expects a directory, not a file. Add a trailing slash.") - # Remove extension if present sanitized = ctx.user_input[:ext_match.start()] if ext_match else ctx.user_input check_path = Path(sanitized) - # Check parent directory exists - if check_path.parent != Path('.') and not _parent_exists(check_path): - return _result(False, None, "Parent folder does not exist.") + if err := _check_parent_exists(check_path): + return err - # Warn if directory already exists if not ctx.autocorrect and _folder_exists(check_path): msg = "Directory already exists. Stripped extension for DIFFUSERS." if ext_match else "Directory already exists." return _result(True, sanitized if ext_match else None, msg, 'warning') - # Extension was stripped if ext_match: return _result(True, sanitized, f"Stripped extension for {ctx.output_format}.") return _result(True, None, "") def _validate_single_file_path(ctx: ValidationContext) -> ValidationResult: - """Validate single-file format path (SAFETENSORS, CKPT).""" + """Validate single-file formats (SAFETENSORS, CKPT).""" required_ext = ctx.output_format.file_extension() - # Check if input looks like a directory is_dir_like = bool(TRAILING_SLASH_RE.search(ctx.user_input)) path = Path(ctx.user_input) - # If no slash but exists as folder, treat as directory if not is_dir_like and '/' not in ctx.user_input and '\\' not in ctx.user_input and _folder_exists(path): is_dir_like = True - # Handle directory-like input if is_dir_like: if ctx.autocorrect: if not _folder_exists(path): @@ -213,13 +186,11 @@ def _validate_single_file_path(ctx: ValidationContext) -> ValidationResult: return _result(True, corrected, f"Appended default filename: {default_name}") return _result(False, None, "Path points to a folder; a filename is required.") - # Check parent directory exists - if path.parent != Path('.') and not _parent_exists(path): - return _result(False, None, "Parent folder does not exist.") + if err := _check_parent_exists(path): + return err base = path.name - # Handle trailing periods if base.endswith('.'): if ctx.autocorrect: new_base = base.rstrip('.') @@ -227,22 +198,18 @@ def _validate_single_file_path(ctx: ValidationContext) -> ValidationResult: return _result(True, new_path, "Removed trailing period(s).") return _result(False, None, "Filename cannot end with a period.") - # Check if already has correct extension if _has_extension(base, required_ext): if path.exists(): return _result(True, None, "WARNING: File already exists and will be overwritten.", 'warning') return _result(True, None, "") - # Check for partial or incorrect extension ext_match = ENDS_WITH_EXT.search(base) if ext_match: current_ext = ext_match.group(0).lower() chars_after_dot = len(current_ext) - 1 start_idx = ext_match.start() - # Extension has at least 3 characters after dot if chars_after_dot >= 3: - # Check if it's a partial match (e.g., ".safe" for ".safetensors") if _is_partial_match(current_ext, required_ext): if ctx.autocorrect: new_base = base[:start_idx] + required_ext @@ -250,18 +217,15 @@ def _validate_single_file_path(ctx: ValidationContext) -> ValidationResult: return _result(True, new_path, f"Completed {current_ext} -> {required_ext}") return _result(False, None, f"Extension is incomplete; expected {required_ext}") else: - # Different extension entirely if ctx.autocorrect: new_base = base[:start_idx] + required_ext new_path = _format_path(path.with_name(new_base), ctx.separator) return _result(True, new_path, f"Replaced {current_ext} with {required_ext}") return _result(False, None, f"Wrong extension; expected {required_ext}") - # Extension too short for auto-correction if 1 <= chars_after_dot < 3 and required_ext.startswith(current_ext): return _result(False, None, f"Extension {current_ext} is too short; Type at least 3 characters for auto-correction to occur.") - # Unrecognized short extension - append required if ctx.autocorrect: corrected = _format_path(Path(path.parent, base + required_ext), ctx.separator) return _result(True, corrected, f"Appended {required_ext} extension.") @@ -317,11 +281,9 @@ def validate_destination( separator='/' if '/' in raw else '\\' ) - # Basic validation checks if basic_check := _validate_basic_input(ctx): return basic_check - # Format-specific validation if output_format == ModelFormat.DIFFUSERS: return _validate_diffusers_path(ctx) else: @@ -333,18 +295,7 @@ def validate_basic_type( nullable: bool, default_val: Any, ) -> ValidationResult: - """ - Validate basic Python types (int, float, bool, str). - - Args: - value: The string value to validate - declared_type: The expected Python type (int, float, bool, str) - nullable: Whether None/empty is allowed - default_val: The default value for this field - - Returns: - ValidationResult with validation outcome - """ + """Validate basic Python types (int, float, bool, str).""" # Handle empty values if value == "": if nullable or (declared_type is str and default_val == ""): @@ -371,18 +322,7 @@ def validate_file_path( valid_extensions: list[str] | None = None, path_type: str = "file", ) -> ValidationResult: - """ - Validate a file or directory path. - - Args: - value: Path string to validate - is_output: Whether this is an output path - valid_extensions: List of valid extensions (with dots), or None for any - path_type: Either "file" or "directory" - - Returns: - ValidationResult with validation outcome - """ + """Validate a file or directory path.""" if not value: return _result(True, None, "") @@ -394,8 +334,8 @@ def validate_file_path( if path_type == "directory": if is_output: - if path.parent != Path('.') and not _parent_exists(path): - return _result(False, None, "Parent folder does not exist.") + if err := _check_parent_exists(path): + return err else: if not _folder_exists(path): return _result(False, None, "Directory does not exist.") @@ -406,8 +346,8 @@ def validate_file_path( return _result(False, None, f"File must have one of these extensions: {', '.join(valid_extensions)}") if is_output: - if path.parent != Path('.') and not _parent_exists(path): - return _result(False, None, "Parent folder does not exist.") + if err := _check_parent_exists(path): + return err else: if not path.exists(): return _result(False, None, "File does not exist") From 2d79fe06f385066908f932e24bd2020e38a31065 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Mon, 10 Nov 2025 03:07:33 +1100 Subject: [PATCH 26/64] Remove dead function and remove CLI save bugfix, out of scope --- modules/modelSaver/BaseModelSaver.py | 13 ------------- .../modelSaver/StableDiffusionXLLoRAModelSaver.py | 5 ----- modules/util/ui/validation.py | 4 ---- 3 files changed, 22 deletions(-) diff --git a/modules/modelSaver/BaseModelSaver.py b/modules/modelSaver/BaseModelSaver.py index 0928225b0..4ad07a49a 100644 --- a/modules/modelSaver/BaseModelSaver.py +++ b/modules/modelSaver/BaseModelSaver.py @@ -11,19 +11,6 @@ class BaseModelSaver(metaclass=ABCMeta): def __init__(self): super().__init__() - def _ensure_correct_extension( - self, - output_model_format: ModelFormat, - output_model_destination: str, - ) -> str: - - if output_model_format.is_single_file(): - expected_extension = output_model_format.file_extension() - if not output_model_destination.endswith(expected_extension): - return output_model_destination + expected_extension - - return output_model_destination - @abstractmethod def save( self, diff --git a/modules/modelSaver/StableDiffusionXLLoRAModelSaver.py b/modules/modelSaver/StableDiffusionXLLoRAModelSaver.py index 972a44aa3..cd38ea052 100644 --- a/modules/modelSaver/StableDiffusionXLLoRAModelSaver.py +++ b/modules/modelSaver/StableDiffusionXLLoRAModelSaver.py @@ -24,11 +24,6 @@ def save( output_model_destination: str, dtype: torch.dtype | None, ): - # Ensure the output destination has the correct file extension - output_model_destination = self._ensure_correct_extension( - output_model_format, - output_model_destination - ) lora_model_saver = StableDiffusionXLLoRASaver() embedding_model_saver = StableDiffusionXLEmbeddingSaver() diff --git a/modules/util/ui/validation.py b/modules/util/ui/validation.py index 24d67d465..a40868d34 100644 --- a/modules/util/ui/validation.py +++ b/modules/util/ui/validation.py @@ -84,10 +84,6 @@ def _check_parent_exists(path: Path) -> ValidationResult | None: def _has_extension(path: str | Path, extension: str) -> bool: return str(path).lower().endswith(extension.lower()) -def _get_extension(path: str | Path) -> str | None: - match = ENDS_WITH_EXT.search(str(path)) - return match.group(0).lower() if match else None - def _is_partial_match(current: str, required: str) -> bool: return len(current) > 3 and required.startswith(current) From 4778222a0edfa84ae1981de08801613c63626fd5 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Mon, 10 Nov 2025 18:38:18 +1100 Subject: [PATCH 27/64] - Make tensorboard use the model output name - Add "Prevent overwrite" behaviour and toggle - If the user writes a generic name, valid prefix will be added if exists - Option to auto prefix output for manually typed files. --- modules/trainer/GenericTrainer.py | 4 +- modules/ui/TrainUI.py | 12 ++- modules/util/config/TrainConfig.py | 4 + modules/util/tensorboard_util.py | 19 +++++ modules/util/ui/components.py | 130 ++++++++++++++++++++--------- modules/util/ui/validation.py | 76 ++++++++++++++++- 6 files changed, 202 insertions(+), 43 deletions(-) diff --git a/modules/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index bfe7eb5e5..d4b3a8e35 100644 --- a/modules/trainer/GenericTrainer.py +++ b/modules/trainer/GenericTrainer.py @@ -29,6 +29,7 @@ from modules.util.enum.TimeUnit import TimeUnit from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.memory_util import TorchMemoryRecorder +from modules.util.tensorboard_util import get_tensorboard_run_name from modules.util.time_util import get_string_timestamp from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -69,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_name = get_tensorboard_run_name(config) + self.tensorboard = SummaryWriter(os.path.join(tensorboard_log_dir, f"{run_name}_{get_string_timestamp()}")) self.model = None self.one_step_trained = False diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index 20fff9d3c..10d16b956 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -322,10 +322,18 @@ def create_general_tab(self, master): tooltip="Use friendly names instead of timestamps for auto-corrected/generated filenames") components.switch(frame, 10, 1, self.ui_state, "use_friendly_names") + components.label(frame, 10, 2, "Prevent Overwrite", + tooltip="Automatically append a number or word to filenames that would overwrite existing files (requires Auto-correct Input)") + components.switch(frame, 10, 3, self.ui_state, "prevent_overwrite") + + components.label(frame, 11, 0, "Auto-prefix", + tooltip="Automatically prefix manually-entered filenames with the Save Filename Prefix (requires Auto-correct Input)") + components.switch(frame, 11, 1, self.ui_state, "auto_prefix") + # device - components.label(frame, 11, 0, "Dataloader Threads", + components.label(frame, 11, 2, "Dataloader Threads", tooltip="Number of threads used for the data loader. Increase if your GPU has room during caching, decrease if it's going out of memory during caching.") - components.entry(frame, 11, 1, self.ui_state, "dataloader_threads") + components.entry(frame, 11, 3, self.ui_state, "dataloader_threads") components.label(frame, 12, 0, "Train Device", tooltip="The device used for training. Can be \"cuda\", \"cuda:0\", \"cuda:1\" etc. Default:\"cuda\". Must be \"cuda\" for multi-GPU training.") diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 858ceb49a..25b7ab887 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -317,6 +317,8 @@ class TrainConfig(BaseConfig): validation_auto_correct: bool validation_show_tooltips: bool use_friendly_names: bool + prevent_overwrite: bool + auto_prefix: bool # multi-GPU multi_gpu: bool @@ -895,6 +897,8 @@ def default_values() -> 'TrainConfig': data.append(("validation_auto_correct", True, bool, False)) data.append(("validation_show_tooltips", True, bool, False)) data.append(("use_friendly_names", False, bool, False)) + data.append(("prevent_overwrite", True, bool, False)) + data.append(("auto_prefix", False, bool, False)) #multi-GPU data.append(("multi_gpu", False, bool, False)) diff --git a/modules/util/tensorboard_util.py b/modules/util/tensorboard_util.py index 0ade3ecda..af69626cb 100644 --- a/modules/util/tensorboard_util.py +++ b/modules/util/tensorboard_util.py @@ -9,6 +9,25 @@ from modules.util.enum.TensorboardMode import TensorboardMode +def get_tensorboard_run_name(config: TrainConfig) -> str: + """Extract run name from output_model_destination, falling back to prefix.""" + if not config.output_model_destination or not config.output_model_destination.strip(): + return config.save_filename_prefix + + path = Path(config.output_model_destination) + name = path.name + + # remove extension + if path.suffix: + name = path.stem + + # fall back to prefix if no name could be extracted + if not name: + return config.save_filename_prefix + + return name + + class TensorboardManager: _instance = None diff --git a/modules/util/ui/components.py b/modules/util/ui/components.py index 0fd953f36..69af04cfc 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -82,15 +82,11 @@ def __init__( self.custom_validator = custom_validator self.validation_state = validation_state - # Get border colors try: self.original_border_color = component.cget("border_color") except Exception: self.original_border_color = "gray50" - self.error_border_color = "#dc3545" - self.warning_border_color = "#ff9500" - self.validation_after_id = None self.revert_after_id = None self.touched = False @@ -119,10 +115,8 @@ def _show_validation_tooltip(self): def validate_value(self, value: str, revert_delay_ms: int | None) -> bool: self._cancel_after(self.revert_after_id) - meta = self.ui_state.get_field_metadata(self.var_name) - # Perform basic type validation basic_result = validate_basic_type( value=value, declared_type=meta.type, @@ -130,11 +124,9 @@ def validate_value(self, value: str, revert_delay_ms: int | None) -> bool: default_val=meta.default ) - # If basic validation fails, handle failure if not basic_result.ok: return self._fail(basic_result.message, revert_delay_ms) - # Custom validation if self.custom_validator: result = self.custom_validator(value) @@ -153,7 +145,7 @@ def _success(self, value: str) -> bool: return True def _warning(self, reason: str, value: str) -> bool: - self.component.configure(border_color=self.warning_border_color) + self.component.configure(border_color="#ff9500") if self.should_show_tooltip(): self.validation_tooltip.show_warning(reason) self.last_valid_value = value @@ -164,7 +156,7 @@ def _warning(self, reason: str, value: str) -> bool: return True def _fail(self, reason: str, revert_delay_ms: int | None) -> bool: - self.component.configure(border_color=self.error_border_color) + self.component.configure(border_color="#dc3545") if self.should_show_tooltip(): self.validation_tooltip.show_error(reason) @@ -239,38 +231,38 @@ def __init__( self.prefix_var = ui_state.get_var(prefix_var_name) self.autocorrect_var = ui_state.get_var("validation_auto_correct") self.friendly_names_var = ui_state.get_var("use_friendly_names") + self.prevent_overwrite_var = ui_state.get_var("prevent_overwrite") + self.auto_prefix_var = ui_state.get_var("auto_prefix") self.state = ValidationState() self._trace_ids: dict[str, str] = {} + self._debounce_after_id = None + self._tk_widget = None + self._last_prefix = self.prefix_var.get() if self.prefix_var else "" def _get_enum_value(self, var, default_enum): with contextlib.suppress(KeyError, ValueError): return type(default_enum)[var.get()] return default_enum - def validate(self, value: str) -> ValidationResult: - + def validate(self, value: str, skip_overwrite_protection: bool = False) -> ValidationResult: value = value.strip() if not value: self.state.clear() return ValidationResult(ok=True, corrected=None, message="", status='success') - output_format = self._get_enum_value(self.format_var, ModelFormat.SAFETENSORS) - training_method = self._get_enum_value(self.method_var, TrainingMethod.FINE_TUNE) - - autocorrect = _safe_bool(self.autocorrect_var) - prefix = self.prefix_var.get() if self.prefix_var else "" - use_friendly_names = _safe_bool(self.friendly_names_var, default=False) - # Use validation logic from validation.py result = validate_destination( value, - output_format, - training_method, - autocorrect=autocorrect, - prefix=prefix, - use_friendly_names=use_friendly_names, - is_output=True + output_format=self._get_enum_value(self.format_var, ModelFormat.SAFETENSORS), + training_method=self._get_enum_value(self.method_var, TrainingMethod.FINE_TUNE), + autocorrect=_safe_bool(self.autocorrect_var), + prefix=self.prefix_var.get() if self.prefix_var else "", + use_friendly_names=_safe_bool(self.friendly_names_var, default=False), + is_output=True, + prevent_overwrite=_safe_bool(self.prevent_overwrite_var, default=False), + auto_prefix=_safe_bool(self.auto_prefix_var, default=False), + skip_overwrite_protection=skip_overwrite_protection, ) if result.status: @@ -283,17 +275,77 @@ def validate(self, value: str) -> ValidationResult: return result + def _remove_old_prefix_if_needed(self): + current_value = self.var.get().strip() + if not current_value or not self._last_prefix: + return + + path = Path(current_value) + filename = path.name + old_prefix_pattern = f"{self._last_prefix}-" + + if filename.startswith(old_prefix_pattern): + new_filename = filename[len(old_prefix_pattern):] + new_value = str(path.parent / new_filename) if path.parent != Path('.') else new_filename + self.var.set(new_value) + + def _schedule_validation(self, skip_overwrite: bool = False): + if not self._tk_widget: + return + + if self._debounce_after_id: + with contextlib.suppress(Exception): + self._tk_widget.after_cancel(self._debounce_after_id) + + self._debounce_after_id = self._tk_widget.after( + COMPONENT_VALIDATION_SETTINGS.debounce_stop_typing_ms, + lambda: self.validate(self.var.get(), skip_overwrite_protection=skip_overwrite) + ) + + def _schedule_prefix_change_validation(self): + if not self._tk_widget: + return + + if self._debounce_after_id: + with contextlib.suppress(Exception): + self._tk_widget.after_cancel(self._debounce_after_id) + + def on_prefix_change(): + self._remove_old_prefix_if_needed() + self._last_prefix = self.prefix_var.get() if self.prefix_var else "" + self.validate(self.var.get(), skip_overwrite_protection=False) + + self._debounce_after_id = self._tk_widget.after( + COMPONENT_VALIDATION_SETTINGS.debounce_stop_typing_ms, + on_prefix_change + ) + def setup_traces(self): - trace_vars = [ - ('format', self.format_var), - ('method', self.method_var), - ('prefix', self.prefix_var), - ('friendly_names', self.friendly_names_var), + trace_config = [ + ('format', self.format_var, False, True, False), + ('method', self.method_var, False, True, False), + ('prefix', self.prefix_var, True, False, True), + ('friendly_names', self.friendly_names_var, False, True, False), + ('prevent_overwrite', self.prevent_overwrite_var, False, True, False), + ('auto_prefix', self.auto_prefix_var, False, True, False), ] - for key, var in trace_vars: - self._trace_ids[key] = var.trace_add("write", lambda *_: self.validate(self.var.get())) + + def make_callback(debounce: bool, skip_overwrite: bool, is_prefix: bool): + if is_prefix: + return lambda *_: self._schedule_prefix_change_validation() + elif debounce: + return lambda *_: self._schedule_validation(skip_overwrite) + return lambda *_: self.validate(self.var.get(), skip_overwrite_protection=skip_overwrite) + + for key, var, debounce, skip_overwrite, is_prefix in trace_config: + self._trace_ids[key] = var.trace_add("write", make_callback(debounce, skip_overwrite, is_prefix)) def cleanup_traces(self): + if self._debounce_after_id and self._tk_widget: + with contextlib.suppress(Exception): + self._tk_widget.after_cancel(self._debounce_after_id) + self._debounce_after_id = None + for key, trace_id in self._trace_ids.items(): var = getattr(self, f"{key}_var", None) if var: @@ -301,6 +353,9 @@ def cleanup_traces(self): var.trace_remove("write", trace_id) self._trace_ids.clear() + def set_widget(self, widget): + self._tk_widget = widget + # UI components @@ -338,11 +393,11 @@ def entry( ) validation_trace_name = var.trace_add("write", validation_handler.debounced_validate) - # Bind focus and input events component.bind("", validation_handler.on_focus_in) component.bind("", validation_handler.on_focus_out) for event in ("", "<>", "<>"): component.bind(event, validation_handler.on_user_input) + if enable_hover_validation and validation_state: show_tooltips_var = ui_state.get_var("validation_show_tooltips") validation_tooltip = validation_handler.validation_tooltip @@ -408,14 +463,12 @@ def path_entry( frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") frame.grid_columnconfigure(0, weight=1) - # Determine if this is an output path if not is_output: meta = ui_state.get_field_metadata(var_name) is_output = getattr(meta, 'is_output', False) var = ui_state.get_var(var_name) - # Set up validator and validation state if use_model_validator: validator = ModelOutputValidator( var=var, ui_state=ui_state, @@ -443,7 +496,6 @@ def simple_path_validator(value: str) -> ValidationResult: custom_validator = simple_path_validator - # Create entry with validation entry_widget = entry( frame, row=0, column=0, ui_state=ui_state, var_name=var_name, custom_validator=custom_validator, @@ -452,9 +504,10 @@ def simple_path_validator(value: str) -> ValidationResult: command=command ) - # Set up validator traces and cleanup if validator: + validator.set_widget(entry_widget) validator.setup_traces() + original_destroy = entry_widget.destroy def new_destroy(): @@ -464,7 +517,6 @@ def new_destroy(): entry_widget.destroy = new_destroy - # Enable drag-and-drop register_drop_target(entry_widget, ui_state, var_name, command) # Browse button @@ -615,7 +667,7 @@ def update_var(): if not deactivate_update_var: for key, value in values: if var.get() == str(value): - if component.winfo_exists(): # the component could already be destroyed + if component.winfo_exists(): component.set(key) if command: command(value) diff --git a/modules/util/ui/validation.py b/modules/util/ui/validation.py index a40868d34..86a110abf 100644 --- a/modules/util/ui/validation.py +++ b/modules/util/ui/validation.py @@ -24,6 +24,8 @@ class ValidationSettings: ONLY_WHITESPACE = re.compile(r"^\s+$") TRAILING_SLASH_RE = re.compile(r"[\\/]$") ENDS_WITH_EXT = re.compile(r"\.[A-Za-z0-9]+$") +INVALID_NAMES = {"", "."} +GENERIC_MODEL_NAMES = {"lora", "embedding", "embeddings", "model", "finetune"} INVALID_CHARS_WIN = set('<>:"|?*') INVALID_NAMES = {"", "."} @@ -87,6 +89,15 @@ def _has_extension(path: str | Path, extension: str) -> bool: def _is_partial_match(current: str, required: str) -> bool: return len(current) > 3 and required.startswith(current) +def _has_prefix(stem: str, prefix: str) -> bool: + """Check if stem already contains the prefix (case-insensitive).""" + if not prefix: + return False + stem_lower, prefix_lower = stem.lower(), prefix.lower() + # Exact match or prefix followed by separator + return (stem_lower == prefix_lower or + stem_lower.startswith((f"{prefix_lower}-", f"{prefix_lower}_"))) + def generate_default_filename( training_method: TrainingMethod, prefix: str = "", @@ -102,6 +113,33 @@ def generate_default_filename( name = f"{prefix}_{timestamp}" if prefix else f"{method_str}_{timestamp}" return f"{name}{extension}" +def make_unique_filename(base_path: Path, use_friendly_names: bool = False) -> Path: + """Generate a unique filename by appending a word or number if the path exists.""" + if not base_path.exists(): + return base_path + + stem = base_path.stem + suffix = base_path.suffix + parent = base_path.parent + + if use_friendly_names: + # Try 5 words before falling back to numbers + for _ in range(5): + word = generate_friendly(1, separator="") + new_path = parent / f"{stem}_{word}{suffix}" + if not new_path.exists(): + return new_path + + # Number-based approach (including friendly word fallback) + counter = 1 + while counter <= 9999: + new_path = parent / f"{stem}_{counter}{suffix}" + if not new_path.exists(): + return new_path + counter += 1 + + return base_path + def get_allowed_formats_for_method(training_method: TrainingMethod) -> set[ModelFormat]: if training_method in (TrainingMethod.EMBEDDING, TrainingMethod.LORA): return {ModelFormat.SAFETENSORS, ModelFormat.CKPT} @@ -118,6 +156,8 @@ class ValidationContext: use_friendly_names: bool is_output: bool separator: str + prevent_overwrite: bool + auto_prefix: bool def _validate_basic_input(ctx: ValidationContext) -> ValidationResult | None: @@ -194,8 +234,33 @@ def _validate_single_file_path(ctx: ValidationContext) -> ValidationResult: return _result(True, new_path, "Removed trailing period(s).") return _result(False, None, "Filename cannot end with a period.") + # Handle generic model names and auto-prefix + stem_lower = path.stem.lower() + + # Check for generic names that should be prefixed + if stem_lower in GENERIC_MODEL_NAMES and ctx.prefix and ctx.autocorrect: + new_stem = f"{ctx.prefix}_{path.stem}" + new_path = _format_path(path.with_name(new_stem + path.suffix), ctx.separator) + return _result(True, new_path, f"Added prefix to generic name '{path.stem}'.", 'warning') + elif stem_lower in GENERIC_MODEL_NAMES and ctx.prefix: + return _result(True, None, f"WARNING: Generic filename '{path.stem}' may cause confusion. Consider adding a prefix.", 'warning') + + # Auto-prefix: add prefix to manually-entered filenames + if ctx.auto_prefix and ctx.prefix and ctx.autocorrect and not _has_prefix(path.stem, ctx.prefix): + new_stem = f"{ctx.prefix}-{path.stem}" + new_path = _format_path(path.with_name(new_stem + path.suffix), ctx.separator) + return _result(True, new_path, f"Added prefix '{ctx.prefix}' to filename.", 'warning') + if _has_extension(base, required_ext): + # Check if file exists and handle overwrite prevention if path.exists(): + if ctx.prevent_overwrite and ctx.autocorrect: + # Generate a unique filename + unique_path = make_unique_filename(path, ctx.use_friendly_names) + if unique_path != path: + corrected = _format_path(unique_path, ctx.separator) + suffix_added = unique_path.stem[len(path.stem):] # Get what was added + return _result(True, corrected, f"File exists. Added '{suffix_added}' to prevent overwrite.") return _result(True, None, "WARNING: File already exists and will be overwritten.", 'warning') return _result(True, None, "") @@ -241,6 +306,9 @@ def validate_destination( prefix: str = "", use_friendly_names: bool = False, is_output: bool = True, + prevent_overwrite: bool = False, + auto_prefix: bool = False, + skip_overwrite_protection: bool = False, ) -> ValidationResult: """ Validate and optionally auto-correct a model output destination path. @@ -252,6 +320,7 @@ def validate_destination( - File extension validation and correction - Parent directory existence - Existing file warnings + - Automatic overwrite prevention Args: raw: Raw user input string @@ -261,6 +330,9 @@ def validate_destination( prefix: Optional prefix for auto-generated filenames use_friendly_names: Use friendly names instead of timestamps is_output: Whether this is an output path (True) or input path (False) + prevent_overwrite: Automatically make filename unique if it would overwrite + auto_prefix: Automatically add prefix to manually-entered filenames + skip_overwrite_protection: Skip overwrite protection (used when validation triggered by settings changes) Returns: ValidationResult with ok status, optional corrected value, and message @@ -274,7 +346,9 @@ def validate_destination( prefix=prefix, use_friendly_names=use_friendly_names, is_output=is_output, - separator='/' if '/' in raw else '\\' + separator='/' if '/' in raw else '\\', + prevent_overwrite=prevent_overwrite and not skip_overwrite_protection, + auto_prefix=auto_prefix, ) if basic_check := _validate_basic_input(ctx): From 33fc76f4d6714ee23dede17eeb878647d8ea47e2 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Tue, 11 Nov 2025 01:27:01 +1100 Subject: [PATCH 28/64] Add logic to validate on load to avoid edge cases --- modules/util/ui/ToolTip.py | 9 ++++++++- modules/util/ui/components.py | 34 +++++++++++++++++++++++++++++----- 2 files changed, 37 insertions(+), 6 deletions(-) diff --git a/modules/util/ui/ToolTip.py b/modules/util/ui/ToolTip.py index a2f4294ef..8a38d5b54 100644 --- a/modules/util/ui/ToolTip.py +++ b/modules/util/ui/ToolTip.py @@ -99,6 +99,11 @@ def _position(self): y = self.widget.winfo_rooty() + self.widget.winfo_height() + 4 self.tw.update_idletasks() + + # Check again after update_idletasks in case tooltip was destroyed + if not self.tw or not self.tw.winfo_exists(): + return + tip_width = self.tw.winfo_reqwidth() tip_height = self.tw.winfo_reqheight() screen_width = self.widget.winfo_screenwidth() @@ -111,7 +116,9 @@ def _position(self): y = max(10, y) - self.tw.wm_geometry(f"+{x}+{y}") + # Final check before setting geometry + if self.tw and self.tw.winfo_exists(): + self.tw.wm_geometry(f"+{x}+{y}") def _on_move(self, _evt=None): if self.tw and self.tw.winfo_exists(): diff --git a/modules/util/ui/components.py b/modules/util/ui/components.py index 69af04cfc..e10241868 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -233,6 +233,7 @@ def __init__( self.friendly_names_var = ui_state.get_var("use_friendly_names") self.prevent_overwrite_var = ui_state.get_var("prevent_overwrite") self.auto_prefix_var = ui_state.get_var("auto_prefix") + self.show_tooltips_var = ui_state.get_var("validation_show_tooltips") self.state = ValidationState() self._trace_ids: dict[str, str] = {} @@ -245,7 +246,7 @@ def _get_enum_value(self, var, default_enum): return type(default_enum)[var.get()] return default_enum - def validate(self, value: str, skip_overwrite_protection: bool = False) -> ValidationResult: + def validate(self, value: str, skip_overwrite_protection: bool = False, readonly: bool = False) -> ValidationResult: value = value.strip() if not value: self.state.clear() @@ -256,12 +257,12 @@ def validate(self, value: str, skip_overwrite_protection: bool = False) -> Valid value, output_format=self._get_enum_value(self.format_var, ModelFormat.SAFETENSORS), training_method=self._get_enum_value(self.method_var, TrainingMethod.FINE_TUNE), - autocorrect=_safe_bool(self.autocorrect_var), + autocorrect=False if readonly else _safe_bool(self.autocorrect_var), prefix=self.prefix_var.get() if self.prefix_var else "", use_friendly_names=_safe_bool(self.friendly_names_var, default=False), is_output=True, - prevent_overwrite=_safe_bool(self.prevent_overwrite_var, default=False), - auto_prefix=_safe_bool(self.auto_prefix_var, default=False), + prevent_overwrite=False if readonly else _safe_bool(self.prevent_overwrite_var, default=False), + auto_prefix=False if readonly else _safe_bool(self.auto_prefix_var, default=False), skip_overwrite_protection=skip_overwrite_protection, ) @@ -270,7 +271,8 @@ def validate(self, value: str, skip_overwrite_protection: bool = False) -> Valid else: self.state.clear() - if result.corrected and result.corrected != value: + # Only apply corrections if not in readonly mode + if not readonly and result.corrected and result.corrected != value: self.var.set(result.corrected) return result @@ -353,8 +355,30 @@ def cleanup_traces(self): var.trace_remove("write", trace_id) self._trace_ids.clear() + def _perform_initial_validation(self): + """Perform initial readonly validation on load without any corrections.""" + current_value = self.var.get().strip() + if not current_value or not self._tk_widget: + return + + # Perform readonly validation - no corrections, just check and show status + self.validate(current_value, skip_overwrite_protection=True, readonly=True) + + # Apply visual feedback based on validation state + if self.state.status == 'error': + self._tk_widget.configure(border_color="#dc3545") + if hasattr(self._tk_widget, '_validation_tooltip'): + self._tk_widget._validation_tooltip.show_error(self.state.message, duration_ms=5000) + elif self.state.status == 'warning': + self._tk_widget.configure(border_color="#ff9500") + if hasattr(self._tk_widget, '_validation_tooltip'): + self._tk_widget._validation_tooltip.show_warning(self.state.message, duration_ms=5000) + def set_widget(self, widget): self._tk_widget = widget + # Perform initial readonly validation if tooltips are enabled + if _safe_bool(self.show_tooltips_var): + self._perform_initial_validation() # UI components From 169c03146721307057445081f60cae9e908175aa Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Tue, 11 Nov 2025 04:17:37 +1100 Subject: [PATCH 29/64] Edgecase fix --- modules/util/ui/components.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/modules/util/ui/components.py b/modules/util/ui/components.py index e10241868..c0ef3e768 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -103,7 +103,10 @@ def _reset_border(self): self.component.configure(border_color=self.original_border_color) def should_show_tooltip(self) -> bool: - return _safe_bool(self.ui_state.get_var("validation_show_tooltips")) + try: + return _safe_bool(self.ui_state.get_var("validation_show_tooltips")) + except (KeyError, AttributeError): + return True def _show_validation_tooltip(self): if not self.should_show_tooltip() or not self.validation_state: @@ -423,11 +426,15 @@ def entry( component.bind(event, validation_handler.on_user_input) if enable_hover_validation and validation_state: - show_tooltips_var = ui_state.get_var("validation_show_tooltips") + try: + show_tooltips_var = ui_state.get_var("validation_show_tooltips") + except (KeyError, AttributeError): + show_tooltips_var = None + validation_tooltip = validation_handler.validation_tooltip def on_hover_enter(_e=None): - if not _safe_bool(show_tooltips_var): + if show_tooltips_var and not _safe_bool(show_tooltips_var): return if validation_state.status == 'error': From 593674035e387d9cb3332e71caa44168685ddb3d Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Tue, 11 Nov 2025 04:35:50 +1100 Subject: [PATCH 30/64] Utilise existing debouncTimer --- modules/util/ui/components.py | 67 +++++++++++++++++------------------ 1 file changed, 33 insertions(+), 34 deletions(-) diff --git a/modules/util/ui/components.py b/modules/util/ui/components.py index c0ef3e768..f0f22dc64 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -11,7 +11,7 @@ from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.path_util import supported_image_extensions from modules.util.ui.ToolTip import ToolTip -from modules.util.ui.ui_utils import register_drop_target +from modules.util.ui.ui_utils import DebounceTimer, register_drop_target from modules.util.ui.UIState import UIState from modules.util.ui.validation import ValidationResult, validate_basic_type, validate_destination, validate_file_path @@ -87,7 +87,16 @@ def __init__( except Exception: self.original_border_color = "gray50" - self.validation_after_id = None + # Replace manual debounce logic with DebounceTimer + self.debounce_timer = DebounceTimer( + widget=component, + delay_ms=COMPONENT_VALIDATION_SETTINGS.debounce_stop_typing_ms, + callback=lambda: self.validate_value( + self.var.get(), + COMPONENT_VALIDATION_SETTINGS.debounced_invalid_revert_ms + ) + ) + self.revert_after_id = None self.touched = False self.last_valid_value = var.get() @@ -179,22 +188,15 @@ def _do_revert(self): def debounced_validate(self, *_): if not self.touched: - self._cancel_after(self.validation_after_id) return self.validation_tooltip.hide() self._reset_border() self._cancel_after(self.revert_after_id) - self._cancel_after(self.validation_after_id) - self.validation_after_id = self.component.after( - COMPONENT_VALIDATION_SETTINGS.debounce_stop_typing_ms, - lambda: self.validate_value( - self.var.get(), - COMPONENT_VALIDATION_SETTINGS.debounced_invalid_revert_ms - ) - ) + # Use DebounceTimer instead of manual after() calls + self.debounce_timer.call() def on_focus_in(self, _e=None): self.touched = False @@ -213,8 +215,8 @@ def on_focus_out(self, _e=None): ) def cleanup(self): - for after_id in (self.validation_after_id, self.revert_after_id): - self._cancel_after(after_id) + self._cancel_after(self.revert_after_id) + # No need to cancel debounce_timer - it handles its own cleanup class ModelOutputValidator: @@ -240,7 +242,9 @@ def __init__( self.state = ValidationState() self._trace_ids: dict[str, str] = {} - self._debounce_after_id = None + + # Replace manual debounce with DebounceTimer + self._debounce_timer = None self._tk_widget = None self._last_prefix = self.prefix_var.get() if self.prefix_var else "" @@ -298,32 +302,32 @@ def _schedule_validation(self, skip_overwrite: bool = False): if not self._tk_widget: return - if self._debounce_after_id: - with contextlib.suppress(Exception): - self._tk_widget.after_cancel(self._debounce_after_id) + if not self._debounce_timer: + self._debounce_timer = DebounceTimer( + widget=self._tk_widget, + delay_ms=COMPONENT_VALIDATION_SETTINGS.debounce_stop_typing_ms, + callback=lambda: self.validate(self.var.get(), skip_overwrite_protection=skip_overwrite) + ) - self._debounce_after_id = self._tk_widget.after( - COMPONENT_VALIDATION_SETTINGS.debounce_stop_typing_ms, - lambda: self.validate(self.var.get(), skip_overwrite_protection=skip_overwrite) - ) + self._debounce_timer.call() def _schedule_prefix_change_validation(self): if not self._tk_widget: return - if self._debounce_after_id: - with contextlib.suppress(Exception): - self._tk_widget.after_cancel(self._debounce_after_id) - def on_prefix_change(): self._remove_old_prefix_if_needed() self._last_prefix = self.prefix_var.get() if self.prefix_var else "" self.validate(self.var.get(), skip_overwrite_protection=False) - self._debounce_after_id = self._tk_widget.after( - COMPONENT_VALIDATION_SETTINGS.debounce_stop_typing_ms, - on_prefix_change - ) + if not self._debounce_timer: + self._debounce_timer = DebounceTimer( + widget=self._tk_widget, + delay_ms=COMPONENT_VALIDATION_SETTINGS.debounce_stop_typing_ms, + callback=on_prefix_change + ) + + self._debounce_timer.call() def setup_traces(self): trace_config = [ @@ -346,11 +350,6 @@ def make_callback(debounce: bool, skip_overwrite: bool, is_prefix: bool): self._trace_ids[key] = var.trace_add("write", make_callback(debounce, skip_overwrite, is_prefix)) def cleanup_traces(self): - if self._debounce_after_id and self._tk_widget: - with contextlib.suppress(Exception): - self._tk_widget.after_cancel(self._debounce_after_id) - self._debounce_after_id = None - for key, trace_id in self._trace_ids.items(): var = getattr(self, f"{key}_var", None) if var: From 134957f4067e2d1592d230cc6304987d27cd4cf0 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Thu, 13 Nov 2025 21:35:23 +0100 Subject: [PATCH 31/64] Remove unnecessary blank line in model saver --- modules/modelSaver/StableDiffusionXLLoRAModelSaver.py | 1 - 1 file changed, 1 deletion(-) diff --git a/modules/modelSaver/StableDiffusionXLLoRAModelSaver.py b/modules/modelSaver/StableDiffusionXLLoRAModelSaver.py index cd38ea052..fe7ebf679 100644 --- a/modules/modelSaver/StableDiffusionXLLoRAModelSaver.py +++ b/modules/modelSaver/StableDiffusionXLLoRAModelSaver.py @@ -24,7 +24,6 @@ def save( output_model_destination: str, dtype: torch.dtype | None, ): - lora_model_saver = StableDiffusionXLLoRASaver() embedding_model_saver = StableDiffusionXLEmbeddingSaver() From 12d634c4942a4fd259bf84b3debae80f106fda31 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 21 Nov 2025 03:27:48 +1100 Subject: [PATCH 32/64] Add missing req --- requirements-global.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements-global.txt b/requirements-global.txt index 8590ea600..cadad561b 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -52,6 +52,7 @@ scalene==1.5.51 # ui customtkinter==5.2.2 tkinterdnd2==0.4.3 +friendlywords=1.2.0 # cloud runpod==1.7.10 From dc900cb1eff9c4f15b2e83d0ecd533625e3bd1b9 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 21 Nov 2025 03:32:52 +1100 Subject: [PATCH 33/64] Make all the enums uppercase --- modules/util/enum/TensorboardMode.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/modules/util/enum/TensorboardMode.py b/modules/util/enum/TensorboardMode.py index aa3274f59..6e92a0926 100644 --- a/modules/util/enum/TensorboardMode.py +++ b/modules/util/enum/TensorboardMode.py @@ -2,6 +2,6 @@ class TensorboardMode(StrEnum): - OFF = "off" - TRAIN_ONLY = "train_only" - ALWAYS_ON = "always_on" + OFF = "OFF" + TRAIN_ONLY = "TRAIN_ONLY" + ALWAYS_ON = "ALWAYS_ON" From 5a06a00b08aecc8446cbf2ebc9a0beff51a27961 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 21 Nov 2025 03:36:21 +1100 Subject: [PATCH 34/64] Remove no longer used code in BaseTrainer --- modules/trainer/BaseTrainer.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/modules/trainer/BaseTrainer.py b/modules/trainer/BaseTrainer.py index 3d9f9e34b..95b78e257 100644 --- a/modules/trainer/BaseTrainer.py +++ b/modules/trainer/BaseTrainer.py @@ -1,4 +1,3 @@ -import subprocess from abc import ABCMeta, abstractmethod from modules.model.BaseModel import BaseModel @@ -21,8 +20,6 @@ class BaseTrainer( metaclass=ABCMeta, ): - tensorboard_subprocess: subprocess.Popen = None - def __init__(self, config: TrainConfig, callbacks: TrainCallbacks, commands: TrainCommands): super().__init__() self.config = config From 74b4f2e20f0f0f1b90cff8bf411cb12f8cd5f167 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 21 Nov 2025 05:02:09 +1100 Subject: [PATCH 35/64] - Align run name inconsistency - Trim more fat --- modules/trainer/CloudTrainer.py | 6 +-- modules/trainer/GenericTrainer.py | 10 ++--- modules/ui/SampleWindow.py | 2 +- modules/ui/TrainUI.py | 4 +- modules/util/config/TrainConfig.py | 19 ++++++++++ modules/util/tensorboard_util.py | 60 ++++-------------------------- 6 files changed, 36 insertions(+), 65 deletions(-) diff --git a/modules/trainer/CloudTrainer.py b/modules/trainer/CloudTrainer.py index 980f0522e..7f8cae406 100644 --- a/modules/trainer/CloudTrainer.py +++ b/modules/trainer/CloudTrainer.py @@ -13,7 +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 TensorboardManager +from modules.util.tensorboard_util import get_tensorboard_manager class CloudTrainer(BaseTrainer): @@ -31,7 +31,7 @@ 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) - TensorboardManager().handle_training_start(config) + get_tensorboard_manager().handle_training_start(config) match config.cloud.type: case CloudType.RUNPOD: @@ -114,7 +114,7 @@ def train(self): def end(self): try: - TensorboardManager().handle_training_end(self.config) + 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") diff --git a/modules/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index 2d875daf8..2af342625 100644 --- a/modules/trainer/GenericTrainer.py +++ b/modules/trainer/GenericTrainer.py @@ -29,7 +29,6 @@ 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_run_name from modules.util.time_util import get_string_timestamp from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -70,8 +69,7 @@ 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) - run_name = get_tensorboard_run_name(config) - self.tensorboard = SummaryWriter(os.path.join(tensorboard_log_dir, f"{run_name}_{get_string_timestamp()}")) + self.tensorboard = SummaryWriter(os.path.join(tensorboard_log_dir, f"{config.get_run_name()}_{get_string_timestamp()}")) self.model = None self.one_step_trained = False @@ -158,7 +156,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) @@ -228,7 +226,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): @@ -842,7 +840,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 diff --git a/modules/ui/SampleWindow.py b/modules/ui/SampleWindow.py index a93b3593f..f49dc050e 100644 --- a/modules/ui/SampleWindow.py +++ b/modules/ui/SampleWindow.py @@ -184,7 +184,7 @@ def __sample(self): progress = self.model.train_progress sample_path = os.path.join( sample_dir, - f"{get_string_timestamp()}-training-sample-{progress.filename_string()}" + f"{self.current_train_config.get_run_name()}_{get_string_timestamp()}-training-sample-{progress.filename_string()}" ) self.model.eval() diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index 7cace5b25..ac4343dcd 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -36,7 +36,7 @@ from modules.util.enum.ModelType import ModelType from modules.util.enum.TensorboardMode import TensorboardMode from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.tensorboard_util import TensorboardManager +from modules.util.tensorboard_util import get_tensorboard_manager from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress from modules.util.ui import components @@ -133,7 +133,7 @@ def __init__(self): self.training_commands = None self.current_workspace_dir = self.train_config.workspace_dir - self.tensorboard_manager = TensorboardManager() + self.tensorboard_manager = get_tensorboard_manager() # Initialize tensorboard if needed if self.train_config.tensorboard_mode == TensorboardMode.ALWAYS_ON: diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 40b842467..491ead643 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -2,6 +2,7 @@ import os import uuid from copy import deepcopy +from pathlib import Path from typing import Any from modules.util.config.BaseConfig import BaseConfig @@ -823,6 +824,24 @@ def all_embedding_configs(self): else: return self.additional_embeddings + def get_run_name(self) -> str: + """Extract run name from output_model_destination, falling back to prefix.""" + if not self.output_model_destination or not self.output_model_destination.strip(): + return self.save_filename_prefix + + path = Path(self.output_model_destination) + name = path.name + + # remove extension + if path.suffix: + name = path.stem + + # fall back to prefix if no name could be extracted + if not name: + return self.save_filename_prefix + + return name + def get_last_backup_path(self) -> str | None: backups_path = os.path.join(self.workspace_dir, "backup") if os.path.exists(backups_path): diff --git a/modules/util/tensorboard_util.py b/modules/util/tensorboard_util.py index af69626cb..6a7725aab 100644 --- a/modules/util/tensorboard_util.py +++ b/modules/util/tensorboard_util.py @@ -9,61 +9,15 @@ from modules.util.enum.TensorboardMode import TensorboardMode -def get_tensorboard_run_name(config: TrainConfig) -> str: - """Extract run name from output_model_destination, falling back to prefix.""" - if not config.output_model_destination or not config.output_model_destination.strip(): - return config.save_filename_prefix - - path = Path(config.output_model_destination) - name = path.name - - # remove extension - if path.suffix: - name = path.stem - - # fall back to prefix if no name could be extracted - if not name: - return config.save_filename_prefix - - return name - - class TensorboardManager: - _instance = None - _lock = threading.Lock() - - def __new__(cls): - if cls._instance is None: - with cls._lock: - if cls._instance is None: - cls._instance = super().__new__(cls) - cls._instance._initialized = False - return cls._instance - def __init__(self): - if self._initialized: - return self._process = None - self._config = None - self._mode = None self._previous_mode = None - self._initialized = True def _is_cloud_tunnel_enabled(self, config: TrainConfig) -> bool: return config.cloud.enabled and config.cloud.tensorboard_tunnel - def should_start(self, config: TrainConfig, is_training: bool = False) -> bool: - if self._is_cloud_tunnel_enabled(config): - return False - - mode = config.tensorboard_mode - if mode == TensorboardMode.OFF: - return False - if mode == TensorboardMode.ALWAYS_ON: - return not is_training - return mode == TensorboardMode.TRAIN_ONLY and is_training - def start(self, config: TrainConfig, mode: TensorboardMode | None = None) -> bool: if self._process is not None: return False @@ -73,8 +27,6 @@ def start(self, config: TrainConfig, mode: TensorboardMode | None = None) -> boo return False self._process = start_filtered_tensorboard(config) - self._config = config - self._mode = effective_mode self._previous_mode = effective_mode return self._process is not None @@ -84,8 +36,6 @@ def stop(self) -> bool: stop_tensorboard(self._process) self._process = None - self._config = None - self._mode = None return True def restart(self, config: TrainConfig) -> bool: @@ -95,9 +45,6 @@ def restart(self, config: TrainConfig) -> bool: def is_running(self) -> bool: return self._process is not None - def get_mode(self) -> TensorboardMode | None: - return self._mode - def handle_training_start(self, config: TrainConfig): if config.tensorboard_mode == TensorboardMode.TRAIN_ONLY: if not self._is_cloud_tunnel_enabled(config): @@ -238,3 +185,10 @@ def stop_tensorboard(process: subprocess.Popen | None): with contextlib.suppress(Exception): if process.stdout: process.stdout.close() + + +_tensorboard_manager = TensorboardManager() + + +def get_tensorboard_manager() -> TensorboardManager: + return _tensorboard_manager From 3ba4cd8b86b31baac1b8126729b17969459b5cfd Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 21 Nov 2025 05:09:40 +1100 Subject: [PATCH 36/64] Move all relevant code under finally block --- modules/ui/TrainUI.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index ac4343dcd..f8ed194aa 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -815,16 +815,15 @@ def __training_thread_function(self): torch.clear_autocast_cache() torch_gc() - if error_caught: - self.on_update_status("Error: check the console for details") - else: - self.on_update_status("Stopped") - self.delete_eta_label() + if error_caught: + self.on_update_status("Error: check the console for details") + else: + self.on_update_status("Stopped") + self.delete_eta_label() - # Use centralized manager - self.tensorboard_manager.handle_training_end(self.train_config) + self.tensorboard_manager.handle_training_end(self.train_config) - self.after(0, self._set_training_button_idle) + self.after(0, self._set_training_button_idle) def start_training(self): if self.training_thread is None: From 3cd211efbe45ae56f4a83f9846ad017affacd009 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 21 Nov 2025 05:25:35 +1100 Subject: [PATCH 37/64] Attempt to trace the state of cloud enabled and tcp tunnel as well. --- modules/ui/TrainUI.py | 20 ++++++++ modules/util/tensorboard_util.py | 82 ++++++++++++++------------------ 2 files changed, 57 insertions(+), 45 deletions(-) diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index f8ed194aa..68790d954 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -106,6 +106,7 @@ def __init__(self): self.train_config = TrainConfig.default_values() self.ui_state = UIState(self, self.train_config) + self.cloud_state = self.ui_state.get_var("cloud") self.grid_rowconfigure(0, weight=0) self.grid_rowconfigure(1, weight=1) @@ -151,6 +152,15 @@ def __init__(self): lambda: self._on_tensorboard_mode_change() ) + self.cloud_enabled_trace_id = self.cloud_state.add_var_trace( + "enabled", + lambda: self._on_cloud_settings_change() + ) + self.cloud_tunnel_trace_id = self.cloud_state.add_var_trace( + "tensorboard_tunnel", + lambda: self._on_cloud_settings_change() + ) + # Persistent profiling window. self.profiling_window = ProfilingWindow(self) @@ -163,6 +173,10 @@ def __close(self): self.ui_state.remove_var_trace("workspace_dir", self.workspace_dir_trace_id) if hasattr(self, 'tensorboard_mode_trace_id'): self.ui_state.remove_var_trace("tensorboard_mode", self.tensorboard_mode_trace_id) + if hasattr(self, 'cloud_enabled_trace_id'): + self.cloud_state.remove_var_trace("enabled", self.cloud_enabled_trace_id) + if hasattr(self, 'cloud_tunnel_trace_id'): + self.cloud_state.remove_var_trace("tensorboard_tunnel", self.cloud_tunnel_trace_id) self.quit() def top_bar(self, master): @@ -890,6 +904,12 @@ def _on_tensorboard_mode_change(self): is_training=self.training_thread is not None ) + def _on_cloud_settings_change(self): + self.tensorboard_manager.handle_mode_change( + self.train_config, + is_training=self.training_thread is not None + ) + def _set_training_button_style(self, mode: str): if not self.training_button: return diff --git a/modules/util/tensorboard_util.py b/modules/util/tensorboard_util.py index 6a7725aab..8b77e6998 100644 --- a/modules/util/tensorboard_util.py +++ b/modules/util/tensorboard_util.py @@ -14,20 +14,52 @@ class TensorboardManager: def __init__(self): self._process = None self._previous_mode = None + self._cloud_tunnel_active = False def _is_cloud_tunnel_enabled(self, config: TrainConfig) -> bool: return config.cloud.enabled and config.cloud.tensorboard_tunnel + def _should_run(self, config: TrainConfig, is_training: bool) -> bool: + if self._is_cloud_tunnel_enabled(config): + return False + + mode = config.tensorboard_mode + if mode == TensorboardMode.OFF: + return False + if mode == TensorboardMode.TRAIN_ONLY: + return is_training + return mode == TensorboardMode.ALWAYS_ON + + def _sync_state(self, config: TrainConfig, is_training: bool): + cloud_tunnel_enabled = self._is_cloud_tunnel_enabled(config) + + if cloud_tunnel_enabled != self._cloud_tunnel_active: + self._cloud_tunnel_active = cloud_tunnel_enabled + if cloud_tunnel_enabled and self.is_running(): + self.stop() + + should_run = self._should_run(config, is_training) + + if should_run and not self.is_running(): + self.start(config) + elif not should_run and self.is_running(): + self.stop() + + self._previous_mode = config.tensorboard_mode + def start(self, config: TrainConfig, mode: TensorboardMode | None = None) -> bool: if self._process is not None: return False effective_mode = mode or config.tensorboard_mode - if effective_mode == TensorboardMode.OFF or self._is_cloud_tunnel_enabled(config): + cloud_tunnel_enabled = self._is_cloud_tunnel_enabled(config) + self._previous_mode = effective_mode + self._cloud_tunnel_active = cloud_tunnel_enabled + + if effective_mode == TensorboardMode.OFF or cloud_tunnel_enabled: return False self._process = start_filtered_tensorboard(config) - self._previous_mode = effective_mode return self._process is not None def stop(self) -> bool: @@ -46,53 +78,13 @@ def is_running(self) -> bool: return self._process is not None def handle_training_start(self, config: TrainConfig): - if config.tensorboard_mode == TensorboardMode.TRAIN_ONLY: - if not self._is_cloud_tunnel_enabled(config): - self.start(config) + self._sync_state(config, is_training=True) def handle_training_end(self, config: TrainConfig): - if config.tensorboard_mode == TensorboardMode.TRAIN_ONLY: - self.stop() - elif config.tensorboard_mode == TensorboardMode.ALWAYS_ON: - if not self._is_cloud_tunnel_enabled(config) and not self.is_running(): - self.start(config) + self._sync_state(config, is_training=False) def handle_mode_change(self, config: TrainConfig, is_training: bool = False): - new_mode = config.tensorboard_mode - - if self._previous_mode is None: - self._previous_mode = new_mode - if (new_mode == TensorboardMode.ALWAYS_ON and not is_training - and not self._is_cloud_tunnel_enabled(config)): - self.start(config) - return - - old_mode = self._previous_mode - - if new_mode == old_mode: - return - - if self._is_cloud_tunnel_enabled(config): - self.stop() - self._previous_mode = new_mode - return - - if is_training: - if new_mode == TensorboardMode.OFF: - self.stop() - elif old_mode == TensorboardMode.OFF and new_mode in (TensorboardMode.TRAIN_ONLY, TensorboardMode.ALWAYS_ON): - self.start(config) - self._previous_mode = new_mode - return - - # Not training - if old_mode == TensorboardMode.ALWAYS_ON: - self.stop() - - if new_mode == TensorboardMode.ALWAYS_ON: - self.start(config) - - self._previous_mode = new_mode + self._sync_state(config, is_training=is_training) def handle_workspace_change(self, config: TrainConfig): if config.tensorboard_mode == TensorboardMode.ALWAYS_ON: From d6146e1e31246d504680c8fcc7d371b76833ecd2 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 21 Nov 2025 05:55:56 +1100 Subject: [PATCH 38/64] Fix scalene and prevent each input binding event handlers to root. Saves about 1.4 seconds or more. --- .gitignore | 3 +++ modules/ui/ProfilingWindow.py | 7 ++++++- modules/util/ui/ToolTip.py | 15 +++++++++------ modules/util/ui/components.py | 4 ---- start-ui.bat | 5 +++-- 5 files changed, 21 insertions(+), 13 deletions(-) diff --git a/.gitignore b/.gitignore index dce73072c..315557641 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,9 @@ debug*.ipynb debug*.py /debug* +profile.json +profile.html +stacks.txt # user data /workspace* diff --git a/modules/ui/ProfilingWindow.py b/modules/ui/ProfilingWindow.py index 8d298abe3..a65af3a3c 100644 --- a/modules/ui/ProfilingWindow.py +++ b/modules/ui/ProfilingWindow.py @@ -44,8 +44,13 @@ def _dump_stack(self): def _end_profiler(self): scalene_profiler.stop() + try: + scalene_profiler.Scalene.output_profile() + self._message_label.configure(text='Profile dumped to profile.json/html') + except Exception as e: + print(f"Error dumping profile: {e}") + self._message_label.configure(text='Error dumping profile') - self._message_label.configure(text='Inactive') self._profile_button.configure(text='Start Profiling') self._profile_button.configure(command=self._start_profiler) diff --git a/modules/util/ui/ToolTip.py b/modules/util/ui/ToolTip.py index 8a38d5b54..60c38cb09 100644 --- a/modules/util/ui/ToolTip.py +++ b/modules/util/ui/ToolTip.py @@ -19,18 +19,14 @@ def __init__(self, widget, text='widget info', x_position=20, wide=False, self.id = None self.tw = None self._after_id = None + self._root = None + self._bound = False if hover_only: self.widget.bind("", self.enter) self.widget.bind("", self.leave) self.widget.bind("", self.leave) - if track_movement: - self._root = widget.winfo_toplevel() - self._root.bind("", self._on_move, add="+") - self._root.bind("", self.hide, add="+") - self._root.bind("", self.hide, add="+") - def enter(self, event=None): self.schedule() @@ -60,6 +56,13 @@ def show(self, text=None, color=None, duration_ms=None, fg_color="#2b2b2b"): self._cancel_scheduled_hide() self.hide() + if self.track_movement and not self._bound: + self._root = self.widget.winfo_toplevel() + self._root.bind("", self._on_move, add="+") + self._root.bind("", self.hide, add="+") + self._root.bind("", self.hide, add="+") + self._bound = True + self.tw = ctk.CTkToplevel(self.widget) self.tw.wm_overrideredirect(True) if self.track_movement: diff --git a/modules/util/ui/components.py b/modules/util/ui/components.py index f0f22dc64..099e9b463 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -369,12 +369,8 @@ def _perform_initial_validation(self): # Apply visual feedback based on validation state if self.state.status == 'error': self._tk_widget.configure(border_color="#dc3545") - if hasattr(self._tk_widget, '_validation_tooltip'): - self._tk_widget._validation_tooltip.show_error(self.state.message, duration_ms=5000) elif self.state.status == 'warning': self._tk_widget.configure(border_color="#ff9500") - if hasattr(self._tk_widget, '_validation_tooltip'): - self._tk_widget._validation_tooltip.show_warning(self.state.message, duration_ms=5000) def set_widget(self, widget): self._tk_widget = widget diff --git a/start-ui.bat b/start-ui.bat index 0f859af59..09dbebe63 100644 --- a/start-ui.bat +++ b/start-ui.bat @@ -32,7 +32,8 @@ if not exist "%VENV_DIR%\Scripts\python.exe" ( goto :end ) set PYTHON="%VENV_DIR%\Scripts\python.exe" -X utf8 -if defined PROFILE (set PYTHON=%PYTHON% -m scalene --off --cpu --gpu --profile-all --no-browser) +set LAUNCH_CMD=%PYTHON% +if defined PROFILE (set LAUNCH_CMD=%PYTHON% -m scalene --off --cpu --gpu --profile-all --no-browser) echo Using Python %PYTHON% REM Disable HF_HUB_DISABLE_XET, buggy; default disables Xet (set to 0 to enable) - https://github.com/Nerogar/OneTrainer/issues/949 @@ -60,7 +61,7 @@ if errorlevel 1 ( :launch echo Starting UI... -%PYTHON% scripts\train_ui.py +%LAUNCH_CMD% scripts\train_ui.py if errorlevel 1 ( echo Error: UI script exited with code %ERRORLEVEL% ) From b12b20def17348de6bd0f59ddb05c4762e6b78a2 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 21 Nov 2025 18:07:59 +1100 Subject: [PATCH 39/64] Update scalene --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index cadad561b..4859d5ac0 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -47,7 +47,7 @@ prodigy-plus-schedule-free==2.0.1 # Prodigy plus optimizer adv_optm==1.1.3 # advanced optimizers # Profiling -scalene==1.5.51 +scalene==1.5.55 # ui customtkinter==5.2.2 From 9a15f6d794b6afc387c3fe01ac393d9ffd5f544c Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 21 Nov 2025 20:50:55 +1100 Subject: [PATCH 40/64] Introduce deferred tab loading. --- modules/ui/ConceptTab.py | 37 +++++++++++-- modules/ui/ConfigList.py | 30 ++++++++--- modules/ui/TopBar.py | 11 +++- modules/ui/TrainUI.py | 112 ++++++++++++++++++++++++++++++++++----- 4 files changed, 164 insertions(+), 26 deletions(-) diff --git a/modules/ui/ConceptTab.py b/modules/ui/ConceptTab.py index a3806c15d..e3a68e701 100644 --- a/modules/ui/ConceptTab.py +++ b/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( "", 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 diff --git a/modules/ui/ConfigList.py b/modules/ui/ConfigList.py index 2dc0b7081..f4d1f7438 100644 --- a/modules/ui/ConfigList.py +++ b/modules/ui/ConfigList.py @@ -53,6 +53,7 @@ def __init__( self.show_toggle_button = show_toggle_button self.is_opening_window = False self._is_current_item_enabled = False + self._deferred_init_done = False self.master.grid_rowconfigure(0, weight=0) self.master.grid_rowconfigure(1, weight=1) @@ -66,29 +67,42 @@ def __init__( self.element_list = None self.configs = [] - self.__load_available_config_names() - + # defer heavy operations to avoid blocking startup self.current_config = getattr(self.train_config, self.attr_name) self.widgets = [] - self.__load_current_config(getattr(self.train_config, self.attr_name)) self.__create_configs_dropdown() components.button(self.top_frame, 0, 1, "Add Config", self.__add_config, tooltip="Adds a new config, which are containers for concepts, which themselves contain your dataset", width=20, padx=5) components.button(self.top_frame, 0, 2, add_button_text, self.__add_element, tooltip=add_button_tooltip, width=30, padx=5) + + # defer loading config names and creating widgets + self.master.after(150, self.__deferred_initialization) else: self.top_frame = ctk.CTkFrame(self.master, fg_color="transparent") self.top_frame.grid(row=0, column=0, sticky="nsew") components.button(self.top_frame, 0, 2, add_button_text, self.__add_element, width=20, padx=5) self.current_config = getattr(self.train_config, self.attr_name) + self.widgets = [] self.element_list = None - self._create_element_list() + # defer widget creation for non-external file configs too + self.master.after(100, self._create_element_list) if show_toggle_button: # tooltips break if you initialize with an empty string, default to a single space self.toggle_button = components.button(self.top_frame, 0, 3, " ", self._toggle, tooltip="Disables/Enables all visible items in the current view", width=30, padx=5) - self._update_toggle_button_text() + if hasattr(self, 'widgets'): + self._update_toggle_button_text() + + def __deferred_initialization(self): + if self._deferred_init_done: + return + self._deferred_init_done = True + + self.__load_available_config_names() + self.__load_current_config(getattr(self.train_config, self.attr_name)) + self.__create_configs_dropdown() # refresh with loaded configs @@ -123,6 +137,10 @@ def _reset_filters(self): # pragma: no cover - default noop def _update_item_enabled_state(self): # Only count items that match current filters + if not hasattr(self, 'widgets') or not self.widgets: + self._is_current_item_enabled = False + return + self._is_current_item_enabled = any( item.ui_state.get_var(self.enable_key).get() for i, item in enumerate(self.widgets) @@ -162,7 +180,7 @@ def __create_configs_dropdown(self): def _create_placeholder(self): placeholder = ctk.CTkLabel( self.element_list, - text="Either click the 'Add Concept' button or drag and drop a directory here to add a new concept", + text="Either click the 'Add Concept/Embedding' button or drag and drop a directory here to add a new concept/embeddings", text_color="gray50" ) placeholder.grid(row=0, column=0, padx=20, pady=20) diff --git a/modules/ui/TopBar.py b/modules/ui/TopBar.py index a3dd7c84e..964027f9b 100644 --- a/modules/ui/TopBar.py +++ b/modules/ui/TopBar.py @@ -43,7 +43,7 @@ def __init__( self.config_ui_state = UIState(master, self.config_ui_data) self.configs = [("", path_util.canonical_join(self.dir, "#.json"))] - self.__load_available_config_names() + self._configs_loaded = False self.current_config = [] @@ -59,6 +59,8 @@ def __init__( self.configs_dropdown = None self.__create_configs_dropdown() + master.after(500, self.__load_available_config_names_deferred) + # remove button # TODO # components.icon_button(self.frame, 0, 2, "-", self.__remove_config) @@ -172,6 +174,13 @@ def __load_available_config_names(self): self.configs.append((name, path)) self.configs.sort() + def __load_available_config_names_deferred(self): + if self._configs_loaded: + return + self.__load_available_config_names() + self._configs_loaded = True + self.__create_configs_dropdown() + def __save_to_file(self, name) -> str: name = path_util.safe_filename(name) path = path_util.canonical_join("training_presets", f"{name}.json") diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index 68790d954..df4ed43fc 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -124,10 +124,19 @@ def __init__(self): self.lora_tab = None self.cloud_tab = None self.additional_embeddings_tab = None + self.concepts_tab = None + self.sampling_tab = None + self.backup_tab = None + self.tools_tab = None self.top_bar_component = self.top_bar(self) + self.update_idletasks() + self.content_frame(self) + self.update_idletasks() + self.bottom_bar(self) + self.update_idletasks() self.training_thread = None self.training_callbacks = None @@ -245,21 +254,93 @@ def content_frame(self, master): self.tabview = ctk.CTkTabview(frame) self.tabview.grid(row=0, column=0, sticky="nsew") + # track which tabs have been created + self._tabs_created = set() + + # only create the general tab immediately for fast startup self.general_tab = self.create_general_tab(self.tabview.add("general")) - self.model_tab = self.create_model_tab(self.tabview.add("model")) - self.data_tab = self.create_data_tab(self.tabview.add("data")) - self.concepts_tab = self.create_concepts_tab(self.tabview.add("concepts")) - self.training_tab = self.create_training_tab(self.tabview.add("training")) - self.sampling_tab = self.create_sampling_tab(self.tabview.add("sampling")) - self.backup_tab = self.create_backup_tab(self.tabview.add("backup")) - self.tools_tab = self.create_tools_tab(self.tabview.add("tools")) - self.additional_embeddings_tab = self.create_additional_embeddings_tab(self.tabview.add("additional embeddings")) - self.cloud_tab = self.create_cloud_tab(self.tabview.add("cloud")) - - self.change_training_method(self.train_config.training_method) + self._tabs_created.add("general") + + # add placeholders for other tabs + for tab in ["model", "data", "concepts", "training", "sampling", "backup", "tools", "additional embeddings", "cloud"]: + self.tabview.add(tab) + + # add dynamic tabs based on current config + if self.train_config.training_method == TrainingMethod.LORA: + self.tabview.add("LoRA") + elif self.train_config.training_method == TrainingMethod.EMBEDDING: + self.tabview.add("embedding") + + # initialize other tab references as None + self.data_tab = None # bind tab change event for lazy loading + self._original_set = self.tabview.set + self.tabview.set = self._lazy_tab_switch + + # defer creation of remaining tabs significantly to show UI first + self.after(500, self._create_remaining_tabs) return frame + def _lazy_tab_switch(self, tab_name): + """Intercept tab switches to create tabs on-demand""" + if tab_name not in self._tabs_created: + self._create_tab_content(tab_name) + + return self._original_set(tab_name) + + def _create_tab_content(self, tab_name): + """Create content for a specific tab on-demand""" + if tab_name in self._tabs_created: + return + + master = self.tabview._tab_dict[tab_name] + creators = { + "model": ("model_tab", self.create_model_tab), + "data": ("data_tab", self.create_data_tab), + "concepts": ("concepts_tab", self.create_concepts_tab), + "training": ("training_tab", self.create_training_tab), + "sampling": ("sampling_tab", self.create_sampling_tab), + "backup": ("backup_tab", self.create_backup_tab), + "tools": ("tools_tab", self.create_tools_tab), + "additional embeddings": ("additional_embeddings_tab", self.create_additional_embeddings_tab), + "cloud": ("cloud_tab", self.create_cloud_tab), + } + + if tab_name in creators: + attr, method = creators[tab_name] + setattr(self, attr, method(master)) + elif tab_name == "LoRA": + self.lora_tab = LoraTab(master, self.train_config, self.ui_state) + elif tab_name == "embedding": + self.embedding_tab(master) + + self._tabs_created.add(tab_name) + + def _create_remaining_tabs(self): + """Create remaining tabs in background after startup""" + tabs = ["model", "training", "data", "concepts", "sampling", "backup", "tools", "additional embeddings", "cloud"] + if self.train_config.training_method == TrainingMethod.LORA: + tabs.insert(1, "LoRA") + elif self.train_config.training_method == TrainingMethod.EMBEDDING: + tabs.insert(1, "embedding") + + self._remaining_tabs_queue = [t for t in tabs if t not in self._tabs_created] + self._create_next_tab_deferred() + + def _create_next_tab_deferred(self): + """Create one tab at a time with delays to keep UI responsive""" + if not hasattr(self, '_remaining_tabs_queue') or not self._remaining_tabs_queue: + # all tabs created, apply training method changes + self.change_training_method(self.train_config.training_method) + return + + tab_name = self._remaining_tabs_queue.pop(0) + self._create_tab_content(tab_name) + + # schedule next tab creation after a delay + if self._remaining_tabs_queue: + self.after(100, self._create_next_tab_deferred) + def create_general_tab(self, master): frame = ctk.CTkScrollableFrame(master, fg_color="transparent") frame.grid_columnconfigure(0, weight=0) @@ -856,9 +937,12 @@ def start_training(self): def save_default(self): self.top_bar_component.save_default() - self.concepts_tab.save_current_config() - self.sampling_tab.save_current_config() - self.additional_embeddings_tab.save_current_config() + if self.concepts_tab: + self.concepts_tab.save_current_config() + if self.sampling_tab: + self.sampling_tab.save_current_config() + if self.additional_embeddings_tab: + self.additional_embeddings_tab.save_current_config() def export_training(self): file_path = filedialog.asksaveasfilename(filetypes=[ From 87cb6b8bdbe8f8d7d6e8f03cfea004ada6d95533 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sat, 29 Nov 2025 13:19:18 +0100 Subject: [PATCH 41/64] Fix syntax for friendlywords version in requirements --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index 4058dafa9..6f2582ec6 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -52,7 +52,7 @@ scalene==1.5.55 # ui customtkinter==5.2.2 tkinterdnd2==0.4.3 -friendlywords=1.2.0 +friendlywords==1.2.0 # cloud runpod==1.7.10 From 339d376773be001b1b10b56ac2c60e6b195db0c8 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Thu, 4 Dec 2025 18:20:45 +1100 Subject: [PATCH 42/64] Fix StrEnum and tweak deprecated warning message --- modules/util/enum/TensorboardMode.py | 5 ++++- requirements-global.txt | 3 ++- scripts/util/version_check.py | 2 +- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/modules/util/enum/TensorboardMode.py b/modules/util/enum/TensorboardMode.py index 6e92a0926..f350e00ab 100644 --- a/modules/util/enum/TensorboardMode.py +++ b/modules/util/enum/TensorboardMode.py @@ -1,4 +1,7 @@ -from enum import StrEnum +try: + from enum import StrEnum +except ImportError: + from strenum import StrEnum class TensorboardMode(StrEnum): diff --git a/requirements-global.txt b/requirements-global.txt index 4058dafa9..f49fc311a 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -16,6 +16,7 @@ scenedetect==0.6.6 accelerate==1.7.0 safetensors==0.5.3 tensorboard==2.19.0 +StrEnum==0.4.15; python_version < "3.11" # For tensoboardMode on 3.10 pytorch-lightning==2.5.1.post0 # diffusion models @@ -52,7 +53,7 @@ scalene==1.5.55 # ui customtkinter==5.2.2 tkinterdnd2==0.4.3 -friendlywords=1.2.0 +friendlywords==1.2.0 # cloud runpod==1.7.10 diff --git a/scripts/util/version_check.py b/scripts/util/version_check.py index c41786d23..1035a5a74 100644 --- a/scripts/util/version_check.py +++ b/scripts/util/version_check.py @@ -41,7 +41,7 @@ def exit_wrong_version(msg, min_ver, too_high_ver): exit_wrong_version("too low", min_ver, too_high_ver) if sys.version_info < deprecated_ver: - sys.stderr.write("Warning: Deprecated Python version found. Update to %s or newer\n" % (tuple_to_str(deprecated_ver))) + sys.stderr.write("Warning: Deprecated Python version found. Its strongly recommended to update to %s or newer\n" % (tuple_to_str(deprecated_ver))) sys.stderr.flush() if sys.version_info >= too_high_ver: From b7de8ada57391b5f9bce096c249255a125cb8412 Mon Sep 17 00:00:00 2001 From: dxqb Date: Tue, 16 Dec 2025 21:17:30 +0100 Subject: [PATCH 43/64] merge fix --- modules/ui/CloudTab.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/ui/CloudTab.py b/modules/ui/CloudTab.py index 33ea64f99..499784405 100644 --- a/modules/ui/CloudTab.py +++ b/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.") From 69aca0c72e726645f3ae52e553f1896b09970258 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 2 Jan 2026 09:57:21 +1100 Subject: [PATCH 44/64] Fix accidental row reuse --- modules/ui/TrainUI.py | 20 ++++++++++---------- modules/util/ui/components.py | 1 - modules/util/ui/ui_utils.py | 1 - 3 files changed, 10 insertions(+), 12 deletions(-) diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index 57ee6976c..cb591d81f 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -441,29 +441,29 @@ def create_general_tab(self, master): tooltip="Multi-GPU: A comma-separated list of device indexes. If empty, all your GPUs are used. With a list such as \"0,1,3,4\" you can omit a GPU, for example an on-board graphics GPU.") components.entry(frame, 13, 3, self.ui_state, "device_indexes") - components.label(frame, 13, 0, "Gradient Reduce Precision", + components.label(frame, 14, 0, "Gradient Reduce Precision", tooltip="WEIGHT_DTYPE: Reduce gradients between GPUs in your weight data type; can be imprecise, but more efficient than float32\n" "WEIGHT_DTYPE_STOCHASTIC: Sum up the gradients in your weight data type, but average them in float32 and stochastically round if your weight data type is bfloat16\n" "FLOAT_32: Reduce gradients in float32\n" "FLOAT_32_STOCHASTIC: Reduce gradients in float32; use stochastic rounding to bfloat16 if your weight data type is bfloat16", wide_tooltip=True) - components.options(frame, 13, 1, [str(x) for x in list(GradientReducePrecision)], self.ui_state, + components.options(frame, 14, 1, [str(x) for x in list(GradientReducePrecision)], self.ui_state, "gradient_reduce_precision") - components.label(frame, 13, 2, "Fused Gradient Reduce", + components.label(frame, 14, 2, "Fused Gradient Reduce", tooltip="Multi-GPU: Gradient synchronisation during the backward pass. Can be more efficient, especially with Async Gradient Reduce") - components.switch(frame, 13, 3, self.ui_state, "fused_gradient_reduce") + components.switch(frame, 14, 3, self.ui_state, "fused_gradient_reduce") - components.label(frame, 14, 0, "Async Gradient Reduce", + components.label(frame, 15, 0, "Async Gradient Reduce", tooltip="Multi-GPU: Asynchroniously start the gradient reduce operations during the backward pass. Can be more efficient, but requires some VRAM.") - components.switch(frame, 14, 1, self.ui_state, "async_gradient_reduce") - components.label(frame, 14, 2, "Buffer size (MB)", + components.switch(frame, 15, 1, self.ui_state, "async_gradient_reduce") + components.label(frame, 15, 2, "Buffer size (MB)", tooltip="Multi-GPU: Maximum VRAM for \"Async Gradient Reduce\", in megabytes. A multiple of this value can be needed if combined with \"Fused Back Pass\" and/or \"Layer offload fraction\"") - components.entry(frame, 14, 3, self.ui_state, "async_gradient_reduce_buffer") + components.entry(frame, 15, 3, self.ui_state, "async_gradient_reduce_buffer") - components.label(frame, 15, 0, "Temp Device", + components.label(frame, 16, 0, "Temp Device", tooltip="The device used to temporarily offload models while they are not used. Default:\"cpu\"") - components.entry(frame, 15, 1, self.ui_state, "temp_device") + components.entry(frame, 16, 1, self.ui_state, "temp_device") frame.pack(fill="both", expand=1) return frame diff --git a/modules/util/ui/components.py b/modules/util/ui/components.py index 0574245bb..f5b227499 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -21,7 +21,6 @@ PAD = 10 -# Common filetypes for dialogs MODEL_FILETYPES = [ ("All Files", "*.*"), ("Diffusers", "model_index.json"), diff --git a/modules/util/ui/ui_utils.py b/modules/util/ui/ui_utils.py index 506d1ae58..5f9dd4d36 100644 --- a/modules/util/ui/ui_utils.py +++ b/modules/util/ui/ui_utils.py @@ -110,7 +110,6 @@ def _drop_enter(event): def _drop_leave(event): return event.action - def _create_drop_handler(entry_widget, ui_state, var_name, command=None): def drop(event): if event.data: From 0f9cb0d70ab6ec3c6a976b448431745eb8aaca9d Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 2 Jan 2026 12:38:49 +1100 Subject: [PATCH 45/64] Make model output dest dragNdrop reject invalid filepaths i.e txt --- modules/util/ui/components.py | 51 ++++++++++++++++++++++++++++++++++- modules/util/ui/ui_utils.py | 19 ++++++++----- 2 files changed, 62 insertions(+), 8 deletions(-) diff --git a/modules/util/ui/components.py b/modules/util/ui/components.py index f5b227499..0ae2cbdb2 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -119,6 +119,9 @@ def should_show_tooltip(self) -> bool: def _show_validation_tooltip(self): if not self.should_show_tooltip() or not self.validation_state: return + # Don't show validation tooltip if a drop rejection is active + if getattr(self.component, '_drop_rejection_active', False): + return if self.validation_state.status == 'error': self.validation_tooltip.show_error(self.validation_state.message, duration_ms=None) elif self.validation_state.status == 'warning': @@ -542,7 +545,53 @@ def new_destroy(): entry_widget.destroy = new_destroy - register_drop_target(entry_widget, ui_state, var_name, command) + drop_validator = None + on_reject = None + if use_model_validator: + valid_model_extensions = {'.safetensors', '.ckpt', '.pt', '.bin'} + valid_extensions_display = ', '.join(sorted(valid_model_extensions)) + + def drop_validator(file_path: str) -> bool: + dropped_path = Path(file_path) + return dropped_path.is_dir() or dropped_path.suffix.lower() in valid_model_extensions + + def on_reject(file_path: str): + tooltip = getattr(entry_widget, '_validation_tooltip', None) + if not tooltip: + return + + entry_widget._drop_rejection_active = True + saved_status = validation_state.status if validation_state else None + saved_message = validation_state.message if validation_state else "" + tooltip.hide() + + rejected_path = Path(file_path) + extension = rejected_path.suffix or "(no extension)" + reject_message = ( + f"Rejected: {rejected_path.name}\n" + f"'{extension}' is not a valid model output format.\n" + f"Expected: {valid_extensions_display} or a directory" + ) + + def show_rejection(): + tooltip.show_error(reject_message, duration_ms=5000) + + def restore_prior_state(): + entry_widget._drop_rejection_active = False + tooltip.hide() + try: + if saved_status and saved_message and entry_widget.focus_get() == entry_widget: + if saved_status == 'error': + tooltip.show_error(saved_message, duration_ms=5000) + elif saved_status == 'warning': + tooltip.show_warning(saved_message, duration_ms=5000) + except Exception: + pass + + entry_widget.after(10, show_rejection) + entry_widget.after(5000, restore_prior_state) + + register_drop_target(entry_widget, ui_state, var_name, command, drop_validator, on_reject) # Browse button def open_dialog(): diff --git a/modules/util/ui/ui_utils.py b/modules/util/ui/ui_utils.py index 5f9dd4d36..c7c08ae23 100644 --- a/modules/util/ui/ui_utils.py +++ b/modules/util/ui/ui_utils.py @@ -110,22 +110,27 @@ def _drop_enter(event): def _drop_leave(event): return event.action -def _create_drop_handler(entry_widget, ui_state, var_name, command=None): +def _create_drop_handler(entry_widget, ui_state, var_name, command=None, drop_validator=None, on_reject=None): def drop(event): if event.data: paths = _parse_dropped_paths(event.data) if paths: - file_path = paths[0] # only take the first path for single-value entries + dropped_path = paths[0] - ui_state.get_var(var_name).set(file_path) + if drop_validator and not drop_validator(dropped_path): + if on_reject: + on_reject(dropped_path) + return event.action + + ui_state.get_var(var_name).set(dropped_path) entry_widget.focus_force() entry_widget.event_generate('') - entry_widget.event_generate('') # mark as touched + entry_widget.event_generate('') entry_widget.event_generate('') if command: - command(file_path) + command(dropped_path) return event.action return drop @@ -155,11 +160,11 @@ def _parse_dropped_paths(event_data: str) -> list[str]: return [p.strip() for p in paths if p.strip()] -def register_drop_target(entry_widget, ui_state, var_name, command=None): +def register_drop_target(entry_widget, ui_state, var_name, command=None, drop_validator=None, on_reject=None): try: entry_widget.drop_target_register(DND_FILES) for event, handler in [('<>', _drop_enter), ('<>', _drop_leave), - ('<>', _create_drop_handler(entry_widget, ui_state, var_name, command))]: + ('<>', _create_drop_handler(entry_widget, ui_state, var_name, command, drop_validator, on_reject))]: entry_widget.dnd_bind(event, handler) except Exception: pass From 2a128cf7eadeaa24a8d611015be7db174cf037ce Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 2 Jan 2026 12:45:46 +1100 Subject: [PATCH 46/64] Modify scalene env var name to OT_PROFILE --- start-ui.bat | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/start-ui.bat b/start-ui.bat index 09dbebe63..6a996d11d 100644 --- a/start-ui.bat +++ b/start-ui.bat @@ -33,7 +33,7 @@ if not exist "%VENV_DIR%\Scripts\python.exe" ( ) set PYTHON="%VENV_DIR%\Scripts\python.exe" -X utf8 set LAUNCH_CMD=%PYTHON% -if defined PROFILE (set LAUNCH_CMD=%PYTHON% -m scalene --off --cpu --gpu --profile-all --no-browser) +if defined OT_PROFILE (set LAUNCH_CMD=%PYTHON% -m scalene --off --cpu --gpu --profile-all --no-browser) echo Using Python %PYTHON% REM Disable HF_HUB_DISABLE_XET, buggy; default disables Xet (set to 0 to enable) - https://github.com/Nerogar/OneTrainer/issues/949 From 1e6923ce1cf2cb39dc4efe308ca4bf012fbd7d7f Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 2 Jan 2026 12:49:44 +1100 Subject: [PATCH 47/64] Remove redundant checks --- modules/ui/ConfigList.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/modules/ui/ConfigList.py b/modules/ui/ConfigList.py index f4d1f7438..47216c21e 100644 --- a/modules/ui/ConfigList.py +++ b/modules/ui/ConfigList.py @@ -92,8 +92,7 @@ def __init__( if show_toggle_button: # tooltips break if you initialize with an empty string, default to a single space self.toggle_button = components.button(self.top_frame, 0, 3, " ", self._toggle, tooltip="Disables/Enables all visible items in the current view", width=30, padx=5) - if hasattr(self, 'widgets'): - self._update_toggle_button_text() + self._update_toggle_button_text() def __deferred_initialization(self): if self._deferred_init_done: @@ -137,7 +136,7 @@ def _reset_filters(self): # pragma: no cover - default noop def _update_item_enabled_state(self): # Only count items that match current filters - if not hasattr(self, 'widgets') or not self.widgets: + if not self.widgets: self._is_current_item_enabled = False return From 66653ddc0b072f2d3c4a97b9a75387d51c31753f Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 2 Jan 2026 13:07:35 +1100 Subject: [PATCH 48/64] Tweak gitignore. --- .gitignore | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 315557641..0a0b43032 100644 --- a/.gitignore +++ b/.gitignore @@ -19,7 +19,7 @@ stacks.txt /training_user_settings /external /update.var -config.json +config*.json secrets.json *.zip From 30b156d1c8c3e2c9f6efa79a195412c3234de2fd Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 2 Jan 2026 13:35:30 +1100 Subject: [PATCH 49/64] Fix CLI tensorboard oversight --- modules/trainer/GenericTrainer.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/modules/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index 39724a007..4fb7a9d8e 100644 --- a/modules/trainer/GenericTrainer.py +++ b/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 @@ -72,6 +73,8 @@ def __init__(self, config: TrainConfig, callbacks: TrainCallbacks, commands: Tra os.makedirs(Path(tensorboard_log_dir).absolute(), exist_ok=True) 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 self.grad_hook_handles = [] @@ -868,6 +871,7 @@ def end(self): if multi.is_master(): self.tensorboard.close() + get_tensorboard_manager().handle_training_end(self.config) for handle in self.grad_hook_handles: handle.remove() From d8ebd21b842f142ee6ed71775dd4094b4ca58419 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Fri, 2 Jan 2026 13:56:05 +1100 Subject: [PATCH 50/64] Make saves and backups use run name --- modules/trainer/GenericTrainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modules/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index 4fb7a9d8e..3436e7a84 100644 --- a/modules/trainer/GenericTrainer.py +++ b/modules/trainer/GenericTrainer.py @@ -432,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. @@ -482,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) From ff8edb9f846d740c7049f46d7fdc7e8266357cf2 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Sun, 11 Jan 2026 15:09:00 +1100 Subject: [PATCH 51/64] Attempts at debugging Xubuntu issues with dnd --- modules/util/ui/ui_utils.py | 48 ++++++++++++++++++++++++++++++++----- 1 file changed, 42 insertions(+), 6 deletions(-) diff --git a/modules/util/ui/ui_utils.py b/modules/util/ui/ui_utils.py index c7c08ae23..e7f708b5b 100644 --- a/modules/util/ui/ui_utils.py +++ b/modules/util/ui/ui_utils.py @@ -3,6 +3,7 @@ import platform import sys import tkinter as tk +import urllib.parse from collections.abc import Callable from pathlib import Path from tkinter import EventType @@ -11,6 +12,21 @@ from customtkinter import CTk, CTkToplevel from tkinterdnd2 import DND_FILES +# Temp debug +print(f"[DnD Debug] Platform: {sys.platform}") +print(f"[DnD Debug] XDG_SESSION_TYPE: {os.environ.get('XDG_SESSION_TYPE', 'not set')}") +print(f"[DnD Debug] WAYLAND_DISPLAY: {os.environ.get('WAYLAND_DISPLAY', 'not set')}") +print(f"[DnD Debug] DISPLAY: {os.environ.get('DISPLAY', 'not set')}") +try: + from tkinterdnd2 import TkinterDnD + _test_root = tk.Tk() + _test_root.withdraw() + _tkdnd_version = TkinterDnD._require(_test_root) + print(f"Debug: tkdnd ver: {_tkdnd_version}") + _test_root.destroy() +except Exception as e: + print(f"Debug: tkdnd load error: {e}") +print("=" * 50) def bind_mousewheel( widget: Any, @@ -135,6 +151,12 @@ def drop(event): return drop def _parse_dropped_paths(event_data: str) -> list[str]: + """Parse dropped path data from drag-and-drop events: + + - Brace-wrapped paths with spaces: {/path/with spaces/file.txt} + - file:// URI format: file:///home/user/file.txt + - URL-encoded characters: %20 for spaces + """ paths = [] current_path = "" in_braces = False @@ -150,7 +172,6 @@ def _parse_dropped_paths(event_data: str) -> list[str]: current_path += char continue - # Common path appending logic for '}' and ' ' if current_path: paths.append(current_path) current_path = "" @@ -158,7 +179,22 @@ def _parse_dropped_paths(event_data: str) -> list[str]: if current_path: paths.append(current_path) - return [p.strip() for p in paths if p.strip()] + + cleaned_paths = [] + for p in paths: + p = p.strip() + if not p: + continue + + if p.startswith('file://'): + p = urllib.parse.unquote(p[7:]) # remove 'file://' and decode %xx + if p.startswith('//'): + p = p[2:] # file:///path -> /path + + cleaned_paths.append(p) + + return cleaned_paths + def register_drop_target(entry_widget, ui_state, var_name, command=None, drop_validator=None, on_reject=None): try: @@ -166,8 +202,8 @@ def register_drop_target(entry_widget, ui_state, var_name, command=None, drop_va for event, handler in [('<>', _drop_enter), ('<>', _drop_leave), ('<>', _create_drop_handler(entry_widget, ui_state, var_name, command, drop_validator, on_reject))]: entry_widget.dnd_bind(event, handler) - except Exception: - pass + except Exception as e: + print(f"Failed to register drop target: {e}") def register_concept_drop_target(widget, drop_callback: Callable[[str], None], allow_multiple: bool = True): def drop_handler(event): @@ -190,8 +226,8 @@ def drop_handler(event): widget.dnd_bind('<>', _drop_enter) widget.dnd_bind('<>', _drop_leave) widget.dnd_bind('<>', drop_handler) - except Exception: - pass + except Exception as e: + print(f"Failed to register concept drop target: {e}") class DebounceTimer: def __init__(self, widget, delay_ms: int, callback: Callable[..., Any]): From b2b44f2a95dcf1fc78f116e2263c6525ad452641 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Sun, 11 Jan 2026 22:09:53 +1100 Subject: [PATCH 52/64] Make the placeholder text implementation better --- modules/ui/AdditionalEmbeddingsTab.py | 3 ++- modules/ui/ConceptTab.py | 3 ++- modules/ui/ConfigList.py | 6 +++++- modules/ui/SamplingTab.py | 3 ++- 4 files changed, 11 insertions(+), 4 deletions(-) diff --git a/modules/ui/AdditionalEmbeddingsTab.py b/modules/ui/AdditionalEmbeddingsTab.py index a8df2e8c4..634528bcd 100644 --- a/modules/ui/AdditionalEmbeddingsTab.py +++ b/modules/ui/AdditionalEmbeddingsTab.py @@ -20,7 +20,8 @@ def __init__(self, master, train_config: TrainConfig, ui_state: UIState): from_external_file=False, add_button_text="add embedding", is_full_width=True, - show_toggle_button=True + show_toggle_button=True, + placeholder_text="No additional embeddings defined. Click 'add embedding' to add one." ) def refresh_ui(self): diff --git a/modules/ui/ConceptTab.py b/modules/ui/ConceptTab.py index e3a68e701..f2e7125b6 100644 --- a/modules/ui/ConceptTab.py +++ b/modules/ui/ConceptTab.py @@ -35,7 +35,8 @@ def __init__(self, master, train_config: TrainConfig, ui_state: UIState): add_button_text="Add Concept", add_button_tooltip="Adds a new concept to the current config.", is_full_width=False, - show_toggle_button=True + show_toggle_button=True, + placeholder_text="You have no concepts defined. Either click the 'Add Concept' button or drag and drop a directory here to add a new concept." ) self._toolbar = None self._toolbar_is_wrapped = False diff --git a/modules/ui/ConfigList.py b/modules/ui/ConfigList.py index 47216c21e..7bf919482 100644 --- a/modules/ui/ConfigList.py +++ b/modules/ui/ConfigList.py @@ -31,6 +31,7 @@ def __init__( add_button_tooltip: str = "", is_full_width: bool = "", show_toggle_button: bool = False, + placeholder_text: str = "", ): self.master = master self.train_config = train_config @@ -43,6 +44,7 @@ def __init__( self.default_config_name = default_config_name self.is_full_width = is_full_width + self.placeholder_text = placeholder_text # From search-concepts self.filters = {"search": "", "type": "ALL", "show_disabled": True} @@ -177,9 +179,11 @@ def __create_configs_dropdown(self): self._update_toggle_button_text() def _create_placeholder(self): + if not self.placeholder_text: + return placeholder = ctk.CTkLabel( self.element_list, - text="Either click the 'Add Concept/Embedding' button or drag and drop a directory here to add a new concept/embeddings", + text=self.placeholder_text, text_color="gray50" ) placeholder.grid(row=0, column=0, padx=20, pady=20) diff --git a/modules/ui/SamplingTab.py b/modules/ui/SamplingTab.py index 64527989c..e71650223 100644 --- a/modules/ui/SamplingTab.py +++ b/modules/ui/SamplingTab.py @@ -22,7 +22,8 @@ def __init__(self, master, train_config: TrainConfig, ui_state: UIState): add_button_text="Add Sample", add_button_tooltip="Add a new sample configuration.", is_full_width=True, - show_toggle_button=True + show_toggle_button=True, + placeholder_text="You have no samples defined. Click 'Add Sample' to add a new sample to your current config." ) def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): From ab07e7b649c6a7b01fc52aa4fa290eddf1f7cb00 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Sat, 17 Jan 2026 15:10:35 +1100 Subject: [PATCH 53/64] Bump tensorboard --- requirements-global.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-global.txt b/requirements-global.txt index f91199c45..bf944d019 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -15,7 +15,7 @@ scenedetect==0.6.6 # pytorch accelerate==1.7.0 safetensors==0.5.3 -tensorboard==2.19.0 +tensorboard==2.20.0 StrEnum==0.4.15; python_version < "3.11" # For tensoboardMode on 3.10 pytorch-lightning==2.5.1.post0 From 1fe32575a813ce652c25e64c5e4dad1079ea2b18 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Sun, 18 Jan 2026 21:54:13 +1100 Subject: [PATCH 54/64] Remove debug, bump av to allow fresh install --- modules/util/ui/ui_utils.py | 21 --------------------- requirements-global.txt | 2 +- 2 files changed, 1 insertion(+), 22 deletions(-) diff --git a/modules/util/ui/ui_utils.py b/modules/util/ui/ui_utils.py index e7f708b5b..b3c70acc6 100644 --- a/modules/util/ui/ui_utils.py +++ b/modules/util/ui/ui_utils.py @@ -12,21 +12,6 @@ from customtkinter import CTk, CTkToplevel from tkinterdnd2 import DND_FILES -# Temp debug -print(f"[DnD Debug] Platform: {sys.platform}") -print(f"[DnD Debug] XDG_SESSION_TYPE: {os.environ.get('XDG_SESSION_TYPE', 'not set')}") -print(f"[DnD Debug] WAYLAND_DISPLAY: {os.environ.get('WAYLAND_DISPLAY', 'not set')}") -print(f"[DnD Debug] DISPLAY: {os.environ.get('DISPLAY', 'not set')}") -try: - from tkinterdnd2 import TkinterDnD - _test_root = tk.Tk() - _test_root.withdraw() - _tkdnd_version = TkinterDnD._require(_test_root) - print(f"Debug: tkdnd ver: {_tkdnd_version}") - _test_root.destroy() -except Exception as e: - print(f"Debug: tkdnd load error: {e}") -print("=" * 50) def bind_mousewheel( widget: Any, @@ -151,12 +136,6 @@ def drop(event): return drop def _parse_dropped_paths(event_data: str) -> list[str]: - """Parse dropped path data from drag-and-drop events: - - - Brace-wrapped paths with spaces: {/path/with spaces/file.txt} - - file:// URI format: file:///home/user/file.txt - - URL-encoded characters: %20 for spaces - """ paths = [] current_path = "" in_braces = False diff --git a/requirements-global.txt b/requirements-global.txt index bf944d019..1738d55b4 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -8,7 +8,7 @@ PyYAML==6.0.2 huggingface-hub==0.34.4 scipy==1.15.3 matplotlib==3.10.3 -av==14.4.0 +av==16.1.0 yt-dlp #no pinned version, frequently updated for compatibility with sites scenedetect==0.6.6 From 7e09675182f8e6b73f43e46c4e647146bd085f1a Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Sun, 18 Jan 2026 23:00:07 +1100 Subject: [PATCH 55/64] Destroy fig and canvas when window is closed --- modules/ui/ConceptWindow.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/modules/ui/ConceptWindow.py b/modules/ui/ConceptWindow.py index 5c0c35d56..babe84c00 100644 --- a/modules/ui/ConceptWindow.py +++ b/modules/ui/ConceptWindow.py @@ -91,6 +91,8 @@ def __init__( self.text_ui_state = text_ui_state self.image_preview_file_index = 0 self.preview_augmentations = ctk.BooleanVar(self, True) + self.bucket_fig = None + self.canvas = None self.title("Concept") self.geometry("800x700") @@ -514,6 +516,10 @@ def __concept_stats_tab(self, master): self.text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" plt.set_loglevel('WARNING') #suppress errors about data type in bar chart + + # Store previous figure to release once we're no longer using it + prev_bucket_fig = self.bucket_fig + self.bucket_fig, self.bucket_ax = plt.subplots(figsize=(7,3)) self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=frame) self.canvas.get_tk_widget().grid(row=19, column=0, columnspan=4, rowspan=2) @@ -531,6 +537,11 @@ def __concept_stats_tab(self, master): self.bucket_ax.xaxis.label.set_color(self.text_color) self.bucket_ax.yaxis.label.set_color(self.text_color) + # close previous figure if any + if prev_bucket_fig is not None: + plt.close(prev_bucket_fig) + prev_bucket_fig = None + #refresh stats - must be after all labels are defined or will give error self.refresh_basic_stats_button = components.button(master=frame, row=0, column=0, text="Refresh Basic", command=lambda: self.__get_concept_stats_threaded(False, 9999), tooltip="Reload basic statistics for the concept directory") @@ -921,5 +932,17 @@ def __auto_update_concept_stats(self): if self.concept.concept_stats["processing_time"] < 0.1: self.__get_concept_stats(True, 2) #do advanced scan automatically if basic took <0.1s + def destroy(self): + if self.canvas is not None: + self.canvas.get_tk_widget().destroy() + self.canvas = None + + if self.bucket_fig is not None: + plt.close(self.bucket_fig) + self.bucket_fig = None + self.bucket_ax = None + + super().destroy() + def __ok(self): self.destroy() From 2e937657b78f99d4b4592cc858cceabb1fd186f2 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Sun, 18 Jan 2026 23:34:23 +1100 Subject: [PATCH 56/64] Allow HF repo style inputs to skip validation --- modules/util/ui/validation.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/modules/util/ui/validation.py b/modules/util/ui/validation.py index 86a110abf..da15dfd1d 100644 --- a/modules/util/ui/validation.py +++ b/modules/util/ui/validation.py @@ -27,6 +27,9 @@ class ValidationSettings: INVALID_NAMES = {"", "."} GENERIC_MODEL_NAMES = {"lora", "embedding", "embeddings", "model", "finetune"} +# org/repo-name (alphanumeric, hyphens, underscores, dots allowed) +HUGGINGFACE_REPO_RE = re.compile(r"^[A-Za-z0-9_.-]+\/[A-Za-z0-9_.-]+$") + INVALID_CHARS_WIN = set('<>:"|?*') INVALID_NAMES = {"", "."} @@ -63,6 +66,29 @@ def _has_invalid_chars(path: str) -> bool: return any(c in INVALID_CHARS_WIN for c in path[start_idx:]) return "\x00" in path +def _is_huggingface_repo(value: str) -> bool: + """Check if valid HF repo after trimming whitespace""" + trimmed = value.strip() + + if " " in trimmed or "\t" in trimmed: + return False + + if trimmed.startswith(("\\\\", "//")): + return False + + if trimmed.startswith("/"): + return False + + if len(trimmed) >= 2 and trimmed[1] == ":" and trimmed[0].isalpha(): + return False + + # HF repos have exactly one slash (owner/repo) + if trimmed.count("/") != 1: + return False + + return bool(HUGGINGFACE_REPO_RE.match(trimmed)) + + def _safe_exists(path: Path, check_dir: bool = True) -> bool: try: return path.is_dir() if check_dir else path.parent.exists() @@ -396,6 +422,10 @@ def validate_file_path( if not value: return _result(True, None, "") + # For input paths (not output), accept HuggingFace repo format + if not is_output and _is_huggingface_repo(value): + return _result(True, None, "") + # Check for invalid characters if _has_invalid_chars(value): return _result(False, None, "Input contains invalid path characters.") From 610520430ccad6c68fb6e69c4f6b47fde6814eaf Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Sun, 18 Jan 2026 23:44:23 +1100 Subject: [PATCH 57/64] Add two HF rules I wasnt aware of. --- modules/util/ui/validation.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/modules/util/ui/validation.py b/modules/util/ui/validation.py index da15dfd1d..6c3329f17 100644 --- a/modules/util/ui/validation.py +++ b/modules/util/ui/validation.py @@ -70,9 +70,15 @@ def _is_huggingface_repo(value: str) -> bool: """Check if valid HF repo after trimming whitespace""" trimmed = value.strip() + if trimmed.__len__() > 96: + return False + if " " in trimmed or "\t" in trimmed: return False + if "—" in trimmed or ".." in trimmed: + return False + if trimmed.startswith(("\\\\", "//")): return False From 2e60c7b618891dd0ec40e71719d3ce0ea3301a34 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Mon, 19 Jan 2026 00:18:14 +1100 Subject: [PATCH 58/64] Allow HF HTTP urls to also skip validation --- modules/util/ui/validation.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/modules/util/ui/validation.py b/modules/util/ui/validation.py index 6c3329f17..ff63fa059 100644 --- a/modules/util/ui/validation.py +++ b/modules/util/ui/validation.py @@ -70,6 +70,9 @@ def _is_huggingface_repo(value: str) -> bool: """Check if valid HF repo after trimming whitespace""" trimmed = value.strip() + if re.match(r"^https://huggingface\.[a-zA-Z]+/", trimmed): + return True + if trimmed.__len__() > 96: return False From 71cde89069556287dc1af4b0476a2435e58aa260 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Mon, 19 Jan 2026 01:06:15 +1100 Subject: [PATCH 59/64] Experimental test to see if FigureCanvasTkAgg is the issue --- modules/ui/ConceptWindow.py | 101 +++++++++++++++++++----------------- 1 file changed, 52 insertions(+), 49 deletions(-) diff --git a/modules/ui/ConceptWindow.py b/modules/ui/ConceptWindow.py index babe84c00..4e5caaf8e 100644 --- a/modules/ui/ConceptWindow.py +++ b/modules/ui/ConceptWindow.py @@ -1,4 +1,5 @@ import fractions +import io import math import os import pathlib @@ -43,9 +44,11 @@ import customtkinter as ctk import huggingface_hub +import matplotlib as mpl from customtkinter import AppearanceModeTracker, ThemeManager + +mpl.use('Agg') # Use non-interactive backend to avoid Linux DnD issues from matplotlib import pyplot as plt -from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg from PIL import Image @@ -91,8 +94,6 @@ def __init__( self.text_ui_state = text_ui_state self.image_preview_file_index = 0 self.preview_augmentations = ctk.BooleanVar(self, True) - self.bucket_fig = None - self.canvas = None self.title("Concept") self.geometry("800x700") @@ -508,39 +509,19 @@ def __concept_stats_tab(self, master): self.small_bucket_label.configure(font=ctk.CTkFont(underline=True)) self.small_bucket_preview = components.label(frame, 18, 1, pad=0, text="-") - #aspect bucketing plot, mostly copied from timestep preview graph + #aspect bucketing plot - using Agg backend to render to image (avoids Linux DnD issues) appearance_mode = AppearanceModeTracker.get_mode() background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) - background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" + self.chart_background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" self.text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" plt.set_loglevel('WARNING') #suppress errors about data type in bar chart - # Store previous figure to release once we're no longer using it - prev_bucket_fig = self.bucket_fig - - self.bucket_fig, self.bucket_ax = plt.subplots(figsize=(7,3)) - self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=frame) - self.canvas.get_tk_widget().grid(row=19, column=0, columnspan=4, rowspan=2) - self.bucket_fig.tight_layout() - self.bucket_fig.subplots_adjust(bottom=0.15) - - self.bucket_fig.set_facecolor(background_color) - self.bucket_ax.set_facecolor(background_color) - self.bucket_ax.spines['bottom'].set_color(self.text_color) - self.bucket_ax.spines['left'].set_color(self.text_color) - self.bucket_ax.spines['top'].set_visible(False) - self.bucket_ax.spines['right'].set_color(self.text_color) - self.bucket_ax.tick_params(axis='x', colors=self.text_color, which="both") - self.bucket_ax.tick_params(axis='y', colors=self.text_color, which="both") - self.bucket_ax.xaxis.label.set_color(self.text_color) - self.bucket_ax.yaxis.label.set_color(self.text_color) - - # close previous figure if any - if prev_bucket_fig is not None: - plt.close(prev_bucket_fig) - prev_bucket_fig = None + # Create a placeholder image for the chart + self.chart_ctk_image = ctk.CTkImage(light_image=Image.new('RGB', (700, 300), self.chart_background_color), size=(700, 300)) + self.chart_label = ctk.CTkLabel(frame, image=self.chart_ctk_image, text="") + self.chart_label.grid(row=19, column=0, columnspan=4, rowspan=2) #refresh stats - must be after all labels are defined or will give error self.refresh_basic_stats_button = components.button(master=frame, row=0, column=0, text="Refresh Basic", command=lambda: self.__get_concept_stats_threaded(False, 9999), @@ -854,17 +835,48 @@ def __update_concept_stats(self): min_bucket_str.strip() self.small_bucket_preview.configure(text=min_bucket_str) - self.bucket_ax.cla() - aspects = [str(x) for x in list(aspect_buckets.keys())] - aspect_ratios = [self.decimal_to_aspect_ratio(x) for x in list(aspect_buckets.keys())] - counts = list(aspect_buckets.values()) - b = self.bucket_ax.bar(aspect_ratios, counts) - self.bucket_ax.bar_label(b, color=self.text_color) - sec = self.bucket_ax.secondary_xaxis(location=-0.1) - sec.spines["bottom"].set_linewidth(0) - sec.set_xticks([0, (len(aspects)-1)/2, len(aspects)-1], labels=["Wide", "Square", "Tall"]) - sec.tick_params('x', length=0) - self.canvas.draw() + self.__update_bucket_chart(aspect_buckets) + + def __update_bucket_chart(self, aspect_buckets: dict): + """Render the aspect bucket chart using Agg backend and update CTkImage.""" + fig, ax = plt.subplots(figsize=(7, 3)) + try: + fig.set_facecolor(self.chart_background_color) + ax.set_facecolor(self.chart_background_color) + ax.spines['bottom'].set_color(self.text_color) + ax.spines['left'].set_color(self.text_color) + ax.spines['top'].set_visible(False) + ax.spines['right'].set_color(self.text_color) + ax.tick_params(axis='x', colors=self.text_color, which="both") + ax.tick_params(axis='y', colors=self.text_color, which="both") + ax.xaxis.label.set_color(self.text_color) + ax.yaxis.label.set_color(self.text_color) + + aspects = [str(x) for x in list(aspect_buckets.keys())] + aspect_ratios = [self.decimal_to_aspect_ratio(x) for x in list(aspect_buckets.keys())] + counts = list(aspect_buckets.values()) + b = ax.bar(aspect_ratios, counts) + ax.bar_label(b, color=self.text_color) + if len(aspects) > 0: + sec = ax.secondary_xaxis(location=-0.1) + sec.spines["bottom"].set_linewidth(0) + sec.set_xticks([0, (len(aspects)-1)/2, len(aspects)-1], labels=["Wide", "Square", "Tall"]) + sec.tick_params('x', length=0) + + fig.tight_layout() + fig.subplots_adjust(bottom=0.15) + + # Render to PIL Image via buffer + buf = io.BytesIO() + fig.savefig(buf, format='png', facecolor=fig.get_facecolor(), edgecolor='none') + buf.seek(0) + pil_image = Image.open(buf).copy() + buf.close() + + # Update the CTkImage + self.chart_ctk_image.configure(light_image=pil_image, size=pil_image.size) + finally: + plt.close(fig) def decimal_to_aspect_ratio(self, value : float): #find closest fraction to decimal aspect value and convert to a:b format @@ -933,15 +945,6 @@ def __auto_update_concept_stats(self): self.__get_concept_stats(True, 2) #do advanced scan automatically if basic took <0.1s def destroy(self): - if self.canvas is not None: - self.canvas.get_tk_widget().destroy() - self.canvas = None - - if self.bucket_fig is not None: - plt.close(self.bucket_fig) - self.bucket_fig = None - self.bucket_ax = None - super().destroy() def __ok(self): From c26003162a4a486725197ee183f598c8f8d9667a Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Mon, 19 Jan 2026 02:26:37 +1100 Subject: [PATCH 60/64] Revert "Experimental test to see if FigureCanvasTkAgg is the issue" This reverts commit 71cde89069556287dc1af4b0476a2435e58aa260. --- modules/ui/ConceptWindow.py | 101 +++++++++++++++++------------------- 1 file changed, 49 insertions(+), 52 deletions(-) diff --git a/modules/ui/ConceptWindow.py b/modules/ui/ConceptWindow.py index 4e5caaf8e..babe84c00 100644 --- a/modules/ui/ConceptWindow.py +++ b/modules/ui/ConceptWindow.py @@ -1,5 +1,4 @@ import fractions -import io import math import os import pathlib @@ -44,11 +43,9 @@ import customtkinter as ctk import huggingface_hub -import matplotlib as mpl from customtkinter import AppearanceModeTracker, ThemeManager - -mpl.use('Agg') # Use non-interactive backend to avoid Linux DnD issues from matplotlib import pyplot as plt +from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg from PIL import Image @@ -94,6 +91,8 @@ def __init__( self.text_ui_state = text_ui_state self.image_preview_file_index = 0 self.preview_augmentations = ctk.BooleanVar(self, True) + self.bucket_fig = None + self.canvas = None self.title("Concept") self.geometry("800x700") @@ -509,19 +508,39 @@ def __concept_stats_tab(self, master): self.small_bucket_label.configure(font=ctk.CTkFont(underline=True)) self.small_bucket_preview = components.label(frame, 18, 1, pad=0, text="-") - #aspect bucketing plot - using Agg backend to render to image (avoids Linux DnD issues) + #aspect bucketing plot, mostly copied from timestep preview graph appearance_mode = AppearanceModeTracker.get_mode() background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) - self.chart_background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" + background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" self.text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" plt.set_loglevel('WARNING') #suppress errors about data type in bar chart - # Create a placeholder image for the chart - self.chart_ctk_image = ctk.CTkImage(light_image=Image.new('RGB', (700, 300), self.chart_background_color), size=(700, 300)) - self.chart_label = ctk.CTkLabel(frame, image=self.chart_ctk_image, text="") - self.chart_label.grid(row=19, column=0, columnspan=4, rowspan=2) + # Store previous figure to release once we're no longer using it + prev_bucket_fig = self.bucket_fig + + self.bucket_fig, self.bucket_ax = plt.subplots(figsize=(7,3)) + self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=frame) + self.canvas.get_tk_widget().grid(row=19, column=0, columnspan=4, rowspan=2) + self.bucket_fig.tight_layout() + self.bucket_fig.subplots_adjust(bottom=0.15) + + self.bucket_fig.set_facecolor(background_color) + self.bucket_ax.set_facecolor(background_color) + self.bucket_ax.spines['bottom'].set_color(self.text_color) + self.bucket_ax.spines['left'].set_color(self.text_color) + self.bucket_ax.spines['top'].set_visible(False) + self.bucket_ax.spines['right'].set_color(self.text_color) + self.bucket_ax.tick_params(axis='x', colors=self.text_color, which="both") + self.bucket_ax.tick_params(axis='y', colors=self.text_color, which="both") + self.bucket_ax.xaxis.label.set_color(self.text_color) + self.bucket_ax.yaxis.label.set_color(self.text_color) + + # close previous figure if any + if prev_bucket_fig is not None: + plt.close(prev_bucket_fig) + prev_bucket_fig = None #refresh stats - must be after all labels are defined or will give error self.refresh_basic_stats_button = components.button(master=frame, row=0, column=0, text="Refresh Basic", command=lambda: self.__get_concept_stats_threaded(False, 9999), @@ -835,48 +854,17 @@ def __update_concept_stats(self): min_bucket_str.strip() self.small_bucket_preview.configure(text=min_bucket_str) - self.__update_bucket_chart(aspect_buckets) - - def __update_bucket_chart(self, aspect_buckets: dict): - """Render the aspect bucket chart using Agg backend and update CTkImage.""" - fig, ax = plt.subplots(figsize=(7, 3)) - try: - fig.set_facecolor(self.chart_background_color) - ax.set_facecolor(self.chart_background_color) - ax.spines['bottom'].set_color(self.text_color) - ax.spines['left'].set_color(self.text_color) - ax.spines['top'].set_visible(False) - ax.spines['right'].set_color(self.text_color) - ax.tick_params(axis='x', colors=self.text_color, which="both") - ax.tick_params(axis='y', colors=self.text_color, which="both") - ax.xaxis.label.set_color(self.text_color) - ax.yaxis.label.set_color(self.text_color) - - aspects = [str(x) for x in list(aspect_buckets.keys())] - aspect_ratios = [self.decimal_to_aspect_ratio(x) for x in list(aspect_buckets.keys())] - counts = list(aspect_buckets.values()) - b = ax.bar(aspect_ratios, counts) - ax.bar_label(b, color=self.text_color) - if len(aspects) > 0: - sec = ax.secondary_xaxis(location=-0.1) - sec.spines["bottom"].set_linewidth(0) - sec.set_xticks([0, (len(aspects)-1)/2, len(aspects)-1], labels=["Wide", "Square", "Tall"]) - sec.tick_params('x', length=0) - - fig.tight_layout() - fig.subplots_adjust(bottom=0.15) - - # Render to PIL Image via buffer - buf = io.BytesIO() - fig.savefig(buf, format='png', facecolor=fig.get_facecolor(), edgecolor='none') - buf.seek(0) - pil_image = Image.open(buf).copy() - buf.close() - - # Update the CTkImage - self.chart_ctk_image.configure(light_image=pil_image, size=pil_image.size) - finally: - plt.close(fig) + self.bucket_ax.cla() + aspects = [str(x) for x in list(aspect_buckets.keys())] + aspect_ratios = [self.decimal_to_aspect_ratio(x) for x in list(aspect_buckets.keys())] + counts = list(aspect_buckets.values()) + b = self.bucket_ax.bar(aspect_ratios, counts) + self.bucket_ax.bar_label(b, color=self.text_color) + sec = self.bucket_ax.secondary_xaxis(location=-0.1) + sec.spines["bottom"].set_linewidth(0) + sec.set_xticks([0, (len(aspects)-1)/2, len(aspects)-1], labels=["Wide", "Square", "Tall"]) + sec.tick_params('x', length=0) + self.canvas.draw() def decimal_to_aspect_ratio(self, value : float): #find closest fraction to decimal aspect value and convert to a:b format @@ -945,6 +933,15 @@ def __auto_update_concept_stats(self): self.__get_concept_stats(True, 2) #do advanced scan automatically if basic took <0.1s def destroy(self): + if self.canvas is not None: + self.canvas.get_tk_widget().destroy() + self.canvas = None + + if self.bucket_fig is not None: + plt.close(self.bucket_fig) + self.bucket_fig = None + self.bucket_ax = None + super().destroy() def __ok(self): From f434a760ac9df1f7e4a3cafe1477230bb6789008 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Mon, 19 Jan 2026 04:30:47 +1100 Subject: [PATCH 61/64] Move statistics to default Agg backend instead of interactive one --- modules/ui/ConceptWindow.py | 127 +++++++++++++++++++++++----------- modules/util/concept_stats.py | 8 +-- 2 files changed, 92 insertions(+), 43 deletions(-) diff --git a/modules/ui/ConceptWindow.py b/modules/ui/ConceptWindow.py index babe84c00..c359a11ea 100644 --- a/modules/ui/ConceptWindow.py +++ b/modules/ui/ConceptWindow.py @@ -1,4 +1,5 @@ import fractions +import io import math import os import pathlib @@ -45,7 +46,6 @@ import huggingface_hub from customtkinter import AppearanceModeTracker, ThemeManager from matplotlib import pyplot as plt -from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg from PIL import Image @@ -92,7 +92,9 @@ def __init__( self.image_preview_file_index = 0 self.preview_augmentations = ctk.BooleanVar(self, True) self.bucket_fig = None - self.canvas = None + self._destroyed = False + self._last_bucket_data = None + self._last_chart_update = 0 self.title("Concept") self.geometry("800x700") @@ -521,10 +523,11 @@ def __concept_stats_tab(self, master): prev_bucket_fig = self.bucket_fig self.bucket_fig, self.bucket_ax = plt.subplots(figsize=(7,3)) - self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=frame) - self.canvas.get_tk_widget().grid(row=19, column=0, columnspan=4, rowspan=2) + self.bucket_image = ctk.CTkImage(light_image=Image.new('RGB', (700, 300)), size=(700, 300)) + self.bucket_label = ctk.CTkLabel(frame, text="", image=self.bucket_image) + self.bucket_label.grid(row=19, column=0, columnspan=4, rowspan=2) self.bucket_fig.tight_layout() - self.bucket_fig.subplots_adjust(bottom=0.15) + self.bucket_fig.subplots_adjust(bottom=0.18, left=0.1) self.bucket_fig.set_facecolor(background_color) self.bucket_ax.set_facecolor(background_color) @@ -537,6 +540,8 @@ def __concept_stats_tab(self, master): self.bucket_ax.xaxis.label.set_color(self.text_color) self.bucket_ax.yaxis.label.set_color(self.text_color) + self.__draw_empty_bucket_chart() + # close previous figure if any if prev_bucket_fig is not None: plt.close(prev_bucket_fig) @@ -756,6 +761,9 @@ def __get_preview_image(self): return image, filename_output, prompt_output def __update_concept_stats(self): + if self._destroyed: + return + #file size self.file_size_preview.configure(text=str(int(self.concept.concept_stats["file_size"]/1048576)) + " MB") self.processing_time.configure(text=str(round(self.concept.concept_stats["processing_time"], 2)) + " s") @@ -841,37 +849,72 @@ def __update_concept_stats(self): #aspect bucketing aspect_buckets = self.concept.concept_stats["aspect_buckets"] - if len(aspect_buckets) != 0 and max(val for val in aspect_buckets.values()) > 0: #check aspect_bucket data exists and is not all zero - min_val = min(val for val in aspect_buckets.values() if val > 0) #smallest nonzero values - if max(val for val in aspect_buckets.values()) > min_val: #check if any buckets larger than min_val exist - if all images are same aspect then there won't be - min_val2 = min(val for val in aspect_buckets.values() if (val > 0 and val != min_val)) #second smallest bucket - else: - min_val2 = min_val #if no second smallest bucket exists set to min_val - min_aspect_buckets = {key: val for key,val in aspect_buckets.items() if val in (min_val, min_val2)} - min_bucket_str = "" - for key, val in min_aspect_buckets.items(): - min_bucket_str += f'aspect {self.decimal_to_aspect_ratio(key)} : {val} img\n' - min_bucket_str.strip() - self.small_bucket_preview.configure(text=min_bucket_str) + if len(aspect_buckets) == 0 or max(aspect_buckets.values()) == 0: + self.small_bucket_preview.configure(text="-") + self.__draw_empty_bucket_chart() + return + + min_val = min(val for val in aspect_buckets.values() if val > 0) #smallest nonzero values + if max(aspect_buckets.values()) > min_val: #check if any buckets larger than min_val exist - if all images are same aspect then there won't be + min_val2 = min(val for val in aspect_buckets.values() if (val > 0 and val != min_val)) #second smallest bucket + else: + min_val2 = min_val #if no second smallest bucket exists set to min_val + min_aspect_buckets = {key: val for key, val in aspect_buckets.items() if val in (min_val, min_val2)} + min_bucket_str = "" + for key, val in min_aspect_buckets.items(): + min_bucket_str += f'aspect {self.decimal_to_aspect_ratio(key)} : {val} img\n' + self.small_bucket_preview.configure(text=min_bucket_str.strip()) + + current_time = time.perf_counter() + bucket_tuple = tuple(sorted(aspect_buckets.items())) + if (current_time - self._last_chart_update < 3.0 and + self._last_bucket_data == bucket_tuple): + return + + self._last_bucket_data = bucket_tuple + self._last_chart_update = current_time self.bucket_ax.cla() - aspects = [str(x) for x in list(aspect_buckets.keys())] - aspect_ratios = [self.decimal_to_aspect_ratio(x) for x in list(aspect_buckets.keys())] + aspect_ratios = [self.decimal_to_aspect_ratio(x) for x in aspect_buckets] counts = list(aspect_buckets.values()) b = self.bucket_ax.bar(aspect_ratios, counts) self.bucket_ax.bar_label(b, color=self.text_color) - sec = self.bucket_ax.secondary_xaxis(location=-0.1) - sec.spines["bottom"].set_linewidth(0) - sec.set_xticks([0, (len(aspects)-1)/2, len(aspects)-1], labels=["Wide", "Square", "Tall"]) - sec.tick_params('x', length=0) - self.canvas.draw() + self.bucket_ax.set_xlabel("Aspect Ratio (width:height)", color=self.text_color) + self.bucket_ax.set_ylabel("Image Count", color=self.text_color) - def decimal_to_aspect_ratio(self, value : float): + if len(aspect_ratios) > 1: + sec = self.bucket_ax.secondary_xaxis(location=-0.1) + sec.spines["bottom"].set_linewidth(0) + sec.set_xticks([0, (len(aspect_ratios)-1)/2, len(aspect_ratios)-1], labels=["Wide", "Square", "Tall"]) + sec.tick_params('x', length=0) + + self.__render_bucket_chart() + + def decimal_to_aspect_ratio(self, value: float): #find closest fraction to decimal aspect value and convert to a:b format aspect_fraction = fractions.Fraction(value).limit_denominator(16) aspect_string = f'{aspect_fraction.denominator}:{aspect_fraction.numerator}' return aspect_string + def __draw_empty_bucket_chart(self): + self.bucket_ax.cla() + self.bucket_ax.set_xlabel("Aspect Ratio (width:height)", color=self.text_color) + self.bucket_ax.set_ylabel("Image Count", color=self.text_color) + self.bucket_ax.set_xlim(-0.5, 4.5) + self.bucket_ax.set_ylim(0, 10) + self.bucket_ax.text(2, 5, "Run Advanced Scan\nto view aspect buckets", + ha='center', va='center', color=self.text_color, + fontsize=10, style='italic') + self.__render_bucket_chart() + + def __render_bucket_chart(self): + buf = io.BytesIO() + self.bucket_fig.savefig(buf, format='png', facecolor=self.bucket_fig.get_facecolor()) + buf.seek(0) + img = Image.open(buf) + self.bucket_image.configure(light_image=img, size=(img.width, img.height)) + buf.close() + def __get_concept_stats(self, advanced_checks: bool, wait_time: float): if not os.path.isdir(self.concept.path): print(f"Unable to get statistics for invalid concept path: {self.concept.path}") @@ -890,30 +933,35 @@ def __get_concept_stats(self, advanced_checks: bool, wait_time: float): stats_dict = concept_stats.init_concept_stats(advanced_checks) for path in subfolders: - if self.cancel_scan_flag.is_set() or time.perf_counter() - start_time > wait_time: + if self._destroyed or self.cancel_scan_flag.is_set() or time.perf_counter() - start_time > wait_time: break stats_dict = concept_stats.folder_scan(path, stats_dict, advanced_checks, self.concept, start_time, wait_time, self.cancel_scan_flag) - if self.concept.include_subdirectories and not self.cancel_scan_flag.is_set(): #add all subfolders of current directory to for loop + if self.concept.include_subdirectories and not self.cancel_scan_flag.is_set() and not self._destroyed: #add all subfolders of current directory to for loop subfolders.extend([f for f in os.scandir(path) if f.is_dir()]) self.concept.concept_stats = stats_dict - #update GUI approx every half second - if time.perf_counter() > (last_update + 0.5): + #update GUI approx every second + if time.perf_counter() > (last_update + 1): last_update = time.perf_counter() self.concept_stats_tab.after(0, self.__update_concept_stats) self.cancel_scan_flag.clear() - self.concept_stats_tab.after(0, self.__enable_scan_buttons) - self.concept_stats_tab.after(0, self.__update_concept_stats) + if not self._destroyed: + self.concept_stats_tab.after(0, self.__enable_scan_buttons) + self.concept_stats_tab.after(0, self.__update_concept_stats) def __get_concept_stats_threaded(self, advanced_checks : bool, waittime : float): self.scan_thread = threading.Thread(target=self.__get_concept_stats, args=[advanced_checks, waittime], daemon=True) self.scan_thread.start() def __disable_scan_buttons(self): + if self._destroyed: + return self.refresh_basic_stats_button.configure(state="disabled") self.refresh_advanced_stats_button.configure(state="disabled") def __enable_scan_buttons(self): + if self._destroyed: + return self.refresh_basic_stats_button.configure(state="normal") self.refresh_advanced_stats_button.configure(state="normal") @@ -921,21 +969,22 @@ def __cancel_concept_stats(self): self.cancel_scan_flag.set() def __auto_update_concept_stats(self): + if self._destroyed: + return try: - self.__update_concept_stats() #load stats from config if available, else raises KeyError + self.after(0, self.__update_concept_stats) #load stats from config if available, else raises KeyError if self.concept.concept_stats["file_size"] == 0: #force rescan if empty raise KeyError except KeyError: concept_path = self.get_concept_path(self.concept.path) - if concept_path: - self.__get_concept_stats(False, 2) #force rescan if config is empty, timeout of 2 sec - if self.concept.concept_stats["processing_time"] < 0.1: - self.__get_concept_stats(True, 2) #do advanced scan automatically if basic took <0.1s + if concept_path and not self._destroyed: + # Run advanced scan directly with short timeout - avoids scanning twice + # since we'd do advanced anyway if basic was fast + self.__get_concept_stats(True, 2) def destroy(self): - if self.canvas is not None: - self.canvas.get_tk_widget().destroy() - self.canvas = None + self._destroyed = True + self.cancel_scan_flag.set() if self.bucket_fig is not None: plt.close(self.bucket_fig) diff --git a/modules/util/concept_stats.py b/modules/util/concept_stats.py index f0abb30a6..3e1df0961 100644 --- a/modules/util/concept_stats.py +++ b/modules/util/concept_stats.py @@ -93,7 +93,7 @@ def folder_scan(dir, stats_dict : dict, advanced_checks : bool, conceptconfig : file_list = [f for f in os.scandir(dir) if f.is_file()] #this may take time on large directories if advanced_checks: aspect_ratio_list = list(stats_dict["aspect_buckets"].keys()) - file_list_str = [x.path for x in file_list] #seems faster to check list of strings for matching files than list of path objects + file_list_set = {x.path for x in file_list} #use set for O(1) lookup instead of O(n) list search stats_dict["directory_count"] += 1 for path in file_list: @@ -107,10 +107,10 @@ def folder_scan(dir, stats_dict : dict, advanced_checks : bool, conceptconfig : stats_dict["file_size"] += path.stat().st_size if advanced_checks: #check if image has a corresponding mask/caption in the same directory - if (basename + "-masklabel.png") in file_list_str: + if (basename + "-masklabel.png") in file_list_set: stats_dict["paired_masks"] += 1 stats_dict["image_with_mask_count"] += 1 - if (basename + ".txt") in file_list_str: + if (basename + ".txt") in file_list_set: stats_dict["paired_captions"] += 1 stats_dict["image_with_caption_count"] += 1 with open(basename + ".txt", "r") as captionfile: @@ -155,7 +155,7 @@ def folder_scan(dir, stats_dict : dict, advanced_checks : bool, conceptconfig : # if (basename + "-masklabel.png") in file_list_str: # stats_dict["paired_masks"] += 1 # stats_dict["video_with_mask_count"] += 1 - if (basename + ".txt") in file_list_str: + if (basename + ".txt") in file_list_set: stats_dict["paired_captions"] += 1 stats_dict["video_with_caption_count"] += 1 with open(basename + ".txt", "r") as captionfile: From a145528ab92601032eb2bafce6326187fcb6f9b3 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Mon, 19 Jan 2026 04:48:31 +1100 Subject: [PATCH 62/64] Skip scanning for invalid concepts --- modules/ui/ConceptWindow.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/modules/ui/ConceptWindow.py b/modules/ui/ConceptWindow.py index c359a11ea..950f61461 100644 --- a/modules/ui/ConceptWindow.py +++ b/modules/ui/ConceptWindow.py @@ -764,6 +764,10 @@ def __update_concept_stats(self): if self._destroyed: return + # Check if concept_stats is properly initialized + if not self.concept.concept_stats.get("file_size"): + return + #file size self.file_size_preview.configure(text=str(int(self.concept.concept_stats["file_size"]/1048576)) + " MB") self.processing_time.configure(text=str(round(self.concept.concept_stats["processing_time"], 2)) + " s") @@ -971,11 +975,16 @@ def __cancel_concept_stats(self): def __auto_update_concept_stats(self): if self._destroyed: return + + # Skip scanning for empty/invalid concept paths + if not self.concept.path or not self.concept.path.strip(): + return + try: self.after(0, self.__update_concept_stats) #load stats from config if available, else raises KeyError if self.concept.concept_stats["file_size"] == 0: #force rescan if empty raise KeyError - except KeyError: + except (KeyError, AttributeError, TypeError): concept_path = self.get_concept_path(self.concept.path) if concept_path and not self._destroyed: # Run advanced scan directly with short timeout - avoids scanning twice From 79498d82b4b011d56f43a9883cb573496c915ec6 Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Wed, 4 Feb 2026 20:08:08 +1100 Subject: [PATCH 63/64] Disable DnD on linux due to unstable behaviour --- modules/ui/TrainUI.py | 18 +++++++++++++++--- modules/util/ui/ui_utils.py | 7 +++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index cb591d81f..13c5a2191 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -47,7 +47,16 @@ import customtkinter as ctk from customtkinter import AppearanceModeTracker -from tkinterdnd2 import TkinterDnD + +# tkinterdnd2 is unstable on Linux - conditionally import and provide no-op mixin +# https://github.com/Eliav2/tkinterdnd2/issues/12#issuecomment-3776598066 +if platform.system() != "Linux": + from tkinterdnd2 import TkinterDnD + _DnDMixin = TkinterDnD.DnDWrapper +else: + TkinterDnD = None + class _DnDMixin: + """No-op mixin when DnD is disabled on Linux""" # chunk for forcing Windows to ignore DPI scaling when moving between monitors # fixes the long standing transparency bug https://github.com/Nerogar/OneTrainer/issues/90 @@ -56,7 +65,7 @@ # https://learn.microsoft.com/en-us/windows/win32/hidpi/setting-the-default-dpi-awareness-for-a-process#setting-default-awareness-programmatically ctypes.windll.shcore.SetProcessDpiAwareness(1) # PROCESS_SYSTEM_DPI_AWARE -class TrainUI(ctk.CTk, TkinterDnD.DnDWrapper): +class TrainUI(ctk.CTk, _DnDMixin): set_step_progress: Callable[[int, int], None] set_epoch_progress: Callable[[int, int], None] @@ -65,6 +74,8 @@ class TrainUI(ctk.CTk, TkinterDnD.DnDWrapper): training_callbacks: TrainCallbacks | None training_commands: TrainCommands | None + _DND_ENABLED = TkinterDnD is not None + _TRAIN_BUTTON_STYLES = { "idle": { "text": "Start Training", @@ -93,7 +104,8 @@ class TrainUI(ctk.CTk, TkinterDnD.DnDWrapper): def __init__(self): super().__init__() - self.TkdndVersion = TkinterDnD._require(self) + if self._DND_ENABLED: + self.TkdndVersion = TkinterDnD._require(self) self.title("OneTrainer") self.geometry("1100x740") diff --git a/modules/util/ui/ui_utils.py b/modules/util/ui/ui_utils.py index b3c70acc6..cb490b38b 100644 --- a/modules/util/ui/ui_utils.py +++ b/modules/util/ui/ui_utils.py @@ -176,6 +176,10 @@ def _parse_dropped_paths(event_data: str) -> list[str]: def register_drop_target(entry_widget, ui_state, var_name, command=None, drop_validator=None, on_reject=None): + # tkinterdnd2 is unstable on Linux https://github.com/Eliav2/tkinterdnd2/issues/12#issuecomment-3776598066 + if sys.platform == "linux": + return + try: entry_widget.drop_target_register(DND_FILES) for event, handler in [('<>', _drop_enter), ('<>', _drop_leave), @@ -185,6 +189,9 @@ def register_drop_target(entry_widget, ui_state, var_name, command=None, drop_va print(f"Failed to register drop target: {e}") def register_concept_drop_target(widget, drop_callback: Callable[[str], None], allow_multiple: bool = True): + if sys.platform == "linux": + return + def drop_handler(event): if not event.data: return event.action From 935dc0ef18522468023b3de8e8a1f850378ddfdc Mon Sep 17 00:00:00 2001 From: O-J1 <18110006+O-J1@users.noreply.github.com> Date: Wed, 4 Feb 2026 23:37:26 +1100 Subject: [PATCH 64/64] Add undo/redo. Remove revert behaviour. Prevent training startining with invalid values. Logic tweaks. --- modules/ui/TrainUI.py | 9 ++ modules/ui/VideoToolUI.py | 2 +- modules/util/ui/ToolTip.py | 11 ++ modules/util/ui/UIState.py | 37 +++++ modules/util/ui/components.py | 246 +++++++++++++++++++++++++--------- modules/util/ui/ui_utils.py | 18 +++ modules/util/ui/validation.py | 24 +++- 7 files changed, 277 insertions(+), 70 deletions(-) diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index 13c5a2191..572de5a9d 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -40,6 +40,7 @@ from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress from modules.util.ui import components +from modules.util.ui.components import ValidationErrorDialog from modules.util.ui.ui_utils import set_window_icon from modules.util.ui.UIState import UIState @@ -931,6 +932,14 @@ def __training_thread_function(self): def start_training(self): if self.training_thread is None: + self.ui_state.flush_all_validations() + + # Check for validation errors before starting + invalid_fields = self.ui_state.get_invalid_fields() + if invalid_fields: + ValidationErrorDialog(self, invalid_fields) + return + self.save_default() self._set_training_button_running() diff --git a/modules/ui/VideoToolUI.py b/modules/ui/VideoToolUI.py index 96970fdb6..caaee6ff3 100644 --- a/modules/ui/VideoToolUI.py +++ b/modules/ui/VideoToolUI.py @@ -281,7 +281,7 @@ def __video_download_tab(self, master): components.label(frame, 3, 0, "Additional Args", tooltip="Any additional arguments to pass to yt-dlp, for example '--restrict-filenames --force-overwrite'. \ Default args will hide most terminal outputs.") - self.download_args_entry = ctk.CTkTextbox(frame, width=220, height=90, border_width=2) + self.download_args_entry = ctk.CTkTextbox(frame, width=220, height=90, border_width=2, undo=True, autoseparators=True, maxundo=-1) self.download_args_entry.grid(row=3, column=1, rowspan=2, sticky="w", padx=5, pady=5) self.download_args_entry.insert(index="1.0", text="--quiet --no-warnings --progress") components.button(frame, 3, 2, "yt-dlp info", diff --git a/modules/util/ui/ToolTip.py b/modules/util/ui/ToolTip.py index 60c38cb09..dd2af0f1a 100644 --- a/modules/util/ui/ToolTip.py +++ b/modules/util/ui/ToolTip.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import contextlib import customtkinter as ctk @@ -5,6 +7,7 @@ class ToolTip: MAX_TOOLTIP_LENGTH = 350 + _active_tooltip: ToolTip | None = None # currently visible tooltip def __init__(self, widget, text='widget info', x_position=20, wide=False, hover_only=True, track_movement=False): @@ -53,9 +56,14 @@ def show(self, text=None, color=None, duration_ms=None, fg_color="#2b2b2b"): if text: self.text = text[:self.MAX_TOOLTIP_LENGTH] + ("..." if len(text) > self.MAX_TOOLTIP_LENGTH else "") + if ToolTip._active_tooltip is not None and ToolTip._active_tooltip is not self: + ToolTip._active_tooltip.hide() + self._cancel_scheduled_hide() self.hide() + ToolTip._active_tooltip = self + if self.track_movement and not self._bound: self._root = self.widget.winfo_toplevel() self._root.bind("", self._on_move, add="+") @@ -134,6 +142,9 @@ def hide(self, _evt=None): self.tw.destroy() self.tw = None + if ToolTip._active_tooltip is self: + ToolTip._active_tooltip = None + def show_error(self, message: str, duration_ms: int | None = 7000): self.show(message, color="#ff6b6b", duration_ms=duration_ms) diff --git a/modules/util/ui/UIState.py b/modules/util/ui/UIState.py index 3f7fc3579..331dc0d07 100644 --- a/modules/util/ui/UIState.py +++ b/modules/util/ui/UIState.py @@ -13,6 +13,8 @@ class UIState: __vars: dict[str, Any] __var_traces: dict[str, dict[int, Callable[[], None]]] __latest_var_trace_id: int + __validation_states: dict[str, Any] # Maps field names to ValidationState objects + __validation_handlers: dict[str, Any] # Maps field names to validation handler objects with debounce timers def __init__(self, master, obj): self.master = master @@ -21,6 +23,8 @@ def __init__(self, master, obj): self.__var_types: dict[str, type] = {} self.__var_nullables: dict[str, bool] = {} self.__var_defaults: dict[str, Any] = {} + self.__validation_states = {} + self.__validation_handlers = {} self.__vars = self.__create_vars(obj) self.__var_traces = {name: {} for name in self.__vars} @@ -311,3 +315,36 @@ def get_field_metadata(self, name: str) -> "UIState.VarMeta": state.__var_nullables.get(leaf, False), state.__var_defaults.get(leaf, None), ) + + # validation state tracking + def register_validation_state(self, field_name: str, validation_state: Any) -> None: + if validation_state is not None: + self.__validation_states[field_name] = validation_state + + def unregister_validation_state(self, field_name: str) -> None: + self.__validation_states.pop(field_name, None) + + def get_invalid_fields(self) -> list[tuple[str, str]]: + return [ + (name, getattr(s, 'message', 'Invalid value')) + for name, s in self.__validation_states.items() + if getattr(s, 'status', None) == 'error' + ] + + # validation handler tracking (for flushing debounced validations) + def register_validation_handler(self, field_name: str, handler: Any) -> None: + if handler is not None: + self.__validation_handlers[field_name] = handler + + def unregister_validation_handler(self, field_name: str) -> None: + self.__validation_handlers.pop(field_name, None) + + def flush_all_validations(self) -> None: + """Flush all pending debounced validations to ensure validation state is current.""" + for handler in self.__validation_handlers.values(): + # Flush EntryValidationHandler debounce timers + if hasattr(handler, 'debounce_timer') and handler.debounce_timer: + handler.debounce_timer.flush() + # Flush ModelOutputValidator debounce timers + if hasattr(handler, '_debounce_timer') and handler._debounce_timer: + handler._debounce_timer.flush() diff --git a/modules/util/ui/components.py b/modules/util/ui/components.py index 0ae2cbdb2..0648ee50d 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -3,7 +3,7 @@ from collections.abc import Callable from dataclasses import dataclass from pathlib import Path -from tkinter import filedialog +from tkinter import filedialog, messagebox from typing import Any, Literal from modules.util.enum.ModelFormat import ModelFormat @@ -45,12 +45,97 @@ def _wrap_dropdown_destroy(component): @dataclass(frozen=True) class ComponentValidationSettings: debounce_stop_typing_ms: int = 1700 - debounced_invalid_revert_ms: int = 1000 - focusout_invalid_revert_ms: int = 1200 COMPONENT_VALIDATION_SETTINGS = ComponentValidationSettings() + +class EntryUndoHandler: + """Provides undo/redo (Ctrl+Z / Ctrl+Shift+Z) for CTkEntry widgets.""" + + def __init__(self, component: ctk.CTkEntry, var: tk.Variable, debounce_ms: int = 500, maxundo: int = -1): + self.component = component + self.var = var + self._undo: list[str] = [var.get()] + self._redo: list[str] = [] + self._last = self._undo[0] + self._in_undo = False + self._debounce_id = None + self._debounce_ms = debounce_ms + self._maxundo = maxundo + + self._trace_id = var.trace_add("write", self._on_change) + for key in ("", ""): + component.bind(key, self._do_undo) + for key in ("", "", "", ""): + component.bind(key, self._do_redo) + + def _on_change(self, *_): + if self._in_undo: + return + current = self.var.get() + if current == self._last: + return + + self._redo.clear() + self._last = current + + # Debounce: cancel pending, schedule new snapshot + if self._debounce_id: + self.component.after_cancel(self._debounce_id) + self._debounce_id = self.component.after(self._debounce_ms, self._snapshot) + + def _snapshot(self): + self._debounce_id = None + if self._undo[-1] != self._last: + self._undo.append(self._last) + # Enforce maxundo limit (-1 means unlimited) + if self._maxundo > 0 and len(self._undo) > self._maxundo + 1: + self._undo.pop(0) + + def _do_undo(self, _=None): + if len(self._undo) <= 1: + return "break" + self._in_undo = True + if self._undo[-1] != self._last: + self._undo.append(self._last) + self._redo.append(self._undo.pop()) + self._last = self._undo[-1] + self.var.set(self._last) + self._in_undo = False + return "break" + + def _do_redo(self, _=None): + if not self._redo: + return "break" + self._in_undo = True + state = self._redo.pop() + self._undo.append(state) + self._last = state + self.var.set(state) + self._in_undo = False + return "break" + + def cleanup(self): + if self._debounce_id: + with contextlib.suppress(Exception): + self.component.after_cancel(self._debounce_id) + with contextlib.suppress(Exception): + self.var.trace_remove("write", self._trace_id) + + +class ValidationErrorDialog(ctk.CTkToplevel): + def __init__(self, parent, invalid_fields: list[tuple[str, str]], *args, **kwargs): + + lines = ["Cannot start training. Please fix the following errors:\n"] + for field_name, message in invalid_fields: + display_name = field_name.replace("_", " ").title() + lines.append(f"• {display_name}: {message}") + + error_text = "\n".join(lines) + + messagebox.showerror("Validation Errors", error_text, parent=parent) + @dataclass class ValidationState: status: Literal['error', 'warning'] | None = None @@ -90,23 +175,14 @@ def __init__( self.debounce_timer = DebounceTimer( widget=component, delay_ms=COMPONENT_VALIDATION_SETTINGS.debounce_stop_typing_ms, - callback=lambda: self.validate_value( - self.var.get(), - COMPONENT_VALIDATION_SETTINGS.debounced_invalid_revert_ms - ) + callback=lambda: self.validate_value(self.var.get()) ) - self.revert_after_id = None self.touched = False - self.last_valid_value = var.get() self.validation_tooltip = ToolTip(component, text="", hover_only=False, track_movement=True, wide=True) component._validation_tooltip = self.validation_tooltip - def _cancel_after(self, after_id): - with contextlib.suppress(Exception): - self.component.after_cancel(after_id) - def _reset_border(self): self.component.configure(border_color=self.original_border_color) @@ -127,8 +203,7 @@ def _show_validation_tooltip(self): elif self.validation_state.status == 'warning': self.validation_tooltip.show_warning(self.validation_state.message, duration_ms=None) - def validate_value(self, value: str, revert_delay_ms: int | None) -> bool: - self._cancel_after(self.revert_after_id) + def validate_value(self, value: str) -> bool: meta = self.ui_state.get_field_metadata(self.var_name) basic_result = validate_basic_type( @@ -139,7 +214,7 @@ def validate_value(self, value: str, revert_delay_ms: int | None) -> bool: ) if not basic_result.ok: - return self._fail(basic_result.message, revert_delay_ms) + return self._fail(basic_result.message) if self.custom_validator: result = self.custom_validator(value) @@ -148,28 +223,28 @@ def validate_value(self, value: str, revert_delay_ms: int | None) -> bool: if result.status == 'warning': return self._warning(result.message, value) elif not result.ok: - return self._fail(result.message, revert_delay_ms) + return self._fail(result.message) return self._success(value) def _success(self, value: str) -> bool: self._reset_border() self.validation_tooltip.hide() - self.last_valid_value = value + if self.validation_state: + self.validation_state.clear() return True def _warning(self, reason: str, value: str) -> bool: self.component.configure(border_color="#ff9500") if self.should_show_tooltip(): self.validation_tooltip.show_warning(reason) - self.last_valid_value = value if self.validation_state: self.validation_state.set_status('warning', reason) return True - def _fail(self, reason: str, revert_delay_ms: int | None) -> bool: + def _fail(self, reason: str) -> bool: self.component.configure(border_color="#dc3545") if self.should_show_tooltip(): self.validation_tooltip.show_error(reason) @@ -177,17 +252,8 @@ def _fail(self, reason: str, revert_delay_ms: int | None) -> bool: if self.validation_state: self.validation_state.set_status('error', reason) - if revert_delay_ms is not None: - self.revert_after_id = self.component.after(revert_delay_ms, self._do_revert) - else: - self._do_revert() return False - def _do_revert(self): - self.var.set(self.last_valid_value) - self._reset_border() - self.validation_tooltip.hide() - def debounced_validate(self, *_): if not self.touched: return @@ -195,9 +261,6 @@ def debounced_validate(self, *_): self.validation_tooltip.hide() self._reset_border() - self._cancel_after(self.revert_after_id) - - # Use DebounceTimer instead of manual after() calls self.debounce_timer.call() def on_focus_in(self, _e=None): @@ -211,14 +274,14 @@ def on_focus_out(self, _e=None): self.validation_tooltip.hide() if self.touched: - self.validate_value( - self.var.get(), - COMPONENT_VALIDATION_SETTINGS.focusout_invalid_revert_ms - ) + self.validate_value(self.var.get()) def cleanup(self): - self._cancel_after(self.revert_after_id) - # No need to cancel debounce_timer - it handles its own cleanup + self.debounce_timer.cancel() + with contextlib.suppress(Exception): + self.validation_tooltip.hide() + self.debounce_timer = None + self.validation_tooltip = None class ModelOutputValidator: @@ -250,41 +313,56 @@ def __init__( self._tk_widget = None self._last_prefix = self.prefix_var.get() if self.prefix_var else "" + self._in_validation = False + self._applying_correction = False + def _get_enum_value(self, var, default_enum): with contextlib.suppress(KeyError, ValueError): return type(default_enum)[var.get()] return default_enum def validate(self, value: str, skip_overwrite_protection: bool = False, readonly: bool = False) -> ValidationResult: + # Guard recursion + if self._in_validation or self._applying_correction: + return ValidationResult(ok=True, corrected=None, message="", status='success') + value = value.strip() if not value: self.state.clear() return ValidationResult(ok=True, corrected=None, message="", status='success') - # Use validation logic from validation.py - result = validate_destination( - value, - output_format=self._get_enum_value(self.format_var, ModelFormat.SAFETENSORS), - training_method=self._get_enum_value(self.method_var, TrainingMethod.FINE_TUNE), - autocorrect=False if readonly else _safe_bool(self.autocorrect_var), - prefix=self.prefix_var.get() if self.prefix_var else "", - use_friendly_names=_safe_bool(self.friendly_names_var, default=False), - is_output=True, - prevent_overwrite=False if readonly else _safe_bool(self.prevent_overwrite_var, default=False), - auto_prefix=False if readonly else _safe_bool(self.auto_prefix_var, default=False), - skip_overwrite_protection=skip_overwrite_protection, - ) + self._in_validation = True + try: + # Use validation logic from validation.py + result = validate_destination( + value, + output_format=self._get_enum_value(self.format_var, ModelFormat.SAFETENSORS), + training_method=self._get_enum_value(self.method_var, TrainingMethod.FINE_TUNE), + autocorrect=False if readonly else _safe_bool(self.autocorrect_var), + prefix=self.prefix_var.get() if self.prefix_var else "", + use_friendly_names=_safe_bool(self.friendly_names_var, default=False), + is_output=True, + prevent_overwrite=False if readonly else _safe_bool(self.prevent_overwrite_var, default=False), + auto_prefix=False if readonly else _safe_bool(self.auto_prefix_var, default=False), + skip_overwrite_protection=skip_overwrite_protection, + ) - if result.status: - self.state.set_status(result.status, result.message) - else: - self.state.clear() + if result.status: + self.state.set_status(result.status, result.message) + else: + self.state.clear() - # Only apply corrections if not in readonly mode - if not readonly and result.corrected and result.corrected != value: - self.var.set(result.corrected) + # Only apply corrections if not in readonly mode + if not readonly and result.corrected and result.corrected != value: + self._applying_correction = True + try: + self.var.set(result.corrected) + finally: + self._applying_correction = False - return result + return result + finally: + self._in_validation = False def _remove_old_prefix_if_needed(self): current_value = self.var.get().strip() @@ -342,11 +420,19 @@ def setup_traces(self): ] def make_callback(debounce: bool, skip_overwrite: bool, is_prefix: bool): - if is_prefix: - return lambda *_: self._schedule_prefix_change_validation() - elif debounce: - return lambda *_: self._schedule_validation(skip_overwrite) - return lambda *_: self.validate(self.var.get(), skip_overwrite_protection=skip_overwrite) + def callback(*_): + if self._applying_correction: + return + + if is_prefix: + self._schedule_prefix_change_validation() + elif debounce: + self._schedule_validation(skip_overwrite) + else: + if self._debounce_timer: + self._debounce_timer.cancel() + self.validate(self.var.get(), skip_overwrite_protection=skip_overwrite) + return callback for key, var, debounce, skip_overwrite, is_prefix in trace_config: self._trace_ids[key] = var.trace_add("write", make_callback(debounce, skip_overwrite, is_prefix)) @@ -397,6 +483,7 @@ def entry( custom_validator: Callable[[str], ValidationResult] | None = None, validation_state: ValidationState = None, enable_hover_validation: bool = False, + undo: bool = True, ): var = ui_state.get_var(var_name) @@ -407,6 +494,14 @@ def entry( component = ctk.CTkEntry(master, textvariable=var, width=width) component.grid(row=row, column=column, padx=PAD, pady=PAD, sticky=sticky) + # set up undo/redo handler if enabled, its enabled by default + undo_handler = EntryUndoHandler(component, var) if undo else None + + if validation_state is None: + validation_state = ValidationState() + + ui_state.register_validation_state(var_name, validation_state) # register global state + validation_handler = EntryValidationHandler( component=component, var=var, @@ -415,6 +510,7 @@ def entry( custom_validator=custom_validator, validation_state=validation_state, ) + ui_state.register_validation_handler(var_name, validation_handler) # register for debounce flushing validation_trace_name = var.trace_add("write", validation_handler.debounced_validate) component.bind("", validation_handler.on_focus_in) @@ -453,9 +549,16 @@ def new_destroy(): component._textvariable.trace_remove("write", component._textvariable_callback_name) component._textvariable_callback_name = "" + if undo_handler: + undo_handler.cleanup() + validation_handler.cleanup() var.trace_remove("write", validation_trace_name) + # Unregister validation state and handler + ui_state.unregister_validation_state(var_name) + ui_state.unregister_validation_handler(var_name) + if command is not None and trace_id is not None: ui_state.remove_var_trace(var_name, trace_id) @@ -508,7 +611,7 @@ def path_entry( custom_validator = validator.validate else: validator = None - validation_state = None + validation_state = ValidationState() # Create state for simple path validation def simple_path_validator(value: str) -> ValidationResult: if value and path_type == "directory": @@ -520,10 +623,18 @@ def simple_path_validator(value: str) -> ValidationResult: trimmed_path = f"{parent_str}{sep}{trimmed_name}" if parent_str != '.' else trimmed_name var.set(trimmed_path) value = trimmed_path - return validate_file_path(value, is_output, valid_extensions, path_type) + result = validate_file_path(value, is_output, valid_extensions, path_type) + + if result.ok: + validation_state.clear() + else: + validation_state.set_status('error', result.message) + return result custom_validator = simple_path_validator + ui_state.register_validation_state(var_name, validation_state) + entry_widget = entry( frame, row=0, column=0, ui_state=ui_state, var_name=var_name, custom_validator=custom_validator, @@ -535,12 +646,15 @@ def simple_path_validator(value: str) -> ValidationResult: if validator: validator.set_widget(entry_widget) validator.setup_traces() + ui_state.register_validation_handler(f"{var_name}_model_validator", validator) # register for debounce flushing original_destroy = entry_widget.destroy def new_destroy(): if validator: validator.cleanup_traces() + ui_state.unregister_validation_handler(f"{var_name}_model_validator") + ui_state.unregister_validation_state(var_name) original_destroy() entry_widget.destroy = new_destroy diff --git a/modules/util/ui/ui_utils.py b/modules/util/ui/ui_utils.py index cb490b38b..5dbfd89f0 100644 --- a/modules/util/ui/ui_utils.py +++ b/modules/util/ui/ui_utils.py @@ -233,3 +233,21 @@ def fire(): with contextlib.suppress(tk.TclError): self._after_id = self.widget.after(self.delay_ms, fire) + + def cancel(self): + """Cancel any pending debounced callback.""" + if self._after_id: + with contextlib.suppress(tk.TclError): + self.widget.after_cancel(self._after_id) + self._after_id = None + + def flush(self): + """Execute any pending callback immediately and cancel the scheduled timer.""" + if self._after_id: + with contextlib.suppress(tk.TclError): + self.widget.after_cancel(self._after_id) + after_id = self._after_id + self._after_id = None + # Only execute if pending callback + if after_id: + self.callback() diff --git a/modules/util/ui/validation.py b/modules/util/ui/validation.py index ff63fa059..da6cbb85b 100644 --- a/modules/util/ui/validation.py +++ b/modules/util/ui/validation.py @@ -26,7 +26,7 @@ class ValidationSettings: ENDS_WITH_EXT = re.compile(r"\.[A-Za-z0-9]+$") INVALID_NAMES = {"", "."} GENERIC_MODEL_NAMES = {"lora", "embedding", "embeddings", "model", "finetune"} - +TIMESTAMP_PATTERN = re.compile(r"[_-](\d{8}_\d{6}|\d{8})(?=[_.-]|$)") # org/repo-name (alphanumeric, hyphens, underscores, dots allowed) HUGGINGFACE_REPO_RE = re.compile(r"^[A-Za-z0-9_.-]+\/[A-Za-z0-9_.-]+$") @@ -133,6 +133,14 @@ def _has_prefix(stem: str, prefix: str) -> bool: return (stem_lower == prefix_lower or stem_lower.startswith((f"{prefix_lower}-", f"{prefix_lower}_"))) +def _has_timestamp(stem: str) -> bool: + return bool(TIMESTAMP_PATTERN.search(stem)) + +def _extension_is_substring(current_ext: str, required_ext: str) -> bool: + current_lower = current_ext.lower() + required_lower = required_ext.lower() + return current_lower != required_lower and required_lower.startswith(current_lower) + def generate_default_filename( training_method: TrainingMethod, prefix: str = "", @@ -305,6 +313,14 @@ def _validate_single_file_path(ctx: ValidationContext) -> ValidationResult: chars_after_dot = len(current_ext) - 1 start_idx = ext_match.start() + # Check if current extension is a partial match of required (e.g., .safetensor vs .safetensors) + if _extension_is_substring(current_ext, required_ext): + if ctx.autocorrect: + new_base = base[:start_idx] + required_ext + new_path = _format_path(path.with_name(new_base), ctx.separator) + return _result(True, new_path, f"Completed {current_ext} -> {required_ext}") + return _result(False, None, f"Extension is incomplete; expected {required_ext}") + if chars_after_dot >= 3: if _is_partial_match(current_ext, required_ext): if ctx.autocorrect: @@ -323,8 +339,10 @@ def _validate_single_file_path(ctx: ValidationContext) -> ValidationResult: return _result(False, None, f"Extension {current_ext} is too short; Type at least 3 characters for auto-correction to occur.") if ctx.autocorrect: - corrected = _format_path(Path(path.parent, base + required_ext), ctx.separator) - return _result(True, corrected, f"Appended {required_ext} extension.") + # Replace unrecognized extension with required + new_base = base[:start_idx] + required_ext + new_path = _format_path(path.with_name(new_base), ctx.separator) + return _result(True, new_path, f"Replaced {current_ext} with {required_ext}") return _result(False, None, f"Filename must end with {required_ext}") # No extension - append required extension