diff --git a/.gitignore b/.gitignore index dce73072c..0a0b43032 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,9 @@ debug*.ipynb debug*.py /debug* +profile.json +profile.html +stacks.txt # user data /workspace* @@ -16,7 +19,7 @@ debug*.py /training_user_settings /external /update.var -config.json +config*.json secrets.json *.zip diff --git a/modules/trainer/BaseTrainer.py b/modules/trainer/BaseTrainer.py index bef29dd78..95b78e257 100644 --- a/modules/trainer/BaseTrainer.py +++ b/modules/trainer/BaseTrainer.py @@ -1,6 +1,3 @@ -import os -import subprocess -import sys from abc import ABCMeta, abstractmethod from modules.model.BaseModel import BaseModel @@ -23,8 +20,6 @@ class BaseTrainer( metaclass=ABCMeta, ): - tensorboard_subprocess: subprocess.Popen - def __init__(self, config: TrainConfig, callbacks: TrainCallbacks, commands: TrainCommands): super().__init__() self.config = config @@ -80,23 +75,3 @@ def create_model_sampler(self, model: BaseModel) -> BaseModelSampler: self.config.model_type, self.config.training_method ) - def _start_tensorboard(self): - tensorboard_executable = os.path.join(os.path.dirname(sys.executable), "tensorboard") - tensorboard_log_dir = os.path.join(self.config.workspace_dir, "tensorboard") - - tensorboard_args = [ - tensorboard_executable, - "--logdir", - tensorboard_log_dir, - "--port", - str(self.config.tensorboard_port), - "--samples_per_plugin=images=100,scalars=10000", - ] - - if self.config.tensorboard_expose: - tensorboard_args.append("--bind_all") - - self.tensorboard_subprocess = subprocess.Popen(tensorboard_args) - - def _stop_tensorboard(self): - self.tensorboard_subprocess.kill() diff --git a/modules/trainer/CloudTrainer.py b/modules/trainer/CloudTrainer.py index daa595892..7f8cae406 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 get_tensorboard_manager class CloudTrainer(BaseTrainer): @@ -29,8 +30,8 @@ def __init__(self, config: TrainConfig, callbacks: TrainCallbacks, commands: Tra tensorboard_log_dir = os.path.join(config.workspace_dir, "tensorboard") os.makedirs(Path(tensorboard_log_dir).absolute(), exist_ok=True) - if config.tensorboard and not config.cloud.tensorboard_tunnel and not config.tensorboard_always_on: - super()._start_tensorboard() + + get_tensorboard_manager().handle_training_start(config) match config.cloud.type: case CloudType.RUNPOD: @@ -113,8 +114,7 @@ def train(self): def end(self): try: - if self.config.tensorboard and not self.config.cloud.tensorboard_tunnel and not self.config.tensorboard_always_on: - super()._stop_tensorboard() + get_tensorboard_manager().handle_training_end(self.config) if self.config.cloud.delete_workspace and not self.error_caught and not self.commands.get_stop_command(): self.callbacks.on_update_status("Deleting remote workspace") diff --git a/modules/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index 13d1a7607..3436e7a84 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 @@ -70,9 +71,9 @@ def __init__(self, config: TrainConfig, callbacks: TrainCallbacks, commands: Tra if multi.is_master(): tensorboard_log_dir = os.path.join(config.workspace_dir, "tensorboard") os.makedirs(Path(tensorboard_log_dir).absolute(), exist_ok=True) - self.tensorboard = SummaryWriter(os.path.join(tensorboard_log_dir, f"{config.save_filename_prefix}{get_string_timestamp()}")) - if config.tensorboard and not config.tensorboard_always_on: - super()._start_tensorboard() + self.tensorboard = SummaryWriter(os.path.join(tensorboard_log_dir, f"{config.get_run_name()}_{get_string_timestamp()}")) + + get_tensorboard_manager().handle_training_start(config) self.model = None self.one_step_trained = False @@ -164,7 +165,7 @@ def start(self): def __save_config_to_workspace(self): path = path_util.canonical_join(self.config.workspace_dir, "config") os.makedirs(Path(path).absolute(), exist_ok=True) - path = path_util.canonical_join(path, f"{self.config.save_filename_prefix}{get_string_timestamp()}.json") + path = path_util.canonical_join(path, f"{self.config.get_run_name()}_{get_string_timestamp()}.json") with open(path, "w") as f: json.dump(self.config.to_pack_dict(secrets=False), f, indent=4) @@ -234,7 +235,7 @@ def __sample_loop( sample_path = os.path.join( sample_dir, - f"{self.config.save_filename_prefix}{get_string_timestamp()}-training-sample-{train_progress.filename_string()}" + f"{self.config.get_run_name()}_{get_string_timestamp()}-training-sample-{train_progress.filename_string()}" ) def on_sample_default(sampler_output: ModelSamplerOutput): @@ -431,7 +432,7 @@ def __backup(self, train_progress: TrainProgress, print_msg: bool = True, print_ self.callbacks.on_update_status("Creating backup") - backup_name = f"{get_string_timestamp()}-backup-{train_progress.filename_string()}" + backup_name = f"{get_string_timestamp()}-{self.config.get_run_name()}-backup-{train_progress.filename_string()}" backup_path = os.path.join(self.config.workspace_dir, "backup", backup_name) # Special case for schedule-free optimizers. @@ -481,7 +482,7 @@ def __save(self, train_progress: TrainProgress, print_msg: bool = True, print_cb save_path = os.path.join( self.config.workspace_dir, "save", - f"{self.config.save_filename_prefix}{get_string_timestamp()}-save-{train_progress.filename_string()}{self.config.output_model_format.file_extension()}" + f"{self.config.save_filename_prefix}{get_string_timestamp()}-{self.config.get_run_name()}-save-{train_progress.filename_string()}{self.config.output_model_format.file_extension()}" ) if print_msg: print_cb("Saving " + save_path) @@ -851,7 +852,7 @@ def end(self): if os.path.isdir(self.config.output_model_destination) and self.config.output_model_format.is_single_file(): save_path = os.path.join( self.config.output_model_destination, - f"{self.config.save_filename_prefix}{get_string_timestamp()}{self.config.output_model_format.file_extension()}" + f"{self.config.get_run_name()}_{get_string_timestamp()}{self.config.output_model_format.file_extension()}" ) else: save_path = self.config.output_model_destination @@ -870,9 +871,7 @@ def end(self): if multi.is_master(): self.tensorboard.close() - - if self.config.tensorboard and not self.config.tensorboard_always_on: - super()._stop_tensorboard() + get_tensorboard_manager().handle_training_end(self.config) for handle in self.grad_hook_handles: handle.remove() diff --git a/modules/ui/AdditionalEmbeddingsTab.py b/modules/ui/AdditionalEmbeddingsTab.py index d6cef6ca8..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): @@ -89,7 +90,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/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.") diff --git a/modules/ui/ConceptTab.py b/modules/ui/ConceptTab.py index a3806c15d..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 @@ -179,13 +180,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 +226,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/ConceptWindow.py b/modules/ui/ConceptWindow.py index 60b195c75..950f61461 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 @@ -91,6 +91,10 @@ 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._destroyed = False + self._last_bucket_data = None + self._last_chart_update = 0 self.title("Concept") self.geometry("800x700") @@ -147,14 +151,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': @@ -514,11 +518,16 @@ 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) + 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) @@ -531,6 +540,13 @@ 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) + 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") @@ -745,6 +761,13 @@ def __get_preview_image(self): return image, filename_output, prompt_output 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") @@ -830,37 +853,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}") @@ -879,30 +937,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") @@ -910,16 +973,34 @@ def __cancel_concept_stats(self): self.cancel_scan_flag.set() 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.__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: + except (KeyError, AttributeError, TypeError): 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): + self._destroyed = True + self.cancel_scan_flag.set() + + 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() diff --git a/modules/ui/ConfigList.py b/modules/ui/ConfigList.py index 75d69252a..7bf919482 100644 --- a/modules/ui/ConfigList.py +++ b/modules/ui/ConfigList.py @@ -9,7 +9,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 @@ -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} @@ -53,6 +55,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,30 +69,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() + 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 + @abstractmethod @@ -123,6 +138,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 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) @@ -159,6 +178,16 @@ 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=self.placeholder_text, + text_color="gray50" + ) + placeholder.grid(row=0, column=0, padx=20, pady=20) + def _create_element_list(self, **filters): if not self.from_external_file: self.current_config = getattr(self.train_config, self.attr_name) @@ -183,6 +212,9 @@ def _initialize_all_widgets(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, @@ -207,6 +239,21 @@ def _update_widget_visibility(self): else: widget.grid_remove() + 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): + return + + new_element = self.create_new_element() + + if hasattr(new_element, 'path'): + new_element.path = dir_path + + self.__add_element_to_list(new_element) + + ui_utils.register_concept_drop_target(self.master, self.__handle_dir_drop) + def __load_available_config_names(self): if os.path.isdir(self.config_dir): for path in os.listdir(self.config_dir): @@ -214,6 +261,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: @@ -230,9 +278,14 @@ def __create_config(self, name: str): def __add_config(self): dialogs.StringInputDialog(self.master, "name", "Name", self.__create_config) - def __add_element(self): - new_element = self.create_new_element() + def __add_element_to_list(self, new_element): + # placeholder removal + if len(self.current_config) == 0: + for child in self.element_list.winfo_children(): + child.destroy() + self.current_config.append(new_element) + # incremental insertion if widgets already initialized, else fall back to full rebuild if self.widgets_initialized and self.element_list is not None: i = len(self.current_config) - 1 @@ -250,6 +303,10 @@ def __add_element(self): self._create_element_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): new_element = copy.deepcopy(self.current_config[clone_i]) @@ -274,17 +331,24 @@ def __clone_element(self, clone_i, modify_element_fun=None): def __remove_element(self, remove_i): self.current_config.pop(remove_i) + if self.widgets_initialized and 0 <= remove_i < len(self.widgets): removed = self.widgets.pop(remove_i) with contextlib.suppress(tk.TclError, AttributeError): removed.destroy() - # Reindex remaining widgets - for idx, widget in enumerate(self.widgets): - widget.i = idx - self._update_widget_visibility() + + # Create placeholder if config is empty + if len(self.current_config) == 0: + self._create_placeholder() + else: + # Reindex remaining widgets + for idx, widget in enumerate(self.widgets): + widget.i = idx + self._update_widget_visibility() else: self.widgets_initialized = False self._create_element_list() + self.save_current_config() def __load_current_config(self, filename): diff --git a/modules/ui/ConvertModelUI.py b/modules/ui/ConvertModelUI.py index 837b27614..51940301b 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 902f2c8c2..bae748c1f 100644 --- a/modules/ui/LoraTab.py +++ b/modules/ui/LoraTab.py @@ -64,7 +64,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 278369e2b..c4e3d40c5 100644 --- a/modules/ui/ModelTab.py +++ b/modules/ui/ModelTab.py @@ -355,7 +355,7 @@ def __create_base_dtype_components(self, frame, row: int) -> int: # base model components.label(frame, row, 0, "Base Model", tooltip="Filename, directory or Hugging Face repository of the base model") - components.file_entry( + components.path_entry( frame, row, 1, self.ui_state, "base_model_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x ) @@ -400,7 +400,7 @@ def __create_base_components( # prior model components.label(frame, row, 0, "Prior Model", tooltip="Filename, directory or Hugging Face repository of the prior model") - components.file_entry( + components.path_entry( frame, row, 1, self.ui_state, "prior.model_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x ) @@ -418,7 +418,7 @@ def __create_base_components( # transformer model components.label(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( frame, row, 1, self.ui_state, "transformer.model_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x ) @@ -527,7 +527,7 @@ def __create_base_components( # text encoder 4 weight dtype components.label(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( 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 ) @@ -543,8 +543,8 @@ def __create_base_components( if has_vae: # base model components.label(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.") - components.file_entry( + 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.path_entry( frame, row, 1, self.ui_state, "vae.model_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x ) @@ -563,7 +563,7 @@ def __create_effnet_encoder_components(self, frame, row: int): # effnet encoder model components.label(frame, row, 0, "Effnet Encoder Model", tooltip="Filename, directory or Hugging Face repository of the effnet encoder model") - components.file_entry( + components.path_entry( frame, row, 1, self.ui_state, "effnet_encoder.model_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x ) @@ -587,7 +587,7 @@ def __create_decoder_components( # decoder model components.label(frame, row, 0, "Decoder Model", tooltip="Filename, directory or Hugging Face repository of the decoder model") - components.file_entry( + components.path_entry( frame, row, 1, self.ui_state, "decoder.model_name", path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x ) @@ -629,8 +629,16 @@ def __create_output_components( ) -> int: # output model destination components.label(frame, row, 0, "Model Output Destination", - tooltip="Filename or directory where the output model is saved") - components.file_entry(frame, row, 1, self.ui_state, "output_model_destination", is_output=True) + tooltip="Filename or directory where the *final* epoch output model is saved") + + components.path_entry( + 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", + ) # output data type components.label(frame, row, 3, "Output Data Type", @@ -656,7 +664,9 @@ def __create_output_components( components.label(frame, row, 0, "Output Format", tooltip="Format to use when saving the output model") - components.options_kv(frame, row, 1, formats, self.ui_state, "output_model_format") + components.options_kv( + frame, row, 1, formats, self.ui_state, "output_model_format" + ) # include config components.label(frame, row, 3, "Include Config", 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/ui/SampleFrame.py b/modules/ui/SampleFrame.py index 2f35a23a2..1d17d63b9 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, 6, 2, "mask image path:", tooltip="The mask used when inpainting.") - components.file_entry(bottom_frame, 6, 3, self.ui_state, "mask_image_path", + components.path_entry(bottom_frame, 6, 3, self.ui_state, "mask_image_path", allow_model_files=False, allow_image_files=True, ) diff --git a/modules/ui/SampleWindow.py b/modules/ui/SampleWindow.py index 6db4654b0..f1818c6a0 100644 --- a/modules/ui/SampleWindow.py +++ b/modules/ui/SampleWindow.py @@ -189,7 +189,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/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): diff --git a/modules/ui/TopBar.py b/modules/ui/TopBar.py index c53ea4160..13637f0f2 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) @@ -174,6 +176,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 584b68f81..572de5a9d 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -1,10 +1,8 @@ +import contextlib import ctypes import datetime import json -import os import platform -import subprocess -import sys import threading import time import traceback @@ -36,10 +34,13 @@ from modules.util.enum.GradientReducePrecision import GradientReducePrecision 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 get_tensorboard_manager 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 @@ -48,6 +49,16 @@ import customtkinter as ctk from customtkinter import AppearanceModeTracker +# 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 if platform.system() == "Windows": @@ -55,7 +66,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, _DnDMixin): set_step_progress: Callable[[int, int], None] set_epoch_progress: Callable[[int, int], None] @@ -64,6 +75,8 @@ class TrainUI(ctk.CTk): training_callbacks: TrainCallbacks | None training_commands: TrainCommands | None + _DND_ENABLED = TkinterDnD is not None + _TRAIN_BUTTON_STYLES = { "idle": { "text": "Start Training", @@ -92,6 +105,8 @@ class TrainUI(ctk.CTk): def __init__(self): super().__init__() + if self._DND_ENABLED: + self.TkdndVersion = TkinterDnD._require(self) self.title("OneTrainer") self.geometry("1100x740") @@ -104,6 +119,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) @@ -121,20 +137,51 @@ 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 self.training_commands = None - self.always_on_tensorboard_subprocess = None self.current_workspace_dir = self.train_config.workspace_dir - self._check_start_always_on_tensorboard() + self.tensorboard_manager = get_tensorboard_manager() + + # Initialize tensorboard if needed + 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 + 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", + 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) @@ -143,9 +190,15 @@ 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'): + 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): @@ -214,21 +267,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) @@ -239,12 +364,13 @@ 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.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") + 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", @@ -263,24 +389,24 @@ 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.path_entry(frame, 4, 3, self.ui_state, "debug_dir", is_output=True, path_type="directory") # tensorboard - components.label(frame, 6, 0, "Tensorboard", - tooltip="Starts the Tensorboard Web UI during training") - components.switch(frame, 6, 1, self.ui_state, "tensorboard") - - components.label(frame, 6, 2, "Always-On Tensorboard", - tooltip="Keep Tensorboard accessible even when not training. Useful for monitoring completed training sessions.") - components.switch(frame, 6, 3, self.ui_state, "tensorboard_always_on", command=self._on_always_on_tensorboard_toggle) - - 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, 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, 7, 3, self.ui_state, "tensorboard_port") + components.entry(frame, 7, 1, self.ui_state, "tensorboard_port") + 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, 7, 3, self.ui_state, "tensorboard_expose") # validation components.label(frame, 8, 0, "Validation", @@ -291,45 +417,66 @@ 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") + + 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, 10, 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, 10, 1, self.ui_state, "dataloader_threads") + components.entry(frame, 11, 3, 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, "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 @@ -479,7 +626,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 ) @@ -526,7 +673,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 ) @@ -740,6 +887,10 @@ def open_sample_ui(self): def __training_thread_function(self): error_caught = False + # 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, on_update_status=self.on_update_status, @@ -758,37 +909,40 @@ 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 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() + self.tensorboard_manager.handle_training_end(self.train_config) - # 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) - - if self.train_config.tensorboard_always_on and not self.always_on_tensorboard_subprocess: - 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.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() - 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() torch_gc() @@ -801,9 +955,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=[ @@ -830,70 +987,30 @@ 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: - self._start_always_on_tensorboard() - - def _start_always_on_tensorboard(self): - if self.always_on_tensorboard_subprocess: - self._stop_always_on_tensorboard() - - tensorboard_executable = os.path.join(os.path.dirname(sys.executable), "tensorboard") - tensorboard_log_dir = os.path.join(self.train_config.workspace_dir, "tensorboard") + 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) - os.makedirs(Path(tensorboard_log_dir).absolute(), exist_ok=True) - - tensorboard_args = [ - tensorboard_executable, - "--logdir", - tensorboard_log_dir, - "--port", - str(self.train_config.tensorboard_port), - "--samples_per_plugin=images=100,scalars=10000", - ] + def _on_workspace_dir_change_actual(self): + if (new := self.train_config.workspace_dir) == self.current_workspace_dir: + return + self.current_workspace_dir = new + self.tensorboard_manager.handle_workspace_change(self.train_config) - if self.train_config.tensorboard_expose: - tensorboard_args.append("--bind_all") + def _on_tensorboard_mode_change(self): + self.tensorboard_manager.handle_mode_change( + self.train_config, + is_training=self.training_thread is not None + ) - try: - self.always_on_tensorboard_subprocess = subprocess.Popen(tensorboard_args) - except Exception: - self.always_on_tensorboard_subprocess = None - - def _stop_always_on_tensorboard(self): - 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 - - 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 - - if self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: - self._start_always_on_tensorboard() - - def _on_workspace_dir_change_trace(self, *args): - 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() + 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: 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/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: 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 ddaee4b89..d21089f05 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 @@ -23,6 +24,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 @@ -360,7 +362,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 @@ -370,6 +372,13 @@ 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 + prevent_overwrite: bool + auto_prefix: bool + # multi-GPU multi_gpu: bool device_indexes: str @@ -569,7 +578,7 @@ class TrainConfig(BaseConfig): def __init__(self, data: list[(str, Any, type, bool)]): super().__init__( data, - config_version=10, + config_version=11, config_migrations={ 0: self.__migration_0, 1: self.__migration_1, @@ -581,6 +590,7 @@ def __init__(self, data: list[(str, Any, type, bool)]): 7: self.__migration_7, 8: self.__migration_8, 9: self.__migration_9, + 10: self.__migration_10, } ) @@ -761,6 +771,7 @@ 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: @@ -800,6 +811,21 @@ def replace_dtype(part: str): return migrated_data + def __migration_10(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) + + 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, @@ -878,6 +904,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): @@ -945,7 +989,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.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)) @@ -955,6 +999,13 @@ 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)) + data.append(("prevent_overwrite", True, bool, False)) + data.append(("auto_prefix", False, bool, False)) + #multi-GPU data.append(("multi_gpu", False, bool, False)) data.append(("device_indexes", "", str, False)) diff --git a/modules/util/enum/TensorboardMode.py b/modules/util/enum/TensorboardMode.py new file mode 100644 index 000000000..f350e00ab --- /dev/null +++ b/modules/util/enum/TensorboardMode.py @@ -0,0 +1,10 @@ +try: + from enum import StrEnum +except ImportError: + from strenum import StrEnum + + +class TensorboardMode(StrEnum): + OFF = "OFF" + TRAIN_ONLY = "TRAIN_ONLY" + ALWAYS_ON = "ALWAYS_ON" diff --git a/modules/util/tensorboard_util.py b/modules/util/tensorboard_util.py new file mode 100644 index 000000000..8b77e6998 --- /dev/null +++ b/modules/util/tensorboard_util.py @@ -0,0 +1,186 @@ +import contextlib +import os +import subprocess +import sys +import threading +from pathlib import Path + +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.TensorboardMode import TensorboardMode + + +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 + 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) + return self._process is not None + + def stop(self) -> bool: + if self._process is None: + return False + + stop_tensorboard(self._process) + self._process = None + return True + + def restart(self, config: TrainConfig) -> bool: + self.stop() + return self.start(config) + + def is_running(self) -> bool: + return self._process is not None + + def handle_training_start(self, config: TrainConfig): + self._sync_state(config, is_training=True) + + def handle_training_end(self, config: TrainConfig): + self._sync_state(config, is_training=False) + + def handle_mode_change(self, config: TrainConfig, is_training: bool = False): + self._sync_state(config, is_training=is_training) + + def handle_workspace_change(self, config: TrainConfig): + if config.tensorboard_mode == TensorboardMode.ALWAYS_ON: + 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]: + 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, filtering the annoying warnings Google doesnt plan to fix""" + 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 + + 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): + if not process: + return + + try: + print("Stopping Tensorboard...") + process.terminate() + + 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): + if process.stdout: + process.stdout.close() + + +_tensorboard_manager = TensorboardManager() + + +def get_tensorboard_manager() -> TensorboardManager: + return _tensorboard_manager diff --git a/modules/util/ui/ToolTip.py b/modules/util/ui/ToolTip.py index be80e2765..dd2af0f1a 100644 --- a/modules/util/ui/ToolTip.py +++ b/modules/util/ui/ToolTip.py @@ -1,56 +1,170 @@ +from __future__ import annotations + +import contextlib + import customtkinter as ctk class ToolTip: - """ - create a tooltip for a given widget - """ + MAX_TOOLTIP_LENGTH = 350 + _active_tooltip: ToolTip | None = None # currently visible tooltip - 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 + self._root = None + self._bound = False + + if hover_only: + self.widget.bind("", self.enter) + self.widget.bind("", self.leave) + self.widget.bind("", self.leave) 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 _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 "") + + 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="+") + self._root.bind("", self.hide, add="+") + self._root.bind("", self.hide, add="+") + self._bound = True - 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) - 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) + if self.track_movement: + self.tw.attributes("-topmost", True) + + frame = ctk.CTkFrame(self.tw, fg_color=fg_color, corner_radius=6) + frame.pack() + + 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 - def hidetip(self): - tw = self.tw + 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() + + # 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() + 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) + + # 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(): + self._position() + + def hide(self, _evt=None): + self._cancel_scheduled_hide() + + if self.tw and self.tw.winfo_exists(): + self.tw.destroy() self.tw = None - if tw: - tw.destroy() + + 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) + + 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/UIState.py b/modules/util/ui/UIState.py index 0c858094b..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} @@ -77,22 +81,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 @@ -306,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 961d213a0..0648ee50d 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/components.py @@ -1,12 +1,19 @@ import contextlib +import tkinter as tk from collections.abc import Callable -from tkinter import filedialog -from typing import Any +from dataclasses import dataclass +from pathlib import Path +from tkinter import filedialog, messagebox +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 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 import customtkinter as ctk from customtkinter.windows.widgets.scaling import CTkScalingBaseClass @@ -14,29 +21,453 @@ PAD = 10 +MODEL_FILETYPES = [ + ("All Files", "*.*"), + ("Diffusers", "model_index.json"), + ("Checkpoint", "*.ckpt *.pt *.bin"), + ("Safetensors", "*.safetensors"), +] -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) +def _safe_bool(var, default=True): + try: + return bool(var.get()) + except Exception: + return default - 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 _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)) -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 +@dataclass(frozen=True) +class ComponentValidationSettings: + debounce_stop_typing_ms: int = 1700 + + +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 + message: str = "" + + def clear(self): + self.status, self.message = None, "" + + def set_status(self, status: Literal['error', 'warning'], message: str): + self.status, self.message = status, message + + +class EntryValidationHandler: + + 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 + + try: + self.original_border_color = component.cget("border_color") + except Exception: + self.original_border_color = "gray50" + + # 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()) + ) + + self.touched = False + + self.validation_tooltip = ToolTip(component, text="", hover_only=False, track_movement=True, wide=True) + component._validation_tooltip = self.validation_tooltip + + def _reset_border(self): + self.component.configure(border_color=self.original_border_color) + + def should_show_tooltip(self) -> bool: + 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: + 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': + self.validation_tooltip.show_warning(self.validation_state.message, duration_ms=None) + + def validate_value(self, value: str) -> bool: + meta = self.ui_state.get_field_metadata(self.var_name) + + basic_result = validate_basic_type( + value=value, + declared_type=meta.type, + nullable=meta.nullable, + default_val=meta.default + ) + + if not basic_result.ok: + return self._fail(basic_result.message) + + 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) + + return self._success(value) + + def _success(self, value: str) -> bool: + self._reset_border() + self.validation_tooltip.hide() + 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) + + if self.validation_state: + self.validation_state.set_status('warning', reason) + + return True + + def _fail(self, reason: str) -> bool: + self.component.configure(border_color="#dc3545") + if self.should_show_tooltip(): + self.validation_tooltip.show_error(reason) + + if self.validation_state: + self.validation_state.set_status('error', reason) + + return False + + def debounced_validate(self, *_): + if not self.touched: + return + + self.validation_tooltip.hide() + self._reset_border() + + self.debounce_timer.call() + + def on_focus_in(self, _e=None): + self.touched = False + self._show_validation_tooltip() + + def on_user_input(self, _e=None): + self.touched = True + + def on_focus_out(self, _e=None): + self.validation_tooltip.hide() + + if self.touched: + self.validate_value(self.var.get()) + + def cleanup(self): + self.debounce_timer.cancel() + with contextlib.suppress(Exception): + self.validation_tooltip.hide() + self.debounce_timer = None + self.validation_tooltip = None + + +class ModelOutputValidator: + + 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.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] = {} + + # 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 "" + + 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') + + 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() + + # 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 + finally: + self._in_validation = False + + 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 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_timer.call() + + def _schedule_prefix_change_validation(self): + if not self._tk_widget: + return + + 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) + + 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 = [ + ('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), + ] + + def make_callback(debounce: bool, skip_overwrite: bool, is_prefix: bool): + 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)) + + def cleanup_traces(self): + for key, trace_id in self._trace_ids.items(): + var = getattr(self, f"{key}_var", None) + if var: + with contextlib.suppress(Exception): + 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") + elif self.state.status == 'warning': + self._tk_widget.configure(border_color="#ff9500") + + 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 def entry( master, @@ -49,7 +480,12 @@ def entry( 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, + undo: bool = True, ): + var = ui_state.get_var(var_name) trace_id = None if command: @@ -58,134 +494,71 @@ def entry( 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" + # set up undo/redo handler if enabled, its enabled by default + undo_handler = EntryUndoHandler(component, var) if undo else None - error_border_color = "#dc3545" + if validation_state is None: + validation_state = ValidationState() - validation_after_id = None - revert_after_id = None - touched = False - - DEBOUNCE_STOP_TYPING_MS = 1500 - DEBOUNCED_INVALID_REVERT_MS = 1000 - FOCUSOUT_INVALID_REVERT_MS = 1200 - - last_valid_value = var.get() - - 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 - - if revert_after_id: - with contextlib.suppress(Exception): - component.after_cancel(revert_after_id) - revert_after_id = None + ui_state.register_validation_state(var_name, validation_state) # register global state - def success(): - nonlocal last_valid_value - component.configure(border_color=original_border_color) - last_valid_value = value - return True - - def do_revert(): - var.set(last_valid_value) - component.configure(border_color=original_border_color) - - 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 + validation_handler = EntryValidationHandler( + component=component, + var=var, + var_name=var_name, + ui_state=ui_state, + 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) - if value == "": - if nullable: - return success() - if declared_type is str: - if default_val == "": - return success() - return fail("Value required") + 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: 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") - 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: - with contextlib.suppress(Exception): - component.after_cancel(validation_after_id) - validation_after_id = None - return - if revert_after_id: - with contextlib.suppress(Exception): - component.after_cancel(revert_after_id) - revert_after_id = None - if 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) - ) + show_tooltips_var = ui_state.get_var("validation_show_tooltips") + except (KeyError, AttributeError): + show_tooltips_var = None - validation_trace_name = var.trace_add("write", debounced_validate) + validation_tooltip = validation_handler.validation_tooltip - def on_focus_in(_e=None): - nonlocal touched - touched = False + def on_hover_enter(_e=None): + if show_tooltips_var and not _safe_bool(show_tooltips_var): + return - def on_user_input(_e=None): - nonlocal touched - touched = True + 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_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) + def on_hover_leave(_e=None): + validation_tooltip.hide() - 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) + component.bind("", on_hover_enter, add="+") + component.bind("", on_hover_leave, add="+") 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) + 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) @@ -199,80 +572,190 @@ def new_destroy(): return component -def file_entry( - master, row, column, ui_state: UIState, var_name: str, - 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, - command: Callable[[str], None] = None, + 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") - frame.grid_columnconfigure(0, weight=1) - entry(frame,row=0, column=0, ui_state=ui_state, var_name=var_name) - - def __open_dialog(): - 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: - file_path = filedialog.asksaveasfilename(filetypes=filetypes) - else: - file_path = filedialog.askopenfilename(filetypes=filetypes) - - if file_path: - if path_modifier: - file_path = path_modifier(file_path) + if not is_output: + meta = ui_state.get_field_metadata(var_name) + is_output = getattr(meta, 'is_output', False) - ui_state.get_var(var_name).set(file_path) + var = ui_state.get_var(var_name) - if command: - command(file_path) + 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 = ValidationState() # Create state for simple path validation + + 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 + 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 - 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") + custom_validator = simple_path_validator - return frame + 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, + validation_state=validation_state, + enable_hover_validation=bool(validator), + command=command + ) -def dir_entry(master, row, column, ui_state: UIState, var_name: str, command: Callable[[str], None] = None): - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") + 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 - frame.grid_columnconfigure(0, weight=1) + original_destroy = entry_widget.destroy - entry(frame, row=0, column=0, ui_state=ui_state, var_name=var_name) + 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() - def __open_dialog(): - dir_path = filedialog.askdirectory() + entry_widget.destroy = new_destroy + + 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(): + if path_type == "directory": + selected_path = filedialog.askdirectory() + else: + filetypes = [("All Files", "*.*")] + if allow_model_files: + filetypes.extend(MODEL_FILETYPES[1:]) + if allow_image_files: + filetypes.append(("Image", ' '.join(f"*.{x}" for x in supported_image_extensions()))) - if dir_path: - ui_state.get_var(var_name).set(dir_path) + 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(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 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) + ) + 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) + + +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") @@ -426,18 +909,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 @@ -458,18 +930,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} @@ -477,10 +938,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 @@ -496,14 +955,14 @@ 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(): 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) @@ -512,18 +971,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 @@ -544,19 +992,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 diff --git a/modules/util/ui/ui_utils.py b/modules/util/ui/ui_utils.py index ba09e1038..5dbfd89f0 100644 --- a/modules/util/ui/ui_utils.py +++ b/modules/util/ui/ui_utils.py @@ -1,13 +1,16 @@ import contextlib +import os import platform import sys import tkinter as tk +import urllib.parse from collections.abc import Callable from pathlib import Path from tkinter import EventType from typing import Any from customtkinter import CTk, CTkToplevel +from tkinterdnd2 import DND_FILES def bind_mousewheel( @@ -101,6 +104,117 @@ 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, drop_validator=None, on_reject=None): + def drop(event): + if event.data: + paths = _parse_dropped_paths(event.data) + if paths: + dropped_path = paths[0] + + 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('') + entry_widget.event_generate('') + + if command: + command(dropped_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 + elif char == ' ' and not in_braces: + pass # Will append below + else: + current_path += char + continue + + if current_path: + paths.append(current_path) + current_path = "" + + if current_path: + paths.append(current_path) + + + 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): + # 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), + ('<>', _create_drop_handler(entry_widget, ui_state, var_name, command, drop_validator, on_reject))]: + entry_widget.dnd_bind(event, handler) + 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): + if sys.platform == "linux": + return + + def drop_handler(event): + if not event.data: + return event.action + + paths = _parse_dropped_paths(event.data) + + for path in paths: + 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 + + 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 as e: + print(f"Failed to register concept drop target: {e}") + class DebounceTimer: def __init__(self, widget, delay_ms: int, callback: Callable[..., Any]): self.widget = widget @@ -119,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 new file mode 100644 index 000000000..da6cbb85b --- /dev/null +++ b/modules/util/ui/validation.py @@ -0,0 +1,482 @@ +"""Input validation utilities for UI components.""" + +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() + +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"} +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_.-]+$") + +INVALID_CHARS_WIN = set('<>:"|?*') +INVALID_NAMES = {"", "."} + +ValidationStatus = Literal['success', 'warning', 'error'] + +@dataclass +class ValidationResult: + 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' + +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: + 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 _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 + + if " " in trimmed or "\t" in trimmed: + return False + + if "—" in trimmed or ".." 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() + 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: + 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: + return str(path).lower().endswith(extension.lower()) + +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 _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 = "", + extension: str = ".safetensors", + use_friendly_names: bool = False, +) -> str: + 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 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} + return {ModelFormat.SAFETENSORS, ModelFormat.CKPT, ModelFormat.DIFFUSERS} + +@dataclass +class ValidationContext: + raw: str + user_input: str + output_format: ModelFormat + training_method: TrainingMethod + autocorrect: bool + prefix: str + use_friendly_names: bool + is_output: bool + separator: str + prevent_overwrite: bool + auto_prefix: bool + + +def _validate_basic_input(ctx: ValidationContext) -> ValidationResult | None: + """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.") + + if not ctx.user_input or ctx.user_input in INVALID_NAMES: + return _result(False, "", "Input cleared: empty or invalid.") + + if _has_invalid_chars(ctx.user_input): + return _result(False, None, "Input contains invalid path characters.") + + 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.") + + 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 (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.") + + sanitized = ctx.user_input[:ext_match.start()] if ext_match else ctx.user_input + check_path = Path(sanitized) + + if err := _check_parent_exists(check_path): + return err + + 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') + + 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 formats (SAFETENSORS, CKPT).""" + required_ext = ctx.output_format.file_extension() + + is_dir_like = bool(TRAILING_SLASH_RE.search(ctx.user_input)) + path = Path(ctx.user_input) + + 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 + + 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.") + + if err := _check_parent_exists(path): + return err + + base = path.name + + 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.") + + # 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, "") + + 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() + + # 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: + 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: + 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}") + + 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.") + + if ctx.autocorrect: + # 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 + 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, + prevent_overwrite: bool = False, + auto_prefix: bool = False, + skip_overwrite_protection: bool = False, +) -> 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 + - Automatic overwrite prevention + + 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) + 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 + """ + 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 '\\', + prevent_overwrite=prevent_overwrite and not skip_overwrite_protection, + auto_prefix=auto_prefix, + ) + + if basic_check := _validate_basic_input(ctx): + return basic_check + + 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).""" + # 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.""" + 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.") + + path = Path(value) + + if path_type == "directory": + if is_output: + if err := _check_parent_exists(path): + return err + 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 err := _check_parent_exists(path): + return err + else: + if not path.exists(): + return _result(False, None, "File does not exist") + + return _result(True, None, "") diff --git a/requirements-global.txt b/requirements-global.txt index 76fbdaafa..1738d55b4 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -8,14 +8,15 @@ 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 # 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 # diffusion models @@ -48,10 +49,12 @@ adv_optm==1.4.0 # advanced optimizers -e git+https://github.com/KellerJordan/Muon.git@f90a42b#egg=muon-optimizer # Profiling -scalene==1.5.51 +scalene==1.5.55 # ui customtkinter==5.2.2 +tkinterdnd2==0.4.3 +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: diff --git a/start-ui.bat b/start-ui.bat index 0f859af59..6a996d11d 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 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 @@ -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% )